mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge branch 'litellm_internal_staging' into feature/improve-gigachat-provider
This commit is contained in:
commit
e7fc00d724
530 changed files with 24877 additions and 3598 deletions
19
.cargo/config.toml
Normal file
19
.cargo/config.toml
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
[http]
|
||||
# CI has seen transient crates.io failures from libcurl's HTTP/2 multiplexing
|
||||
# during `maturin` metadata resolution. Disable multiplexing and retry more
|
||||
# aggressively so editable `uv sync` builds are not failed by one flaky frame.
|
||||
multiplexing = false
|
||||
|
||||
[net]
|
||||
retry = 5
|
||||
|
||||
# PyO3 cdylib (`litellm-python-bridge`) links against the host interpreter's
|
||||
# symbols, which are not present at link time when building an extension module.
|
||||
# On macOS, tell the linker to resolve undefined `_Py*` symbols dynamically at
|
||||
# load time (the standard pyo3 extension-module flag) so the cdylib links without
|
||||
# a libpython on the link line.
|
||||
[target.x86_64-apple-darwin]
|
||||
rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"]
|
||||
|
||||
[target.aarch64-apple-darwin]
|
||||
rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"]
|
||||
|
|
@ -190,6 +190,8 @@ jobs:
|
|||
working_directory: ~/project
|
||||
environment:
|
||||
UV_PYTHON: "3.11"
|
||||
CARGO_HTTP_MULTIPLEXING: "false"
|
||||
CARGO_NET_RETRY: "5"
|
||||
steps:
|
||||
- checkout
|
||||
- run:
|
||||
|
|
@ -205,6 +207,24 @@ jobs:
|
|||
environment:
|
||||
UV_HTTP_TIMEOUT: "300"
|
||||
command: |
|
||||
$rustupInit = Join-Path $env:TEMP "rustup-init.exe"
|
||||
$rustupVersion = "1.28.2"
|
||||
$rustupUrl = "https://static.rust-lang.org/rustup/archive/$rustupVersion/x86_64-pc-windows-msvc/rustup-init.exe"
|
||||
Invoke-WebRequest -Uri $rustupUrl -OutFile $rustupInit
|
||||
$rustupExpected = "88d8258dcf6ae4f7a80c7d1088e1f36fa7025a1cfd1343731b4ee6f385121fc0"
|
||||
$rustupActual = (Get-FileHash -Path $rustupInit -Algorithm SHA256).Hash.ToLower()
|
||||
if ($rustupActual -ne $rustupExpected) {
|
||||
throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual"
|
||||
}
|
||||
& $rustupInit -y --profile minimal --default-toolchain stable
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
exit $LASTEXITCODE
|
||||
}
|
||||
Remove-Item $rustupInit
|
||||
$cargoBin = Join-Path $HOME ".cargo\bin"
|
||||
$env:Path = "$cargoBin;$env:Path"
|
||||
rustc --version
|
||||
cargo --version
|
||||
$installer = Join-Path $env:TEMP "uv-install.ps1"
|
||||
Invoke-WebRequest -Uri https://astral.sh/uv/0.10.9/install.ps1 -OutFile $installer
|
||||
$expected = "d43ffff8d28e7d1e7d1831a212465f12b24a43c7f87f386e95e2a5915aee5d7d"
|
||||
|
|
@ -222,7 +242,20 @@ jobs:
|
|||
if (-not (Select-String -Path $PROFILE -SimpleMatch $uvBin -Quiet)) {
|
||||
Add-Content -Path $PROFILE -Value "`$env:Path = `"$uvBin;`$env:Path`""
|
||||
}
|
||||
uv sync --frozen --group dev --python 3.11
|
||||
if (-not (Select-String -Path $PROFILE -SimpleMatch $cargoBin -Quiet)) {
|
||||
Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`""
|
||||
}
|
||||
for ($attempt = 1; $attempt -le 5; $attempt++) {
|
||||
Write-Host "uv sync attempt $attempt/5"
|
||||
uv sync --frozen --group dev --python 3.11
|
||||
if ($LASTEXITCODE -eq 0) {
|
||||
break
|
||||
}
|
||||
if ($attempt -eq 5) {
|
||||
exit $LASTEXITCODE
|
||||
}
|
||||
Start-Sleep -Seconds 15
|
||||
}
|
||||
- run:
|
||||
name: Run Windows-specific test
|
||||
command: |
|
||||
|
|
@ -232,6 +265,9 @@ jobs:
|
|||
environment:
|
||||
UV_HTTP_TIMEOUT: "300"
|
||||
command: |
|
||||
$env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path"
|
||||
cargo --version
|
||||
Get-ChildItem -Path "litellm\rust_bridge" -Filter "_native*" -File -ErrorAction SilentlyContinue | Remove-Item -Force
|
||||
uv build --wheel --out-dir dist
|
||||
uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py
|
||||
|
||||
|
|
@ -1020,7 +1056,9 @@ jobs:
|
|||
name: Run tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
|
||||
TEST_FILES=$(printf "%s\n%s\n" \
|
||||
"$(circleci tests glob "tests/ocr_tests/**/test_*.py")" \
|
||||
"tests/test_litellm/ocr/test_rust_bridge.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
|
|
|
|||
|
|
@ -49,6 +49,8 @@ build/
|
|||
*.egg-info/
|
||||
.DS_Store
|
||||
**/node_modules
|
||||
litellm-rust/target/
|
||||
litellm/rust_bridge/_native*.so
|
||||
*.log
|
||||
.env
|
||||
.env.local
|
||||
|
|
|
|||
32
.github/pull_request_template.md
vendored
32
.github/pull_request_template.md
vendored
|
|
@ -1,17 +1,17 @@
|
|||
## Relevant issues
|
||||
|
||||
<!-- e.g. "Fixes #000" -->
|
||||
<!-- e.g., "Fixes #000" -->
|
||||
|
||||
## Linear ticket
|
||||
|
||||
<!-- if you are an internal contributor (e.g., your username is postfixed with -berri or -berriai), add "Resolves " followed by the Linear ticket e.g. "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
|
||||
<!-- if you are an internal contributor (e.g., your username is postfixed with -berri or -berriai), add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
|
||||
|
||||
## Pre-Submission checklist
|
||||
|
||||
**Please complete all items before asking a LiteLLM maintainer to review your PR**
|
||||
|
||||
- [ ] I have added meaningful tests
|
||||
- [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code)
|
||||
- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests)
|
||||
- [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem
|
||||
- [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review
|
||||
|
||||
|
|
@ -19,29 +19,13 @@
|
|||
|
||||
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slack (#pr-review)](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA).
|
||||
|
||||
## CI (LiteLLM team)
|
||||
|
||||
> **CI status guideline:**
|
||||
>
|
||||
> - 50-55 passing tests: main is stable with minor issues.
|
||||
> - 45-49 passing tests: acceptable but needs attention
|
||||
> - <= 40 passing tests: unstable; be careful with your merges and assess the risk.
|
||||
|
||||
- [ ] **Branch creation CI run**
|
||||
Link:
|
||||
|
||||
- [ ] **CI run for the last commit**
|
||||
Link:
|
||||
|
||||
- [ ] **Merge / cherry-pick CI run**
|
||||
Links:
|
||||
|
||||
## Screenshots / Proof of Fix
|
||||
|
||||
<!-- Include screenshots, screen recordings, or log output demonstrating that your changes work as expected.
|
||||
For bug fixes: show reproduction before the fix and passing behavior after.
|
||||
For new features: show the feature working end-to-end.
|
||||
For UI changes: include before/after screenshots. -->
|
||||
<!-- Include screenshots, screen recordings, or command (e.g., curl) + output demonstrating that your changes work as expected
|
||||
The proof must be completely e2e with no mocks, using, for example, actual LLM calls costing real $. `pytest` commands are not enough
|
||||
For bug fixes: show reproduction before the fix and passing behavior after
|
||||
For new features: show the feature working end-to-end
|
||||
For UI changes: include before/after screenshots -->
|
||||
|
||||
## Type
|
||||
|
||||
|
|
|
|||
31
.github/scripts/uv_sync_with_retries.sh
vendored
Executable file
31
.github/scripts/uv_sync_with_retries.sh
vendored
Executable file
|
|
@ -0,0 +1,31 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
max_attempts="${UV_SYNC_MAX_ATTEMPTS:-5}"
|
||||
delay_seconds="${UV_SYNC_RETRY_DELAY_SECONDS:-15}"
|
||||
|
||||
export CARGO_HTTP_MULTIPLEXING="${CARGO_HTTP_MULTIPLEXING:-false}"
|
||||
export CARGO_NET_RETRY="${CARGO_NET_RETRY:-5}"
|
||||
|
||||
if [[ "$#" -eq 0 ]]; then
|
||||
echo "usage: $0 <uv sync args...>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
for attempt in $(seq 1 "${max_attempts}"); do
|
||||
echo "uv sync attempt ${attempt}/${max_attempts}"
|
||||
status=0
|
||||
if uv sync "$@"; then
|
||||
exit 0
|
||||
else
|
||||
status=$?
|
||||
fi
|
||||
|
||||
if [[ "${attempt}" -eq "${max_attempts}" ]]; then
|
||||
echo "uv sync failed after ${max_attempts} attempts" >&2
|
||||
exit "${status}"
|
||||
fi
|
||||
|
||||
echo "uv sync failed; retrying in ${delay_seconds}s..."
|
||||
sleep "${delay_seconds}"
|
||||
done
|
||||
2
.github/workflows/_test-unit-base.yml
vendored
2
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -73,7 +73,7 @@ jobs:
|
|||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
|
|
|
|||
2
.github/workflows/check-ui-api-types.yml
vendored
2
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -46,7 +46,7 @@ jobs:
|
|||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Install backend dependencies
|
||||
run: uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- "uv.lock"
|
||||
|
|
|
|||
4
.github/workflows/guard-main-branch.yml
vendored
4
.github/workflows/guard-main-branch.yml
vendored
|
|
@ -31,12 +31,12 @@ jobs:
|
|||
echo "PR head repo: $HEAD_REPO"
|
||||
echo "PR head branch: $HEAD_REF"
|
||||
if [ "$HEAD_REPO" != "$BASE_REPO" ]; then
|
||||
echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the 'litellm_oss_branch' branch instead."
|
||||
echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the 'litellm_oss_staging' branch instead."
|
||||
exit 1
|
||||
fi
|
||||
if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then
|
||||
echo "Allowed source branch."
|
||||
exit 0
|
||||
fi
|
||||
echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_oss_branch' instead."
|
||||
echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_oss_staging' instead."
|
||||
exit 1
|
||||
|
|
|
|||
65
.github/workflows/image-scan.yml
vendored
Normal file
65
.github/workflows/image-scan.yml
vendored
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
name: Image Scan
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- docker/Dockerfile.non_root
|
||||
- uv.lock
|
||||
- ui/litellm-dashboard/package-lock.json
|
||||
- .github/workflows/image-scan.yml
|
||||
schedule:
|
||||
- cron: "41 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
image-scan:
|
||||
name: image-scan
|
||||
runs-on: ubuntu-latest
|
||||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Download Grype v0.114.0
|
||||
run: |
|
||||
curl -fsSL --retry 3 -o "$RUNNER_TEMP/grype.tar.gz" \
|
||||
https://github.com/anchore/grype/releases/download/v0.114.0/grype_0.114.0_linux_amd64.tar.gz
|
||||
echo "edda0968d8827daab01d32b3cd7de192ae0915005e7bbfcfef9e68e79bc43343 $RUNNER_TEMP/grype.tar.gz" | sha256sum -c -
|
||||
tar xzf "$RUNNER_TEMP/grype.tar.gz" -C "$RUNNER_TEMP" grype
|
||||
chmod +x "$RUNNER_TEMP/grype"
|
||||
|
||||
# Dockerfile.non_root is the rootless variant we ship. The other
|
||||
# Dockerfiles share the same wolfi base and apk set, so OS-layer coverage
|
||||
# is the same; matrix-scan if those variants ever diverge.
|
||||
- name: Build runtime image
|
||||
run: docker build -f docker/Dockerfile.non_root -t litellm-image-scan:${{ github.sha }} .
|
||||
|
||||
# Scans the whole shipped artifact: OS/apk plus every language package
|
||||
# baked into the image, including ones no lockfile declares (e.g. prisma's
|
||||
# vendored node engine) that osv-scan cannot see. osv-scan stays the fast
|
||||
# source-level gate; this is the customer's-eye-view backstop. Credential-
|
||||
# free OSS, run as a pinned, checksum-verified binary; no GitHub Action
|
||||
# dependency and no vendor SaaS callout.
|
||||
- name: Scan image for fixable HIGH/CRITICAL CVEs
|
||||
run: |
|
||||
"$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \
|
||||
--only-fixed \
|
||||
--fail-on high \
|
||||
--output table
|
||||
2
.github/workflows/mutation-test.yml
vendored
2
.github/workflows/mutation-test.yml
vendored
|
|
@ -55,7 +55,7 @@ jobs:
|
|||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
|
|
|
|||
2
.github/workflows/osv-scan.yml
vendored
2
.github/workflows/osv-scan.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
schedule:
|
||||
- cron: "23 6 * * *"
|
||||
|
|
|
|||
2
.github/workflows/test-code-quality.yml
vendored
2
.github/workflows/test-code-quality.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
15
.github/workflows/test-linting.yml
vendored
15
.github/workflows/test-linting.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
@ -50,11 +50,16 @@ jobs:
|
|||
run: |
|
||||
uv sync --frozen
|
||||
|
||||
- name: Check Black formatting
|
||||
- name: Check ruff format
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
cd litellm
|
||||
uv run --no-sync black --check --exclude '/enterprise/' .
|
||||
cd ..
|
||||
git diff --name-only "$BASE_SHA"...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
|
||||
if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
|
||||
echo "No changed litellm Python files to check with ruff format."
|
||||
exit 0
|
||||
fi
|
||||
xargs uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt"
|
||||
|
||||
- name: Debug - Check file state
|
||||
run: |
|
||||
|
|
|
|||
4
.github/workflows/test-litellm-ui-build.yml
vendored
4
.github/workflows/test-litellm-ui-build.yml
vendored
|
|
@ -7,7 +7,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
jobs:
|
||||
|
|
@ -111,4 +111,4 @@ jobs:
|
|||
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
|
||||
run: |
|
||||
npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true
|
||||
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json
|
||||
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json --check eslint-metrics.json
|
||||
|
|
|
|||
4
.github/workflows/test-mcp.yml
vendored
4
.github/workflows/test-mcp.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
@ -39,7 +39,7 @@ jobs:
|
|||
- name: Install dependencies
|
||||
run: |
|
||||
uv lock --check
|
||||
uv sync --frozen --group proxy-dev --extra proxy --extra semantic-router
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router
|
||||
|
||||
- name: Run MCP tests
|
||||
run: |
|
||||
|
|
|
|||
2
.github/workflows/test-model-map.yaml
vendored
2
.github/workflows/test-model-map.yaml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
2
.github/workflows/test-rust.yml
vendored
2
.github/workflows/test-rust.yml
vendored
|
|
@ -9,7 +9,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
|
|
|
|||
2
.github/workflows/test-semgrep.yml
vendored
2
.github/workflows/test-semgrep.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
2
.github/workflows/test-unit-core-utils.yml
vendored
2
.github/workflows/test-unit-core-utils.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
@ -54,7 +54,7 @@ jobs:
|
|||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
2
.github/workflows/test-unit-integrations.yml
vendored
2
.github/workflows/test-unit-integrations.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
2
.github/workflows/test-unit-misc.yml
vendored
2
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
2
.github/workflows/test-unit-proxy-auth.yml
vendored
2
.github/workflows/test-unit-proxy-auth.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
|
|||
3
.github/workflows/test-unit-proxy-infra.yml
vendored
3
.github/workflows/test-unit-proxy-infra.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
@ -29,6 +29,7 @@ jobs:
|
|||
tests/test_litellm/proxy/_experimental
|
||||
tests/test_litellm/proxy/experimental
|
||||
tests/test_litellm/proxy/common_utils
|
||||
tests/test_litellm/proxy/logging_endpoints
|
||||
tests/test_litellm/proxy/test_*.py
|
||||
workers: 2
|
||||
reruns: 2
|
||||
|
|
|
|||
4
.github/workflows/test-unit-proxy-legacy.yml
vendored
4
.github/workflows/test-unit-proxy-legacy.yml
vendored
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
@ -71,7 +71,7 @@ jobs:
|
|||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
2
.github/workflows/test_server_root_path.yml
vendored
2
.github/workflows/test_server_root_path.yml
vendored
|
|
@ -7,7 +7,7 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
jobs:
|
||||
|
|
|
|||
11
.gitignore
vendored
11
.gitignore
vendored
|
|
@ -9,12 +9,17 @@ litellm/proxy/myenv/*
|
|||
litellm_uuid.txt
|
||||
__pycache__/
|
||||
*.pyc
|
||||
|
||||
# Rust bridge build artifacts (compiled, platform-specific; regenerated by maturin/cargo)
|
||||
litellm/rust_bridge/_native*.so
|
||||
litellm/rust_bridge/_native*.pyd
|
||||
litellm-rust/target/
|
||||
|
||||
bun.lockb
|
||||
**/.DS_Store
|
||||
.aider*
|
||||
litellm_results.jsonl
|
||||
secrets.toml
|
||||
.gitignore
|
||||
litellm/proxy/litellm_secrets.toml
|
||||
litellm/proxy/api_log.json
|
||||
.idea/
|
||||
|
|
@ -36,7 +41,6 @@ litellm/tests/dynamo*.log
|
|||
.vscode/settings.json
|
||||
litellm/proxy/log.txt
|
||||
proxy_server_config_@.yaml
|
||||
.gitignore
|
||||
proxy_server_config_2.yaml
|
||||
litellm/proxy/secret_managers/credentials.json
|
||||
hosted_config.yaml
|
||||
|
|
@ -123,3 +127,6 @@ crash.*.log
|
|||
# and should be committed.
|
||||
.vscode
|
||||
.pin_list.txt
|
||||
|
||||
# pytest coverage data
|
||||
.coverage
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ If you ever make public-facing PR descriptions, comments, issues, commit message
|
|||
- don't use "—". Instead, reach for ";", ".", etc.
|
||||
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
|
||||
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
|
||||
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: unless there's a sentence immediately after, don't add a "."
|
||||
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
|
||||
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
|
||||
|
||||
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
|
||||
|
|
@ -54,7 +54,7 @@ When working on a PR, keep the PR description in sync with new commits being mad
|
|||
|
||||
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
|
||||
|
||||
Do not put names of customers or customer company names in code, PRs, and issues. The codebase is public
|
||||
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
|
||||
|
||||
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ RUN apk add --no-cache \
|
|||
gcc \
|
||||
python3 \
|
||||
python3-dev \
|
||||
rust \
|
||||
openssl \
|
||||
openssl-dev \
|
||||
nodejs \
|
||||
|
|
|
|||
16
Makefile
16
Makefile
|
|
@ -20,13 +20,13 @@ help:
|
|||
@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 format - Apply Black code formatting"
|
||||
@echo " make format-check - Check Black code formatting (matches CI)"
|
||||
@echo " make lint - Run all linting (Ruff, basedpyright, Black check, circular imports, import safety)"
|
||||
@echo " make format - Apply ruff format code formatting"
|
||||
@echo " make format-check - Check ruff format code formatting (matches CI)"
|
||||
@echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)"
|
||||
@echo " make lint-ruff - Run Ruff linting only"
|
||||
@echo " make lint-basedpyright - Run basedpyright strict, gated by per-rule error counts"
|
||||
@echo " make lint-basedpyright-budget-update - Re-capture the basedpyright per-rule budget (ratchet)"
|
||||
@echo " make lint-black - Check Black formatting (matches CI)"
|
||||
@echo " make lint-format - Check ruff format formatting (matches CI)"
|
||||
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its ceiling"
|
||||
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
|
||||
@echo " make lint-ruff-budget-update - Re-capture per-rule baselines in ruff-strict-budget.json (ratchet)"
|
||||
|
|
@ -82,11 +82,13 @@ install-hooks:
|
|||
./scripts/install_git_hooks.sh
|
||||
|
||||
# Formatting
|
||||
# 88-column wrap matches the Black width the whole repo is formatted to; ruff.toml's
|
||||
# global line-length is 120 (for E501/isort), so 88 is forced here.
|
||||
format: install-dev
|
||||
cd litellm && $(UV_RUN) black . && cd ..
|
||||
cd litellm && $(UV_RUN) ruff format --line-length 88 --exclude '/enterprise/' . && cd ..
|
||||
|
||||
format-check: install-dev
|
||||
cd litellm && $(UV_RUN) black --check . && cd ..
|
||||
cd litellm && $(UV_RUN) ruff format --check --line-length 88 --exclude '/enterprise/' . && cd ..
|
||||
|
||||
# Linting targets
|
||||
lint-ruff: install-dev
|
||||
|
|
@ -131,7 +133,7 @@ lint-basedpyright: install-dev
|
|||
lint-basedpyright-budget-update: install-dev
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --update
|
||||
|
||||
lint-black: format-check
|
||||
lint-format: format-check
|
||||
|
||||
lint-ruff-budget: install-dev
|
||||
$(UV_RUN) python scripts/ruff_strict_gate.py
|
||||
|
|
|
|||
|
|
@ -84,6 +84,8 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/active/callbacks",
|
||||
"/callbacks",
|
||||
"/team_callback",
|
||||
# Rust data-plane gateway → proxy control-plane API (logging today, auth later)
|
||||
"/v1/rust_control_plane/",
|
||||
# Alerting / email / IP allowlist
|
||||
"/alerting/",
|
||||
"/email/",
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@
|
|||
},
|
||||
"reportMatchNotExhaustive": {
|
||||
"baseline": 1,
|
||||
"slack": 3
|
||||
"slack": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"baseline": 3933,
|
||||
|
|
|
|||
|
|
@ -3,6 +3,9 @@ codecov:
|
|||
notify:
|
||||
wait_for_ci: false # post as soon as expected uploads arrive, don't wait on CI
|
||||
|
||||
ignore:
|
||||
- "litellm-rust/**"
|
||||
|
||||
# Uploads are flagged per workflow/shard (GHA) or "circleci". carryforward makes
|
||||
# a re-upload of a flag replace its prior session instead of accumulating a
|
||||
# conflicting one, and lets a commit reuse a flag from its parent when that flag
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ RUN for i in 1 2 3; do \
|
|||
python3 \
|
||||
python3-dev \
|
||||
gcc \
|
||||
rust \
|
||||
bash \
|
||||
coreutils \
|
||||
curl \
|
||||
|
|
|
|||
|
|
@ -13,7 +13,9 @@ from litellm import Router, verbose_logger
|
|||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_file_metadata,
|
||||
)
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.llms.base_llm.managed_resources.isolation import (
|
||||
build_list_page,
|
||||
|
|
@ -981,9 +983,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
target_model_names_list: List[str],
|
||||
) -> OpenAIFileObject:
|
||||
## GET THE FILE TYPE FROM THE CREATE FILE REQUEST
|
||||
file_data = extract_file_data(create_file_request["file"])
|
||||
|
||||
file_type = file_data["content_type"]
|
||||
_, file_type = extract_file_metadata(create_file_request["file"])
|
||||
|
||||
output_file_id = file_objects[0].id
|
||||
model_id = file_objects[0]._hidden_params.get("model_id")
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.43"
|
||||
version = "0.1.44"
|
||||
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.43"
|
||||
version = "0.1.44"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
17
litellm-rust/AGENTS.md
Normal file
17
litellm-rust/AGENTS.md
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
# AGENTS.md
|
||||
|
||||
litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers.
|
||||
|
||||
## Crates
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
|
||||
|
||||
Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
|
||||
|
|
@ -2,27 +2,32 @@
|
|||
|
||||
This file defines the rules for Rust work in LiteLLM.
|
||||
|
||||
## Crates (exactly three — see AGENTS.md)
|
||||
|
||||
`litellm-core` describes work; `litellm-ai-gateway` executes it; `litellm-python-bridge`
|
||||
exposes it to the Python SDK. A crate is a **layer**, not a route — add modules, not crates.
|
||||
|
||||
## Core Boundary
|
||||
|
||||
The `core` and `providers` crates describe work; hosts execute work.
|
||||
`litellm-core` is the pure translation layer; the `litellm-ai-gateway` host executes work.
|
||||
|
||||
Route-level Rust structure mirrors LiteLLM's Python responsibilities:
|
||||
- `core/src/<route>/` owns the route contract, shared types, and provider
|
||||
template traits. For OCR, this means `core/src/ocr`.
|
||||
- `providers/src/<provider>/<route>/transformation.rs` owns the
|
||||
- `core/src/providers/<provider>/<route>/transformation.rs` owns the
|
||||
provider-specific transform. For Mistral OCR, this means
|
||||
`providers/src/mistral/ocr/transformation.rs`.
|
||||
- Future network execution belongs in a host/transport layer such as
|
||||
`llm_http_handler`, not inside `core` or `providers`.
|
||||
`core/src/providers/mistral/ocr/transformation.rs`.
|
||||
- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`),
|
||||
never inside `core`.
|
||||
|
||||
Allowed in `core` and `providers`:
|
||||
Allowed in `core`:
|
||||
- Pure request transforms
|
||||
- Pure response transforms
|
||||
- Pure stream chunk normalization
|
||||
- Shared data types and validation errors
|
||||
- Deterministic token/cost helper logic
|
||||
|
||||
Not allowed in `core` or `providers`:
|
||||
Not allowed in `core`:
|
||||
- Network calls
|
||||
- Environment variable or secret reads
|
||||
- Filesystem access
|
||||
|
|
@ -72,6 +77,20 @@ such as `ai-gateway`, router hosts, or standalone servers:
|
|||
- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is
|
||||
impossible by construction and documented.
|
||||
|
||||
## Constants
|
||||
|
||||
Magic numbers and fixed strings go in a crate-level `constants.rs`, never
|
||||
hardcoded inline — the Rust mirror of Python's `litellm/constants.py`.
|
||||
|
||||
- Each crate that needs them has `src/constants.rs` (declared `mod constants;`);
|
||||
import from it (`use crate::constants::...`). Don't scatter `const` values at
|
||||
the top of feature modules.
|
||||
- An env-overridable tunable still lives in `constants.rs` as its `DEFAULT_*`
|
||||
value; the env read (with fallback to that default) happens at the host/config
|
||||
resolution layer, not in `core`/`providers`.
|
||||
- Exception: a value that is purely local to one function and has no meaning
|
||||
elsewhere may stay inline, but prefer `constants.rs` when in doubt.
|
||||
|
||||
## Checks
|
||||
|
||||
Run these before pushing Rust changes. The same checks run in GitHub Actions
|
||||
|
|
@ -80,7 +99,9 @@ for changes under `litellm-rust/`.
|
|||
```bash
|
||||
cd litellm-rust
|
||||
cargo fmt --check
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
# the ai-gateway binary + server code is behind the `server` feature
|
||||
cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings
|
||||
cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings
|
||||
cargo test --workspace
|
||||
```
|
||||
|
||||
|
|
|
|||
192
litellm-rust/Cargo.lock
generated
192
litellm-rust/Cargo.lock
generated
|
|
@ -206,12 +206,24 @@ dependencies = [
|
|||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||
|
||||
[[package]]
|
||||
name = "form_urlencoded"
|
||||
version = "1.2.2"
|
||||
|
|
@ -221,6 +233,21 @@ dependencies = [
|
|||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-executor",
|
||||
"futures-io",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-channel"
|
||||
version = "0.3.32"
|
||||
|
|
@ -237,12 +264,34 @@ version = "0.3.32"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
|
||||
|
||||
[[package]]
|
||||
name = "futures-executor"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-io"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-sink"
|
||||
version = "0.3.32"
|
||||
|
|
@ -261,8 +310,10 @@ version = "0.3.32"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"futures-macro",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"memchr",
|
||||
|
|
@ -307,6 +358,31 @@ dependencies = [
|
|||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.4.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155"
|
||||
dependencies = [
|
||||
"atomic-waker",
|
||||
"bytes",
|
||||
"fnv",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"http",
|
||||
"indexmap",
|
||||
"slab",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.17.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.5.0"
|
||||
|
|
@ -368,6 +444,7 @@ dependencies = [
|
|||
"bytes",
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
"httparse",
|
||||
|
|
@ -521,6 +598,16 @@ dependencies = [
|
|||
"icu_properties",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indexmap"
|
||||
version = "2.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
|
||||
dependencies = [
|
||||
"equivalent",
|
||||
"hashbrown",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indoc"
|
||||
version = "2.0.7"
|
||||
|
|
@ -544,9 +631,9 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
|
|||
|
||||
[[package]]
|
||||
name = "js-sys"
|
||||
version = "0.3.102"
|
||||
version = "0.3.103"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31"
|
||||
checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"futures-util",
|
||||
|
|
@ -564,14 +651,18 @@ name = "litellm-ai-gateway"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"base64",
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"litellm-core",
|
||||
"litellm-providers",
|
||||
"pyo3",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"subtle",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -582,29 +673,19 @@ dependencies = [
|
|||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-providers"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"litellm-core",
|
||||
"reqwest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-ai-gateway",
|
||||
"litellm-core",
|
||||
"litellm-providers",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -738,6 +819,19 @@ dependencies = [
|
|||
"unindent",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-async-runtimes"
|
||||
version = "0.23.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e"
|
||||
dependencies = [
|
||||
"futures",
|
||||
"once_cell",
|
||||
"pin-project-lite",
|
||||
"pyo3",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-build-config"
|
||||
version = "0.23.5"
|
||||
|
|
@ -923,6 +1017,7 @@ dependencies = [
|
|||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
|
|
@ -942,12 +1037,14 @@ dependencies = [
|
|||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams",
|
||||
"web-sys",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
|
@ -1140,6 +1237,17 @@ dependencies = [
|
|||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures",
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "shlex"
|
||||
version = "2.0.1"
|
||||
|
|
@ -1334,6 +1442,19 @@ dependencies = [
|
|||
"tungstenite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-util"
|
||||
version = "0.7.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.3"
|
||||
|
|
@ -1506,9 +1627,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a"
|
||||
checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"once_cell",
|
||||
|
|
@ -1519,9 +1640,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-futures"
|
||||
version = "0.4.75"
|
||||
version = "0.4.76"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280"
|
||||
checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d"
|
||||
dependencies = [
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
|
|
@ -1529,9 +1650,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d"
|
||||
checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1"
|
||||
dependencies = [
|
||||
"quote",
|
||||
"wasm-bindgen-macro-support",
|
||||
|
|
@ -1539,9 +1660,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro-support"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd"
|
||||
checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e"
|
||||
dependencies = [
|
||||
"bumpalo",
|
||||
"proc-macro2",
|
||||
|
|
@ -1552,18 +1673,31 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-shared"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f"
|
||||
checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.102"
|
||||
name = "wasm-streams"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d"
|
||||
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.103"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141"
|
||||
dependencies = [
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
|
|
|
|||
|
|
@ -1,9 +1,8 @@
|
|||
[workspace]
|
||||
members = [
|
||||
"crates/core",
|
||||
"crates/providers",
|
||||
"crates/python-bridge",
|
||||
"crates/ai-gateway",
|
||||
"crates/python-bridge",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
|
|
@ -14,15 +13,18 @@ repository = "https://github.com/BerriAI/litellm"
|
|||
|
||||
[workspace.dependencies]
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-providers = { path = "crates/providers" }
|
||||
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
|
||||
axum = "0.7"
|
||||
pyo3 = "0.23.5"
|
||||
pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] }
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
sha2 = "0.10"
|
||||
subtle = "2"
|
||||
thiserror = "2.0"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
base64 = "0.22"
|
||||
|
|
|
|||
|
|
@ -7,6 +7,16 @@ continues to own auth, configuration, network I/O, retries, routing, logging,
|
|||
callbacks, spend tracking, and customer plugins until each Rust path has parity
|
||||
coverage and production evidence.
|
||||
|
||||
## Crates
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
## Layout
|
||||
|
||||
```text
|
||||
|
|
|
|||
12
litellm-rust/crates/ai-gateway/ARCHITECTURE.md
Normal file
12
litellm-rust/crates/ai-gateway/ARCHITECTURE.md
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
# ai-gateway architecture
|
||||
|
||||
The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is an
|
||||
API callback: it POSTs each finished session to the LiteLLM proxy, which records
|
||||
spend and runs the usual callbacks.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
C[client] <--> G[Rust ai-gateway<br/>LLM inference]
|
||||
G <--> O[OpenAI realtime]
|
||||
G -. spend tracking callback .-> P[litellm proxy]
|
||||
```
|
||||
|
|
@ -5,22 +5,39 @@ edition.workspace = true
|
|||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "litellm_ai_gateway"
|
||||
|
||||
[[bin]]
|
||||
name = "litellm-ai-gateway"
|
||||
path = "src/main.rs"
|
||||
required-features = ["server"]
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
litellm-providers.workspace = true
|
||||
axum = { workspace = true, features = ["ws"] }
|
||||
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
|
||||
# Python proxy callbacks API.
|
||||
reqwest.workspace = true
|
||||
# `sync` powers the bounded mpsc channel the realtime logger drains.
|
||||
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
|
||||
tokio-tungstenite.workspace = true
|
||||
futures-util.workspace = true
|
||||
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time"] }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
subtle.workspace = true
|
||||
base64.workspace = true
|
||||
axum = { workspace = true, features = ["ws"], optional = true }
|
||||
serde.workspace = true
|
||||
subtle = { workspace = true, optional = true }
|
||||
# sha2 hashes the master key into user_api_key_hash (matches the proxy's
|
||||
# SHA-256 hash_token) so the plaintext credential never enters a log payload.
|
||||
sha2 = { workspace = true, optional = true }
|
||||
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
server = ["dep:axum", "dep:subtle", "dep:sha2"]
|
||||
# Build the gateway's config from the proxy YAML via an embedded Python
|
||||
# interpreter (links libpython; requires `litellm` importable at runtime).
|
||||
python-config = ["dep:pyo3"]
|
||||
|
||||
[dev-dependencies]
|
||||
futures-channel = "0.3"
|
||||
|
|
|
|||
|
|
@ -4,9 +4,22 @@ A minimal Axum service that fronts OpenAI's realtime API. Clients open a
|
|||
WebSocket to `GET /v1/realtime`; the gateway authenticates, selects a deployment,
|
||||
dials OpenAI upstream, and splices the two sockets frame-by-frame.
|
||||
|
||||
## Crates
|
||||
|
||||
`litellm-rust` is exactly three crates (a crate is a **layer**, not a route):
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under `io/`) plus the Axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
|
||||
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
|
||||
- **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil`
|
||||
- **Request logs:** POSTed to a LiteLLM proxy at `/v1/rust_control_plane/logs` (see [Request logging](#request-logging))
|
||||
|
||||
> **Realtime serving is pure Rust.** Python is used at **load time only** — to
|
||||
> read the config once at boot. The realtime hot path never touches Python.
|
||||
|
|
@ -53,6 +66,7 @@ overridden at deploy time (e.g. a Render secret file mounted at the same path).
|
|||
| `OPENAI_API_KEY` | yes | — | Upstream OpenAI key. Referenced by config.yaml as `os.environ/OPENAI_API_KEY` for the gateway→OpenAI dial. |
|
||||
| `HOST` | no | `127.0.0.1` | **Set to `0.0.0.0` in any container/deploy** or external traffic is refused. |
|
||||
| `PORT` | no | `4001` | Listen port. Render and most PaaS inject this automatically. |
|
||||
| `LITELLM_PROXY_BASE_URL` | no | `http://localhost:4000` | LiteLLM proxy that request logs are POSTed to. See [Request logging](#request-logging). |
|
||||
|
||||
> Secrets (`LITELLM_MASTER_KEY`, `OPENAI_API_KEY`) are never baked into the image
|
||||
> or `render.yaml` — inject them at deploy time only.
|
||||
|
|
@ -71,6 +85,18 @@ This mode links no libpython and needs no config file, but it only supports one
|
|||
hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
|
||||
stand-in only for the leanest possible build.
|
||||
|
||||
## Request logging
|
||||
|
||||
The gateway runs no spend logic. When a session ends it builds one
|
||||
`StandardLoggingPayload` and POSTs it to `{LITELLM_PROXY_BASE_URL}/v1/rust_control_plane/logs`
|
||||
(admin-only, bearer = `LITELLM_MASTER_KEY`), and the proxy replays it through its
|
||||
normal callbacks (spend logs, Langfuse, etc.). The POST is non-blocking: a bounded
|
||||
channel drained by a background worker, dropping with a counter if the proxy is
|
||||
down. It sends one payload per session. Both env vars are in the table above.
|
||||
|
||||
Worker tuning, rarely needed: `LITELLM_LOG_CHANNEL_CAPACITY` (4096),
|
||||
`LITELLM_LOG_BATCH_SIZE` (256), `LITELLM_LOG_FLUSH_INTERVAL_MS` (500).
|
||||
|
||||
## Build & run with Docker
|
||||
|
||||
The image is built `--features python-config` and installs litellm **from this
|
||||
|
|
|
|||
|
|
@ -12,10 +12,29 @@ use axum::extract::FromRequestParts;
|
|||
use axum::http::header::AUTHORIZATION;
|
||||
use axum::http::request::Parts;
|
||||
use axum::http::StatusCode;
|
||||
use sha2::{Digest, Sha256};
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// SHA-256 hex digest of a token — the exact transform the Python proxy applies
|
||||
/// (`litellm.proxy.utils.hash_token`).
|
||||
///
|
||||
/// STRICT REQUIREMENT: a raw key (`LITELLM_MASTER_KEY`, a virtual key, …) must
|
||||
/// **never** leave this gateway in a log payload. Spend logs and every callback
|
||||
/// integration receive `user_api_key_hash`, so that field must be this hash, not
|
||||
/// the credential. Hashing here also means the value matches the key's hash in
|
||||
/// `LiteLLM_SpendLogs.api_key`, so realtime spend joins with the rest of LiteLLM.
|
||||
pub fn hash_token(token: &str) -> String {
|
||||
let digest = Sha256::digest(token.as_bytes());
|
||||
let mut hex = String::with_capacity(digest.len() * 2);
|
||||
for byte in digest {
|
||||
use std::fmt::Write;
|
||||
let _ = write!(hex, "{byte:02x}");
|
||||
}
|
||||
hex
|
||||
}
|
||||
|
||||
/// Extractor that requires the configured master key as a bearer token.
|
||||
///
|
||||
/// Rejections: `500` when no master key is configured (permanent
|
||||
|
|
@ -52,3 +71,23 @@ impl FromRequestParts<AppState> for RequireMasterKey {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::hash_token;
|
||||
|
||||
#[test]
|
||||
fn hash_token_matches_python_sha256_hexdigest() {
|
||||
// Must equal hashlib.sha256("sk-1234".encode()).hexdigest() — the value
|
||||
// the proxy stores in LiteLLM_SpendLogs.api_key.
|
||||
assert_eq!(
|
||||
hash_token("sk-1234"),
|
||||
"88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
|
||||
);
|
||||
// 64 lowercase hex chars, and never the raw input.
|
||||
let h = hash_token("sk-secret");
|
||||
assert_eq!(h.len(), 64);
|
||||
assert!(h.chars().all(|c| c.is_ascii_hexdigit()));
|
||||
assert_ne!(h, "sk-secret");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
30
litellm-rust/crates/ai-gateway/src/constants.rs
Normal file
30
litellm-rust/crates/ai-gateway/src/constants.rs
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
//! Crate-level constants for the ai-gateway.
|
||||
//!
|
||||
//! Per `litellm-rust/CLAUDE.md`, magic numbers and fixed strings live here
|
||||
//! (the Rust mirror of Python's `litellm/constants.py`), not inline in feature
|
||||
//! modules. Env-overridable tunables keep their `DEFAULT_*` value here; the env
|
||||
//! read + fallback happens at the host/config layer.
|
||||
|
||||
/// Default LiteLLM control-plane base URL for request-log egress when
|
||||
/// `LITELLM_PROXY_BASE_URL` is unset.
|
||||
pub(crate) const DEFAULT_PROXY_BASE_URL: &str = "http://localhost:4000";
|
||||
|
||||
/// The logs ingest path appended to the proxy base. Not a tunable; it is the
|
||||
/// proxy's API contract (the rust-control-plane router on the Python proxy).
|
||||
pub(crate) const RUST_CONTROL_PLANE_LOGS_PATH: &str = "/v1/rust_control_plane/logs";
|
||||
|
||||
/// Default bounded channel depth for the log-egress worker.
|
||||
/// Override: `LITELLM_LOG_CHANNEL_CAPACITY`.
|
||||
pub(crate) const DEFAULT_CHANNEL_CAPACITY: usize = 4096;
|
||||
|
||||
/// Default max records POSTed per request to the control plane.
|
||||
/// Override: `LITELLM_LOG_BATCH_SIZE`.
|
||||
pub(crate) const DEFAULT_MAX_BATCH_SIZE: usize = 256;
|
||||
|
||||
/// Default partial-batch flush cadence, in ms.
|
||||
/// Override: `LITELLM_LOG_FLUSH_INTERVAL_MS`.
|
||||
pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
|
||||
|
||||
/// Provider attributed to realtime sessions in the logging payload.
|
||||
#[cfg(feature = "server")]
|
||||
pub(crate) const DEFAULT_PROVIDER: &str = "openai";
|
||||
127
litellm-rust/crates/ai-gateway/src/integrations/README.md
Normal file
127
litellm-rust/crates/ai-gateway/src/integrations/README.md
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
# LiteLLM Rust integrations
|
||||
|
||||
This directory contains Rust-native equivalents of LiteLLM integration hooks.
|
||||
The first supported surfaces are terminal custom loggers and pre/during-call
|
||||
custom guardrails.
|
||||
|
||||
## File layout
|
||||
|
||||
Every integration is a folder:
|
||||
|
||||
- `mod.rs` contains the implementation, trait, runner, or adapter
|
||||
- `types.rs` contains the integration-local request, response, error, and future
|
||||
types
|
||||
|
||||
Do not add new flat integration files such as `custom_logger.rs`. Shared wire
|
||||
contracts that are used by multiple integrations can stay in
|
||||
`integrations/types.rs`.
|
||||
|
||||
Call ordering and lifecycle timing live in `litellm-core/src/call_lifecycle`.
|
||||
Call-type modules, such as OCR, adapt their request and response shapes into
|
||||
that generic lifecycle runner.
|
||||
|
||||
## CustomLogger
|
||||
|
||||
Implement `CustomLogger` when Rust code needs to observe terminal success or
|
||||
failure events. Method names intentionally match Python `CustomLogger` names.
|
||||
|
||||
```rust
|
||||
use litellm_ai_gateway::integrations::custom_logger::{
|
||||
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
|
||||
};
|
||||
|
||||
struct RecordingLogger;
|
||||
|
||||
impl CustomLogger for RecordingLogger {
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
response_obj: &'a CallbackValue,
|
||||
timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
let model = &model_call_details.model;
|
||||
let provider = &model_call_details.custom_llm_provider;
|
||||
let call_type = model_call_details.call_type.to_string();
|
||||
let request_id = model_call_details.request_id.as_deref();
|
||||
let response_object = &response_obj.object;
|
||||
let duration = timing.end_time - timing.start_time;
|
||||
let standard_payload = model_call_details.standard_logging_payload.as_ref();
|
||||
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
response_obj: Option<&'a CallbackValue>,
|
||||
timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
let error = model_call_details.failure_error.as_ref();
|
||||
let response_object = response_obj.map(|value| value.object.as_str());
|
||||
let duration = timing.end_time - timing.start_time;
|
||||
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Use `CustomLoggerRunner` to fan out terminal events to configured loggers. The
|
||||
runner is a no-op when no loggers are configured, which is the expected fast
|
||||
path for requests without callbacks.
|
||||
|
||||
## CustomGuardrail
|
||||
|
||||
Implement `CustomGuardrail` when Rust code needs to run pre-call or native
|
||||
during-call checks. Method names intentionally match Python `CustomGuardrail`
|
||||
entrypoints inherited from Python `CustomLogger`.
|
||||
|
||||
```rust
|
||||
use litellm_ai_gateway::integrations::custom_guardrail::{
|
||||
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailEventHook,
|
||||
GuardrailFuture, GuardrailRequest,
|
||||
};
|
||||
|
||||
struct BlocklistedPromptGuardrail;
|
||||
|
||||
impl CustomGuardrail for BlocklistedPromptGuardrail {
|
||||
fn guardrail_name(&self) -> &str {
|
||||
"blocklisted-prompt"
|
||||
}
|
||||
|
||||
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
|
||||
&[GuardrailEventHook::PreCall]
|
||||
}
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move {
|
||||
if request.data.to_string().contains("blocked phrase") {
|
||||
return Ok(GuardrailDecision::Block(
|
||||
litellm_ai_gateway::integrations::custom_guardrail::GuardrailError::blocked(
|
||||
"blocked phrase detected",
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(GuardrailDecision::Allow(request))
|
||||
})
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Use `CustomGuardrailRunner::run_pre_call` for `pre_call` guardrails and
|
||||
`CustomGuardrailRunner::run_during_call` for `during_call` guardrails. A
|
||||
`GuardrailDecision::Mask` continues with modified request data.
|
||||
`GuardrailDecision::Block` short-circuits the provider call.
|
||||
|
||||
## Current boundary
|
||||
|
||||
These are Rust-only primitives. Python callback and guardrail adapters are a
|
||||
separate layer that should implement these Rust traits instead of changing the
|
||||
runner interfaces.
|
||||
|
|
@ -0,0 +1,468 @@
|
|||
//! Rust mirror of Python `CustomGuardrail` entrypoints used by the proxy.
|
||||
//!
|
||||
//! This module is intentionally Rust-only: Python/PyO3 adapters are a later
|
||||
//! layer that should implement this trait rather than changing the runner.
|
||||
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::integrations::custom_logger::{
|
||||
CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
|
||||
};
|
||||
|
||||
pub mod types;
|
||||
|
||||
pub use types::{
|
||||
GuardrailContext, GuardrailDecision, GuardrailDispatchReport, GuardrailError,
|
||||
GuardrailEventHook, GuardrailFuture, GuardrailRequest,
|
||||
};
|
||||
|
||||
pub trait CustomGuardrail: Send + Sync {
|
||||
fn guardrail_name(&self) -> &str;
|
||||
|
||||
fn supported_event_hooks(&self) -> &[GuardrailEventHook];
|
||||
|
||||
/// Python 1:1 name: `async_pre_call_hook(user_api_key_dict, cache, data, call_type)`.
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move { Ok(GuardrailDecision::Allow(request)) })
|
||||
}
|
||||
|
||||
/// Python 1:1 name: `async_moderation_hook(data, user_api_key_dict, call_type)`.
|
||||
fn async_moderation_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move { Ok(GuardrailDecision::Allow(request)) })
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CustomGuardrailRunner {
|
||||
guardrails: Vec<Arc<dyn CustomGuardrail>>,
|
||||
}
|
||||
|
||||
impl CustomGuardrailRunner {
|
||||
pub fn new(guardrails: Vec<Arc<dyn CustomGuardrail>>) -> Self {
|
||||
Self { guardrails }
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.guardrails.is_empty()
|
||||
}
|
||||
|
||||
pub async fn run_pre_call(
|
||||
&self,
|
||||
context: &GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
|
||||
self.run_hook(GuardrailEventHook::PreCall, context, request)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn run_during_call(
|
||||
&self,
|
||||
context: &GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
|
||||
self.run_hook(GuardrailEventHook::DuringCall, context, request)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn run_before_provider<F, Fut, T>(
|
||||
&self,
|
||||
event_hook: GuardrailEventHook,
|
||||
context: &GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
provider: F,
|
||||
) -> Result<T, GuardrailError>
|
||||
where
|
||||
F: FnOnce(GuardrailRequest) -> Fut,
|
||||
Fut: Future<Output = Result<T, GuardrailError>>,
|
||||
{
|
||||
let (request, _) = self.run_hook(event_hook, context, request).await?;
|
||||
provider(request).await
|
||||
}
|
||||
|
||||
pub async fn run_pre_call_with_failure_logging(
|
||||
&self,
|
||||
context: &GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
logger_runner: &CustomLoggerRunner,
|
||||
model_call_details: &ModelCallDetails,
|
||||
timing: CallbackTiming,
|
||||
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
|
||||
match self.run_pre_call(context, request).await {
|
||||
Ok(result) => Ok(result),
|
||||
Err(error) => {
|
||||
let failure_details = model_call_details.clone().with_failure_error(LoggingError {
|
||||
message: error.message.clone(),
|
||||
kind: error.kind.clone(),
|
||||
});
|
||||
let response_obj = CallbackValue::new(
|
||||
"guardrail_error",
|
||||
serde_json::json!({
|
||||
"message": error.message,
|
||||
"kind": error.kind,
|
||||
}),
|
||||
);
|
||||
logger_runner
|
||||
.async_log_failure_event(&failure_details, Some(&response_obj), timing)
|
||||
.await;
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_hook(
|
||||
&self,
|
||||
event_hook: GuardrailEventHook,
|
||||
context: &GuardrailContext,
|
||||
mut request: GuardrailRequest,
|
||||
) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> {
|
||||
if self.guardrails.is_empty() {
|
||||
return Ok((request, GuardrailDispatchReport::default()));
|
||||
}
|
||||
|
||||
let mut report = GuardrailDispatchReport::default();
|
||||
for guardrail in &self.guardrails {
|
||||
if !self.should_run(guardrail.as_ref(), event_hook, context) {
|
||||
continue;
|
||||
}
|
||||
|
||||
report.invoked += 1;
|
||||
let decision = match event_hook {
|
||||
GuardrailEventHook::PreCall => {
|
||||
guardrail
|
||||
.async_pre_call_hook(context, request.clone())
|
||||
.await?
|
||||
}
|
||||
GuardrailEventHook::DuringCall => {
|
||||
guardrail
|
||||
.async_moderation_hook(context, request.clone())
|
||||
.await?
|
||||
}
|
||||
};
|
||||
match decision.into_request() {
|
||||
Ok(next_request) => request = next_request,
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
Ok((request, report))
|
||||
}
|
||||
|
||||
fn should_run(
|
||||
&self,
|
||||
guardrail: &dyn CustomGuardrail,
|
||||
event_hook: GuardrailEventHook,
|
||||
context: &GuardrailContext,
|
||||
) -> bool {
|
||||
let supports_hook = guardrail.supported_event_hooks().contains(&event_hook);
|
||||
let selected = context.selected_guardrails.is_empty()
|
||||
|| context
|
||||
.selected_guardrails
|
||||
.iter()
|
||||
.any(|name| name == guardrail.guardrail_name());
|
||||
supports_hook && selected
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::integrations::custom_logger::{CallType, CallbackValue, CustomLogger, LogFuture};
|
||||
use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
|
||||
use serde_json::json;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Clone)]
|
||||
enum TestDecision {
|
||||
Allow,
|
||||
Mask,
|
||||
Block,
|
||||
}
|
||||
|
||||
struct RecordingCustomGuardrail {
|
||||
name: String,
|
||||
hooks: Vec<GuardrailEventHook>,
|
||||
decision: TestDecision,
|
||||
calls: Mutex<Vec<&'static str>>,
|
||||
}
|
||||
|
||||
impl RecordingCustomGuardrail {
|
||||
fn new(name: &str, hooks: Vec<GuardrailEventHook>, decision: TestDecision) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
hooks,
|
||||
decision,
|
||||
calls: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn calls(&self) -> Vec<&'static str> {
|
||||
self.calls.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
fn decision(&self, mut request: GuardrailRequest) -> GuardrailDecision {
|
||||
match self.decision {
|
||||
TestDecision::Allow => GuardrailDecision::Allow(request),
|
||||
TestDecision::Mask => {
|
||||
request.data["masked"] = json!(true);
|
||||
GuardrailDecision::Mask(request)
|
||||
}
|
||||
TestDecision::Block => {
|
||||
GuardrailDecision::Block(GuardrailError::blocked("blocked by guardrail"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CustomGuardrail for RecordingCustomGuardrail {
|
||||
fn guardrail_name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
|
||||
&self.hooks
|
||||
}
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.calls.lock().unwrap().push("async_pre_call_hook");
|
||||
Ok(self.decision(request))
|
||||
})
|
||||
}
|
||||
|
||||
fn async_moderation_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a GuardrailContext,
|
||||
request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.calls.lock().unwrap().push("async_moderation_hook");
|
||||
Ok(self.decision(request))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_call_dispatches_to_async_pre_call_hook() {
|
||||
let guardrail = Arc::new(RecordingCustomGuardrail::new(
|
||||
"pre",
|
||||
vec![GuardrailEventHook::PreCall],
|
||||
TestDecision::Allow,
|
||||
));
|
||||
let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]);
|
||||
let context =
|
||||
GuardrailContext::new(CallType::Ocr).with_selected_guardrails(vec!["pre".to_string()]);
|
||||
let request = GuardrailRequest::new(json!({"messages": ["hello"]}));
|
||||
|
||||
let (result, report) = runner
|
||||
.run_pre_call(&context, request)
|
||||
.await
|
||||
.expect("guardrail allows request");
|
||||
|
||||
assert_eq!(report.invoked, 1);
|
||||
assert_eq!(result.data["messages"], json!(["hello"]));
|
||||
assert_eq!(guardrail.calls(), vec!["async_pre_call_hook"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn during_call_dispatches_to_async_moderation_hook() {
|
||||
let guardrail = Arc::new(RecordingCustomGuardrail::new(
|
||||
"during",
|
||||
vec![GuardrailEventHook::DuringCall],
|
||||
TestDecision::Allow,
|
||||
));
|
||||
let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]);
|
||||
let context = GuardrailContext::new(CallType::Completion)
|
||||
.with_selected_guardrails(vec!["during".to_string()]);
|
||||
let request = GuardrailRequest::new(json!({"prompt": "hello"}));
|
||||
|
||||
let (_result, report) = runner
|
||||
.run_during_call(&context, request)
|
||||
.await
|
||||
.expect("guardrail allows request");
|
||||
|
||||
assert_eq!(report.invoked, 1);
|
||||
assert_eq!(guardrail.calls(), vec!["async_moderation_hook"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mask_decision_continues_with_updated_request() {
|
||||
let guardrail = Arc::new(RecordingCustomGuardrail::new(
|
||||
"masker",
|
||||
vec![GuardrailEventHook::PreCall],
|
||||
TestDecision::Mask,
|
||||
));
|
||||
let runner = CustomGuardrailRunner::new(vec![guardrail]);
|
||||
let context = GuardrailContext::new(CallType::Ocr);
|
||||
let request = GuardrailRequest::new(json!({"document": "secret"}));
|
||||
|
||||
let (result, report) = runner
|
||||
.run_pre_call(&context, request)
|
||||
.await
|
||||
.expect("mask continues");
|
||||
|
||||
assert_eq!(report.invoked, 1);
|
||||
assert_eq!(result.data["masked"], json!(true));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn block_decision_short_circuits_and_logs_failure() {
|
||||
struct RecordingFailureLogger {
|
||||
errors: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl CustomLogger for RecordingFailureLogger {
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: Option<&'a CallbackValue>,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.errors.lock().unwrap().push(
|
||||
model_call_details
|
||||
.failure_error
|
||||
.as_ref()
|
||||
.map(|error| error.kind.clone())
|
||||
.unwrap_or_default(),
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
let guardrail = Arc::new(RecordingCustomGuardrail::new(
|
||||
"blocker",
|
||||
vec![GuardrailEventHook::PreCall],
|
||||
TestDecision::Block,
|
||||
));
|
||||
let guardrail_runner = CustomGuardrailRunner::new(vec![guardrail]);
|
||||
let logger = Arc::new(RecordingFailureLogger {
|
||||
errors: Mutex::new(Vec::new()),
|
||||
});
|
||||
let logger_runner = CustomLoggerRunner::new(vec![logger.clone()]);
|
||||
let context = GuardrailContext::new(CallType::Ocr);
|
||||
let details = ModelCallDetails::from_standard_logging_payload(StandardLoggingPayload {
|
||||
id: "req_ocr".to_string(),
|
||||
litellm_call_id: "req_ocr".to_string(),
|
||||
call_type: "ocr".to_string(),
|
||||
model: "mistral-ocr-latest".to_string(),
|
||||
custom_llm_provider: "mistral".to_string(),
|
||||
response_cost: 0.0,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
total_tokens: 0,
|
||||
start_time: 1.0,
|
||||
end_time: 1.0,
|
||||
stream: false,
|
||||
metadata: StandardLoggingMetadata::default(),
|
||||
messages: None,
|
||||
});
|
||||
|
||||
let err = guardrail_runner
|
||||
.run_pre_call_with_failure_logging(
|
||||
&context,
|
||||
GuardrailRequest::new(json!({"document": "bad"})),
|
||||
&logger_runner,
|
||||
&details,
|
||||
CallbackTiming::new(1.0, 2.0),
|
||||
)
|
||||
.await
|
||||
.expect_err("guardrail blocks request");
|
||||
|
||||
assert_eq!(err.kind, "GuardrailBlocked");
|
||||
assert_eq!(
|
||||
logger.errors.lock().unwrap().as_slice(),
|
||||
["GuardrailBlocked"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn block_decision_short_circuits_later_guardrails_and_provider_work() {
|
||||
let blocking_guardrail = Arc::new(RecordingCustomGuardrail::new(
|
||||
"blocker",
|
||||
vec![GuardrailEventHook::PreCall],
|
||||
TestDecision::Block,
|
||||
));
|
||||
let later_guardrail = Arc::new(RecordingCustomGuardrail::new(
|
||||
"later",
|
||||
vec![GuardrailEventHook::PreCall],
|
||||
TestDecision::Allow,
|
||||
));
|
||||
let runner =
|
||||
CustomGuardrailRunner::new(vec![blocking_guardrail.clone(), later_guardrail.clone()]);
|
||||
let provider_called = Arc::new(Mutex::new(false));
|
||||
let provider_called_for_closure = provider_called.clone();
|
||||
|
||||
let result = runner
|
||||
.run_before_provider(
|
||||
GuardrailEventHook::PreCall,
|
||||
&GuardrailContext::new(CallType::Completion),
|
||||
GuardrailRequest::new(json!({"prompt": "blocked"})),
|
||||
move |_request| async move {
|
||||
*provider_called_for_closure.lock().unwrap() = true;
|
||||
Ok("provider response")
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
assert_eq!(blocking_guardrail.calls(), vec!["async_pre_call_hook"]);
|
||||
assert_eq!(later_guardrail.calls(), Vec::<&'static str>::new());
|
||||
assert!(!*provider_called.lock().unwrap());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_before_provider_returns_provider_guardrail_error_directly() {
|
||||
let guardrail = Arc::new(RecordingCustomGuardrail::new(
|
||||
"allow",
|
||||
vec![GuardrailEventHook::PreCall],
|
||||
TestDecision::Allow,
|
||||
));
|
||||
let runner = CustomGuardrailRunner::new(vec![guardrail]);
|
||||
|
||||
let result = runner
|
||||
.run_before_provider(
|
||||
GuardrailEventHook::PreCall,
|
||||
&GuardrailContext::new(CallType::Completion),
|
||||
GuardrailRequest::new(json!({"prompt": "allowed"})),
|
||||
|_request| async move {
|
||||
Err::<&'static str, GuardrailError>(GuardrailError::blocked(
|
||||
"provider-side guardrail error",
|
||||
))
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let err = result.expect_err("provider error is returned directly");
|
||||
assert_eq!(err.kind, "GuardrailBlocked");
|
||||
assert_eq!(err.message, "provider-side guardrail error");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_guardrails_fast_path_dispatches_nothing() {
|
||||
let runner = CustomGuardrailRunner::new(Vec::new());
|
||||
let context = GuardrailContext::new(CallType::Ocr);
|
||||
let request = GuardrailRequest::new(json!({"document": "ok"}));
|
||||
|
||||
let (result, report) = runner
|
||||
.run_pre_call(&context, request)
|
||||
.await
|
||||
.expect("no guardrails allow request");
|
||||
|
||||
assert!(runner.is_empty());
|
||||
assert_eq!(report, GuardrailDispatchReport::default());
|
||||
assert_eq!(result.data["document"], json!("ok"));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,110 @@
|
|||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::integrations::custom_logger::CallType;
|
||||
|
||||
pub type GuardrailFuture<'a> =
|
||||
Pin<Box<dyn Future<Output = Result<GuardrailDecision, GuardrailError>> + Send + 'a>>;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum GuardrailEventHook {
|
||||
PreCall,
|
||||
DuringCall,
|
||||
}
|
||||
|
||||
impl GuardrailEventHook {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::PreCall => "pre_call",
|
||||
Self::DuringCall => "during_call",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct GuardrailError {
|
||||
pub message: String,
|
||||
pub kind: String,
|
||||
}
|
||||
|
||||
impl GuardrailError {
|
||||
pub fn blocked(message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
message: message.into(),
|
||||
kind: "GuardrailBlocked".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for GuardrailError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}: {}", self.kind, self.message)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for GuardrailError {}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct GuardrailContext {
|
||||
pub call_type: CallType,
|
||||
pub selected_guardrails: Vec<String>,
|
||||
pub metadata: HashMap<String, Value>,
|
||||
pub user_api_key_hash: Option<String>,
|
||||
pub user_api_key_user_id: Option<String>,
|
||||
pub user_api_key_team_id: Option<String>,
|
||||
pub trace_parent: Option<String>,
|
||||
}
|
||||
|
||||
impl GuardrailContext {
|
||||
pub fn new(call_type: CallType) -> Self {
|
||||
Self {
|
||||
call_type,
|
||||
selected_guardrails: Vec::new(),
|
||||
metadata: HashMap::new(),
|
||||
user_api_key_hash: None,
|
||||
user_api_key_user_id: None,
|
||||
user_api_key_team_id: None,
|
||||
trace_parent: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_selected_guardrails(mut self, selected_guardrails: Vec<String>) -> Self {
|
||||
self.selected_guardrails = selected_guardrails;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct GuardrailRequest {
|
||||
pub data: Value,
|
||||
}
|
||||
|
||||
impl GuardrailRequest {
|
||||
pub fn new(data: Value) -> Self {
|
||||
Self { data }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum GuardrailDecision {
|
||||
Allow(GuardrailRequest),
|
||||
Mask(GuardrailRequest),
|
||||
Block(GuardrailError),
|
||||
}
|
||||
|
||||
impl GuardrailDecision {
|
||||
pub(super) fn into_request(self) -> Result<GuardrailRequest, GuardrailError> {
|
||||
match self {
|
||||
Self::Allow(request) | Self::Mask(request) => Ok(request),
|
||||
Self::Block(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct GuardrailDispatchReport {
|
||||
pub invoked: usize,
|
||||
}
|
||||
|
|
@ -0,0 +1,317 @@
|
|||
//! The `CustomLogger` trait — the Rust mirror of Python
|
||||
//! `litellm/integrations/custom_logger.py::CustomLogger`.
|
||||
//!
|
||||
//! The Python-named async terminal methods are the public Rust callback shape.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub mod types;
|
||||
|
||||
pub use types::{
|
||||
CallType, CallbackDispatchReport, CallbackTiming, CallbackValue, LogError, LogFuture,
|
||||
LoggingError, ModelCallDetails,
|
||||
};
|
||||
|
||||
pub trait CustomLogger: Send + Sync {
|
||||
/// Python 1:1 name: `async_log_success_event(model_call_details, response_obj, start_time, end_time)`.
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
_model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: &'a CallbackValue,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
/// Python 1:1 name: `async_log_failure_event(model_call_details, response_obj, start_time, end_time)`.
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
_model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: Option<&'a CallbackValue>,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CustomLoggerRunner {
|
||||
loggers: Vec<Arc<dyn CustomLogger>>,
|
||||
}
|
||||
|
||||
impl CustomLoggerRunner {
|
||||
pub fn new(loggers: Vec<Arc<dyn CustomLogger>>) -> Self {
|
||||
Self { loggers }
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.loggers.is_empty()
|
||||
}
|
||||
|
||||
pub async fn async_log_success_event(
|
||||
&self,
|
||||
model_call_details: &ModelCallDetails,
|
||||
response_obj: &CallbackValue,
|
||||
timing: CallbackTiming,
|
||||
) -> CallbackDispatchReport {
|
||||
if self.loggers.is_empty() {
|
||||
return CallbackDispatchReport::default();
|
||||
}
|
||||
|
||||
let mut report = CallbackDispatchReport::default();
|
||||
for logger in &self.loggers {
|
||||
report.invoked += 1;
|
||||
if let Err(err) = logger
|
||||
.async_log_success_event(model_call_details, response_obj, timing)
|
||||
.await
|
||||
{
|
||||
report.dropped += 1;
|
||||
eprintln!("litellm-ai-gateway: async_log_success_event dropped: {err}");
|
||||
}
|
||||
}
|
||||
report
|
||||
}
|
||||
|
||||
pub async fn async_log_failure_event(
|
||||
&self,
|
||||
model_call_details: &ModelCallDetails,
|
||||
response_obj: Option<&CallbackValue>,
|
||||
timing: CallbackTiming,
|
||||
) -> CallbackDispatchReport {
|
||||
if self.loggers.is_empty() {
|
||||
return CallbackDispatchReport::default();
|
||||
}
|
||||
|
||||
let mut report = CallbackDispatchReport::default();
|
||||
for logger in &self.loggers {
|
||||
report.invoked += 1;
|
||||
if let Err(err) = logger
|
||||
.async_log_failure_event(model_call_details, response_obj, timing)
|
||||
.await
|
||||
{
|
||||
report.dropped += 1;
|
||||
eprintln!("litellm-ai-gateway: async_log_failure_event dropped: {err}");
|
||||
}
|
||||
}
|
||||
report
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
|
||||
use serde_json::json;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
struct RecordedEvent {
|
||||
hook: &'static str,
|
||||
model: String,
|
||||
provider: String,
|
||||
call_type: String,
|
||||
request_id: Option<String>,
|
||||
litellm_call_id: Option<String>,
|
||||
user_id: Option<String>,
|
||||
response_object: Option<String>,
|
||||
error_kind: Option<String>,
|
||||
start_time: f64,
|
||||
end_time: f64,
|
||||
standard_logging_model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingCustomLogger {
|
||||
events: Mutex<Vec<RecordedEvent>>,
|
||||
}
|
||||
|
||||
impl RecordingCustomLogger {
|
||||
fn events(&self) -> Vec<RecordedEvent> {
|
||||
self.events.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl CustomLogger for RecordingCustomLogger {
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
response_obj: &'a CallbackValue,
|
||||
timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push(RecordedEvent {
|
||||
hook: "async_log_success_event",
|
||||
model: model_call_details.model.clone(),
|
||||
provider: model_call_details.custom_llm_provider.clone(),
|
||||
call_type: model_call_details.call_type.to_string(),
|
||||
request_id: model_call_details.request_id.clone(),
|
||||
litellm_call_id: model_call_details.litellm_call_id.clone(),
|
||||
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
|
||||
response_object: Some(response_obj.object.clone()),
|
||||
error_kind: None,
|
||||
start_time: timing.start_time,
|
||||
end_time: timing.end_time,
|
||||
standard_logging_model: model_call_details
|
||||
.standard_logging_payload
|
||||
.as_ref()
|
||||
.map(|payload| payload.model.clone()),
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
response_obj: Option<&'a CallbackValue>,
|
||||
timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push(RecordedEvent {
|
||||
hook: "async_log_failure_event",
|
||||
model: model_call_details.model.clone(),
|
||||
provider: model_call_details.custom_llm_provider.clone(),
|
||||
call_type: model_call_details.call_type.to_string(),
|
||||
request_id: model_call_details.request_id.clone(),
|
||||
litellm_call_id: model_call_details.litellm_call_id.clone(),
|
||||
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
|
||||
response_object: response_obj.map(|value| value.object.clone()),
|
||||
error_kind: model_call_details
|
||||
.failure_error
|
||||
.as_ref()
|
||||
.map(|error| error.kind.clone()),
|
||||
start_time: timing.start_time,
|
||||
end_time: timing.end_time,
|
||||
standard_logging_model: model_call_details
|
||||
.standard_logging_payload
|
||||
.as_ref()
|
||||
.map(|payload| payload.model.clone()),
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn payload(call_type: &str, model: &str, provider: &str) -> StandardLoggingPayload {
|
||||
StandardLoggingPayload {
|
||||
id: format!("req_{call_type}"),
|
||||
litellm_call_id: format!("call_{call_type}"),
|
||||
call_type: call_type.to_string(),
|
||||
model: model.to_string(),
|
||||
custom_llm_provider: provider.to_string(),
|
||||
response_cost: 0.25,
|
||||
prompt_tokens: 3,
|
||||
completion_tokens: 4,
|
||||
total_tokens: 7,
|
||||
start_time: 10.0,
|
||||
end_time: 11.5,
|
||||
stream: false,
|
||||
metadata: StandardLoggingMetadata {
|
||||
user_api_key_hash: Some("hash".to_string()),
|
||||
user_api_key_user_id: Some("user".to_string()),
|
||||
user_api_key_team_id: Some("team".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
messages: Some(json!([{"role": "user", "content": "read this"}])),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rust_custom_logger_reads_success_payload_for_ocr() {
|
||||
let logger = Arc::new(RecordingCustomLogger::default());
|
||||
let runner = CustomLoggerRunner::new(vec![logger.clone()]);
|
||||
let details = ModelCallDetails::from_standard_logging_payload(payload(
|
||||
"ocr",
|
||||
"mistral-ocr-latest",
|
||||
"mistral",
|
||||
));
|
||||
let response = CallbackValue::new("ocr", json!({"pages": [{"markdown": "ok"}]}));
|
||||
let report = runner
|
||||
.async_log_success_event(&details, &response, CallbackTiming::new(10.0, 11.5))
|
||||
.await;
|
||||
|
||||
assert_eq!(report.invoked, 1);
|
||||
assert_eq!(report.dropped, 0);
|
||||
assert_eq!(
|
||||
logger.events(),
|
||||
vec![RecordedEvent {
|
||||
hook: "async_log_success_event",
|
||||
model: "mistral-ocr-latest".to_string(),
|
||||
provider: "mistral".to_string(),
|
||||
call_type: "ocr".to_string(),
|
||||
request_id: Some("req_ocr".to_string()),
|
||||
litellm_call_id: Some("call_ocr".to_string()),
|
||||
user_id: Some("user".to_string()),
|
||||
response_object: Some("ocr".to_string()),
|
||||
error_kind: None,
|
||||
start_time: 10.0,
|
||||
end_time: 11.5,
|
||||
standard_logging_model: Some("mistral-ocr-latest".to_string()),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rust_custom_logger_reads_failure_payload_for_non_ocr_call_type() {
|
||||
let logger = Arc::new(RecordingCustomLogger::default());
|
||||
let runner = CustomLoggerRunner::new(vec![logger.clone()]);
|
||||
let details = ModelCallDetails::from_standard_logging_payload(payload(
|
||||
"acompletion",
|
||||
"gpt-4.1-mini",
|
||||
"openai",
|
||||
))
|
||||
.with_failure_error(LoggingError {
|
||||
message: "provider failed".to_string(),
|
||||
kind: "ProviderError".to_string(),
|
||||
});
|
||||
let response = CallbackValue::new("error", json!({"message": "provider failed"}));
|
||||
let report = runner
|
||||
.async_log_failure_event(&details, Some(&response), CallbackTiming::new(2.0, 3.0))
|
||||
.await;
|
||||
|
||||
assert_eq!(report.invoked, 1);
|
||||
assert_eq!(report.dropped, 0);
|
||||
assert_eq!(
|
||||
logger.events(),
|
||||
vec![RecordedEvent {
|
||||
hook: "async_log_failure_event",
|
||||
model: "gpt-4.1-mini".to_string(),
|
||||
provider: "openai".to_string(),
|
||||
call_type: "acompletion".to_string(),
|
||||
request_id: Some("req_acompletion".to_string()),
|
||||
litellm_call_id: Some("call_acompletion".to_string()),
|
||||
user_id: Some("user".to_string()),
|
||||
response_object: Some("error".to_string()),
|
||||
error_kind: Some("ProviderError".to_string()),
|
||||
start_time: 2.0,
|
||||
end_time: 3.0,
|
||||
standard_logging_model: Some("gpt-4.1-mini".to_string()),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_callback_fast_path_dispatches_nothing() {
|
||||
let runner = CustomLoggerRunner::new(Vec::new());
|
||||
let details = ModelCallDetails::new("mistral-ocr-latest", "mistral", CallType::Ocr);
|
||||
let response = CallbackValue::new("ocr", json!({}));
|
||||
|
||||
let report = runner
|
||||
.async_log_success_event(&details, &response, CallbackTiming::new(1.0, 1.5))
|
||||
.await;
|
||||
|
||||
assert!(runner.is_empty());
|
||||
assert_eq!(report, CallbackDispatchReport::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_standard_logging_payload_keeps_top_level_fields_in_sync() {
|
||||
let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion)
|
||||
.with_standard_logging_payload(payload("ocr", "mistral-ocr-latest", "mistral"));
|
||||
|
||||
assert_eq!(details.model, "mistral-ocr-latest");
|
||||
assert_eq!(details.custom_llm_provider, "mistral");
|
||||
assert_eq!(details.call_type, CallType::Ocr);
|
||||
assert_eq!(details.request_id, Some("req_ocr".to_string()));
|
||||
assert_eq!(details.litellm_call_id, Some("call_ocr".to_string()));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,194 @@
|
|||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload};
|
||||
|
||||
pub type LogFuture<'a> = Pin<Box<dyn Future<Output = Result<(), LogError>> + Send + 'a>>;
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct CallbackDispatchReport {
|
||||
pub invoked: usize,
|
||||
pub dropped: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum CallType {
|
||||
Ocr,
|
||||
Realtime,
|
||||
Completion,
|
||||
Acompletion,
|
||||
ChatCompletion,
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl CallType {
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
Self::Ocr => "ocr",
|
||||
Self::Realtime => "realtime",
|
||||
Self::Completion => "completion",
|
||||
Self::Acompletion => "acompletion",
|
||||
Self::ChatCompletion => "chat_completion",
|
||||
Self::Other(value) => value.as_str(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for CallType {
|
||||
fn from(value: &str) -> Self {
|
||||
match value {
|
||||
"ocr" => Self::Ocr,
|
||||
"realtime" => Self::Realtime,
|
||||
"completion" => Self::Completion,
|
||||
"acompletion" => Self::Acompletion,
|
||||
"chat_completion" => Self::ChatCompletion,
|
||||
other => Self::Other(other.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for CallType {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct CallbackTiming {
|
||||
pub start_time: f64,
|
||||
pub end_time: f64,
|
||||
}
|
||||
|
||||
impl CallbackTiming {
|
||||
pub fn new(start_time: f64, end_time: f64) -> Self {
|
||||
Self {
|
||||
start_time,
|
||||
end_time,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct CallbackValue {
|
||||
pub object: String,
|
||||
pub value: Value,
|
||||
}
|
||||
|
||||
impl CallbackValue {
|
||||
pub fn new(object: impl Into<String>, value: Value) -> Self {
|
||||
Self {
|
||||
object: object.into(),
|
||||
value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ModelCallDetails {
|
||||
pub model: String,
|
||||
pub custom_llm_provider: String,
|
||||
pub call_type: CallType,
|
||||
pub metadata: StandardLoggingMetadata,
|
||||
pub extra_metadata: HashMap<String, Value>,
|
||||
pub request_id: Option<String>,
|
||||
pub litellm_call_id: Option<String>,
|
||||
pub response_cost: Option<f64>,
|
||||
pub standard_logging_payload: Option<StandardLoggingPayload>,
|
||||
pub failure_error: Option<LoggingError>,
|
||||
}
|
||||
|
||||
impl ModelCallDetails {
|
||||
pub fn new(
|
||||
model: impl Into<String>,
|
||||
custom_llm_provider: impl Into<String>,
|
||||
call_type: CallType,
|
||||
) -> Self {
|
||||
Self {
|
||||
model: model.into(),
|
||||
custom_llm_provider: custom_llm_provider.into(),
|
||||
call_type,
|
||||
metadata: StandardLoggingMetadata::default(),
|
||||
extra_metadata: HashMap::new(),
|
||||
request_id: None,
|
||||
litellm_call_id: None,
|
||||
response_cost: None,
|
||||
standard_logging_payload: None,
|
||||
failure_error: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_standard_logging_payload(payload: StandardLoggingPayload) -> Self {
|
||||
let request_id = Some(payload.id.clone());
|
||||
let litellm_call_id = Some(payload.litellm_call_id.clone());
|
||||
let response_cost = Some(payload.response_cost);
|
||||
let metadata = payload.metadata.clone();
|
||||
Self {
|
||||
model: payload.model.clone(),
|
||||
custom_llm_provider: payload.custom_llm_provider.clone(),
|
||||
call_type: CallType::from(payload.call_type.as_str()),
|
||||
metadata,
|
||||
extra_metadata: HashMap::new(),
|
||||
request_id,
|
||||
litellm_call_id,
|
||||
response_cost,
|
||||
standard_logging_payload: Some(payload),
|
||||
failure_error: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_standard_logging_payload(mut self, payload: StandardLoggingPayload) -> Self {
|
||||
self.model = payload.model.clone();
|
||||
self.custom_llm_provider = payload.custom_llm_provider.clone();
|
||||
self.call_type = CallType::from(payload.call_type.as_str());
|
||||
self.request_id = Some(payload.id.clone());
|
||||
self.litellm_call_id = Some(payload.litellm_call_id.clone());
|
||||
self.response_cost = Some(payload.response_cost);
|
||||
self.metadata = payload.metadata.clone();
|
||||
self.standard_logging_payload = Some(payload);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_failure_error(mut self, error: LoggingError) -> Self {
|
||||
self.failure_error = Some(error);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct LoggingError {
|
||||
pub message: String,
|
||||
pub kind: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct LogError {
|
||||
pub message: String,
|
||||
pub kind: String,
|
||||
}
|
||||
|
||||
impl LogError {
|
||||
pub fn channel_full() -> Self {
|
||||
Self {
|
||||
message: "logging channel is full; dropping record".to_string(),
|
||||
kind: "ChannelFull".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn channel_closed() -> Self {
|
||||
Self {
|
||||
message: "logging channel is closed; worker has shut down".to_string(),
|
||||
kind: "ChannelClosed".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for LogError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}: {}", self.kind, self.message)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for LogError {}
|
||||
|
|
@ -0,0 +1,197 @@
|
|||
//! A `CustomLogger` that ships finished events to the LiteLLM Python proxy's
|
||||
//! `/v1/rust_control_plane/logs` endpoint.
|
||||
//!
|
||||
//! The callback path is non-blocking: `async_log_success_event` /
|
||||
//! `async_log_failure_event`
|
||||
//! build a `LogRecord` and `try_send` it onto a bounded channel, returning a
|
||||
//! `LogError` (never panicking, never awaiting) if the channel is full or the
|
||||
//! worker has gone away. A spawned background worker drains the channel, batches
|
||||
//! records into `{"records":[...]}`, and POSTs them to the proxy with a pooled
|
||||
//! `reqwest::Client`.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use reqwest::Client;
|
||||
use tokio::sync::mpsc::{self, Receiver, Sender};
|
||||
use tokio::time::interval;
|
||||
|
||||
use crate::constants::{DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH};
|
||||
use crate::integrations::custom_logger::{
|
||||
CallbackTiming, CallbackValue, CustomLogger, LogError, LogFuture, LoggingError,
|
||||
ModelCallDetails,
|
||||
};
|
||||
use types::{CallbackLogsRequest, EgressTunables, LogRecord};
|
||||
|
||||
pub mod types;
|
||||
|
||||
/// Ships realtime logging events to the LiteLLM Python proxy.
|
||||
pub struct LiteLLMPythonProxyAPILogger {
|
||||
sink: Sender<LogRecord>,
|
||||
}
|
||||
|
||||
impl LiteLLMPythonProxyAPILogger {
|
||||
/// Spawn the background worker and return a logger handle. `base` is the
|
||||
/// proxy base URL (no trailing path); `master_key` is sent as a bearer token.
|
||||
pub fn start(base: String, master_key: String) -> Arc<Self> {
|
||||
let tunables = EgressTunables::from_env();
|
||||
let (sink, receiver) = mpsc::channel::<LogRecord>(tunables.channel_capacity);
|
||||
let url = format!(
|
||||
"{}{}",
|
||||
base.trim_end_matches('/'),
|
||||
RUST_CONTROL_PLANE_LOGS_PATH
|
||||
);
|
||||
let client = Client::new();
|
||||
tokio::spawn(worker_loop(
|
||||
receiver,
|
||||
client,
|
||||
url,
|
||||
master_key,
|
||||
tunables.max_batch_size,
|
||||
tunables.flush_interval,
|
||||
));
|
||||
Arc::new(Self { sink })
|
||||
}
|
||||
|
||||
/// Build a logger from the environment: `LITELLM_PROXY_BASE_URL` (default
|
||||
/// `http://localhost:4000`) and `LITELLM_MASTER_KEY`.
|
||||
///
|
||||
/// `LITELLM_PROXY_BASE_URL` is treated as the full base and the route is
|
||||
/// appended verbatim, so if the proxy runs under a `SERVER_ROOT_PATH`
|
||||
/// (e.g. served at `https://host/litellm`), include it in the base
|
||||
/// (`LITELLM_PROXY_BASE_URL=https://host/litellm`) and the POST lands at
|
||||
/// `https://host/litellm/v1/rust_control_plane/logs`.
|
||||
pub fn from_env() -> Arc<Self> {
|
||||
let base = std::env::var("LITELLM_PROXY_BASE_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or_else(|| DEFAULT_PROXY_BASE_URL.to_string());
|
||||
let key = std::env::var("LITELLM_MASTER_KEY").unwrap_or_default();
|
||||
Self::start(base, key)
|
||||
}
|
||||
|
||||
fn enqueue(&self, record: LogRecord) -> Result<(), LogError> {
|
||||
self.sink.try_send(record).map_err(|err| match err {
|
||||
mpsc::error::TrySendError::Full(_) => LogError::channel_full(),
|
||||
mpsc::error::TrySendError::Closed(_) => LogError::channel_closed(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl CustomLogger for LiteLLMPythonProxyAPILogger {
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: &'a CallbackValue,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
if let Some(payload) = &model_call_details.standard_logging_payload {
|
||||
self.enqueue(LogRecord {
|
||||
status: "success".to_string(),
|
||||
payload: payload.clone(),
|
||||
error: None,
|
||||
})?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: Option<&'a CallbackValue>,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
if let Some(payload) = &model_call_details.standard_logging_payload {
|
||||
let fallback_error;
|
||||
let error = match &model_call_details.failure_error {
|
||||
Some(error) => error,
|
||||
None => {
|
||||
fallback_error = LoggingError {
|
||||
message: "callback failure event".to_string(),
|
||||
kind: "CallbackFailure".to_string(),
|
||||
};
|
||||
&fallback_error
|
||||
}
|
||||
};
|
||||
self.enqueue(LogRecord {
|
||||
status: "failure".to_string(),
|
||||
payload: payload.clone(),
|
||||
error: Some(format!("{}: {}", error.kind, error.message)),
|
||||
})?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Drain the channel, batching records and POSTing them to the proxy. Exits when
|
||||
/// the channel is closed (all senders dropped) and drained.
|
||||
async fn worker_loop(
|
||||
mut receiver: Receiver<LogRecord>,
|
||||
client: Client,
|
||||
url: String,
|
||||
master_key: String,
|
||||
max_batch_size: usize,
|
||||
flush_interval: Duration,
|
||||
) {
|
||||
let mut ticker = interval(flush_interval);
|
||||
let mut batch: Vec<LogRecord> = Vec::with_capacity(max_batch_size);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
maybe_record = receiver.recv() => {
|
||||
match maybe_record {
|
||||
Some(record) => {
|
||||
batch.push(record);
|
||||
if batch.len() >= max_batch_size {
|
||||
flush(&client, &url, &master_key, &mut batch).await;
|
||||
}
|
||||
}
|
||||
None => {
|
||||
// Channel closed: flush remaining and exit.
|
||||
flush(&client, &url, &master_key, &mut batch).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = ticker.tick() => {
|
||||
flush(&client, &url, &master_key, &mut batch).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST the current batch (if any), clearing it. Errors are logged, not fatal.
|
||||
async fn flush(client: &Client, url: &str, master_key: &str, batch: &mut Vec<LogRecord>) {
|
||||
if batch.is_empty() {
|
||||
return;
|
||||
}
|
||||
let records = std::mem::take(batch)
|
||||
.into_iter()
|
||||
.map(LogRecord::into_callback_record)
|
||||
.collect();
|
||||
let body = CallbackLogsRequest { records };
|
||||
|
||||
let response = client
|
||||
.post(url)
|
||||
.bearer_auth(master_key)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await;
|
||||
|
||||
match response {
|
||||
Ok(resp) if resp.status().is_success() => {}
|
||||
Ok(resp) => {
|
||||
eprintln!(
|
||||
"litellm-ai-gateway: callback logs POST returned {} to {url}",
|
||||
resp.status()
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!("litellm-ai-gateway: callback logs POST failed to {url}: {err}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,72 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::constants::{
|
||||
DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE,
|
||||
};
|
||||
use crate::integrations::types::StandardLoggingPayload;
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct CallbackLogsRequest {
|
||||
pub records: Vec<CallbackLogRecord>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct CallbackLogRecord {
|
||||
pub status: String,
|
||||
pub standard_logging_payload: StandardLoggingPayload,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct LogRecord {
|
||||
pub status: String,
|
||||
pub payload: StandardLoggingPayload,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
impl LogRecord {
|
||||
pub fn into_callback_record(self) -> CallbackLogRecord {
|
||||
CallbackLogRecord {
|
||||
status: self.status,
|
||||
standard_logging_payload: self.payload,
|
||||
error: self.error,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct EgressTunables {
|
||||
pub channel_capacity: usize,
|
||||
pub max_batch_size: usize,
|
||||
pub flush_interval: Duration,
|
||||
}
|
||||
|
||||
impl EgressTunables {
|
||||
pub fn from_env() -> Self {
|
||||
Self {
|
||||
channel_capacity: env_positive(
|
||||
"LITELLM_LOG_CHANNEL_CAPACITY",
|
||||
DEFAULT_CHANNEL_CAPACITY,
|
||||
),
|
||||
max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE),
|
||||
flush_interval: Duration::from_millis(env_positive(
|
||||
"LITELLM_LOG_FLUSH_INTERVAL_MS",
|
||||
DEFAULT_FLUSH_INTERVAL_MS,
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn env_positive<T>(name: &str, default: T) -> T
|
||||
where
|
||||
T: std::str::FromStr + PartialOrd + From<u8>,
|
||||
{
|
||||
let zero = T::from(0u8);
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<T>().ok())
|
||||
.filter(|n| *n > zero)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
12
litellm-rust/crates/ai-gateway/src/integrations/mod.rs
Normal file
12
litellm-rust/crates/ai-gateway/src/integrations/mod.rs
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
//! Pure-Rust logging integrations. Names map 1:1 to Python
|
||||
//! `litellm/integrations/`:
|
||||
//! - [`custom_guardrail::CustomGuardrail`] — the guardrail callback trait
|
||||
//! - [`custom_logger::CustomLogger`] — the callback trait
|
||||
//! - [`litellm_python_proxy_api::LiteLLMPythonProxyAPILogger`] — ships events
|
||||
//! to the Python proxy's `/v1/rust_control_plane/logs` endpoint
|
||||
//! - [`types`] — the typed `StandardLoggingPayload` wire contract
|
||||
|
||||
pub mod custom_guardrail;
|
||||
pub mod custom_logger;
|
||||
pub mod litellm_python_proxy_api;
|
||||
pub mod types;
|
||||
83
litellm-rust/crates/ai-gateway/src/integrations/types.rs
Normal file
83
litellm-rust/crates/ai-gateway/src/integrations/types.rs
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
//! Typed payloads for the LiteLLM `/v1/callbacks/logs` realtime-logging contract.
|
||||
//!
|
||||
//! Field names below are the EXACT JSON keys the Python replay path + spend-logs
|
||||
//! builder read. Note the deliberate mix:
|
||||
//! - `startTime` / `endTime` are camelCase (epoch f64 seconds)
|
||||
//! - `response_cost` / `prompt_tokens` / etc. are snake_case
|
||||
//!
|
||||
//! Mirrors Python `litellm/integrations/` + the proxy `CallbackLogsRequest`
|
||||
//! contract 1:1.
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Cumulative token usage for a realtime session.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub struct Usage {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
}
|
||||
|
||||
/// Cost-attribution metadata threaded from the authenticated request.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct RequestMetadata {
|
||||
pub user_api_key_hash: Option<String>,
|
||||
pub user_api_key_user_id: Option<String>,
|
||||
pub user_api_key_team_id: Option<String>,
|
||||
}
|
||||
|
||||
/// The self-describing payload. Field names are the EXACT JSON keys the Python
|
||||
/// replay path + spend-logs builder read.
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
pub struct StandardLoggingPayload {
|
||||
pub id: String,
|
||||
pub litellm_call_id: String,
|
||||
|
||||
/// e.g. "realtime", "acompletion". Falls back to "acompletion" if absent.
|
||||
pub call_type: String,
|
||||
|
||||
pub model: String,
|
||||
pub custom_llm_provider: String,
|
||||
|
||||
/// Spend ($) written to LiteLLM_SpendLogs.spend.
|
||||
pub response_cost: f64,
|
||||
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
|
||||
/// EPOCH SECONDS as float — camelCase keys, NOT snake_case.
|
||||
#[serde(rename = "startTime")]
|
||||
pub start_time: f64,
|
||||
#[serde(rename = "endTime")]
|
||||
pub end_time: f64,
|
||||
|
||||
pub stream: bool,
|
||||
|
||||
pub metadata: StandardLoggingMetadata,
|
||||
|
||||
/// Optional; stored as request input on the spend log row.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub messages: Option<Value>,
|
||||
}
|
||||
|
||||
/// Cost-attribution keys. The replayer maps these into litellm_params.metadata,
|
||||
/// which the spend-logs builder reads to set user / team_id / organization_id.
|
||||
#[derive(Clone, Debug, Serialize, Default)]
|
||||
pub struct StandardLoggingMetadata {
|
||||
pub user_api_key_hash: Option<String>, // -> SpendLogs.api_key
|
||||
pub user_api_key_user_id: Option<String>, // -> SpendLogs.user
|
||||
pub user_api_key_team_id: Option<String>, // -> SpendLogs.team_id
|
||||
|
||||
// Optional but read by the builder; include when known:
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub user_api_key_alias: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub user_api_key_org_id: Option<String>, // -> SpendLogs.organization_id
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub user_api_key_end_user_id: Option<String>, // -> SpendLogs.end_user
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub spend_logs_metadata: Option<HashMap<String, Value>>,
|
||||
}
|
||||
|
|
@ -1,5 +1,3 @@
|
|||
pub mod mistral;
|
||||
pub mod ocr;
|
||||
pub mod openai;
|
||||
pub mod realtime;
|
||||
pub mod realtime_pool;
|
||||
1
litellm-rust/crates/ai-gateway/src/io/ocr.rs
Normal file
1
litellm-rust/crates/ai-gateway/src/io/ocr.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub use crate::ocr::{ocr, OcrRequest};
|
||||
|
|
@ -1,13 +1,13 @@
|
|||
//! End-to-end OpenAI realtime invocation.
|
||||
//!
|
||||
//! The host-facing entry point, mirroring `providers::ocr::run_ocr`: open the
|
||||
//! WebSocket to OpenAI, then splice a client realtime stream to the upstream,
|
||||
//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms.
|
||||
//! The host-facing entry point opens the WebSocket to OpenAI, then splices a
|
||||
//! client realtime stream to the upstream, driving typed events through the pure
|
||||
//! `OPENAI_REALTIME_CONFIG` transforms.
|
||||
//! Network, auth header, key resolution, and wire (de)serialization live here so
|
||||
//! the `transformation` module stays pure and typed.
|
||||
//!
|
||||
//! The dial and splice steps are factored out ([`dial_upstream`], [`splice`]) so
|
||||
//! the connection pool ([`crate::realtime_pool`]) can pre-establish an upstream,
|
||||
//! the connection pool ([`crate::io::realtime_pool`]) can pre-establish an upstream,
|
||||
//! buffer its `session.created`, and later hand the live socket to the same
|
||||
//! splice loop a fresh dial uses.
|
||||
|
||||
|
|
@ -26,7 +26,7 @@ use tokio_tungstenite::tungstenite::http::HeaderValue;
|
|||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
|
||||
|
||||
use crate::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
|
||||
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
|
||||
|
||||
/// Environment variable holding the OpenAI API key (last-resort fallback).
|
||||
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
|
||||
|
|
@ -126,6 +126,9 @@ pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<Realt
|
|||
/// `session.created` here; the fresh-dial path passes `None` and lets the upstream
|
||||
/// deliver it). Then a single select loop forwards both directions through the
|
||||
/// transforms until either side closes or the idle timeout fires.
|
||||
/// `observe` is invoked on **upstream→client** events only (the trusted side that
|
||||
/// carries `session.created` and `response.done` usage) — never on client events,
|
||||
/// so a client cannot fabricate usage into its own logs.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn splice<In, Out>(
|
||||
model: &str,
|
||||
|
|
@ -133,6 +136,7 @@ pub(crate) async fn splice<In, Out>(
|
|||
mut upstream_rx: UpstreamRx,
|
||||
prelude: Option<RealtimeEvent>,
|
||||
idle_timeout: Option<Duration>,
|
||||
mut observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
mut client_in: In,
|
||||
mut client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
|
|
@ -165,6 +169,10 @@ where
|
|||
// client -> upstream
|
||||
client_event = client_in.next() => {
|
||||
let Some(event) = client_event else { break }; // client disconnected
|
||||
// NOTE: do NOT observe client events. session.created / response.done
|
||||
// (carrying usage) are server→client events; observing the client arm
|
||||
// would let an authenticated client POST a fabricated response.done and
|
||||
// inflate its own spend log. Logging observes upstream events only.
|
||||
for outbound in config.transform_realtime_request(&event, model)?.events {
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
|
|
@ -181,6 +189,7 @@ where
|
|||
Message::Text(text) => {
|
||||
let event: RealtimeEvent = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
observe(&event);
|
||||
for outbound in config.transform_realtime_response(&event, model)?.events {
|
||||
client_out
|
||||
.send(outbound)
|
||||
|
|
@ -207,11 +216,13 @@ where
|
|||
/// framework-agnostic; the gateway adapts its axum socket to these. This is the
|
||||
/// fresh-dial path: dial, then splice. The pool's warm-handoff path skips the dial
|
||||
/// and calls [`splice`] directly with a buffered `session.created`.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn realtime<In, Out>(
|
||||
model: &str,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
idle_timeout: Option<Duration>,
|
||||
observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
|
|
@ -229,19 +240,22 @@ where
|
|||
upstream_rx,
|
||||
None,
|
||||
idle_timeout,
|
||||
observe,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Splice a pre-warmed upstream (taken from [`crate::realtime_pool`]) to the
|
||||
/// Splice a pre-warmed upstream (taken from [`crate::io::realtime_pool`]) to the
|
||||
/// client. Relays the buffered `session.created` first, then splices exactly like
|
||||
/// the fresh-dial path — so a warm session is indistinguishable from a fresh one.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn realtime_warm<In, Out>(
|
||||
model: &str,
|
||||
handoff: crate::realtime_pool::WarmHandoff,
|
||||
handoff: crate::io::realtime_pool::WarmHandoff,
|
||||
idle_timeout: Option<Duration>,
|
||||
observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
|
|
@ -256,6 +270,7 @@ where
|
|||
handoff.rx,
|
||||
Some(handoff.session_created),
|
||||
idle_timeout,
|
||||
observe,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
|
|
@ -281,7 +296,7 @@ mod tests {
|
|||
|
||||
/// Live end-to-end check against OpenAI. Ignored by default (CI never runs
|
||||
/// it); run explicitly with `OPENAI_API_KEY` set:
|
||||
/// `cargo test -p litellm-providers realtime_invokes_openai -- --ignored --nocapture`
|
||||
/// `cargo test -p litellm-ai-gateway --features server realtime_invokes_openai -- --ignored --nocapture`
|
||||
#[tokio::test]
|
||||
#[ignore = "hits the live OpenAI realtime API; needs OPENAI_API_KEY"]
|
||||
async fn realtime_invokes_openai_and_responds() {
|
||||
|
|
@ -303,6 +318,7 @@ mod tests {
|
|||
Some(&key_owned),
|
||||
None,
|
||||
None,
|
||||
|_| {},
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
|
|
@ -6,7 +6,7 @@
|
|||
//! sockets **already connected and already past `session.created`** so a connect
|
||||
//! can be served from a warm socket and the handshake is off the critical path.
|
||||
//!
|
||||
//! Layering: this stays in `providers` (axum-free) next to the dial/splice it
|
||||
//! Layering: this lives in the gateway's `io` module next to the dial/splice it
|
||||
//! reuses. The gateway holds an `Arc<RealtimePool>` in its state and asks for a
|
||||
//! warm socket per connect; on a miss it fresh-dials exactly as before. The pool
|
||||
//! is a latency optimization, never a correctness dependency — see the gateway's
|
||||
|
|
@ -31,7 +31,7 @@ use futures_util::StreamExt;
|
|||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::CoreResult;
|
||||
|
||||
use crate::realtime::{
|
||||
use crate::io::realtime::{
|
||||
dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs,
|
||||
};
|
||||
|
||||
37
litellm-rust/crates/ai-gateway/src/lib.rs
Normal file
37
litellm-rust/crates/ai-gateway/src/lib.rs
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
//! LiteLLM AI Gateway library.
|
||||
//!
|
||||
//! Two layers, split by feature so the Python `cdylib` can depend on the I/O
|
||||
//! without pulling in the HTTP server:
|
||||
//!
|
||||
//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks,
|
||||
//! and provider I/O. Always available — no feature required.
|
||||
//! - [`io`]: compatibility exports and realtime WebSocket splice helpers.
|
||||
//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
|
||||
//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
|
||||
//! binary turns on. The `python-config` feature additionally pulls in [`python`]
|
||||
//! for the load-time config reader.
|
||||
|
||||
pub mod io;
|
||||
pub mod ocr;
|
||||
|
||||
/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
|
||||
/// the `python-config` reader, so it is available without either feature.
|
||||
pub mod gil;
|
||||
|
||||
#[cfg(feature = "server")]
|
||||
pub mod auth;
|
||||
#[cfg(feature = "server")]
|
||||
pub mod routes;
|
||||
#[cfg(feature = "server")]
|
||||
pub mod state;
|
||||
|
||||
// Realtime request logging. Only the server serves realtime, so these are
|
||||
// `server`-gated; `io::realtime` exposes the generic `observe` hook while the
|
||||
// collector and callback fan-out live here.
|
||||
mod constants;
|
||||
pub mod integrations;
|
||||
#[cfg(feature = "server")]
|
||||
mod realtime;
|
||||
|
||||
#[cfg(feature = "python-config")]
|
||||
pub mod python;
|
||||
|
|
@ -1,22 +1,25 @@
|
|||
//! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router.
|
||||
//!
|
||||
//! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment
|
||||
//! (simple-shuffle) → `providers::realtime::realtime()` invokes OpenAI. The
|
||||
//! (simple-shuffle) → `io::realtime::realtime()` invokes OpenAI. The
|
||||
//! server owns transport + config; routing lives in the `router` crate.
|
||||
|
||||
mod auth;
|
||||
mod gil;
|
||||
#[cfg(feature = "python-config")]
|
||||
mod python;
|
||||
mod routes;
|
||||
mod state;
|
||||
//!
|
||||
//! The binary requires the `server` feature (declared in `Cargo.toml` via
|
||||
//! `required-features`), so cargo skips it unless that feature is on. Everything
|
||||
//! the binary needs lives in the library (`litellm_ai_gateway`); `main` just
|
||||
//! wires startup.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool};
|
||||
use litellm_ai_gateway::routes;
|
||||
use litellm_ai_gateway::state::AppState;
|
||||
use litellm_core::router::{Deployment, LiteLLMParams, Router};
|
||||
use litellm_providers::realtime_pool::{upstream_key, PoolConfig, RealtimePool};
|
||||
|
||||
use crate::state::AppState;
|
||||
use litellm_ai_gateway::integrations::custom_logger::CustomLogger;
|
||||
use litellm_ai_gateway::integrations::litellm_python_proxy_api::LiteLLMPythonProxyAPILogger;
|
||||
#[cfg(feature = "python-config")]
|
||||
use litellm_ai_gateway::python;
|
||||
|
||||
/// Bind to localhost by default so the gateway is not a public, unauthenticated
|
||||
/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`).
|
||||
|
|
@ -38,6 +41,12 @@ async fn main() {
|
|||
);
|
||||
}
|
||||
|
||||
// Spawn the realtime-logging worker (drains a channel → POSTs batches to the
|
||||
// Python proxy's /v1/callbacks/logs). Built here so the spawn lands on the
|
||||
// tokio runtime. `from_env` reads LITELLM_PROXY_BASE_URL + LITELLM_MASTER_KEY.
|
||||
let proxy_logger = LiteLLMPythonProxyAPILogger::from_env();
|
||||
let loggers: Vec<Arc<dyn CustomLogger>> = vec![proxy_logger];
|
||||
|
||||
let router = Arc::new(build_router());
|
||||
|
||||
// Build the pre-warmed realtime pool and register each deployment's upstream
|
||||
|
|
@ -61,6 +70,7 @@ async fn main() {
|
|||
let state = AppState {
|
||||
router,
|
||||
master_key,
|
||||
loggers: Arc::new(loggers),
|
||||
realtime_pool,
|
||||
};
|
||||
|
||||
|
|
|
|||
14
litellm-rust/crates/ai-gateway/src/ocr/client.rs
Normal file
14
litellm-rust/crates/ai-gateway/src/ocr/client.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
const OCR_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
pub(super) fn http_client() -> &'static reqwest::Client {
|
||||
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
|
||||
.build()
|
||||
.expect("failed to build reqwest client")
|
||||
})
|
||||
}
|
||||
447
litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs
Normal file
447
litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs
Normal file
|
|
@ -0,0 +1,447 @@
|
|||
use std::net::IpAddr;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use base64::Engine;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::CoreResult;
|
||||
use reqwest::Url;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use litellm_core::providers::azure_ai::ocr::transformation::{
|
||||
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
|
||||
};
|
||||
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
|
||||
use litellm_core::providers::vertex_ai::ocr::transformation::{
|
||||
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
|
||||
};
|
||||
|
||||
use super::client::http_client;
|
||||
|
||||
const ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
|
||||
const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0;
|
||||
const MAX_SAFE_FETCH_REDIRECTS: usize = 10;
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
|
||||
pub(super) fn ocr_provider_config(
|
||||
provider: &str,
|
||||
model: &str,
|
||||
) -> Option<&'static dyn OcrProviderConfig> {
|
||||
match provider {
|
||||
"mistral" => Some(&MISTRAL_OCR_CONFIG),
|
||||
"azure_ai" if is_azure_document_intelligence_model(model) => {
|
||||
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
|
||||
}
|
||||
"azure_ai" => Some(&AZURE_AI_OCR_CONFIG),
|
||||
"vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG),
|
||||
"vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_azure_document_intelligence_model(model: &str) -> bool {
|
||||
let model = model.to_ascii_lowercase();
|
||||
model.contains("doc-intelligence") || model.contains("documentintelligence")
|
||||
}
|
||||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"OCR extra_headers.{key} must be a string, got {}",
|
||||
litellm_core::error::json_type_name(&value)
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
|
||||
headers
|
||||
.iter()
|
||||
.any(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
fn document_url_field(document: &Value) -> CoreResult<Option<(&str, &str)>> {
|
||||
let Some(object) = document.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(doc_type) = object.get("type").and_then(Value::as_str) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let field = match doc_type {
|
||||
"document_url" => "document_url",
|
||||
"image_url" => "image_url",
|
||||
_ => return Ok(None),
|
||||
};
|
||||
let Some(url) = object.get(field).and_then(Value::as_str) else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some((field, url)))
|
||||
}
|
||||
|
||||
fn is_url_requiring_fetch(url: &str) -> bool {
|
||||
!url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://"))
|
||||
}
|
||||
|
||||
fn max_document_download_bytes() -> u64 {
|
||||
let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB")
|
||||
.ok()
|
||||
.and_then(|value| value.parse::<f64>().ok())
|
||||
.unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB);
|
||||
(max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64
|
||||
}
|
||||
|
||||
fn is_blocked_ip(ip: IpAddr) -> bool {
|
||||
match ip {
|
||||
IpAddr::V4(ip) => {
|
||||
ip.is_private()
|
||||
|| ip.is_loopback()
|
||||
|| ip.is_link_local()
|
||||
|| ip.is_broadcast()
|
||||
|| ip.is_multicast()
|
||||
|| ip.is_unspecified()
|
||||
}
|
||||
IpAddr::V6(ip) => {
|
||||
let first_segment = ip.segments()[0];
|
||||
let is_unique_local = (first_segment & 0xfe00) == 0xfc00;
|
||||
let is_link_local = (first_segment & 0xffc0) == 0xfe80;
|
||||
ip.is_loopback()
|
||||
|| ip.is_unspecified()
|
||||
|| ip.is_multicast()
|
||||
|| is_unique_local
|
||||
|| is_link_local
|
||||
|| ip
|
||||
.to_ipv4_mapped()
|
||||
.or_else(|| ip.to_ipv4())
|
||||
.map(|v4| is_blocked_ip(IpAddr::V4(v4)))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn blocked_url_error(url: &Url) -> CoreError {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"OCR document URL rejected by SSRF protection: {url}"
|
||||
))
|
||||
}
|
||||
|
||||
async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
|
||||
if !matches!(url.scheme(), "http" | "https") {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
|
||||
let host = url.host_str().ok_or_else(|| blocked_url_error(url))?;
|
||||
if let Ok(ip) = host.parse::<IpAddr>() {
|
||||
if is_blocked_ip(ip) {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let port = url
|
||||
.port_or_known_default()
|
||||
.ok_or_else(|| blocked_url_error(url))?;
|
||||
let addresses = tokio::net::lookup_host((host, port))
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
let mut saw_address = false;
|
||||
for address in addresses {
|
||||
saw_address = true;
|
||||
if is_blocked_ip(address.ip()) {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
}
|
||||
if !saw_address {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult<Url> {
|
||||
let location = response
|
||||
.headers()
|
||||
.get(reqwest::header::LOCATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("OCR document redirect missing Location header".to_string())
|
||||
})?;
|
||||
url.join(location)
|
||||
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}")))
|
||||
}
|
||||
|
||||
async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> {
|
||||
let client = reqwest::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
let mut current_url = Url::parse(url)
|
||||
.map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
|
||||
|
||||
for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
|
||||
validate_safe_fetch_url(¤t_url).await?;
|
||||
let response = client
|
||||
.get(current_url.clone())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
if !response.status().is_redirection() {
|
||||
return Ok((current_url, response));
|
||||
}
|
||||
current_url = redirect_location(&response, ¤t_url)?;
|
||||
}
|
||||
|
||||
Err(CoreError::InvalidRequest(
|
||||
"Too many redirects while fetching OCR document URL".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> {
|
||||
if max_bytes == 0 {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
|
||||
)));
|
||||
}
|
||||
if content_length > max_bytes {
|
||||
let size_mb = content_length as f64 / (1024.0 * 1024.0);
|
||||
let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_response_with_limit(
|
||||
mut response: reqwest::Response,
|
||||
url: &Url,
|
||||
) -> CoreResult<Vec<u8>> {
|
||||
let max_bytes = max_document_download_bytes();
|
||||
if let Some(content_length) = response.content_length() {
|
||||
enforce_download_size(content_length, max_bytes, url)?;
|
||||
} else {
|
||||
enforce_download_size(0, max_bytes, url)?;
|
||||
}
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
let mut bytes_downloaded: u64 = 0;
|
||||
while let Some(chunk) = response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?
|
||||
{
|
||||
bytes_downloaded += chunk.len() as u64;
|
||||
enforce_download_size(bytes_downloaded, max_bytes, url)?;
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult<Value> {
|
||||
let Some((field, url)) = document_url_field(&document)? else {
|
||||
return Ok(document);
|
||||
};
|
||||
if !is_url_requiring_fetch(url) {
|
||||
return Ok(document);
|
||||
}
|
||||
|
||||
let (final_url, response) = safe_get_document_url(url).await?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&body),
|
||||
});
|
||||
}
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.split(';').next())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("application/octet-stream")
|
||||
.to_string();
|
||||
let bytes = read_response_with_limit(response, &final_url).await?;
|
||||
let data_uri = format!(
|
||||
"data:{content_type};base64,{}",
|
||||
BASE64_STANDARD.encode(bytes)
|
||||
);
|
||||
|
||||
let mut transformed = document
|
||||
.as_object()
|
||||
.cloned()
|
||||
.ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?;
|
||||
transformed.insert(field.to_string(), Value::String(data_uri));
|
||||
Ok(Value::Object(transformed))
|
||||
}
|
||||
|
||||
fn same_origin(left: &str, right: &str) -> bool {
|
||||
let Ok(left) = reqwest::Url::parse(left) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(right) = reqwest::Url::parse(right) else {
|
||||
return false;
|
||||
};
|
||||
left.scheme() == right.scheme()
|
||||
&& left.host_str() == right.host_str()
|
||||
&& left.port_or_known_default() == right.port_or_known_default()
|
||||
}
|
||||
|
||||
fn retry_after_secs(response: &reqwest::Response) -> u64 {
|
||||
response
|
||||
.headers()
|
||||
.get(reqwest::header::RETRY_AFTER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
.unwrap_or(2)
|
||||
}
|
||||
|
||||
fn operation_status(response_json: &Value) -> CoreResult<&str> {
|
||||
let status = response_json
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("status"))?;
|
||||
match status {
|
||||
"succeeded" => Ok("succeeded"),
|
||||
"running" | "notStarted" => Ok("running"),
|
||||
"failed" => {
|
||||
let message = response_json
|
||||
.get("error")
|
||||
.and_then(|error| error.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("Unknown error");
|
||||
Err(CoreError::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed: {message}"
|
||||
)))
|
||||
}
|
||||
other => Err(CoreError::InvalidResponse(format!(
|
||||
"Unknown operation status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn poll_document_intelligence(
|
||||
operation_url: &str,
|
||||
original_url: &str,
|
||||
headers: &[(String, String)],
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Value> {
|
||||
if !same_origin(operation_url, original_url) {
|
||||
return Err(CoreError::InvalidResponse(
|
||||
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
let timeout = timeout.unwrap_or(Duration::from_secs(
|
||||
AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS,
|
||||
));
|
||||
loop {
|
||||
if start.elapsed() > timeout {
|
||||
return Err(CoreError::Network(format!(
|
||||
"Azure Document Intelligence operation polling timed out after {} seconds",
|
||||
timeout.as_secs()
|
||||
)));
|
||||
}
|
||||
|
||||
let mut request_builder = http_client().get(operation_url);
|
||||
for (key, value) in headers {
|
||||
if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
}
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
let retry_after = retry_after_secs(&response);
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
|
||||
})?;
|
||||
if operation_status(&response_json)? == "succeeded" {
|
||||
return Ok(response_json);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_secs(retry_after)).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn blocks_private_and_metadata_ips() {
|
||||
assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("10.0.0.1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("169.254.169.254".parse().unwrap()));
|
||||
assert!(is_blocked_ip("::1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("fd00::1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("fe80::1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap()));
|
||||
assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap()));
|
||||
assert!(!is_blocked_ip("8.8.8.8".parse().unwrap()));
|
||||
assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn convert_document_url_rejects_loopback_fetch() {
|
||||
let error = convert_document_url_to_data_uri(json!({
|
||||
"type": "image_url",
|
||||
"image_url": "http://127.0.0.1/image.png"
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
CoreError::InvalidRequest(message)
|
||||
if message.contains("SSRF protection")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn convert_document_url_leaves_data_uri_untouched() {
|
||||
let document = json!({
|
||||
"type": "image_url",
|
||||
"image_url": "data:image/png;base64,abcd"
|
||||
});
|
||||
|
||||
let transformed = convert_document_url_to_data_uri(document.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(transformed, document);
|
||||
}
|
||||
}
|
||||
71
litellm-rust/crates/ai-gateway/src/ocr/handler.rs
Normal file
71
litellm-rust/crates/ai-gateway/src/ocr/handler.rs
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrResponseHandling;
|
||||
use litellm_core::CoreResult;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::common_utils::{poll_document_intelligence, truncate_error_body};
|
||||
use super::types::ProviderOcrRequest;
|
||||
|
||||
pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
&& status.as_u16() == 202
|
||||
{
|
||||
let operation_url = response
|
||||
.headers()
|
||||
.get("operation-location")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse(
|
||||
"Azure Document Intelligence returned 202 but no Operation-Location header found"
|
||||
.to_string(),
|
||||
)
|
||||
})?;
|
||||
let response_json = poll_document_intelligence(
|
||||
&operation_url,
|
||||
&request.url,
|
||||
&request.upstream_headers,
|
||||
request.timeout,
|
||||
)
|
||||
.await?;
|
||||
return Ok(request
|
||||
.config
|
||||
.transform_ocr_response(&request.model, response_json)?
|
||||
.into_json());
|
||||
}
|
||||
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let response_json: Value = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
|
||||
|
||||
Ok(request
|
||||
.config
|
||||
.transform_ocr_response(&request.model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
329
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
Normal file
329
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrAuthStrategy;
|
||||
use litellm_core::CoreResult;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use super::common_utils::{
|
||||
convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers,
|
||||
};
|
||||
use super::types::{PreparedOcrRequest, ProviderOcrRequest};
|
||||
use crate::integrations::custom_guardrail::{
|
||||
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
|
||||
};
|
||||
use crate::integrations::custom_logger::{
|
||||
CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
|
||||
};
|
||||
use crate::integrations::types::{
|
||||
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
|
||||
};
|
||||
|
||||
pub(crate) struct OcrLifecycleHooks {
|
||||
logger_runner: CustomLoggerRunner,
|
||||
guardrail_runner: CustomGuardrailRunner,
|
||||
request_metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
|
||||
type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
|
||||
|
||||
impl OcrLifecycleHooks {
|
||||
pub(crate) fn new(
|
||||
logger_runner: CustomLoggerRunner,
|
||||
guardrail_runner: CustomGuardrailRunner,
|
||||
request_metadata: RequestMetadata,
|
||||
) -> Self {
|
||||
Self {
|
||||
logger_runner,
|
||||
guardrail_runner,
|
||||
request_metadata,
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_pre_call_guardrails(
|
||||
&self,
|
||||
request: PreparedOcrRequest,
|
||||
) -> CoreResult<PreparedOcrRequest> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
|
||||
let context = guardrail_context(&self.request_metadata);
|
||||
let guardrail_request = GuardrailRequest::new(json!({
|
||||
"model": request.model,
|
||||
"custom_llm_provider": request.custom_llm_provider,
|
||||
"document": request.document,
|
||||
"optional_params": request.optional_params,
|
||||
}));
|
||||
let (guardrail_request, _) = self
|
||||
.guardrail_runner
|
||||
.run_pre_call(&context, guardrail_request)
|
||||
.await
|
||||
.map_err(guardrail_error_to_core_error)?;
|
||||
let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
|
||||
Ok(PreparedOcrRequest {
|
||||
document,
|
||||
optional_params,
|
||||
..request
|
||||
})
|
||||
}
|
||||
|
||||
async fn prepare_provider_request(
|
||||
&self,
|
||||
request: PreparedOcrRequest,
|
||||
) -> CoreResult<ProviderOcrRequest> {
|
||||
let config = ocr_provider_config(&request.custom_llm_provider, &request.model)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
let headers = string_headers(request.extra_headers)?;
|
||||
let auth_strategy = config.auth_strategy();
|
||||
let api_key = (!has_header(&headers, auth_strategy.header_name()))
|
||||
.then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup))
|
||||
.transpose()?;
|
||||
let url = config.complete_url(
|
||||
request.api_base.as_deref(),
|
||||
&request.model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_ocr_params(&request.optional_params);
|
||||
let model = request.model.clone();
|
||||
let custom_llm_provider = request.custom_llm_provider.clone();
|
||||
let document = if config.requires_data_uri_document() {
|
||||
convert_document_url_to_data_uri(request.document).await?
|
||||
} else {
|
||||
request.document
|
||||
};
|
||||
let body = config
|
||||
.transform_ocr_request(&request.model, document, filtered_params)?
|
||||
.data;
|
||||
let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref());
|
||||
let body = self
|
||||
.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
|
||||
.await?;
|
||||
Ok(ProviderOcrRequest {
|
||||
model,
|
||||
config,
|
||||
url,
|
||||
body,
|
||||
upstream_headers,
|
||||
timeout: request.timeout,
|
||||
})
|
||||
}
|
||||
|
||||
async fn run_during_call_guardrails(
|
||||
&self,
|
||||
model: &str,
|
||||
custom_llm_provider: &str,
|
||||
url: &str,
|
||||
body: Value,
|
||||
) -> CoreResult<Value> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(body);
|
||||
}
|
||||
|
||||
let context = guardrail_context(&self.request_metadata);
|
||||
let guardrail_request = GuardrailRequest::new(json!({
|
||||
"model": model,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"url": url,
|
||||
"body": body,
|
||||
}));
|
||||
let (guardrail_request, _) = self
|
||||
.guardrail_runner
|
||||
.run_during_call(&context, guardrail_request)
|
||||
.await
|
||||
.map_err(guardrail_error_to_core_error)?;
|
||||
parse_ocr_during_call_guardrail_request(guardrail_request)
|
||||
}
|
||||
|
||||
fn standard_logging_payload(
|
||||
&self,
|
||||
context: &CallLifecycleContext,
|
||||
timing: &CallLifecycleTiming,
|
||||
) -> StandardLoggingPayload {
|
||||
StandardLoggingPayload {
|
||||
id: context.litellm_call_id.clone(),
|
||||
litellm_call_id: context.litellm_call_id.clone(),
|
||||
call_type: context.call_type.clone(),
|
||||
model: context.model.clone(),
|
||||
custom_llm_provider: context.custom_llm_provider.clone(),
|
||||
response_cost: 0.0,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
total_tokens: 0,
|
||||
start_time: timing.start_time,
|
||||
end_time: timing.end_time,
|
||||
stream: false,
|
||||
metadata: StandardLoggingMetadata {
|
||||
user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
|
||||
user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
|
||||
user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
|
||||
..Default::default()
|
||||
},
|
||||
messages: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLifecycleHooks {
|
||||
type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
|
||||
type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>;
|
||||
type SuccessFuture<'a> = OcrLogFuture<'a>;
|
||||
type FailureFuture<'a> = OcrLogFuture<'a>;
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
request: PreparedOcrRequest,
|
||||
) -> Self::PreCallFuture<'a> {
|
||||
Box::pin(async move { self.run_pre_call_guardrails(request).await })
|
||||
}
|
||||
|
||||
fn async_during_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
request: PreparedOcrRequest,
|
||||
) -> Self::DuringCallFuture<'a> {
|
||||
Box::pin(async move { self.prepare_provider_request(request).await })
|
||||
}
|
||||
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
response: &'a Value,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::SuccessFuture<'a> {
|
||||
Box::pin(async move {
|
||||
if self.logger_runner.is_empty() {
|
||||
return;
|
||||
}
|
||||
let response_obj = CallbackValue::new("ocr", response.clone());
|
||||
self.logger_runner
|
||||
.async_log_success_event(
|
||||
&ModelCallDetails::from_standard_logging_payload(
|
||||
self.standard_logging_payload(context, timing),
|
||||
),
|
||||
&response_obj,
|
||||
CallbackTiming::new(timing.start_time, timing.end_time),
|
||||
)
|
||||
.await;
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
error: &'a CoreError,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
if self.logger_runner.is_empty() {
|
||||
return;
|
||||
}
|
||||
let logging_error = LoggingError {
|
||||
message: error.to_string(),
|
||||
kind: core_error_kind(error).to_string(),
|
||||
};
|
||||
let response_obj = CallbackValue::new(
|
||||
"error",
|
||||
json!({
|
||||
"message": logging_error.message,
|
||||
"kind": logging_error.kind,
|
||||
}),
|
||||
);
|
||||
self.logger_runner
|
||||
.async_log_failure_event(
|
||||
&ModelCallDetails::from_standard_logging_payload(
|
||||
self.standard_logging_payload(context, timing),
|
||||
)
|
||||
.with_failure_error(logging_error),
|
||||
Some(&response_obj),
|
||||
CallbackTiming::new(timing.start_time, timing.end_time),
|
||||
)
|
||||
.await;
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn upstream_headers(
|
||||
headers: &[(String, String)],
|
||||
auth_strategy: OcrAuthStrategy,
|
||||
api_key: Option<&str>,
|
||||
) -> Vec<(String, String)> {
|
||||
api_key
|
||||
.map(|api_key| match auth_strategy {
|
||||
OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
|
||||
OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()),
|
||||
})
|
||||
.into_iter()
|
||||
.chain(headers.iter().cloned())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
|
||||
GuardrailContext {
|
||||
call_type: CallType::Ocr,
|
||||
selected_guardrails: Vec::new(),
|
||||
metadata: std::collections::HashMap::new(),
|
||||
user_api_key_hash: metadata.user_api_key_hash.clone(),
|
||||
user_api_key_user_id: metadata.user_api_key_user_id.clone(),
|
||||
user_api_key_team_id: metadata.user_api_key_team_id.clone(),
|
||||
trace_parent: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_ocr_pre_call_guardrail_request(
|
||||
request: GuardrailRequest,
|
||||
) -> CoreResult<(Value, Map<String, Value>)> {
|
||||
let Value::Object(mut data) = request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"OCR pre_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
let document = data.remove("document").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string())
|
||||
})?;
|
||||
let optional_params = match data.remove("optional_params") {
|
||||
Some(Value::Object(params)) => params,
|
||||
Some(_) => {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"OCR pre_call guardrail optional_params must be an object".to_string(),
|
||||
))
|
||||
}
|
||||
None => Map::new(),
|
||||
};
|
||||
Ok((document, optional_params))
|
||||
}
|
||||
|
||||
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult<Value> {
|
||||
let Value::Object(mut data) = request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"OCR during_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
data.remove("body").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
|
||||
CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
|
||||
}
|
||||
|
||||
fn core_error_kind(error: &CoreError) -> &'static str {
|
||||
match error {
|
||||
CoreError::Auth(_) => "AuthError",
|
||||
CoreError::InvalidProvider(_) => "InvalidProvider",
|
||||
CoreError::InvalidRequest(_) => "InvalidRequest",
|
||||
CoreError::InvalidType { .. } => "InvalidType",
|
||||
CoreError::MissingField(_) => "MissingField",
|
||||
CoreError::Http { .. } => "HttpError",
|
||||
CoreError::InvalidResponse(_) => "InvalidResponse",
|
||||
CoreError::Network(_) => "NetworkError",
|
||||
CoreError::Routing(_) => "RoutingError",
|
||||
}
|
||||
}
|
||||
25
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
Normal file
25
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
use litellm_core::call_lifecycle::CallLifecycle;
|
||||
use litellm_core::CoreResult;
|
||||
use serde_json::Value;
|
||||
|
||||
mod client;
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod hooks;
|
||||
mod prepare;
|
||||
mod types;
|
||||
|
||||
pub use types::OcrRequest;
|
||||
|
||||
use handler::execute_ocr_provider_call;
|
||||
use prepare::{prepare_ocr_call, PreparedOcrCall};
|
||||
|
||||
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
|
||||
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
|
||||
CallLifecycle::default()
|
||||
.run_request(request, &hooks, execute_ocr_provider_call)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
57
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
Normal file
57
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
|
||||
|
||||
use super::hooks::OcrLifecycleHooks;
|
||||
use super::types::{OcrRequest, PreparedOcrRequest};
|
||||
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
|
||||
use crate::integrations::custom_logger::CustomLoggerRunner;
|
||||
|
||||
pub(crate) struct PreparedOcrCall {
|
||||
pub(crate) request: PreparedOcrRequest,
|
||||
pub(crate) hooks: OcrLifecycleHooks,
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
|
||||
let call_id = request
|
||||
.litellm_call_id
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(new_ocr_call_id);
|
||||
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
|
||||
.unwrap_or(CustomLlmProvider {
|
||||
model: request.model,
|
||||
custom_llm_provider: "mistral",
|
||||
});
|
||||
let model = provider_info.model.to_string();
|
||||
let custom_llm_provider = provider_info.custom_llm_provider.to_string();
|
||||
|
||||
PreparedOcrCall {
|
||||
request: PreparedOcrRequest {
|
||||
model,
|
||||
custom_llm_provider,
|
||||
litellm_call_id: call_id,
|
||||
document: request.document,
|
||||
api_key: request.api_key.map(str::to_string),
|
||||
api_base: request.api_base.map(str::to_string),
|
||||
extra_headers: request.extra_headers,
|
||||
optional_params: request.optional_params,
|
||||
timeout: request.timeout,
|
||||
},
|
||||
hooks: OcrLifecycleHooks::new(
|
||||
CustomLoggerRunner::new(request.callbacks),
|
||||
CustomGuardrailRunner::new(request.guardrails),
|
||||
request.request_metadata,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn new_ocr_call_id() -> String {
|
||||
static COUNTER: AtomicU64 = AtomicU64::new(1);
|
||||
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||
let timestamp = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_nanos())
|
||||
.unwrap_or(0);
|
||||
format!("ocr-{timestamp}-{sequence}")
|
||||
}
|
||||
610
litellm-rust/crates/ai-gateway/src/ocr/tests.rs
Normal file
610
litellm-rust/crates/ai-gateway/src/ocr/tests.rs
Normal file
|
|
@ -0,0 +1,610 @@
|
|||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrResponseHandling;
|
||||
use serde_json::{json, Map, Value};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body};
|
||||
use super::{ocr, OcrRequest};
|
||||
use crate::integrations::custom_guardrail::{
|
||||
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
|
||||
GuardrailFuture, GuardrailRequest,
|
||||
};
|
||||
use crate::integrations::custom_logger::{
|
||||
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
|
||||
};
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
|
||||
async fn read_http_headers(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
loop {
|
||||
let n = socket.read(&mut buffer).await.expect("reads request");
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
if request.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
String::from_utf8(request).expect("request is utf8")
|
||||
}
|
||||
|
||||
async fn read_http_request(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
let header_end = loop {
|
||||
let n = socket.read(&mut buffer).await.expect("reads request");
|
||||
if n == 0 {
|
||||
break request.len();
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
|
||||
break position + 4;
|
||||
}
|
||||
};
|
||||
let headers = String::from_utf8_lossy(&request[..header_end]);
|
||||
let content_length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().ok())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or(0);
|
||||
while request.len().saturating_sub(header_end) < content_length {
|
||||
let n = socket.read(&mut buffer).await.expect("reads body");
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
}
|
||||
String::from_utf8(request).expect("request is utf8")
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
struct RecordedLogEvent {
|
||||
hook: &'static str,
|
||||
model: String,
|
||||
call_type: String,
|
||||
user_id: Option<String>,
|
||||
response_object: Option<String>,
|
||||
error_kind: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingOcrLogger {
|
||||
events: Mutex<Vec<RecordedLogEvent>>,
|
||||
}
|
||||
|
||||
impl RecordingOcrLogger {
|
||||
fn events(&self) -> Vec<RecordedLogEvent> {
|
||||
self.events.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl CustomLogger for RecordingOcrLogger {
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
response_obj: &'a CallbackValue,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push(RecordedLogEvent {
|
||||
hook: "async_log_success_event",
|
||||
model: model_call_details.model.clone(),
|
||||
call_type: model_call_details.call_type.to_string(),
|
||||
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
|
||||
response_object: Some(response_obj.object.clone()),
|
||||
error_kind: None,
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
response_obj: Option<&'a CallbackValue>,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push(RecordedLogEvent {
|
||||
hook: "async_log_failure_event",
|
||||
model: model_call_details.model.clone(),
|
||||
call_type: model_call_details.call_type.to_string(),
|
||||
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
|
||||
response_object: response_obj.map(|value| value.object.clone()),
|
||||
error_kind: model_call_details
|
||||
.failure_error
|
||||
.as_ref()
|
||||
.map(|error| error.kind.clone()),
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct RecordingOcrGuardrail {
|
||||
hooks: Vec<GuardrailEventHook>,
|
||||
events: Mutex<Vec<&'static str>>,
|
||||
block_pre_call: bool,
|
||||
}
|
||||
|
||||
impl RecordingOcrGuardrail {
|
||||
fn new(hooks: Vec<GuardrailEventHook>) -> Self {
|
||||
Self {
|
||||
hooks,
|
||||
events: Mutex::new(Vec::new()),
|
||||
block_pre_call: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn blocking_pre_call() -> Self {
|
||||
Self {
|
||||
hooks: vec![GuardrailEventHook::PreCall],
|
||||
events: Mutex::new(Vec::new()),
|
||||
block_pre_call: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn events(&self) -> Vec<&'static str> {
|
||||
self.events.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl CustomGuardrail for RecordingOcrGuardrail {
|
||||
fn guardrail_name(&self) -> &str {
|
||||
"recording-ocr-guardrail"
|
||||
}
|
||||
|
||||
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
|
||||
&self.hooks
|
||||
}
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a GuardrailContext,
|
||||
mut request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("async_pre_call_hook");
|
||||
if self.block_pre_call {
|
||||
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
|
||||
"blocked before provider",
|
||||
)));
|
||||
}
|
||||
request.data["document"]["guarded_pre"] = json!(true);
|
||||
Ok(GuardrailDecision::Mask(request))
|
||||
})
|
||||
}
|
||||
|
||||
fn async_moderation_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a GuardrailContext,
|
||||
mut request: GuardrailRequest,
|
||||
) -> GuardrailFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("async_moderation_hook");
|
||||
request.data["body"]["guarded_during"] = json!(true);
|
||||
Ok(GuardrailDecision::Mask(request))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_passes_short_strings_through() {
|
||||
let body = "Unauthorized";
|
||||
assert_eq!(truncate_error_body(body), "Unauthorized");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_caps_long_payloads() {
|
||||
let body = "x".repeat(306);
|
||||
let truncated = truncate_error_body(&body);
|
||||
|
||||
assert!(truncated.ends_with("... (truncated)"));
|
||||
let prefix_chars = truncated
|
||||
.strip_suffix("... (truncated)")
|
||||
.expect("truncated marker present")
|
||||
.chars()
|
||||
.count();
|
||||
assert_eq!(prefix_chars, 256);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_does_not_split_multibyte_chars() {
|
||||
let body = "é".repeat(266);
|
||||
let truncated = truncate_error_body(&body);
|
||||
assert!(truncated.is_char_boundary(truncated.len()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ocr_dispatch_supports_migrated_providers() {
|
||||
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
|
||||
assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409")
|
||||
.expect("azure ai config resolves")
|
||||
.requires_data_uri_document());
|
||||
assert_eq!(
|
||||
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
|
||||
.expect("document intelligence config resolves")
|
||||
.response_handling(),
|
||||
OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
);
|
||||
assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
|
||||
.expect("vertex deepseek config resolves")
|
||||
.supported_ocr_params()
|
||||
.contains(&"temperature"));
|
||||
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_headers_accepts_string_values() {
|
||||
let headers = json!({
|
||||
"x-trace-id": "trace-1"
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
|
||||
assert_eq!(
|
||||
string_headers(Some(headers)).expect("string headers accepted"),
|
||||
vec![("x-trace-id".to_string(), "trace-1".to_string())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_header_detection_is_case_insensitive() {
|
||||
let headers = vec![
|
||||
("x-trace-id".to_string(), "trace-1".to_string()),
|
||||
("authorization".to_string(), "Bearer sk-test".to_string()),
|
||||
];
|
||||
|
||||
assert!(has_header(&headers, "authorization"));
|
||||
|
||||
let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
|
||||
assert!(has_header(&headers, "authorization"));
|
||||
|
||||
let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
|
||||
assert!(!has_header(&headers, "authorization"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts one request");
|
||||
let request = read_http_request(&mut socket).await;
|
||||
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
response_body.len(),
|
||||
response_body
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
request
|
||||
});
|
||||
|
||||
let logger = Arc::new(RecordingOcrLogger::default());
|
||||
let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![
|
||||
GuardrailEventHook::PreCall,
|
||||
GuardrailEventHook::DuringCall,
|
||||
]));
|
||||
let response = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("sk-test"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("mistral"),
|
||||
extra_headers: None,
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
callbacks: vec![logger.clone()],
|
||||
guardrails: vec![guardrail.clone()],
|
||||
request_metadata: RequestMetadata {
|
||||
user_api_key_user_id: Some("user-1".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
litellm_call_id: Some("ocr-call-1"),
|
||||
})
|
||||
.await
|
||||
.expect("ocr request succeeds");
|
||||
|
||||
assert_eq!(response["pages"][0]["markdown"], "ok");
|
||||
assert_eq!(
|
||||
guardrail.events(),
|
||||
vec!["async_pre_call_hook", "async_moderation_hook"]
|
||||
);
|
||||
assert_eq!(
|
||||
logger.events(),
|
||||
vec![RecordedLogEvent {
|
||||
hook: "async_log_success_event",
|
||||
model: "mistral-ocr-latest".to_string(),
|
||||
call_type: "ocr".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
response_object: Some("ocr".to_string()),
|
||||
error_kind: None,
|
||||
}]
|
||||
);
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
|
||||
assert!(request.contains(r#""guarded_during":true"#), "{request}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts one request");
|
||||
let _request = read_http_request(&mut socket).await;
|
||||
let response_body = "provider failed";
|
||||
let response = format!(
|
||||
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
response_body.len(),
|
||||
response_body
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
});
|
||||
|
||||
let logger = Arc::new(RecordingOcrLogger::default());
|
||||
let err = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("sk-test"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("mistral"),
|
||||
extra_headers: None,
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
callbacks: vec![logger.clone()],
|
||||
guardrails: Vec::new(),
|
||||
request_metadata: RequestMetadata::default(),
|
||||
litellm_call_id: Some("ocr-call-2"),
|
||||
})
|
||||
.await
|
||||
.expect_err("provider error propagates");
|
||||
|
||||
assert!(matches!(err, CoreError::Http { status: 500, .. }));
|
||||
server.await.expect("server task completes");
|
||||
assert_eq!(
|
||||
logger.events(),
|
||||
vec![RecordedLogEvent {
|
||||
hook: "async_log_failure_event",
|
||||
model: "mistral-ocr-latest".to_string(),
|
||||
call_type: "ocr".to_string(),
|
||||
user_id: None,
|
||||
response_object: Some("error".to_string()),
|
||||
error_kind: Some("HttpError".to_string()),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
let logger = Arc::new(RecordingOcrLogger::default());
|
||||
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call());
|
||||
|
||||
let err = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("sk-test"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("mistral"),
|
||||
extra_headers: None,
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_millis(100)),
|
||||
callbacks: vec![logger.clone()],
|
||||
guardrails: vec![guardrail.clone()],
|
||||
request_metadata: RequestMetadata::default(),
|
||||
litellm_call_id: Some("ocr-call-3"),
|
||||
})
|
||||
.await
|
||||
.expect_err("guardrail blocks request");
|
||||
|
||||
assert!(matches!(err, CoreError::InvalidRequest(_)));
|
||||
assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]);
|
||||
assert_eq!(
|
||||
logger.events(),
|
||||
vec![RecordedLogEvent {
|
||||
hook: "async_log_failure_event",
|
||||
model: "mistral-ocr-latest".to_string(),
|
||||
call_type: "ocr".to_string(),
|
||||
user_id: None,
|
||||
response_object: Some("error".to_string()),
|
||||
error_kind: Some("InvalidRequest".to_string()),
|
||||
}]
|
||||
);
|
||||
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
|
||||
assert!(accepted.is_err(), "provider socket should not be touched");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts one request");
|
||||
let request = read_http_headers(&mut socket).await;
|
||||
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
response_body.len(),
|
||||
response_body
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
request
|
||||
});
|
||||
|
||||
let mut headers = Map::new();
|
||||
headers.insert(
|
||||
"Authorization".to_string(),
|
||||
Value::String("Bearer sk-from-python".to_string()),
|
||||
);
|
||||
headers.insert(
|
||||
"x-trace-id".to_string(),
|
||||
Value::String("trace-1".to_string()),
|
||||
);
|
||||
|
||||
let response = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("sk-for-rust-fallback"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("mistral"),
|
||||
extra_headers: Some(headers),
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
callbacks: Vec::new(),
|
||||
guardrails: Vec::new(),
|
||||
request_metadata: RequestMetadata::default(),
|
||||
litellm_call_id: None,
|
||||
})
|
||||
.await
|
||||
.expect("ocr request succeeds");
|
||||
|
||||
assert_eq!(response["pages"][0]["markdown"], "ok");
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
let authorization_count = request
|
||||
.lines()
|
||||
.filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
|
||||
.count();
|
||||
assert_eq!(authorization_count, 1, "{request}");
|
||||
assert!(
|
||||
request.contains("authorization: Bearer sk-from-python")
|
||||
|| request.contains("Authorization: Bearer sk-from-python"),
|
||||
"{request}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn document_intelligence_poll_uses_resolved_subscription_key() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
let operation_url = format!("http://{addr}/operations/1");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
|
||||
let post_request = read_http_headers(&mut post_socket).await;
|
||||
let post_response = format!(
|
||||
"HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
|
||||
);
|
||||
post_socket
|
||||
.write_all(post_response.as_bytes())
|
||||
.await
|
||||
.expect("writes post response");
|
||||
|
||||
let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
|
||||
let poll_request = read_http_headers(&mut poll_socket).await;
|
||||
let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
|
||||
let poll_response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
response_body.len(),
|
||||
response_body
|
||||
);
|
||||
poll_socket
|
||||
.write_all(poll_response.as_bytes())
|
||||
.await
|
||||
.expect("writes poll response");
|
||||
(post_request, poll_request)
|
||||
});
|
||||
|
||||
let response = ocr(OcrRequest {
|
||||
model: "doc-intelligence/prebuilt-read",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("di-key"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("azure_ai"),
|
||||
extra_headers: None,
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
callbacks: Vec::new(),
|
||||
guardrails: Vec::new(),
|
||||
request_metadata: RequestMetadata::default(),
|
||||
litellm_call_id: None,
|
||||
})
|
||||
.await
|
||||
.expect("document intelligence request succeeds");
|
||||
|
||||
assert_eq!(response["pages"][0]["markdown"], "ok");
|
||||
|
||||
let (post_request, poll_request) = server.await.expect("server task completes");
|
||||
assert!(
|
||||
post_request
|
||||
.to_ascii_lowercase()
|
||||
.contains("ocp-apim-subscription-key: di-key"),
|
||||
"{post_request}"
|
||||
);
|
||||
assert!(
|
||||
poll_request
|
||||
.to_ascii_lowercase()
|
||||
.contains("ocp-apim-subscription-key: di-key"),
|
||||
"{poll_request}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_headers_rejects_non_string_values() {
|
||||
let headers = json!({
|
||||
"x-retry-count": 3
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
|
||||
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidRequest(
|
||||
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
57
litellm-rust/crates/ai-gateway/src/ocr/types.rs
Normal file
57
litellm-rust/crates/ai-gateway/src/ocr/types.rs
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::integrations::custom_guardrail::CustomGuardrail;
|
||||
use crate::integrations::custom_logger::CustomLogger;
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
|
||||
pub struct OcrRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub document: Value,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub timeout: Option<Duration>,
|
||||
pub callbacks: Vec<Arc<dyn CustomLogger>>,
|
||||
pub guardrails: Vec<Arc<dyn CustomGuardrail>>,
|
||||
pub request_metadata: RequestMetadata,
|
||||
pub litellm_call_id: Option<&'a str>,
|
||||
}
|
||||
|
||||
pub(crate) struct PreparedOcrRequest {
|
||||
pub(crate) model: String,
|
||||
pub(crate) custom_llm_provider: String,
|
||||
pub(crate) litellm_call_id: String,
|
||||
pub(crate) document: Value,
|
||||
pub(crate) api_key: Option<String>,
|
||||
pub(crate) api_base: Option<String>,
|
||||
pub(crate) extra_headers: Option<Map<String, Value>>,
|
||||
pub(crate) optional_params: Map<String, Value>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl CallLifecycleRequest for PreparedOcrRequest {
|
||||
fn lifecycle_context(&self) -> CallLifecycleContext {
|
||||
CallLifecycleContext::new(
|
||||
"ocr",
|
||||
self.model.clone(),
|
||||
self.custom_llm_provider.clone(),
|
||||
self.litellm_call_id.clone(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct ProviderOcrRequest {
|
||||
pub(crate) model: String,
|
||||
pub(crate) config: &'static dyn OcrProviderConfig,
|
||||
pub(crate) url: String,
|
||||
pub(crate) body: Value,
|
||||
pub(crate) upstream_headers: Vec<(String, String)>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
4
litellm-rust/crates/ai-gateway/src/realtime/mod.rs
Normal file
4
litellm-rust/crates/ai-gateway/src/realtime/mod.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
//! Realtime logging collector. Observes the realtime event stream and emits a
|
||||
//! `StandardLoggingPayload` to the registered callbacks on session close.
|
||||
|
||||
pub mod streaming;
|
||||
388
litellm-rust/crates/ai-gateway/src/realtime/streaming.rs
Normal file
388
litellm-rust/crates/ai-gateway/src/realtime/streaming.rs
Normal file
|
|
@ -0,0 +1,388 @@
|
|||
//! `RealTimeStreaming` — the realtime logging collector.
|
||||
//!
|
||||
//! Mirrors Python `litellm.realtime_api.main.RealTimeStreaming`: it observes the
|
||||
//! event stream in O(1) (never buffering frames), accumulating just the fields
|
||||
//! the spend log needs (model, id, cumulative usage), then on session close
|
||||
//! builds a `StandardLoggingPayload` and fans it out to every registered
|
||||
//! `CustomLogger`.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::constants::DEFAULT_PROVIDER;
|
||||
use crate::integrations::custom_logger::{
|
||||
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
|
||||
};
|
||||
use crate::integrations::types::{
|
||||
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, Usage,
|
||||
};
|
||||
|
||||
/// Current wall-clock time as epoch seconds (float), matching the Python
|
||||
/// `startTime`/`endTime` contract.
|
||||
fn epoch_seconds() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs_f64())
|
||||
.unwrap_or(0.0)
|
||||
}
|
||||
|
||||
/// Status of a finished realtime session, mapped to the callback record status.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum SessionStatus {
|
||||
Success,
|
||||
Failure,
|
||||
}
|
||||
|
||||
/// Accumulates realtime session state and emits a logging payload on close.
|
||||
pub struct RealTimeStreaming {
|
||||
callbacks: Vec<Arc<dyn CustomLogger>>,
|
||||
/// REQUEST-ID RULE: the SpendLogs `request_id` == the OpenAI realtime session
|
||||
/// id (`sess_…`), captured from `session.created`. Both `id` and
|
||||
/// `litellm_call_id` are set to that value so the Python writer logs the same
|
||||
/// id regardless of which field it reads. The gateway-generated `rt-…` id
|
||||
/// (the constructor seed) is only a fallback for sessions that fail before
|
||||
/// `session.created` arrives.
|
||||
litellm_call_id: String,
|
||||
/// See the request-id rule above — mirrors `litellm_call_id`.
|
||||
id: String,
|
||||
model: String,
|
||||
custom_llm_provider: String,
|
||||
usage: Usage,
|
||||
response_cost: f64,
|
||||
start_time: f64,
|
||||
end_time: f64,
|
||||
metadata: RequestMetadata,
|
||||
/// Count of logging callbacks that failed to enqueue (non-fatal).
|
||||
dropped: u64,
|
||||
}
|
||||
|
||||
impl RealTimeStreaming {
|
||||
/// Create a collector for one session. `litellm_call_id` is the gateway's
|
||||
/// per-connection id; `model` is the requested model (a sane default until
|
||||
/// `session.created` reports the upstream model).
|
||||
pub fn new(
|
||||
callbacks: Vec<Arc<dyn CustomLogger>>,
|
||||
litellm_call_id: String,
|
||||
model: String,
|
||||
metadata: RequestMetadata,
|
||||
) -> Self {
|
||||
let now = epoch_seconds();
|
||||
Self {
|
||||
callbacks,
|
||||
id: litellm_call_id.clone(),
|
||||
litellm_call_id,
|
||||
model,
|
||||
custom_llm_provider: DEFAULT_PROVIDER.to_string(),
|
||||
usage: Usage::default(),
|
||||
response_cost: 0.0,
|
||||
start_time: now,
|
||||
end_time: now,
|
||||
metadata,
|
||||
dropped: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of logging callbacks that failed to enqueue so far (test/observ.).
|
||||
#[allow(dead_code)]
|
||||
pub fn dropped(&self) -> u64 {
|
||||
self.dropped
|
||||
}
|
||||
|
||||
/// Observe one realtime event. O(1): updates accumulated state only; never
|
||||
/// buffers frames. Safe to call on every event in either direction.
|
||||
pub fn observe(&mut self, event: &RealtimeEvent) {
|
||||
match event.event_type.as_str() {
|
||||
"session.created" | "session.updated" => self.on_session(event),
|
||||
"response.done" => self.on_response_done(event),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// `session.created` / `session.updated` → capture upstream id + model.
|
||||
/// Per the request-id rule, the OpenAI session id becomes BOTH `id` and
|
||||
/// `litellm_call_id`, replacing the gateway-generated fallback.
|
||||
fn on_session(&mut self, event: &RealtimeEvent) {
|
||||
let session = event.data.get("session").and_then(Value::as_object);
|
||||
if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str) {
|
||||
if !id.is_empty() {
|
||||
self.id = id.to_string();
|
||||
self.litellm_call_id = id.to_string();
|
||||
}
|
||||
}
|
||||
if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str) {
|
||||
if !model.is_empty() {
|
||||
self.model = model.to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `response.done` → add this response's usage to the cumulative totals.
|
||||
fn on_response_done(&mut self, event: &RealtimeEvent) {
|
||||
let usage = event
|
||||
.data
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|r| r.get("usage"))
|
||||
.and_then(Value::as_object);
|
||||
let Some(usage) = usage else { return };
|
||||
|
||||
let input = usage.get("input_tokens").and_then(Value::as_u64);
|
||||
let output = usage.get("output_tokens").and_then(Value::as_u64);
|
||||
let total = usage.get("total_tokens").and_then(Value::as_u64);
|
||||
|
||||
if let Some(input) = input {
|
||||
self.usage.prompt_tokens += input;
|
||||
}
|
||||
if let Some(output) = output {
|
||||
self.usage.completion_tokens += output;
|
||||
}
|
||||
// Prefer the upstream-reported total; otherwise derive it.
|
||||
match total {
|
||||
Some(total) => self.usage.total_tokens += total,
|
||||
None => {
|
||||
self.usage.total_tokens += input.unwrap_or(0) + output.unwrap_or(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the per-session response cost ($). Cost computation is Python-side in
|
||||
/// the proxy; the gateway forwards 0.0 by default and lets the proxy price.
|
||||
/// Public API (exercised in tests) for the future path where the gateway
|
||||
/// prices realtime sessions itself.
|
||||
#[allow(dead_code)]
|
||||
pub fn set_response_cost(&mut self, cost: f64) {
|
||||
self.response_cost = cost;
|
||||
}
|
||||
|
||||
/// Build the `StandardLoggingPayload` from accumulated state.
|
||||
pub fn build_payload(&self) -> StandardLoggingPayload {
|
||||
StandardLoggingPayload {
|
||||
id: self.id.clone(),
|
||||
litellm_call_id: self.litellm_call_id.clone(),
|
||||
call_type: "realtime".to_string(),
|
||||
model: self.model.clone(),
|
||||
custom_llm_provider: self.custom_llm_provider.clone(),
|
||||
response_cost: self.response_cost,
|
||||
prompt_tokens: self.usage.prompt_tokens,
|
||||
completion_tokens: self.usage.completion_tokens,
|
||||
total_tokens: self.usage.total_tokens,
|
||||
start_time: self.start_time,
|
||||
end_time: self.end_time,
|
||||
stream: true,
|
||||
metadata: StandardLoggingMetadata {
|
||||
user_api_key_hash: self.metadata.user_api_key_hash.clone(),
|
||||
user_api_key_user_id: self.metadata.user_api_key_user_id.clone(),
|
||||
user_api_key_team_id: self.metadata.user_api_key_team_id.clone(),
|
||||
..Default::default()
|
||||
},
|
||||
messages: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Finish the session: stamp the end time and fan the payload out to every
|
||||
/// callback. On a logger enqueue error we bump a non-fatal counter (the
|
||||
/// realtime session has already ended; a dropped log must never propagate).
|
||||
pub async fn log_messages(&mut self, status: SessionStatus) {
|
||||
self.end_time = epoch_seconds();
|
||||
let payload = self.build_payload();
|
||||
let timing = CallbackTiming::new(payload.start_time, payload.end_time);
|
||||
let runner = CustomLoggerRunner::new(self.callbacks.clone());
|
||||
|
||||
match status {
|
||||
SessionStatus::Success => {
|
||||
let response = CallbackValue::new("realtime", serde_json::Value::Null);
|
||||
let report = runner
|
||||
.async_log_success_event(
|
||||
&ModelCallDetails::from_standard_logging_payload(payload),
|
||||
&response,
|
||||
timing,
|
||||
)
|
||||
.await;
|
||||
self.dropped += report.dropped as u64;
|
||||
}
|
||||
SessionStatus::Failure => {
|
||||
let error = LoggingError {
|
||||
message: "realtime session ended in failure".to_string(),
|
||||
kind: "RealtimeSessionError".to_string(),
|
||||
};
|
||||
let response = CallbackValue::new(
|
||||
"error",
|
||||
serde_json::json!({
|
||||
"message": error.message,
|
||||
"kind": error.kind,
|
||||
}),
|
||||
);
|
||||
let report = runner
|
||||
.async_log_failure_event(
|
||||
&ModelCallDetails::from_standard_logging_payload(payload)
|
||||
.with_failure_error(error),
|
||||
Some(&response),
|
||||
timing,
|
||||
)
|
||||
.await;
|
||||
self.dropped += report.dropped as u64;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::integrations::custom_logger::LogError;
|
||||
use crate::integrations::custom_logger::LogFuture;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
fn event(raw: &str) -> RealtimeEvent {
|
||||
serde_json::from_str(raw).expect("valid event json")
|
||||
}
|
||||
|
||||
/// A test logger that records the last payload it saw.
|
||||
#[derive(Default)]
|
||||
struct CapturingLogger {
|
||||
calls: AtomicU64,
|
||||
last_model: std::sync::Mutex<Option<String>>,
|
||||
last_total_tokens: AtomicU64,
|
||||
}
|
||||
|
||||
impl CustomLogger for CapturingLogger {
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: &'a CallbackValue,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async move {
|
||||
let payload = model_call_details
|
||||
.standard_logging_payload
|
||||
.as_ref()
|
||||
.expect("standard logging payload");
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
*self.last_model.lock().unwrap() = Some(payload.model.clone());
|
||||
self.last_total_tokens
|
||||
.store(payload.total_tokens, Ordering::SeqCst);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn observe_accumulates_model_and_tokens_then_logs() {
|
||||
let logger = Arc::new(CapturingLogger::default());
|
||||
let callbacks: Vec<Arc<dyn CustomLogger>> = vec![logger.clone()];
|
||||
let mut streaming = RealTimeStreaming::new(
|
||||
callbacks,
|
||||
"call_abc".to_string(),
|
||||
"gpt-realtime".to_string(),
|
||||
RequestMetadata {
|
||||
user_api_key_hash: Some("hash123".to_string()),
|
||||
user_api_key_user_id: Some("user-1".to_string()),
|
||||
user_api_key_team_id: Some("team-1".to_string()),
|
||||
},
|
||||
);
|
||||
|
||||
streaming.observe(&event(
|
||||
r#"{"type":"session.created","session":{"id":"sess_001","model":"gpt-realtime-2025"}}"#,
|
||||
));
|
||||
streaming.observe(&event(
|
||||
r#"{"type":"response.done","response":{"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}"#,
|
||||
));
|
||||
// A second response.done accumulates.
|
||||
streaming.observe(&event(
|
||||
r#"{"type":"response.done","response":{"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}}"#,
|
||||
));
|
||||
|
||||
let payload = streaming.build_payload();
|
||||
assert_eq!(payload.model, "gpt-realtime-2025");
|
||||
// Request-id rule: session.created's id becomes BOTH id and
|
||||
// litellm_call_id (replacing the "call_abc" gateway fallback), so the
|
||||
// SpendLogs request_id is always the OpenAI session id.
|
||||
assert_eq!(payload.id, "sess_001");
|
||||
assert_eq!(payload.litellm_call_id, "sess_001");
|
||||
assert_eq!(payload.prompt_tokens, 13);
|
||||
assert_eq!(payload.completion_tokens, 7);
|
||||
assert_eq!(payload.total_tokens, 20);
|
||||
assert_eq!(payload.response_cost, 0.0);
|
||||
assert_eq!(payload.call_type, "realtime");
|
||||
assert_eq!(payload.custom_llm_provider, "openai");
|
||||
assert_eq!(
|
||||
payload.metadata.user_api_key_hash.as_deref(),
|
||||
Some("hash123")
|
||||
);
|
||||
|
||||
streaming.log_messages(SessionStatus::Success).await;
|
||||
assert_eq!(logger.calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(
|
||||
logger.last_model.lock().unwrap().as_deref(),
|
||||
Some("gpt-realtime-2025")
|
||||
);
|
||||
assert_eq!(logger.last_total_tokens.load(Ordering::SeqCst), 20);
|
||||
assert_eq!(streaming.dropped(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn payload_serializes_with_camelcase_times_and_realtime_call_type() {
|
||||
let mut streaming = RealTimeStreaming::new(
|
||||
Vec::new(),
|
||||
"call_xyz".to_string(),
|
||||
"gpt-realtime".to_string(),
|
||||
RequestMetadata::default(),
|
||||
);
|
||||
streaming.observe(&event(
|
||||
r#"{"type":"response.done","response":{"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}"#,
|
||||
));
|
||||
streaming.set_response_cost(0.0042);
|
||||
let payload = streaming.build_payload();
|
||||
let json = serde_json::to_string(&payload).expect("serialize payload");
|
||||
|
||||
assert!(json.contains("\"startTime\""), "missing startTime: {json}");
|
||||
assert!(json.contains("\"endTime\""), "missing endTime: {json}");
|
||||
assert!(
|
||||
json.contains("\"call_type\":\"realtime\""),
|
||||
"missing call_type realtime: {json}"
|
||||
);
|
||||
assert!(
|
||||
json.contains("\"response_cost\""),
|
||||
"missing response_cost: {json}"
|
||||
);
|
||||
assert_eq!(payload.response_cost, 0.0042);
|
||||
}
|
||||
|
||||
/// A logger whose enqueue always fails should bump the dropped counter, not
|
||||
/// panic or propagate.
|
||||
#[tokio::test]
|
||||
async fn failing_logger_bumps_dropped_counter() {
|
||||
struct FailingLogger;
|
||||
impl CustomLogger for FailingLogger {
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
_model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: &'a CallbackValue,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async { Err(LogError::channel_full()) })
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
_model_call_details: &'a ModelCallDetails,
|
||||
_response_obj: Option<&'a CallbackValue>,
|
||||
_timing: CallbackTiming,
|
||||
) -> LogFuture<'a> {
|
||||
Box::pin(async { Err(LogError::channel_closed()) })
|
||||
}
|
||||
}
|
||||
let callbacks: Vec<Arc<dyn CustomLogger>> = vec![Arc::new(FailingLogger)];
|
||||
let mut streaming = RealTimeStreaming::new(
|
||||
callbacks,
|
||||
"call_1".to_string(),
|
||||
"gpt-realtime".to_string(),
|
||||
RequestMetadata::default(),
|
||||
);
|
||||
streaming.log_messages(SessionStatus::Success).await;
|
||||
assert_eq!(streaming.dropped(), 1);
|
||||
}
|
||||
}
|
||||
|
|
@ -6,8 +6,11 @@
|
|||
|
||||
mod service;
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::StatusCode;
|
||||
|
|
@ -17,12 +20,29 @@ use axum::Router;
|
|||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router as ModelRouter;
|
||||
use litellm_providers::realtime_pool::RealtimePool;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::auth::RequireMasterKey;
|
||||
use crate::integrations::custom_logger::CustomLogger;
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
use crate::realtime::streaming::{RealTimeStreaming, SessionStatus};
|
||||
use crate::state::AppState;
|
||||
|
||||
/// Process-local monotonic counter, mixed into the per-session call id so two
|
||||
/// sessions opened in the same nanosecond still get distinct ids.
|
||||
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
/// Generate a per-connection `litellm_call_id`. No external uuid dep: epoch
|
||||
/// nanos + a process-local sequence is unique enough for log correlation.
|
||||
fn new_call_id() -> String {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0);
|
||||
let seq = CALL_SEQ.fetch_add(1, Ordering::Relaxed);
|
||||
format!("rt-{nanos:x}-{seq:x}")
|
||||
}
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route("/v1/realtime", get(handle))
|
||||
|
|
@ -57,26 +77,63 @@ async fn handle(
|
|||
|
||||
let router = state.router.clone();
|
||||
let pool = state.realtime_pool.clone();
|
||||
let loggers = state.loggers.clone();
|
||||
let master_key = state.master_key.clone();
|
||||
let model = query.model;
|
||||
Ok(ws.on_upgrade(move |socket| bridge(socket, router, pool, model)))
|
||||
Ok(ws.on_upgrade(move |socket| bridge(socket, router, pool, loggers, master_key, model)))
|
||||
}
|
||||
|
||||
/// Adapt the axum socket (text frames) to the typed-event `Stream`/`Sink` the
|
||||
/// service wants, keeping axum types out of `service`.
|
||||
///
|
||||
/// This is also the realtime-logging seam: every upstream→client event (the
|
||||
/// direction carrying `session.created` and `response.done` with usage) is fed
|
||||
/// to a [`RealTimeStreaming`] collector via the splice's `observe` callback. The
|
||||
/// observe is O(1) and never buffers frames. When the splice returns (any of the
|
||||
/// three break paths — client disconnect, upstream close, idle timeout), we flush
|
||||
/// one logging payload to the registered callbacks.
|
||||
async fn bridge(
|
||||
socket: WebSocket,
|
||||
router: Arc<ModelRouter>,
|
||||
pool: Arc<RealtimePool>,
|
||||
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
|
||||
master_key: Option<Arc<str>>,
|
||||
model: String,
|
||||
) {
|
||||
let (ws_sink, ws_stream) = socket.split();
|
||||
|
||||
// Attribute the spend log to the key that authenticated this session (the
|
||||
// master key — the gateway is master-key auth). A non-null user_api_key_hash
|
||||
// is required for the Python spend logger to write a SpendLogs row.
|
||||
//
|
||||
// SECURITY: hash the key — never send the raw credential. This field fans out
|
||||
// to spend logs and every callback integration; the SHA-256 (matching the
|
||||
// proxy's hash_token) keeps the plaintext master key out of all of them while
|
||||
// still matching the key's hash in LiteLLM_SpendLogs.
|
||||
let metadata = RequestMetadata {
|
||||
user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token),
|
||||
..RequestMetadata::default()
|
||||
};
|
||||
|
||||
// Owned by THIS task only. The splice observes it via a synchronous `&mut`
|
||||
// callback (below), so there is no Arc/Mutex/atomic on the per-frame hot
|
||||
// path — just a monomorphized FnMut mutating stack-local fields. This is
|
||||
// what lets observe scale: 10K concurrent sessions = 10K independent
|
||||
// collectors, zero cross-task synchronization.
|
||||
let mut collector = RealTimeStreaming::new(
|
||||
loggers.as_ref().clone(),
|
||||
new_call_id(),
|
||||
model.clone(),
|
||||
metadata,
|
||||
);
|
||||
|
||||
let client_in = ws_stream.filter_map(|message| async move {
|
||||
match message {
|
||||
Ok(Message::Text(text)) => serde_json::from_str::<RealtimeEvent>(&text).ok(),
|
||||
_ => None,
|
||||
}
|
||||
});
|
||||
// Plain forwarding sink — no observe here anymore.
|
||||
let client_out = ws_sink.with(|event: RealtimeEvent| async move {
|
||||
Ok::<Message, axum::Error>(Message::Text(
|
||||
serde_json::to_string(&event).unwrap_or_default(),
|
||||
|
|
@ -84,5 +141,26 @@ async fn bridge(
|
|||
});
|
||||
|
||||
futures_util::pin_mut!(client_in, client_out);
|
||||
let _ = service::run(&router, &pool, &model, None, client_in, client_out).await;
|
||||
|
||||
// The observe closure borrows `&mut collector` for the duration of the
|
||||
// splice; the borrow ends when `run` returns, freeing the collector for the
|
||||
// single post-session `log_messages` flush. `run` picks a pooled (warm) or
|
||||
// fresh upstream — observe fires on the upstream arm either way.
|
||||
let result = service::run(
|
||||
&router,
|
||||
&pool,
|
||||
&model,
|
||||
None,
|
||||
|event: &RealtimeEvent| collector.observe(event),
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await;
|
||||
|
||||
let status = if result.is_ok() {
|
||||
SessionStatus::Success
|
||||
} else {
|
||||
SessionStatus::Failure
|
||||
};
|
||||
collector.log_messages(status).await;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
//! Business logic: select a deployment with the (pure) core router, then call the
|
||||
//! provider splice. The seam between `core::router` (selection only) and
|
||||
//! `providers` (the actual WebSocket I/O).
|
||||
//! `io` (the actual WebSocket I/O).
|
||||
//!
|
||||
//! On connect we try a pre-warmed upstream from the pool (handshake already paid,
|
||||
//! `session.created` buffered) and relay it instantly. On a pool miss or dead warm
|
||||
|
|
@ -9,12 +9,12 @@
|
|||
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::io::realtime_pool::{upstream_key, RealtimePool};
|
||||
use futures_util::{Sink, Stream};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_providers::realtime_pool::{upstream_key, RealtimePool};
|
||||
|
||||
/// Select a deployment for `model` and splice the client stream to the provider.
|
||||
///
|
||||
|
|
@ -26,6 +26,7 @@ pub async fn run<In, Out>(
|
|||
pool: &RealtimePool,
|
||||
model: &str,
|
||||
idle_timeout: Option<Duration>,
|
||||
observe: impl FnMut(&RealtimeEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
|
|
@ -52,10 +53,11 @@ where
|
|||
params.api_base.as_deref(),
|
||||
) {
|
||||
if let Some(handoff) = pool.take(&key) {
|
||||
return litellm_providers::realtime::realtime_warm(
|
||||
return crate::io::realtime::realtime_warm(
|
||||
provider_model,
|
||||
handoff,
|
||||
idle_timeout,
|
||||
observe,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
|
|
@ -64,11 +66,12 @@ where
|
|||
}
|
||||
|
||||
// Cold path: fresh dial (the original behavior).
|
||||
litellm_providers::realtime::realtime(
|
||||
crate::io::realtime::realtime(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
idle_timeout,
|
||||
observe,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use litellm_core::router::Router;
|
||||
use litellm_providers::realtime_pool::RealtimePool;
|
||||
|
||||
use crate::integrations::custom_logger::CustomLogger;
|
||||
|
||||
/// Shared application state handed to every route handler.
|
||||
#[derive(Clone)]
|
||||
|
|
@ -10,6 +12,8 @@ pub struct AppState {
|
|||
/// The gateway master key. Any caller presenting it as a bearer token may
|
||||
/// invoke the gateway. `None` → auth not configured (routes fail closed).
|
||||
pub master_key: Option<Arc<str>>,
|
||||
/// Logging callbacks fanned out at the end of each realtime session.
|
||||
pub loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
|
||||
/// Pre-warmed upstream realtime connection pool. Disabled
|
||||
/// (`RealtimePool::disabled()`) when `REALTIME_POOL_SIZE=0`, in which case
|
||||
/// every realtime connect fresh-dials exactly as before.
|
||||
|
|
|
|||
3
litellm-rust/crates/core/AGENTS.md
Normal file
3
litellm-rust/crates/core/AGENTS.md
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
litellm-core is the PURE translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. No network, no I/O, no env reads.
|
||||
|
||||
Routes (ocr, realtime) and providers (mistral, openai) are modules, not crates.
|
||||
|
|
@ -10,3 +10,6 @@ rand.workspace = true
|
|||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
|
||||
|
|
|
|||
167
litellm-rust/crates/core/src/call_lifecycle/README.md
Normal file
167
litellm-rust/crates/core/src/call_lifecycle/README.md
Normal file
|
|
@ -0,0 +1,167 @@
|
|||
# Call lifecycle
|
||||
|
||||
`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call
|
||||
types migrated to Rust. It owns lifecycle ordering, phase timing, and trace
|
||||
observer calls. It must not know about OCR, chat, messages, responses,
|
||||
completions, provider auth, request transforms, or response normalization.
|
||||
|
||||
Call-type modules own their domain behavior. For example, OCR owns document
|
||||
payloads, OCR provider transforms, safe document fetch, guardrail payload shape,
|
||||
callback payload shape, and provider HTTP execution.
|
||||
|
||||
## Runtime order
|
||||
|
||||
Every wrapped call runs in this order:
|
||||
|
||||
1. `async_pre_call_hook`
|
||||
2. `async_during_call_hook`
|
||||
3. provider call
|
||||
4. `async_log_success_event` or `async_log_failure_event`
|
||||
|
||||
`async_pre_call_hook` receives the initial LiteLLM request shape. It is where
|
||||
pre-call custom guardrails run.
|
||||
|
||||
`async_during_call_hook` converts the initial request into the provider-ready
|
||||
request. It is where provider config selection, parameter mapping, auth/header
|
||||
resolution, request transforms, and during-call guardrails belong.
|
||||
|
||||
The provider call receives only the provider-ready request. It should execute
|
||||
I/O and call the provider response transform.
|
||||
|
||||
Success and failure callbacks receive `CallLifecycleTiming`. Callback failures
|
||||
must not replace the original provider or guardrail result.
|
||||
|
||||
## Trace contract
|
||||
|
||||
The lifecycle runner records:
|
||||
|
||||
- full call start and end time
|
||||
- `pre_call` phase timing
|
||||
- `during_call` phase timing
|
||||
- `provider_call` phase timing
|
||||
- `success_callback` phase timing
|
||||
- `failure_callback` phase timing
|
||||
|
||||
`CallLifecycleObserver` receives phase start and end events. The default
|
||||
observer is a no-op. Future OTEL support should implement this observer instead
|
||||
of editing OCR, chat, messages, responses, completions, or provider modules.
|
||||
|
||||
## Required shape
|
||||
|
||||
Each migrated call type should use this folder shape:
|
||||
|
||||
```text
|
||||
litellm-rust/crates/ai-gateway/src/<call_type>/
|
||||
mod.rs # thin public entrypoint
|
||||
types.rs # public request, prepared request, provider request, response types
|
||||
prepare.rs # model/provider/callback/guardrail setup
|
||||
hooks.rs # CallLifecycleHooks implementation
|
||||
handler.rs # provider I/O and response normalization
|
||||
tests.rs # call-type lifecycle and handler tests
|
||||
```
|
||||
|
||||
Provider transforms can live in `litellm-rust/crates/core/src/providers/...`.
|
||||
Shared call-type helpers can live beside the call type, but generic lifecycle
|
||||
code stays in this folder.
|
||||
|
||||
## Core API
|
||||
|
||||
The prepared request implements `CallLifecycleRequest`:
|
||||
|
||||
```rust
|
||||
impl CallLifecycleRequest for PreparedMessagesRequest {
|
||||
fn lifecycle_context(&self) -> CallLifecycleContext {
|
||||
CallLifecycleContext::new(
|
||||
"messages",
|
||||
self.model.clone(),
|
||||
self.custom_llm_provider.clone(),
|
||||
self.litellm_call_id.clone(),
|
||||
)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
The call-type hooks implement `CallLifecycleHooks`:
|
||||
|
||||
```rust
|
||||
impl CallLifecycleHooks<
|
||||
PreparedMessagesRequest,
|
||||
ProviderMessagesRequest,
|
||||
MessagesResponse,
|
||||
> for MessagesLifecycleHooks {
|
||||
fn async_pre_call_hook(...) {
|
||||
// run pre-call custom guardrails against the LiteLLM request shape
|
||||
}
|
||||
|
||||
fn async_during_call_hook(...) {
|
||||
// map params, validate env, transform request, run during-call guardrails
|
||||
}
|
||||
|
||||
fn async_log_success_event(...) {
|
||||
// call async_log_success_event on configured custom loggers
|
||||
}
|
||||
|
||||
fn async_log_failure_event(...) {
|
||||
// call async_log_failure_event without swallowing the original error
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
The public entrypoint stays thin:
|
||||
|
||||
```rust
|
||||
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<MessagesResponse> {
|
||||
let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?;
|
||||
|
||||
CallLifecycle::default()
|
||||
.run_request(request, &hooks, execute_messages_provider_call)
|
||||
.await
|
||||
}
|
||||
```
|
||||
|
||||
Use `run_request` for new call types. Keep `run` available only for specialized
|
||||
tests or existing code that already has a `CallLifecycleContext`.
|
||||
|
||||
## Adding a new call type
|
||||
|
||||
1. Add `<call_type>/types.rs`
|
||||
|
||||
Define the public request accepted by the bridge, the prepared request used by
|
||||
the lifecycle runner, and the provider request consumed by the handler.
|
||||
|
||||
2. Implement `CallLifecycleRequest`
|
||||
|
||||
Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`.
|
||||
Do not put provider-specific logic here.
|
||||
|
||||
3. Add `<call_type>/prepare.rs`
|
||||
|
||||
Resolve model/provider once, generate or preserve `litellm_call_id`, construct
|
||||
callback and guardrail runners, and return `Prepared<CallType>Call`.
|
||||
|
||||
4. Add `<call_type>/hooks.rs`
|
||||
|
||||
Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction,
|
||||
provider config selection, param mapping, request transform, during-call
|
||||
guardrail payload construction, and callback payload construction here.
|
||||
|
||||
5. Add `<call_type>/handler.rs`
|
||||
|
||||
Execute the provider request and normalize the provider response. Do not repeat
|
||||
provider-specific transforms here; call the provider config.
|
||||
|
||||
6. Add tests
|
||||
|
||||
Cover hook order, success callback payload, failure callback payload, pre-call
|
||||
guardrail blocking before provider I/O, during-call body mutation, and provider
|
||||
error mapping.
|
||||
|
||||
## Review checklist
|
||||
|
||||
- Core lifecycle has no call-type or provider-specific branches
|
||||
- Public call-type entrypoint only prepares and calls `run_request`
|
||||
- Provider behavior lives behind provider config/transformation code
|
||||
- Hook method names map to the Python custom logger and guardrail concepts
|
||||
- Phase timing is recorded once in lifecycle, not separately per call type
|
||||
- Callback failures never hide the original provider or guardrail error
|
||||
- Tests prove the provider socket is not touched when pre-call guardrails block
|
||||
414
litellm-rust/crates/core/src/call_lifecycle/mod.rs
Normal file
414
litellm-rust/crates/core/src/call_lifecycle/mod.rs
Normal file
|
|
@ -0,0 +1,414 @@
|
|||
use std::future::Future;
|
||||
use std::time::{Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::{CoreError, CoreResult};
|
||||
|
||||
pub mod types;
|
||||
|
||||
pub use types::{
|
||||
CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest,
|
||||
CallLifecycleTiming,
|
||||
};
|
||||
|
||||
pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
|
||||
type PreCallFuture<'a>: Future<Output = CoreResult<InitialReq>> + Send + 'a
|
||||
where
|
||||
Self: 'a,
|
||||
InitialReq: 'a,
|
||||
ProviderReq: 'a,
|
||||
Resp: 'a;
|
||||
|
||||
type DuringCallFuture<'a>: Future<Output = CoreResult<ProviderReq>> + Send + 'a
|
||||
where
|
||||
Self: 'a,
|
||||
InitialReq: 'a,
|
||||
ProviderReq: 'a,
|
||||
Resp: 'a;
|
||||
|
||||
type SuccessFuture<'a>: Future<Output = ()> + Send + 'a
|
||||
where
|
||||
Self: 'a,
|
||||
Resp: 'a;
|
||||
|
||||
type FailureFuture<'a>: Future<Output = ()> + Send + 'a
|
||||
where
|
||||
Self: 'a;
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
request: InitialReq,
|
||||
) -> Self::PreCallFuture<'a>;
|
||||
|
||||
fn async_during_call_hook<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
request: InitialReq,
|
||||
) -> Self::DuringCallFuture<'a>;
|
||||
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
response: &'a Resp,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::SuccessFuture<'a>;
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
context: &'a CallLifecycleContext,
|
||||
error: &'a CoreError,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a>;
|
||||
}
|
||||
|
||||
pub trait CallLifecycleObserver: Send + Sync {
|
||||
fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {}
|
||||
|
||||
fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct NoopCallLifecycleObserver;
|
||||
|
||||
impl CallLifecycleObserver for NoopCallLifecycleObserver {}
|
||||
|
||||
pub struct CallLifecycle<'a> {
|
||||
observer: &'a dyn CallLifecycleObserver,
|
||||
}
|
||||
|
||||
impl<'a> CallLifecycle<'a> {
|
||||
pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self {
|
||||
Self { observer }
|
||||
}
|
||||
|
||||
pub async fn run_request<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
|
||||
&self,
|
||||
request: InitialReq,
|
||||
hooks: &Hooks,
|
||||
provider_call: ProviderCall,
|
||||
) -> CoreResult<Resp>
|
||||
where
|
||||
InitialReq: CallLifecycleRequest,
|
||||
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
|
||||
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
|
||||
ProviderFuture: Future<Output = CoreResult<Resp>>,
|
||||
{
|
||||
let context = request.lifecycle_context();
|
||||
self.run(context, request, hooks, provider_call).await
|
||||
}
|
||||
|
||||
pub async fn run<InitialReq, ProviderReq, Resp, Hooks, ProviderCall, ProviderFuture>(
|
||||
&self,
|
||||
context: CallLifecycleContext,
|
||||
request: InitialReq,
|
||||
hooks: &Hooks,
|
||||
provider_call: ProviderCall,
|
||||
) -> CoreResult<Resp>
|
||||
where
|
||||
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
|
||||
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
|
||||
ProviderFuture: Future<Output = CoreResult<Resp>>,
|
||||
{
|
||||
let call_start = epoch_seconds();
|
||||
let mut phases = Vec::new();
|
||||
|
||||
let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall);
|
||||
let request = match hooks.async_pre_call_hook(&context, request).await {
|
||||
Ok(request) => {
|
||||
phases.push(self.finish_phase(&context, pre_call));
|
||||
request
|
||||
}
|
||||
Err(error) => {
|
||||
phases.push(self.finish_phase(&context, pre_call));
|
||||
self.log_failure(&context, hooks, &error, call_start, &mut phases)
|
||||
.await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
|
||||
let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall);
|
||||
let provider_request = match hooks.async_during_call_hook(&context, request).await {
|
||||
Ok(request) => {
|
||||
phases.push(self.finish_phase(&context, during_call));
|
||||
request
|
||||
}
|
||||
Err(error) => {
|
||||
phases.push(self.finish_phase(&context, during_call));
|
||||
self.log_failure(&context, hooks, &error, call_start, &mut phases)
|
||||
.await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
|
||||
let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall);
|
||||
let result = provider_call(provider_request).await;
|
||||
phases.push(self.finish_phase(&context, provider_phase));
|
||||
|
||||
match &result {
|
||||
Ok(response) => {
|
||||
let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback);
|
||||
let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
|
||||
hooks
|
||||
.async_log_success_event(&context, response, &timing)
|
||||
.await;
|
||||
phases.push(self.finish_phase(&context, success_phase));
|
||||
}
|
||||
Err(error) => {
|
||||
self.log_failure(&context, hooks, error, call_start, &mut phases)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
async fn log_failure<InitialReq, ProviderReq, Resp, Hooks>(
|
||||
&self,
|
||||
context: &CallLifecycleContext,
|
||||
hooks: &Hooks,
|
||||
error: &CoreError,
|
||||
call_start: f64,
|
||||
phases: &mut Vec<CallLifecyclePhaseTiming>,
|
||||
) where
|
||||
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
|
||||
{
|
||||
let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback);
|
||||
let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone());
|
||||
hooks.async_log_failure_event(context, error, &timing).await;
|
||||
phases.push(self.finish_phase(context, failure_phase));
|
||||
}
|
||||
|
||||
fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart {
|
||||
self.observer.on_phase_start(context, phase);
|
||||
PhaseStart {
|
||||
phase,
|
||||
start_time: epoch_seconds(),
|
||||
started_at: Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
fn finish_phase(
|
||||
&self,
|
||||
context: &CallLifecycleContext,
|
||||
phase_start: PhaseStart,
|
||||
) -> CallLifecyclePhaseTiming {
|
||||
let timing = CallLifecyclePhaseTiming {
|
||||
phase: phase_start.phase,
|
||||
start_time: phase_start.start_time,
|
||||
end_time: epoch_seconds(),
|
||||
duration: phase_start.started_at.elapsed(),
|
||||
};
|
||||
self.observer.on_phase_end(context, &timing);
|
||||
timing
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for CallLifecycle<'static> {
|
||||
fn default() -> Self {
|
||||
static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver;
|
||||
Self::new(&OBSERVER)
|
||||
}
|
||||
}
|
||||
|
||||
struct PhaseStart {
|
||||
phase: CallLifecyclePhase,
|
||||
start_time: f64,
|
||||
started_at: Instant,
|
||||
}
|
||||
|
||||
fn epoch_seconds() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs_f64())
|
||||
.unwrap_or(0.0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Mutex;
|
||||
|
||||
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingHooks {
|
||||
events: Mutex<Vec<&'static str>>,
|
||||
}
|
||||
|
||||
struct RecordingRequest(String);
|
||||
|
||||
impl CallLifecycleRequest for RecordingRequest {
|
||||
fn lifecycle_context(&self) -> CallLifecycleContext {
|
||||
CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1")
|
||||
}
|
||||
}
|
||||
|
||||
impl RecordingHooks {
|
||||
fn events(&self) -> Vec<&'static str> {
|
||||
self.events.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl CallLifecycleHooks<String, String, String> for RecordingHooks {
|
||||
type PreCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
|
||||
type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
|
||||
type SuccessFuture<'a> = BoxFuture<'a, ()>;
|
||||
type FailureFuture<'a> = BoxFuture<'a, ()>;
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
request: String,
|
||||
) -> Self::PreCallFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("pre_call");
|
||||
Ok(format!("{request}:pre"))
|
||||
})
|
||||
}
|
||||
|
||||
fn async_during_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
request: String,
|
||||
) -> Self::DuringCallFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("during_call");
|
||||
Ok(format!("{request}:during"))
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
_response: &'a String,
|
||||
timing: &'a CallLifecycleTiming,
|
||||
) -> Self::SuccessFuture<'a> {
|
||||
Box::pin(async move {
|
||||
assert!(timing.end_time >= timing.start_time);
|
||||
assert_eq!(timing.phases.len(), 3);
|
||||
self.events.lock().unwrap().push("success");
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
_error: &'a CoreError,
|
||||
_timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("failure");
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
|
||||
type PreCallFuture<'a> = BoxFuture<'a, CoreResult<RecordingRequest>>;
|
||||
type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
|
||||
type SuccessFuture<'a> = BoxFuture<'a, ()>;
|
||||
type FailureFuture<'a> = BoxFuture<'a, ()>;
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
request: RecordingRequest,
|
||||
) -> Self::PreCallFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("pre_call");
|
||||
Ok(RecordingRequest(format!("{}:pre", request.0)))
|
||||
})
|
||||
}
|
||||
|
||||
fn async_during_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
request: RecordingRequest,
|
||||
) -> Self::DuringCallFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("during_call");
|
||||
Ok(format!("{}:during", request.0))
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_success_event<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
_response: &'a String,
|
||||
_timing: &'a CallLifecycleTiming,
|
||||
) -> Self::SuccessFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("success");
|
||||
})
|
||||
}
|
||||
|
||||
fn async_log_failure_event<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
_error: &'a CoreError,
|
||||
_timing: &'a CallLifecycleTiming,
|
||||
) -> Self::FailureFuture<'a> {
|
||||
Box::pin(async move {
|
||||
self.events.lock().unwrap().push("failure");
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lifecycle_runs_hooks_around_provider_call() {
|
||||
let hooks = RecordingHooks::default();
|
||||
let response = CallLifecycle::default()
|
||||
.run(
|
||||
CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
|
||||
"request".to_string(),
|
||||
&hooks,
|
||||
|request| async move {
|
||||
assert_eq!(request, "request:pre:during");
|
||||
Ok("response".to_string())
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
assert_eq!(response, "response");
|
||||
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lifecycle_logs_failure_when_provider_fails() {
|
||||
let hooks = RecordingHooks::default();
|
||||
let error = CallLifecycle::default()
|
||||
.run(
|
||||
CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"),
|
||||
"request".to_string(),
|
||||
&hooks,
|
||||
|_request| async move {
|
||||
Err::<String, CoreError>(CoreError::Network("provider down".to_string()))
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("call fails");
|
||||
|
||||
assert_eq!(error, CoreError::Network("provider down".to_string()));
|
||||
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lifecycle_can_run_any_request_with_embedded_context() {
|
||||
let hooks = RecordingHooks::default();
|
||||
let response = CallLifecycle::default()
|
||||
.run_request(
|
||||
RecordingRequest("request".to_string()),
|
||||
&hooks,
|
||||
|request| async move {
|
||||
assert_eq!(request, "request:pre:during");
|
||||
Ok("response".to_string())
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
assert_eq!(response, "response");
|
||||
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]);
|
||||
}
|
||||
}
|
||||
75
litellm-rust/crates/core/src/call_lifecycle/types.rs
Normal file
75
litellm-rust/crates/core/src/call_lifecycle/types.rs
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
use std::time::Duration;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct CallLifecycleContext {
|
||||
pub call_type: String,
|
||||
pub model: String,
|
||||
pub custom_llm_provider: String,
|
||||
pub litellm_call_id: String,
|
||||
}
|
||||
|
||||
impl CallLifecycleContext {
|
||||
pub fn new(
|
||||
call_type: impl Into<String>,
|
||||
model: impl Into<String>,
|
||||
custom_llm_provider: impl Into<String>,
|
||||
litellm_call_id: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
call_type: call_type.into(),
|
||||
model: model.into(),
|
||||
custom_llm_provider: custom_llm_provider.into(),
|
||||
litellm_call_id: litellm_call_id.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait CallLifecycleRequest {
|
||||
fn lifecycle_context(&self) -> CallLifecycleContext;
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum CallLifecyclePhase {
|
||||
PreCall,
|
||||
DuringCall,
|
||||
ProviderCall,
|
||||
SuccessCallback,
|
||||
FailureCallback,
|
||||
}
|
||||
|
||||
impl CallLifecyclePhase {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::PreCall => "pre_call",
|
||||
Self::DuringCall => "during_call",
|
||||
Self::ProviderCall => "provider_call",
|
||||
Self::SuccessCallback => "success_callback",
|
||||
Self::FailureCallback => "failure_callback",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct CallLifecyclePhaseTiming {
|
||||
pub phase: CallLifecyclePhase,
|
||||
pub start_time: f64,
|
||||
pub end_time: f64,
|
||||
pub duration: Duration,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct CallLifecycleTiming {
|
||||
pub start_time: f64,
|
||||
pub end_time: f64,
|
||||
pub phases: Vec<CallLifecyclePhaseTiming>,
|
||||
}
|
||||
|
||||
impl CallLifecycleTiming {
|
||||
pub fn new(start_time: f64, end_time: f64, phases: Vec<CallLifecyclePhaseTiming>) -> Self {
|
||||
Self {
|
||||
start_time,
|
||||
end_time,
|
||||
phases,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -13,6 +13,10 @@ pub enum CoreError {
|
|||
MissingField(&'static str),
|
||||
#[error("invalid response: {0}")]
|
||||
InvalidResponse(String),
|
||||
#[error("invalid provider: {0}")]
|
||||
InvalidProvider(String),
|
||||
#[error("invalid request: {0}")]
|
||||
InvalidRequest(String),
|
||||
#[error("{0}")]
|
||||
Auth(String),
|
||||
#[error("OCR request failed with status {status}: {body}")]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
pub mod call_lifecycle;
|
||||
pub mod error;
|
||||
pub mod ocr;
|
||||
pub mod providers;
|
||||
pub mod realtime;
|
||||
pub mod router;
|
||||
pub mod routing_utils;
|
||||
|
||||
pub use error::{CoreError, CoreResult};
|
||||
|
|
|
|||
|
|
@ -4,7 +4,28 @@ use crate::CoreResult;
|
|||
|
||||
use super::types::{OcrRequestData, OcrResponseData};
|
||||
|
||||
pub trait OcrProviderConfig {
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OcrAuthStrategy {
|
||||
Bearer,
|
||||
Header(&'static str),
|
||||
}
|
||||
|
||||
impl OcrAuthStrategy {
|
||||
pub fn header_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Bearer => "authorization",
|
||||
Self::Header(header_name) => header_name,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OcrResponseHandling {
|
||||
Json,
|
||||
AzureDocumentIntelligencePoll,
|
||||
}
|
||||
|
||||
pub trait OcrProviderConfig: Sync {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str];
|
||||
|
||||
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
|
|
@ -29,4 +50,30 @@ pub trait OcrProviderConfig {
|
|||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData>;
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn auth_strategy(&self) -> OcrAuthStrategy {
|
||||
OcrAuthStrategy::Bearer
|
||||
}
|
||||
|
||||
fn requires_data_uri_document(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn response_handling(&self) -> OcrResponseHandling {
|
||||
OcrResponseHandling::Json
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,520 @@
|
|||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::error::{json_type_name, CoreError, CoreResult};
|
||||
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
|
||||
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96;
|
||||
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages"];
|
||||
|
||||
pub struct AzureAiOcrConfig;
|
||||
pub struct AzureDocumentIntelligenceOcrConfig;
|
||||
|
||||
pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig;
|
||||
pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig =
|
||||
AzureDocumentIntelligenceOcrConfig;
|
||||
|
||||
fn non_empty(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn resolve_value(
|
||||
explicit: Option<&str>,
|
||||
env_name: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
missing_message: &str,
|
||||
) -> CoreResult<String> {
|
||||
non_empty(explicit)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| CoreError::Auth(missing_message.to_string()))
|
||||
}
|
||||
|
||||
pub fn resolve_azure_ai_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_key,
|
||||
AZURE_AI_API_KEY_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params",
|
||||
)
|
||||
}
|
||||
|
||||
pub fn resolve_azure_ai_api_base(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_base,
|
||||
AZURE_AI_API_BASE_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter",
|
||||
)
|
||||
}
|
||||
|
||||
pub fn complete_azure_ai_url(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let base = resolve_azure_ai_api_base(api_base, env_lookup)?;
|
||||
Ok(format!(
|
||||
"{}/providers/mistral/azure/ocr",
|
||||
base.trim_end_matches('/')
|
||||
))
|
||||
}
|
||||
|
||||
pub fn resolve_document_intelligence_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_key,
|
||||
AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter",
|
||||
)
|
||||
}
|
||||
|
||||
pub fn resolve_document_intelligence_endpoint(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_base,
|
||||
AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter",
|
||||
)
|
||||
}
|
||||
|
||||
fn encode_model_id(model: &str) -> String {
|
||||
let model_id = model.rsplit('/').next().unwrap_or(model);
|
||||
model_id
|
||||
.bytes()
|
||||
.flat_map(|byte| match byte {
|
||||
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
|
||||
vec![byte as char]
|
||||
}
|
||||
_ => format!("%{byte:02X}").chars().collect(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn pages_token_is_valid(token: &str) -> bool {
|
||||
let mut parts = token.split('-');
|
||||
let Some(start) = parts.next() else {
|
||||
return false;
|
||||
};
|
||||
if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) {
|
||||
return false;
|
||||
}
|
||||
match parts.next() {
|
||||
None => true,
|
||||
Some(end) => {
|
||||
!end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_pages_param(pages: &Value) -> CoreResult<Option<String>> {
|
||||
match pages {
|
||||
Value::String(value) => {
|
||||
let normalized = value
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
if normalized.split(',').all(pages_token_is_valid) {
|
||||
Ok(Some(normalized))
|
||||
} else {
|
||||
Err(CoreError::InvalidRequest(format!(
|
||||
"Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'."
|
||||
)))
|
||||
}
|
||||
}
|
||||
Value::Array(values) => {
|
||||
if values.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
if values.iter().all(Value::is_i64) {
|
||||
let mut pages = BTreeSet::new();
|
||||
for value in values {
|
||||
let page = value.as_i64().expect("checked is_i64");
|
||||
if page < 0 {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(),
|
||||
));
|
||||
}
|
||||
pages.insert(page + 1);
|
||||
}
|
||||
return Ok(Some(
|
||||
pages
|
||||
.into_iter()
|
||||
.map(|page| page.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(","),
|
||||
));
|
||||
}
|
||||
if values.iter().all(Value::is_string) {
|
||||
let normalized = values
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
if normalized.split(',').all(pages_token_is_valid) {
|
||||
return Ok(Some(normalized));
|
||||
}
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'."
|
||||
)));
|
||||
}
|
||||
Err(CoreError::InvalidRequest(
|
||||
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
|
||||
.to_string(),
|
||||
))
|
||||
}
|
||||
_ => Err(CoreError::InvalidRequest(
|
||||
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
|
||||
.to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn complete_document_intelligence_url(
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?;
|
||||
let mut url = format!(
|
||||
"{}/documentintelligence/documentModels/{}:analyze?api-version={}",
|
||||
endpoint.trim_end_matches('/'),
|
||||
encode_model_id(model),
|
||||
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION
|
||||
);
|
||||
|
||||
if let Some(pages) = optional_params.get("pages") {
|
||||
if let Some(normalized) = normalize_pages_param(pages)? {
|
||||
url.push_str("&pages=");
|
||||
url.push_str(&normalized);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> {
|
||||
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("document.type"))?;
|
||||
let field_name = match doc_type {
|
||||
"document_url" => "document_url",
|
||||
"image_url" => "image_url",
|
||||
other => {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"Invalid document type: {other}. Must be 'document_url' or 'image_url'"
|
||||
)))
|
||||
}
|
||||
};
|
||||
object
|
||||
.get(field_name)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField(field_name))
|
||||
}
|
||||
|
||||
fn extract_base64_from_data_uri(data_uri: &str) -> &str {
|
||||
data_uri
|
||||
.split_once(',')
|
||||
.map(|(_, data)| data)
|
||||
.unwrap_or(data_uri)
|
||||
}
|
||||
|
||||
fn page_markdown(page: &Map<String, Value>) -> String {
|
||||
page.get("lines")
|
||||
.and_then(Value::as_array)
|
||||
.map(|lines| {
|
||||
lines
|
||||
.iter()
|
||||
.filter_map(|line| line.get("content").and_then(Value::as_str))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn page_dimensions(page: &Map<String, Value>) -> Value {
|
||||
let width = page.get("width").and_then(Value::as_f64).unwrap_or(8.5);
|
||||
let height = page.get("height").and_then(Value::as_f64).unwrap_or(11.0);
|
||||
let unit = page.get("unit").and_then(Value::as_str).unwrap_or("inch");
|
||||
let (width, height) = if unit == "inch" {
|
||||
(
|
||||
(width * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64,
|
||||
(height * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64,
|
||||
)
|
||||
} else {
|
||||
(width as i64, height as i64)
|
||||
};
|
||||
json!({
|
||||
"width": width,
|
||||
"height": height,
|
||||
"dpi": AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI,
|
||||
})
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for AzureAiOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_azure_ai_url(api_base, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_azure_ai_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn requires_data_uri_document(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
_model: &str,
|
||||
document: Value,
|
||||
_optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
let document_url = document_url_from_mistral_document(&document)?;
|
||||
let mut data = Map::new();
|
||||
if document_url.starts_with("data:") {
|
||||
data.insert(
|
||||
"base64Source".to_string(),
|
||||
Value::String(extract_base64_from_data_uri(document_url).to_string()),
|
||||
);
|
||||
} else {
|
||||
data.insert(
|
||||
"urlSource".to_string(),
|
||||
Value::String(document_url.to_string()),
|
||||
);
|
||||
}
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
let response = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
let status = response
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("status"))?;
|
||||
if status != "succeeded" {
|
||||
return Err(CoreError::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed with status: {status}"
|
||||
)));
|
||||
}
|
||||
|
||||
let azure_pages = response
|
||||
.get("analyzeResult")
|
||||
.and_then(|result| result.get("pages"))
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
let pages = azure_pages
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.map(|page| {
|
||||
let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1);
|
||||
json!({
|
||||
"index": page_number - 1,
|
||||
"markdown": page_markdown(page),
|
||||
"dimensions": page_dimensions(page),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(OcrResponseData {
|
||||
usage_info: Some(json!({
|
||||
"pages_processed": pages.len(),
|
||||
"doc_size_bytes": null,
|
||||
})),
|
||||
pages,
|
||||
model: model.to_string(),
|
||||
document_annotation: None,
|
||||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_document_intelligence_url(api_base, model, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_document_intelligence_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn auth_strategy(&self) -> OcrAuthStrategy {
|
||||
OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key")
|
||||
}
|
||||
|
||||
fn response_handling(&self) -> OcrResponseHandling {
|
||||
OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn azure_ai_reuses_mistral_body_transform() {
|
||||
let body = AZURE_AI_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"pixtral-12b-2409",
|
||||
json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc"}),
|
||||
serde_json::Map::from_iter([("include_image_base64".to_string(), json!(true))]),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body["model"], "pixtral-12b-2409");
|
||||
assert_eq!(body["include_image_base64"], true);
|
||||
assert_eq!(
|
||||
body["document"]["document_url"],
|
||||
"data:application/pdf;base64,abc"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_intelligence_url_normalizes_zero_based_pages() {
|
||||
let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]);
|
||||
let url = complete_document_intelligence_url(
|
||||
Some("https://example.cognitiveservices.azure.com/"),
|
||||
"azure_ai/doc-intelligence/prebuilt-layout",
|
||||
¶ms,
|
||||
&|_| None,
|
||||
)
|
||||
.expect("url builds");
|
||||
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_intelligence_request_uses_base64_source_for_data_uri() {
|
||||
let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"prebuilt-read",
|
||||
json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc123"}),
|
||||
Map::new(),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body, json!({"base64Source": "abc123"}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_intelligence_response_normalizes_pages() {
|
||||
let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
|
||||
.transform_ocr_response(
|
||||
"prebuilt-layout",
|
||||
json!({
|
||||
"status": "succeeded",
|
||||
"analyzeResult": {
|
||||
"pages": [{
|
||||
"pageNumber": 2,
|
||||
"width": 8.5,
|
||||
"height": 11,
|
||||
"unit": "inch",
|
||||
"lines": [{"content": "hello"}, {"content": "world"}]
|
||||
}]
|
||||
}
|
||||
}),
|
||||
)
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(response.pages[0]["index"], 1);
|
||||
assert_eq!(response.pages[0]["markdown"], "hello\nworld");
|
||||
assert_eq!(response.pages[0]["dimensions"]["width"], 816);
|
||||
assert_eq!(
|
||||
response.usage_info,
|
||||
Some(json!({"pages_processed": 1, "doc_size_bytes": null}))
|
||||
);
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/core/src/providers/mistral/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/mistral/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
use litellm_core::error::{json_type_name, CoreError, CoreResult};
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use crate::error::{json_type_name, CoreError, CoreResult};
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const SUPPORTED_OCR_PARAMS: &[&str] = &[
|
||||
|
|
@ -15,6 +15,7 @@ const SUPPORTED_OCR_PARAMS: &[&str] = &[
|
|||
"extract_footer",
|
||||
"table_format",
|
||||
"confidence_scores_granularity",
|
||||
"include_blocks",
|
||||
"id",
|
||||
];
|
||||
|
||||
|
|
@ -132,6 +133,24 @@ impl OcrProviderConfig for MistralOcrConfig {
|
|||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
Ok(complete_url(api_base))
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_api_key(api_key, env_lookup)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn supported_ocr_params() -> &'static [&'static str] {
|
||||
|
|
@ -175,6 +194,7 @@ mod tests {
|
|||
"extract_footer",
|
||||
"table_format",
|
||||
"confidence_scores_granularity",
|
||||
"include_blocks",
|
||||
"id",
|
||||
]
|
||||
);
|
||||
4
litellm-rust/crates/core/src/providers/mod.rs
Normal file
4
litellm-rust/crates/core/src/providers/mod.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
pub mod azure_ai;
|
||||
pub mod mistral;
|
||||
pub mod openai;
|
||||
pub mod vertex_ai;
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
use litellm_core::realtime::transformation::RealtimeProviderConfig;
|
||||
use litellm_core::realtime::types::{RealtimeEvent, RealtimeTransformResult};
|
||||
use litellm_core::CoreResult;
|
||||
use crate::realtime::transformation::RealtimeProviderConfig;
|
||||
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
|
||||
use crate::CoreResult;
|
||||
|
||||
/// Default OpenAI API base, used when the caller does not override `api_base`.
|
||||
pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com";
|
||||
1
litellm-rust/crates/core/src/providers/vertex_ai/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/vertex_ai/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,435 @@
|
|||
use crate::error::{json_type_name, CoreError, CoreResult};
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
const VERTEX_DEFAULT_LOCATION: &str = "us-central1";
|
||||
const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com";
|
||||
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
|
||||
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
|
||||
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
|
||||
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
|
||||
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
|
||||
|
||||
#[rustfmt::skip]
|
||||
const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[
|
||||
"stream",
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"n",
|
||||
"stop",
|
||||
];
|
||||
|
||||
pub struct VertexAiOcrConfig;
|
||||
pub struct VertexAiDeepSeekOcrConfig;
|
||||
|
||||
pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig;
|
||||
pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig;
|
||||
|
||||
fn string_param<'a>(params: &'a Map<String, Value>, keys: &[&str]) -> Option<&'a str> {
|
||||
keys.iter()
|
||||
.find_map(|key| params.get(*key).and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
pub fn is_deepseek_model(model: &str) -> bool {
|
||||
model.to_ascii_lowercase().contains("deepseek")
|
||||
}
|
||||
|
||||
pub fn resolve_vertex_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
"Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers"
|
||||
.to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn vertex_project(
|
||||
params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
string_param(params, &["vertex_project", "vertex_ai_project"])
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(
|
||||
"Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter"
|
||||
.to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn vertex_location(
|
||||
params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
string_param(params, &["vertex_location", "vertex_ai_location"])
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string())
|
||||
}
|
||||
|
||||
fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String {
|
||||
api_base
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com"))
|
||||
.trim_end_matches('/')
|
||||
.to_string()
|
||||
}
|
||||
|
||||
pub fn complete_vertex_mistral_url(
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let project = vertex_project(optional_params, env_lookup)?;
|
||||
let location = vertex_location(optional_params, env_lookup);
|
||||
let base = vertex_mistral_api_base(api_base, &location);
|
||||
Ok(format!(
|
||||
"{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict"
|
||||
))
|
||||
}
|
||||
|
||||
pub fn complete_vertex_deepseek_url(
|
||||
api_base: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let project = vertex_project(optional_params, env_lookup)?;
|
||||
let location = vertex_location(optional_params, env_lookup);
|
||||
let base = api_base
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE)
|
||||
.trim_end_matches('/');
|
||||
Ok(format!(
|
||||
"{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions"
|
||||
))
|
||||
}
|
||||
|
||||
fn document_content_item(document: &Value) -> CoreResult<Value> {
|
||||
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("document.type"))?;
|
||||
let url_field = match doc_type {
|
||||
"image_url" => "image_url",
|
||||
"document_url" => "document_url",
|
||||
other => {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"Unsupported document type: {other}. Expected 'image_url' or 'document_url'"
|
||||
)))
|
||||
}
|
||||
};
|
||||
let url = object
|
||||
.get(url_field)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField(url_field))?;
|
||||
|
||||
Ok(json!({
|
||||
"type": "image_url",
|
||||
"image_url": url,
|
||||
}))
|
||||
}
|
||||
|
||||
fn deepseek_model_name(model: &str) -> String {
|
||||
if model.starts_with("deepseek-ai/") {
|
||||
model.to_string()
|
||||
} else {
|
||||
format!("deepseek-ai/{model}")
|
||||
}
|
||||
}
|
||||
|
||||
fn first_choice_content(response: &Value) -> CoreResult<Value> {
|
||||
response
|
||||
.get("choices")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|choices| choices.first())
|
||||
.and_then(|choice| choice.get("message"))
|
||||
.and_then(|message| message.get("content"))
|
||||
.cloned()
|
||||
.filter(|content| match content {
|
||||
Value::String(value) => !value.is_empty(),
|
||||
Value::Object(_) => true,
|
||||
_ => false,
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> Value {
|
||||
match content {
|
||||
Value::String(content) => {
|
||||
if content.trim_start().starts_with('{') {
|
||||
serde_json::from_str(&content).unwrap_or_else(|_| {
|
||||
json!({
|
||||
"pages": [{"index": 0, "markdown": content}],
|
||||
"model": model,
|
||||
"usage_info": usage.unwrap_or_else(|| json!({})),
|
||||
})
|
||||
})
|
||||
} else {
|
||||
json!({
|
||||
"pages": [{"index": 0, "markdown": content}],
|
||||
"model": model,
|
||||
"usage_info": usage.unwrap_or_else(|| json!({})),
|
||||
})
|
||||
}
|
||||
}
|
||||
Value::Object(_) => content,
|
||||
other => json!({
|
||||
"pages": [{"index": 0, "markdown": other.to_string()}],
|
||||
"model": model,
|
||||
"usage_info": usage.unwrap_or_else(|| json!({})),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for VertexAiOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_vertex_mistral_url(api_base, model, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_vertex_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn requires_data_uri_document(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
DEEPSEEK_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
let mut data = Map::new();
|
||||
data.insert(
|
||||
"model".to_string(),
|
||||
Value::String(deepseek_model_name(model)),
|
||||
);
|
||||
data.insert(
|
||||
"messages".to_string(),
|
||||
json!([{"role": "user", "content": [document_content_item(&document)?]}]),
|
||||
);
|
||||
for (key, value) in optional_params {
|
||||
if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) {
|
||||
data.insert(key, value);
|
||||
}
|
||||
}
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
let response = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
let usage = response.get("usage").cloned();
|
||||
let content = first_choice_content(&response_json)?;
|
||||
let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model);
|
||||
|
||||
if !ocr_data.get("pages").is_some_and(Value::is_array) {
|
||||
ocr_data = json!({
|
||||
"pages": [{
|
||||
"index": 0,
|
||||
"markdown": match content {
|
||||
Value::String(value) => value,
|
||||
other => other.to_string(),
|
||||
}
|
||||
}],
|
||||
"model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model),
|
||||
"usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})),
|
||||
});
|
||||
}
|
||||
|
||||
let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&ocr_data),
|
||||
})?;
|
||||
let pages = object
|
||||
.get("pages")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let usage_info = object
|
||||
.get("usage_info")
|
||||
.cloned()
|
||||
.or_else(|| response.get("usage").cloned());
|
||||
Ok(OcrResponseData {
|
||||
pages,
|
||||
model: object
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(model)
|
||||
.to_string(),
|
||||
document_annotation: object.get("document_annotation").cloned(),
|
||||
usage_info,
|
||||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_vertex_deepseek_url(api_base, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_vertex_api_key(api_key, env_lookup)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn vertex_mistral_url_uses_project_location_and_model() {
|
||||
let params = Map::from_iter([
|
||||
("vertex_project".to_string(), json!("proj-1")),
|
||||
("vertex_location".to_string(), json!("europe-west4")),
|
||||
]);
|
||||
|
||||
let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None)
|
||||
.expect("url builds");
|
||||
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_mistral_reuses_mistral_body_transform() {
|
||||
let body = VERTEX_AI_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"mistral-ocr-maas",
|
||||
json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}),
|
||||
Map::new(),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body["model"], "mistral-ocr-maas");
|
||||
assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_deepseek_request_uses_ocr_endpoint_shape() {
|
||||
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"deepseek-ocr-maas",
|
||||
json!({"type": "document_url", "document_url": "gs://bucket/doc.pdf"}),
|
||||
Map::from_iter([("temperature".to_string(), json!(0.1))]),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
|
||||
assert_eq!(body["temperature"], 0.1);
|
||||
assert_eq!(
|
||||
body["messages"][0]["content"][0],
|
||||
json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_deepseek_response_wraps_markdown_content() {
|
||||
let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG
|
||||
.transform_ocr_response(
|
||||
"deepseek-ocr-maas",
|
||||
json!({
|
||||
"choices": [{"message": {"content": "# OCR text"}}],
|
||||
"usage": {"prompt_tokens": 1}
|
||||
}),
|
||||
)
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(
|
||||
response.pages,
|
||||
vec![json!({"index": 0, "markdown": "# OCR text"})]
|
||||
);
|
||||
assert_eq!(response.model, "deepseek-ocr-maas");
|
||||
assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1})));
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue