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
ccdc87f80e
1623 changed files with 64386 additions and 23788 deletions
61
.github/workflows/create_daily_oss_branch.yml
vendored
61
.github/workflows/create_daily_oss_branch.yml
vendored
|
|
@ -1,61 +0,0 @@
|
|||
name: Create Daily OSS Branch
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 16 * * 1-5" # 9am PT during daylight saving time, weekdays.
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
date:
|
||||
description: "Branch date in YYYY_MM_DD format. Defaults to today's UTC date."
|
||||
required: false
|
||||
type: string
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
jobs:
|
||||
create-oss-branch:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
|
||||
- name: Create dated OSS branch
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
REQUESTED_DATE: ${{ inputs.date }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
if [ -n "${REQUESTED_DATE}" ]; then
|
||||
if ! echo "${REQUESTED_DATE}" | grep -Eq '^[0-9]{4}_[0-9]{2}_[0-9]{2}$'; then
|
||||
echo "::error::date must use YYYY_MM_DD format, got '${REQUESTED_DATE}'"
|
||||
exit 1
|
||||
fi
|
||||
BRANCH_DATE="${REQUESTED_DATE}"
|
||||
else
|
||||
BRANCH_DATE="$(date -u +'%Y_%m_%d')"
|
||||
fi
|
||||
|
||||
BRANCH_NAME="litellm_oss_daily_${BRANCH_DATE}"
|
||||
echo "Creating branch: ${BRANCH_NAME}"
|
||||
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "github-actions[bot]@users.noreply.github.com"
|
||||
|
||||
git fetch origin main "${BRANCH_NAME}" || true
|
||||
|
||||
if git show-ref --verify --quiet "refs/remotes/origin/${BRANCH_NAME}"; then
|
||||
echo "Branch ${BRANCH_NAME} already exists. Skipping creation."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
git checkout -b "${BRANCH_NAME}" origin/main
|
||||
git push "https://x-access-token:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" "${BRANCH_NAME}"
|
||||
echo "Successfully created and pushed branch: ${BRANCH_NAME}"
|
||||
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 current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) 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 'litellm_internal_staging' 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 the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) 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_internal_staging' instead."
|
||||
exit 1
|
||||
|
|
|
|||
50
.github/workflows/oss_daily_guardrails.yml
vendored
50
.github/workflows/oss_daily_guardrails.yml
vendored
|
|
@ -1,50 +0,0 @@
|
|||
name: OSS Daily Guardrails
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- "litellm_oss_daily_20*"
|
||||
pull_request:
|
||||
branches:
|
||||
- "litellm_oss_daily_20*"
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
oss-safe-checks:
|
||||
name: Run OSS daily safe checks
|
||||
if: startsWith(github.ref_name, 'litellm_oss_daily_20') || startsWith(github.head_ref, 'litellm_oss_daily_20') || startsWith(github.base_ref, 'litellm_oss_daily_20')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Run secret scan test
|
||||
run: |
|
||||
uv run --frozen --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
|
||||
|
||||
- name: Run Ruff
|
||||
run: |
|
||||
uv sync --frozen
|
||||
cd litellm
|
||||
uv run --no-sync ruff check .
|
||||
6
.github/workflows/test-rust.yml
vendored
6
.github/workflows/test-rust.yml
vendored
|
|
@ -61,5 +61,11 @@ jobs:
|
|||
- name: Run Clippy
|
||||
run: cargo clippy --workspace --all-targets --locked -- -D warnings
|
||||
|
||||
- name: Run Clippy with Bedrock auth
|
||||
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
|
||||
|
||||
- name: Run Rust tests
|
||||
run: cargo test --workspace --locked
|
||||
|
||||
- name: Run core tests with Bedrock auth
|
||||
run: cargo test -p litellm-core --features bedrock-auth --locked
|
||||
|
|
|
|||
2
.github/workflows/test-unit-proxy-db.yml
vendored
2
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -5,6 +5,8 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
|
|||
6
.github/workflows/zizmor.yml
vendored
6
.github/workflows/zizmor.yml
vendored
|
|
@ -4,7 +4,11 @@ on:
|
|||
push:
|
||||
branches: [main, litellm_internal_staging]
|
||||
pull_request:
|
||||
branches: [main, litellm_internal_staging]
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
|
|
|
|||
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -15,6 +15,9 @@ litellm/rust_bridge/_native*.so
|
|||
litellm/rust_bridge/_native*.pyd
|
||||
litellm-rust/target/
|
||||
|
||||
# Python package build output
|
||||
dist/
|
||||
|
||||
bun.lockb
|
||||
**/.DS_Store
|
||||
.aider*
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ Same thing for bug fixes. The tests should make it so that this specific bug can
|
|||
|
||||
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
|
||||
|
||||
When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose for internal contributors; external / OSS contributions target the current daily OSS branch instead, named `litellm_oss_daily_YYYY_MM_DD` (a fresh one is cut each weekday, so use the most recent)
|
||||
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
|
||||
|
||||
When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
|
||||
|
||||
|
|
|
|||
|
|
@ -322,7 +322,7 @@ npm run build
|
|||
## Submitting Your PR
|
||||
|
||||
1. **Push your branch**: `git push origin your-feature-branch`
|
||||
2. **Create a PR**: Go to GitHub and open a pull request against the current daily OSS branch, named `litellm_oss_daily_YYYY_MM_DD`. A fresh one is cut each weekday, so pick the most recent from the [branch list](https://github.com/BerriAI/litellm/branches/all?query=litellm_oss_daily). Do not target `main`.
|
||||
2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`.
|
||||
3. **Fill out the PR template**: Provide clear description of changes
|
||||
4. **Wait for review**: Maintainers will review and provide feedback
|
||||
5. **Address feedback**: Make requested changes and push updates
|
||||
|
|
|
|||
29
Dockerfile
29
Dockerfile
|
|
@ -86,7 +86,9 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra semantic-router \
|
||||
--python python3
|
||||
|
||||
RUN prisma generate --schema=./schema.prisma
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
||||
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
|
||||
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
|
@ -100,7 +102,11 @@ USER root
|
|||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
ENV PATH="/app/.venv/bin:${PATH}" \
|
||||
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \
|
||||
PRISMA_CLI_QUERY_ENGINE_TYPE=binary \
|
||||
PRISMA_OFFLINE_MODE=true
|
||||
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
|
|
@ -114,16 +120,19 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
|
|||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
|
||||
# whole /root/.cache drags in the uv build cache (~660 MB, includes a
|
||||
# setuptools wheel that surfaces as a CVE finding even though it's not
|
||||
# on the runtime sys.path).
|
||||
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
|
||||
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
|
||||
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
|
||||
# Prisma CLI + engines are baked under /opt/prisma, a fixed path every
|
||||
# runtime uid can read and that no cache volume mount shadows. The paths are
|
||||
# pinned via PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH and recorded into the
|
||||
# generated client at build time, so `prisma migrate deploy` on a fresh
|
||||
# database needs no npm and no network access (#33650, #24554).
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
|
||||
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
|
||||
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 42
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 10397
|
||||
"limit": 10389
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 11
|
||||
|
|
|
|||
BIN
dist/litellm-1.79.1.tar.gz
vendored
BIN
dist/litellm-1.79.1.tar.gz
vendored
Binary file not shown.
|
|
@ -84,7 +84,9 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra semantic-router \
|
||||
--python python3
|
||||
|
||||
RUN prisma generate --schema=./schema.prisma
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
||||
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
|
||||
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
|
@ -97,7 +99,11 @@ USER root
|
|||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
ENV PATH="/app/.venv/bin:${PATH}" \
|
||||
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \
|
||||
PRISMA_CLI_QUERY_ENGINE_TYPE=binary \
|
||||
PRISMA_OFFLINE_MODE=true
|
||||
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
|
|
@ -111,16 +117,21 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
|
|||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
|
||||
# Only the Prisma subdirs: the whole /root/.cache drags in the uv build cache.
|
||||
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
|
||||
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
|
||||
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
|
||||
# Prisma CLI + engines are baked under /opt/prisma, a fixed path every
|
||||
# runtime uid can read and that no cache volume mount shadows (unlike
|
||||
# /app/.cache or $HOME/.cache in readOnlyRootFilesystem + emptyDir setups).
|
||||
# The paths are pinned via PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH and
|
||||
# recorded into the generated client at build time, so `prisma migrate
|
||||
# deploy` on a fresh database needs no npm and no network access
|
||||
# (#33650, #24554).
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
|
||||
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
|
||||
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
|
|
|
|||
|
|
@ -137,6 +137,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
|
|||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
|
||||
COPY --from=builder /app/.cache /app/.cache
|
||||
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
|
||||
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets
|
||||
|
|
|
|||
|
|
@ -113,6 +113,10 @@ class PagerDutyAlerting(SlackAlerting):
|
|||
user_api_key_spend=_meta.get("user_api_key_spend"),
|
||||
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
|
||||
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
|
||||
user_api_key_user_spend=_meta.get("user_api_key_user_spend"),
|
||||
user_api_key_user_max_budget=_meta.get("user_api_key_user_max_budget"),
|
||||
user_api_key_team_spend=_meta.get("user_api_key_team_spend"),
|
||||
user_api_key_team_max_budget=_meta.get("user_api_key_team_max_budget"),
|
||||
user_api_key_org_id=_meta.get("user_api_key_org_id"),
|
||||
user_api_key_org_alias=_meta.get("user_api_key_org_alias"),
|
||||
user_api_key_team_id=_meta.get("user_api_key_team_id"),
|
||||
|
|
@ -196,6 +200,10 @@ class PagerDutyAlerting(SlackAlerting):
|
|||
if user_api_key_dict.budget_reset_at
|
||||
else None
|
||||
),
|
||||
user_api_key_user_spend=user_api_key_dict.user_spend,
|
||||
user_api_key_user_max_budget=user_api_key_dict.user_max_budget,
|
||||
user_api_key_team_spend=user_api_key_dict.team_spend,
|
||||
user_api_key_team_max_budget=user_api_key_dict.team_max_budget,
|
||||
user_api_key_org_id=user_api_key_dict.org_id,
|
||||
user_api_key_org_alias=user_api_key_dict.organization_alias,
|
||||
user_api_key_team_id=user_api_key_dict.team_id,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.50"
|
||||
version = "0.1.51"
|
||||
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.50"
|
||||
version = "0.1.51"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -46,4 +46,9 @@ Reminders:
|
|||
- gateway.config.proxy_config (rendered into a ConfigMap and mounted at
|
||||
/app/config/config.yaml; gateway reads it via
|
||||
CONFIG_FILE_PATH)
|
||||
- {component}.pdb.{enabled,minAvailable,maxUnavailable} (per-component PodDisruptionBudget; disabled by
|
||||
default — with hpa.minReplicas of 1, minAvailable: 1
|
||||
would block node drains)
|
||||
- {component}.topologySpreadConstraints (standard k8s list, e.g. spread replicas across
|
||||
topology.kubernetes.io/zone)
|
||||
- Enable ingress.enabled=true to dispatch / → ui, gateway data-plane prefixes → gateway, and the catch-all → backend.
|
||||
|
|
|
|||
|
|
@ -295,6 +295,52 @@ harmless no-op for the Job and authoritative for the app pods.
|
|||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
PodDisruptionBudget shared by gateway, backend, and ui.
|
||||
|
||||
Invoke with a dict:
|
||||
(dict "root" $ "component" .Values.gateway "componentName" "gateway"
|
||||
"fullname" (include "litellm.gateway.fullname" .)
|
||||
"selectorLabels" (include "litellm.gateway.selectorLabels" .))
|
||||
|
||||
Renders nothing unless both the component and its `pdb.enabled` are on.
|
||||
Only one of minAvailable / maxUnavailable should be set; if both are,
|
||||
minAvailable wins. If neither is set, falls back to `maxUnavailable: 1` so
|
||||
an enabled-but-unconfigured PDB still permits node drains.
|
||||
|
||||
"Set" means non-nil and non-empty-string, so an explicit 0 (e.g.
|
||||
`maxUnavailable: 0` to forbid all voluntary disruptions) is honored rather
|
||||
than silently replaced by the fallback.
|
||||
*/}}
|
||||
{{- define "litellm.pdb" -}}
|
||||
{{- $root := .root -}}
|
||||
{{- $component := .component -}}
|
||||
{{- $min := $component.pdb.minAvailable -}}
|
||||
{{- $max := $component.pdb.maxUnavailable -}}
|
||||
{{- $minSet := not (or (kindIs "invalid" $min) (eq (printf "%v" $min) "")) -}}
|
||||
{{- $maxSet := not (or (kindIs "invalid" $max) (eq (printf "%v" $max) "")) -}}
|
||||
{{- if and $component.enabled $component.pdb $component.pdb.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ .fullname }}
|
||||
labels:
|
||||
{{- include "litellm.commonLabels" $root | nindent 4 }}
|
||||
app.kubernetes.io/component: {{ .componentName }}
|
||||
spec:
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- .selectorLabels | nindent 6 }}
|
||||
{{- if $minSet }}
|
||||
minAvailable: {{ $min }}
|
||||
{{- else if $maxSet }}
|
||||
maxUnavailable: {{ $max }}
|
||||
{{- else }}
|
||||
maxUnavailable: 1
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Renders `envFrom:` block for a component's `envConfigMaps` / `envSecrets`
|
||||
lists. Each entry is a resource name; the chart wires the whole ConfigMap /
|
||||
|
|
|
|||
|
|
@ -98,4 +98,8 @@ spec:
|
|||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.topologySpreadConstraints }}
|
||||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
6
helm/litellm/templates/backend/poddisruptionbudget.yaml
Normal file
6
helm/litellm/templates/backend/poddisruptionbudget.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
{{- include "litellm.pdb" (dict
|
||||
"root" $
|
||||
"component" .Values.backend
|
||||
"componentName" "backend"
|
||||
"fullname" (include "litellm.backend.fullname" .)
|
||||
"selectorLabels" (include "litellm.backend.selectorLabels" .)) }}
|
||||
|
|
@ -100,4 +100,8 @@ spec:
|
|||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.topologySpreadConstraints }}
|
||||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
6
helm/litellm/templates/gateway/poddisruptionbudget.yaml
Normal file
6
helm/litellm/templates/gateway/poddisruptionbudget.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
{{- include "litellm.pdb" (dict
|
||||
"root" $
|
||||
"component" .Values.gateway
|
||||
"componentName" "gateway"
|
||||
"fullname" (include "litellm.gateway.fullname" .)
|
||||
"selectorLabels" (include "litellm.gateway.selectorLabels" .)) }}
|
||||
|
|
@ -76,4 +76,8 @@ spec:
|
|||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.ui.topologySpreadConstraints }}
|
||||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
6
helm/litellm/templates/ui/poddisruptionbudget.yaml
Normal file
6
helm/litellm/templates/ui/poddisruptionbudget.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
{{- include "litellm.pdb" (dict
|
||||
"root" $
|
||||
"component" .Values.ui
|
||||
"componentName" "ui"
|
||||
"fullname" (include "litellm.ui.fullname" .)
|
||||
"selectorLabels" (include "litellm.ui.selectorLabels" .)) }}
|
||||
188
helm/litellm/tests/pdb_topology_spread_tests.yaml
Normal file
188
helm/litellm/tests/pdb_topology_spread_tests.yaml
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
suite: test pod disruption budgets and topology spread constraints
|
||||
templates:
|
||||
- gateway/poddisruptionbudget.yaml
|
||||
- backend/poddisruptionbudget.yaml
|
||||
- ui/poddisruptionbudget.yaml
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: renders no PDB by default
|
||||
templates:
|
||||
- gateway/poddisruptionbudget.yaml
|
||||
- backend/poddisruptionbudget.yaml
|
||||
- ui/poddisruptionbudget.yaml
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
||||
- it: gateway PDB uses minAvailable and matches the gateway selector labels
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.pdb.enabled: true
|
||||
gateway.pdb.minAvailable: 1
|
||||
asserts:
|
||||
- isKind:
|
||||
of: PodDisruptionBudget
|
||||
- equal:
|
||||
path: apiVersion
|
||||
value: policy/v1
|
||||
- equal:
|
||||
path: metadata.name
|
||||
value: RELEASE-NAME-litellm-gateway
|
||||
- equal:
|
||||
path: spec.minAvailable
|
||||
value: 1
|
||||
- notExists:
|
||||
path: spec.maxUnavailable
|
||||
- equal:
|
||||
path: spec.selector.matchLabels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: gateway
|
||||
|
||||
- it: backend PDB uses maxUnavailable when minAvailable is unset
|
||||
template: backend/poddisruptionbudget.yaml
|
||||
set:
|
||||
backend.pdb.enabled: true
|
||||
backend.pdb.maxUnavailable: 25%
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.maxUnavailable
|
||||
value: 25%
|
||||
- notExists:
|
||||
path: spec.minAvailable
|
||||
- equal:
|
||||
path: spec.selector.matchLabels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: backend
|
||||
|
||||
- it: minAvailable wins when both minAvailable and maxUnavailable are set
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.pdb.enabled: true
|
||||
gateway.pdb.minAvailable: 2
|
||||
gateway.pdb.maxUnavailable: 1
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.minAvailable
|
||||
value: 2
|
||||
- notExists:
|
||||
path: spec.maxUnavailable
|
||||
|
||||
- it: an explicit maxUnavailable 0 is honored instead of the fallback
|
||||
template: backend/poddisruptionbudget.yaml
|
||||
set:
|
||||
backend.pdb.enabled: true
|
||||
backend.pdb.maxUnavailable: 0
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.maxUnavailable
|
||||
value: 0
|
||||
- notExists:
|
||||
path: spec.minAvailable
|
||||
|
||||
- it: an explicit minAvailable 0 is honored and beats a set maxUnavailable
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.pdb.enabled: true
|
||||
gateway.pdb.minAvailable: 0
|
||||
gateway.pdb.maxUnavailable: 1
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.minAvailable
|
||||
value: 0
|
||||
- notExists:
|
||||
path: spec.maxUnavailable
|
||||
|
||||
- it: enabled PDB with neither knob set falls back to maxUnavailable 1
|
||||
template: ui/poddisruptionbudget.yaml
|
||||
set:
|
||||
ui.pdb.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.maxUnavailable
|
||||
value: 1
|
||||
- notExists:
|
||||
path: spec.minAvailable
|
||||
- equal:
|
||||
path: spec.selector.matchLabels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: ui
|
||||
|
||||
- it: renders no PDB for a disabled component even when its pdb is enabled
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.enabled: false
|
||||
gateway.pdb.enabled: true
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
||||
- it: deployments omit topologySpreadConstraints by default
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
asserts:
|
||||
- notExists:
|
||||
path: spec.template.spec.topologySpreadConstraints
|
||||
|
||||
- it: gateway deployment renders configured topologySpreadConstraints
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.topologySpreadConstraints:
|
||||
- maxSkew: 1
|
||||
topologyKey: topology.kubernetes.io/zone
|
||||
whenUnsatisfiable: ScheduleAnyway
|
||||
labelSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/component: gateway
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints
|
||||
value:
|
||||
- maxSkew: 1
|
||||
topologyKey: topology.kubernetes.io/zone
|
||||
whenUnsatisfiable: ScheduleAnyway
|
||||
labelSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/component: gateway
|
||||
|
||||
- it: backend deployment renders configured topologySpreadConstraints
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
backend.topologySpreadConstraints:
|
||||
- maxSkew: 1
|
||||
topologyKey: kubernetes.io/hostname
|
||||
whenUnsatisfiable: DoNotSchedule
|
||||
labelSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/component: backend
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
|
||||
value: kubernetes.io/hostname
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints[0].whenUnsatisfiable
|
||||
value: DoNotSchedule
|
||||
|
||||
- it: ui deployment renders configured topologySpreadConstraints
|
||||
template: ui/deployment.yaml
|
||||
set:
|
||||
ui.topologySpreadConstraints:
|
||||
- maxSkew: 1
|
||||
topologyKey: topology.kubernetes.io/zone
|
||||
whenUnsatisfiable: ScheduleAnyway
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
|
||||
value: topology.kubernetes.io/zone
|
||||
|
|
@ -190,10 +190,28 @@ gateway:
|
|||
maxReplicas: 10
|
||||
targetCPUUtilizationPercentage: 70
|
||||
targetMemoryUtilizationPercentage: 80
|
||||
# PodDisruptionBudget for the gateway pods. Set exactly one of
|
||||
# `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
|
||||
# enabling without either falls back to `maxUnavailable: 1`). Disabled by
|
||||
# default: with the default hpa.minReplicas of 1, a `minAvailable: 1` PDB
|
||||
# would block node drains entirely.
|
||||
pdb:
|
||||
enabled: false
|
||||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
# Standard k8s topologySpreadConstraints for the gateway pods, e.g. to
|
||||
# spread replicas across zones:
|
||||
# - maxSkew: 1
|
||||
# topologyKey: topology.kubernetes.io/zone
|
||||
# whenUnsatisfiable: ScheduleAnyway
|
||||
# labelSelector:
|
||||
# matchLabels:
|
||||
# app.kubernetes.io/component: gateway
|
||||
topologySpreadConstraints: []
|
||||
|
||||
# ---------- backend (UI / management API) ----------
|
||||
backend:
|
||||
|
|
@ -233,10 +251,17 @@ backend:
|
|||
minReplicas: 1
|
||||
maxReplicas: 4
|
||||
targetCPUUtilizationPercentage: 70
|
||||
# Same shape as gateway.pdb.
|
||||
pdb:
|
||||
enabled: false
|
||||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
# Same shape as gateway.topologySpreadConstraints.
|
||||
topologySpreadConstraints: []
|
||||
|
||||
# ---------- ui (Next.js static dashboard) ----------
|
||||
ui:
|
||||
|
|
@ -279,7 +304,14 @@ ui:
|
|||
minReplicas: 1
|
||||
maxReplicas: 3
|
||||
targetCPUUtilizationPercentage: 80
|
||||
# Same shape as gateway.pdb.
|
||||
pdb:
|
||||
enabled: false
|
||||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
# Same shape as gateway.topologySpreadConstraints.
|
||||
topologySpreadConstraints: []
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "issuer" TEXT;
|
||||
|
|
@ -0,0 +1,17 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MCPServerOAuthClient" (
|
||||
"server_id" TEXT NOT NULL,
|
||||
"credentials" JSONB,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_MCPServerOAuthClient_pkey" PRIMARY KEY ("server_id")
|
||||
);
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
|
@ -325,6 +325,7 @@ model LiteLLM_MCPServerTable {
|
|||
command String?
|
||||
args String[] @default([])
|
||||
env Json? @default("{}")
|
||||
issuer String?
|
||||
authorization_url String?
|
||||
token_url String?
|
||||
registration_url String?
|
||||
|
|
@ -395,6 +396,13 @@ model LiteLLM_MCPUserEnvVars {
|
|||
@@index([server_id])
|
||||
}
|
||||
|
||||
model LiteLLM_MCPServerOAuthClient {
|
||||
server_id String @id
|
||||
credentials Json?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
}
|
||||
|
||||
// Generate Tokens for Proxy
|
||||
model LiteLLM_VerificationToken {
|
||||
token String @id
|
||||
|
|
@ -728,6 +736,9 @@ model LiteLLM_DailyUserSpend {
|
|||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
|
|
@ -759,6 +770,9 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
|
|
@ -790,6 +804,9 @@ model LiteLLM_DailyEndUserSpend {
|
|||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
|
|
@ -820,6 +837,9 @@ model LiteLLM_DailyAgentSpend {
|
|||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
|
|
@ -850,6 +870,9 @@ model LiteLLM_DailyTeamSpend {
|
|||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
|
|
@ -882,6 +905,9 @@ model LiteLLM_DailyTagSpend {
|
|||
completion_tokens BigInt @default(0)
|
||||
cache_read_input_tokens BigInt @default(0)
|
||||
cache_creation_input_tokens BigInt @default(0)
|
||||
compression_saved_tokens BigInt @default(0)
|
||||
compression_savings_spend Float @default(0.0)
|
||||
prompt_caching_savings_spend Float @default(0.0)
|
||||
spend Float @default(0.0)
|
||||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.77"
|
||||
version = "0.4.79"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.77"
|
||||
version = "0.4.79"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -6,4 +6,23 @@ Three layers, same for every route (see `ocr` and `realtime` as references):
|
|||
2. **Provider config (pure)** — `crates/providers/src/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
|
||||
3. **HTTP / transport (the host)** — `crates/providers/src/<route>.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O.
|
||||
|
||||
## Coding standards
|
||||
|
||||
Before writing new logic, look for an existing base to extend. When a change is
|
||||
“the same behavior for one more provider/endpoint/integration”, the codebase
|
||||
almost always already has a shared abstraction for it (for example, provider
|
||||
`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
|
||||
helpers in `litellm_core_utils/`, typed request/response models, or factory
|
||||
functions). Find it first with a search, then add the new variant by inheriting
|
||||
from or composing that base, overriding only what genuinely differs (model
|
||||
name, parameter mapping, or auth).
|
||||
|
||||
Never copy an existing implementation and edit it in place, and never hand-roll
|
||||
a parallel version of logic a base already provides. If you catch yourself
|
||||
writing a second copy of a pattern that exists twice already, stop and extract a
|
||||
base instead: put the shared shape in one place and make both call sites thin
|
||||
variants of it. The test for a good abstraction is that adding the next provider
|
||||
is a few declarative lines, not a new file of duplicated flow. Only diverge from
|
||||
the base when behavior is genuinely different, and say so explicitly in the PR.
|
||||
|
||||
**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.
|
||||
|
|
|
|||
|
|
@ -15,3 +15,12 @@ Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-
|
|||
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.
|
||||
|
||||
## Style
|
||||
|
||||
All Rust in `litellm-rust/` follows the official Rust Style Guide:
|
||||
https://doc.rust-lang.org/style-guide/
|
||||
|
||||
`rustfmt` implements its formatting by default, so run `cargo fmt` before committing; CI gates every PR on `cargo fmt --check`. Do not hand-format against rustfmt or add a `rustfmt.toml` that diverges from the default style.
|
||||
|
||||
Beyond formatting, follow the guide's naming and idiom conventions rustfmt cannot auto-apply: `snake_case` items/functions/modules, `UpperCamelCase` types/traits/variants, `SCREAMING_SNAKE_CASE` constants/statics (acronyms as one word, e.g. `HttpClient`), and the import grouping and item ordering it prescribes. See CLAUDE.md for the detailed version.
|
||||
|
|
|
|||
|
|
@ -2,6 +2,25 @@
|
|||
|
||||
This file defines the rules for Rust work in LiteLLM.
|
||||
|
||||
## Provider Coding Standards
|
||||
|
||||
Before writing new logic, look for an existing base to extend. When a change is
|
||||
“the same behavior for one more provider/endpoint/integration”, the codebase
|
||||
almost always already has a shared abstraction for it (for example, provider
|
||||
`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
|
||||
helpers in `litellm_core_utils/`, typed request/response models, or factory
|
||||
functions). Find it first with a search, then add the new variant by inheriting
|
||||
from or composing that base, overriding only what genuinely differs (model
|
||||
name, parameter mapping, or auth).
|
||||
|
||||
Never copy an existing implementation and edit it in place, and never hand-roll
|
||||
a parallel version of logic a base already provides. If you catch yourself
|
||||
writing a second copy of a pattern that exists twice already, stop and extract a
|
||||
base instead: put the shared shape in one place and make both call sites thin
|
||||
variants of it. The test for a good abstraction is that adding the next provider
|
||||
is a few declarative lines, not a new file of duplicated flow. Only diverge from
|
||||
the base when behavior is genuinely different, and say so explicitly in the PR.
|
||||
|
||||
## Crates (exactly three — see AGENTS.md)
|
||||
|
||||
`litellm-core` describes work; `litellm-ai-gateway` executes it; `litellm-python-bridge`
|
||||
|
|
@ -20,6 +39,11 @@ Route-level Rust structure mirrors LiteLLM's Python responsibilities:
|
|||
- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`),
|
||||
never inside `core`.
|
||||
|
||||
Call-hook and lifecycle instrumentation, including phase timing, usage
|
||||
accumulation, and callback payload construction, always lives in `core`.
|
||||
Hosts feed observed events into core and dispatch the completed payloads through
|
||||
their I/O logger; hosts must not own callback orchestration.
|
||||
|
||||
Allowed in `core`:
|
||||
- Pure request transforms
|
||||
- Pure response transforms
|
||||
|
|
@ -38,6 +62,13 @@ Not allowed in `core`:
|
|||
|
||||
Python owns rollout state and fallback while Rust is being introduced. Rust
|
||||
paths must be off by default until parity tests prove equivalence with Python.
|
||||
A new provider/route may instead be implemented rust-only with no Python
|
||||
reference; then the Python interface is a thin dispatch that calls Rust with no
|
||||
fallback, and you state the rust-only choice explicitly in the PR. Either way
|
||||
the Python side stays minimal (it only marshals inputs and calls the Rust
|
||||
interface), never add a per-route feature flag, and never push provider
|
||||
dispatch into `litellm/main.py`; put it in a thin dispatch class under
|
||||
`litellm/llms/<provider>/<route>/`.
|
||||
|
||||
## Production Bar
|
||||
|
||||
|
|
@ -77,6 +108,26 @@ 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.
|
||||
|
||||
## Rust Style Guide
|
||||
|
||||
All Rust in `litellm-rust/` follows the official Rust Style Guide:
|
||||
https://doc.rust-lang.org/style-guide/
|
||||
|
||||
`rustfmt` implements the guide's formatting rules by default, so the mechanical
|
||||
side is enforced for you: run `cargo fmt` before committing and CI gates every
|
||||
PR on `cargo fmt --check` (see Checks). Do not hand-format against rustfmt or add
|
||||
a `rustfmt.toml` that diverges from the default style; the default style *is* the
|
||||
guide.
|
||||
|
||||
The guide also covers conventions rustfmt cannot auto-apply; follow these too:
|
||||
- Naming: `snake_case` for items, functions, and modules; `UpperCamelCase` for
|
||||
types, traits, and enum variants; `SCREAMING_SNAKE_CASE` for constants and
|
||||
statics; acronyms count as one word (`HttpClient`, not `HTTPClient`).
|
||||
- Ordering and grouping the guide prescribes: imports grouped std / external /
|
||||
crate-local, derives before other attributes, and consistent item order.
|
||||
- Idioms the guide recommends over the formatter fighting you (e.g. prefer
|
||||
restructuring an over-long expression rather than forcing an awkward wrap).
|
||||
|
||||
## Constants
|
||||
|
||||
Magic numbers and fixed strings go in a crate-level `constants.rs`, never
|
||||
|
|
|
|||
1373
litellm-rust/Cargo.lock
generated
1373
litellm-rust/Cargo.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -7,7 +7,8 @@ members = [
|
|||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
edition = "2021"
|
||||
edition = "2024"
|
||||
rust-version = "1.88"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/BerriAI/litellm"
|
||||
|
||||
|
|
@ -15,8 +16,8 @@ repository = "https://github.com/BerriAI/litellm"
|
|||
litellm-core = { path = "crates/core" }
|
||||
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"] }
|
||||
pyo3 = "0.29.0"
|
||||
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
|
|
|
|||
|
|
@ -0,0 +1,59 @@
|
|||
# Provider coding standards (litellm-rust)
|
||||
|
||||
Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MISTRAL_OCR_CONFIG`) is the reference; `messages` (`ANTHROPIC_MESSAGES_CONFIG`) is the next port.
|
||||
|
||||
## Provider resolution
|
||||
|
||||
1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string.
|
||||
2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers.
|
||||
|
||||
## Transforms and the base config
|
||||
|
||||
3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src/<route>/transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`).
|
||||
4. Each provider implements that trait as a `const <PROVIDER>_<ROUTE>_CONFIG` in `core/src/providers/<provider>/<route>/transformation.rs`, mirroring the Python provider tree.
|
||||
5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it.
|
||||
6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers.
|
||||
|
||||
## Boundaries
|
||||
|
||||
7. Layers never cross: `core` = pure transforms/types (no network, env, secrets, auth, logging, global mutable state); `ai-gateway` = all I/O, auth headers, HTTP/SSE, lifecycle hooks; `python-bridge` = thin PyO3 adapter.
|
||||
8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers/<provider>/<route>/`; a route is a module, never a new crate.
|
||||
9. Route entry point stays thin: `<route>()` -> `prepare_*` -> `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing. Handlers validate and delegate; no business logic in them.
|
||||
10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Env reads happen only at the host/config layer, with the `DEFAULT_*` fallback defined in `constants.rs`.
|
||||
|
||||
## Types and errors
|
||||
|
||||
11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec<String>` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string.
|
||||
12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input.
|
||||
13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating.
|
||||
14. Early returns over deep nesting; small focused files over god modules.
|
||||
15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test.
|
||||
|
||||
## Safety and data minimization
|
||||
|
||||
16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary.
|
||||
17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer.
|
||||
18. Host I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS.
|
||||
|
||||
## Tests and rollout
|
||||
|
||||
19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity.
|
||||
20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping.
|
||||
21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR.
|
||||
|
||||
## Python bridge (SDK side)
|
||||
|
||||
22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
|
||||
23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
|
||||
24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
|
||||
|
||||
## Checks before push
|
||||
|
||||
25. Run, and keep green:
|
||||
```bash
|
||||
cd litellm-rust
|
||||
cargo fmt --check
|
||||
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
|
||||
```
|
||||
|
|
@ -14,7 +14,7 @@ path = "src/main.rs"
|
|||
required-features = ["server"]
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
litellm-core = { workspace = true, features = ["bedrock-auth"] }
|
||||
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
|
||||
# Python proxy callbacks API.
|
||||
reqwest.workspace = true
|
||||
|
|
@ -41,3 +41,4 @@ python-config = ["dep:pyo3"]
|
|||
|
||||
[dev-dependencies]
|
||||
futures-channel = "0.3"
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
|
|
|
|||
|
|
@ -0,0 +1,48 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub(super) fn audio_transcription_provider_config(
|
||||
provider: &str,
|
||||
) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
|
||||
match provider {
|
||||
"bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn string_headers(
|
||||
headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<BTreeMap<String, String>> {
|
||||
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!(
|
||||
"audio transcription extra_headers.{key} must be a string"
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn has_header(headers: &BTreeMap<String, String>, name: &str) -> bool {
|
||||
headers.keys().any(|key| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
let truncated: String = body.chars().take(256).collect();
|
||||
if truncated.chars().count() == body.chars().count() {
|
||||
truncated
|
||||
} else {
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,89 @@
|
|||
use std::time::SystemTime;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::providers::bedrock::audio_transcription::aws_auth_config;
|
||||
use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::common_utils::truncate_error_body;
|
||||
use super::types::ProviderAudioTranscriptionRequest;
|
||||
use crate::client::http_client;
|
||||
|
||||
pub(crate) async fn execute_audio_transcription_provider_call(
|
||||
request: ProviderAudioTranscriptionRequest,
|
||||
) -> CoreResult<Value> {
|
||||
let body = serde_json::to_vec(&request.body).map_err(|error| {
|
||||
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
|
||||
})?;
|
||||
let mut request_builder = http_client().post(&request.url).body(body.clone());
|
||||
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(|error| CoreError::Network(error.to_string()))?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.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(|error| {
|
||||
CoreError::InvalidResponse(format!("invalid audio response JSON: {error}"))
|
||||
})?;
|
||||
Ok(request
|
||||
.config
|
||||
.transform_transcription_response(&request.model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
|
||||
pub(crate) async fn sign_request(
|
||||
request: &ProviderAudioTranscriptionRequest,
|
||||
optional_params: &serde_json::Map<String, Value>,
|
||||
) -> CoreResult<ProviderAudioTranscriptionRequest> {
|
||||
let env_lookup = environment_lookup;
|
||||
let auth = request
|
||||
.config
|
||||
.auth_strategy(&request.model, optional_params, &env_lookup)?;
|
||||
let body = serde_json::to_vec(&request.body).map_err(|error| {
|
||||
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
|
||||
})?;
|
||||
let mut headers = super::common_utils::string_headers(None)?;
|
||||
headers.insert("Content-Type".to_string(), "application/json".to_string());
|
||||
headers.extend(request.upstream_headers.iter().cloned());
|
||||
match auth {
|
||||
AudioTranscriptionAuth::Bearer => {}
|
||||
AudioTranscriptionAuth::AwsSigV4 { region, .. } => {
|
||||
let credentials =
|
||||
resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup)
|
||||
.await?;
|
||||
headers.extend(sign_bedrock_post(
|
||||
&request.url,
|
||||
&body,
|
||||
&headers,
|
||||
®ion,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?);
|
||||
}
|
||||
}
|
||||
Ok(ProviderAudioTranscriptionRequest {
|
||||
upstream_headers: headers.into_iter().collect(),
|
||||
..request.clone()
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn environment_lookup(key: &str) -> Option<String> {
|
||||
std::env::var(key).ok()
|
||||
}
|
||||
300
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
Normal file
300
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
Normal file
|
|
@ -0,0 +1,300 @@
|
|||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
|
||||
use litellm_core::error::CoreError;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::common_utils::{audio_transcription_provider_config, has_header, string_headers};
|
||||
use super::handler::sign_request;
|
||||
use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
|
||||
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 AudioTranscriptionLifecycleHooks {
|
||||
logger_runner: CustomLoggerRunner,
|
||||
guardrail_runner: CustomGuardrailRunner,
|
||||
request_metadata: RequestMetadata,
|
||||
}
|
||||
|
||||
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
|
||||
type AudioLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
|
||||
|
||||
impl AudioTranscriptionLifecycleHooks {
|
||||
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: PreparedAudioTranscriptionRequest,
|
||||
) -> CoreResult<PreparedAudioTranscriptionRequest> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
let (guardrail_request, _) = self
|
||||
.guardrail_runner
|
||||
.run_pre_call(
|
||||
&guardrail_context(&self.request_metadata),
|
||||
GuardrailRequest::new(json!({
|
||||
"model": request.model,
|
||||
"custom_llm_provider": request.custom_llm_provider,
|
||||
"audio": request.audio,
|
||||
"optional_params": request.optional_params,
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.map_err(guardrail_error_to_core_error)?;
|
||||
let Value::Object(mut data) = guardrail_request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"audio transcription pre_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
let audio = data.remove("audio").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string())
|
||||
})?;
|
||||
let optional_params = match data.remove("optional_params") {
|
||||
Some(Value::Object(value)) => value,
|
||||
Some(_) => {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"audio transcription optional_params must be an object".to_string(),
|
||||
));
|
||||
}
|
||||
None => Map::new(),
|
||||
};
|
||||
Ok(PreparedAudioTranscriptionRequest {
|
||||
audio,
|
||||
optional_params,
|
||||
..request
|
||||
})
|
||||
}
|
||||
|
||||
async fn prepare_provider_request(
|
||||
&self,
|
||||
request: PreparedAudioTranscriptionRequest,
|
||||
) -> CoreResult<ProviderAudioTranscriptionRequest> {
|
||||
let config = audio_transcription_provider_config(&request.custom_llm_provider)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
|
||||
let env_lookup = super::handler::environment_lookup;
|
||||
let headers = string_headers(request.extra_headers)?;
|
||||
let url = config.complete_url(
|
||||
request.api_base.as_deref(),
|
||||
&request.model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_transcription_params(&request.optional_params);
|
||||
let body = config.transform_transcription_request(
|
||||
&request.model,
|
||||
request.audio,
|
||||
filtered_params,
|
||||
)?;
|
||||
let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?;
|
||||
let mut upstream_headers = headers.into_iter().collect::<Vec<_>>();
|
||||
if matches!(auth, AudioTranscriptionAuth::Bearer)
|
||||
&& !has_header(
|
||||
&upstream_headers
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect::<std::collections::BTreeMap<_, _>>(),
|
||||
"authorization",
|
||||
)
|
||||
&& let Some(api_key) = request.api_key.as_deref()
|
||||
{
|
||||
upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
|
||||
}
|
||||
let provider_request = ProviderAudioTranscriptionRequest {
|
||||
model: request.model,
|
||||
config,
|
||||
url,
|
||||
body: body.body,
|
||||
upstream_headers,
|
||||
timeout: request.timeout,
|
||||
};
|
||||
let provider_request = self.run_during_call_guardrails(provider_request).await?;
|
||||
sign_request(&provider_request, &request.optional_params).await
|
||||
}
|
||||
|
||||
async fn run_during_call_guardrails(
|
||||
&self,
|
||||
request: ProviderAudioTranscriptionRequest,
|
||||
) -> CoreResult<ProviderAudioTranscriptionRequest> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
let (guardrail_request, _) = self
|
||||
.guardrail_runner
|
||||
.run_during_call(
|
||||
&guardrail_context(&self.request_metadata),
|
||||
GuardrailRequest::new(json!({
|
||||
"model": request.model,
|
||||
"custom_llm_provider": "bedrock",
|
||||
"url": request.url,
|
||||
"body": request.body,
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.map_err(guardrail_error_to_core_error)?;
|
||||
let Value::Object(mut data) = guardrail_request.data else {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"audio transcription during_call guardrail must return an object".to_string(),
|
||||
));
|
||||
};
|
||||
let body = data.remove("body").ok_or_else(|| {
|
||||
CoreError::InvalidRequest("audio transcription guardrail removed body".to_string())
|
||||
})?;
|
||||
Ok(ProviderAudioTranscriptionRequest { body, ..request })
|
||||
}
|
||||
|
||||
fn 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<PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest, Value>
|
||||
for AudioTranscriptionLifecycleHooks
|
||||
{
|
||||
type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>;
|
||||
type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>;
|
||||
type SuccessFuture<'a> = AudioLogFuture<'a>;
|
||||
type FailureFuture<'a> = AudioLogFuture<'a>;
|
||||
|
||||
fn async_pre_call_hook<'a>(
|
||||
&'a self,
|
||||
_context: &'a CallLifecycleContext,
|
||||
request: PreparedAudioTranscriptionRequest,
|
||||
) -> 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: PreparedAudioTranscriptionRequest,
|
||||
) -> 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;
|
||||
}
|
||||
self.logger_runner
|
||||
.async_log_success_event(
|
||||
&ModelCallDetails::from_standard_logging_payload(
|
||||
self.logging_payload(context, timing),
|
||||
),
|
||||
&CallbackValue::new("audio_transcription", response.clone()),
|
||||
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(),
|
||||
};
|
||||
self.logger_runner
|
||||
.async_log_failure_event(
|
||||
&ModelCallDetails::from_standard_logging_payload(
|
||||
self.logging_payload(context, timing),
|
||||
)
|
||||
.with_failure_error(logging_error.clone()),
|
||||
Some(&CallbackValue::new(
|
||||
"error",
|
||||
json!({"message": logging_error.message, "kind": logging_error.kind}),
|
||||
)),
|
||||
CallbackTiming::new(timing.start_time, timing.end_time),
|
||||
)
|
||||
.await;
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
|
||||
GuardrailContext {
|
||||
call_type: CallType::Other("audio_transcription".to_string()),
|
||||
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 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",
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::call_lifecycle::CallLifecycle;
|
||||
use serde_json::Value;
|
||||
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod hooks;
|
||||
mod prepare;
|
||||
mod types;
|
||||
|
||||
pub use types::AudioTranscriptionRequest;
|
||||
|
||||
use handler::execute_audio_transcription_provider_call;
|
||||
use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
|
||||
|
||||
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult<Value> {
|
||||
let PreparedAudioTranscriptionCall { request, hooks } =
|
||||
prepare_audio_transcription_call(request);
|
||||
CallLifecycle::default()
|
||||
.run_request(request, &hooks, execute_audio_transcription_provider_call)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
@ -0,0 +1,55 @@
|
|||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::hooks::AudioTranscriptionLifecycleHooks;
|
||||
use super::types::{AudioTranscriptionRequest, PreparedAudioTranscriptionRequest};
|
||||
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
|
||||
use crate::integrations::custom_logger::CustomLoggerRunner;
|
||||
|
||||
pub(crate) struct PreparedAudioTranscriptionCall {
|
||||
pub(crate) request: PreparedAudioTranscriptionRequest,
|
||||
pub(crate) hooks: AudioTranscriptionLifecycleHooks,
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_audio_transcription_call(
|
||||
request: AudioTranscriptionRequest<'_>,
|
||||
) -> PreparedAudioTranscriptionCall {
|
||||
let call_id = request
|
||||
.litellm_call_id
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(new_audio_transcription_call_id);
|
||||
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
|
||||
.unwrap_or(CustomLlmProvider {
|
||||
model: request.model,
|
||||
custom_llm_provider: "bedrock",
|
||||
});
|
||||
PreparedAudioTranscriptionCall {
|
||||
request: PreparedAudioTranscriptionRequest {
|
||||
model: provider_info.model.to_string(),
|
||||
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
|
||||
litellm_call_id: call_id,
|
||||
audio: request.audio,
|
||||
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: AudioTranscriptionLifecycleHooks::new(
|
||||
CustomLoggerRunner::new(request.callbacks),
|
||||
CustomGuardrailRunner::new(request.guardrails),
|
||||
request.request_metadata,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn new_audio_transcription_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_or(0, |duration| duration.as_nanos());
|
||||
format!("audio-transcription-{timestamp}-{sequence}")
|
||||
}
|
||||
|
|
@ -0,0 +1,53 @@
|
|||
use std::io::{Read, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::thread;
|
||||
|
||||
use serde_json::{Map, json};
|
||||
|
||||
use super::{AudioTranscriptionRequest, audio_transcription};
|
||||
|
||||
#[tokio::test]
|
||||
async fn bedrock_request_is_signed_and_contains_audio() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
|
||||
let address = listener.local_addr().expect("address");
|
||||
let server = thread::spawn(move || {
|
||||
let (mut stream, _) = listener.accept().expect("connection");
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 16_384];
|
||||
let count = stream.read(&mut buffer).expect("request");
|
||||
request.extend_from_slice(&buffer[..count]);
|
||||
let request = String::from_utf8_lossy(&request);
|
||||
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
|
||||
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
|
||||
assert!(request.contains("x-amz-date:"));
|
||||
assert!(request.contains("\"bytes\":\"AQI=\""));
|
||||
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
|
||||
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
|
||||
stream.write_all(response).expect("response");
|
||||
});
|
||||
|
||||
let optional_params = Map::from_iter([
|
||||
("aws_access_key_id".to_string(), json!("access-key")),
|
||||
("aws_secret_access_key".to_string(), json!("secret-key")),
|
||||
("aws_region_name".to_string(), json!("us-east-1")),
|
||||
]);
|
||||
let api_base = format!("http://{address}");
|
||||
let response = audio_transcription(AudioTranscriptionRequest {
|
||||
model: "mistral.voxtral-mini-3b-2507",
|
||||
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
|
||||
api_key: None,
|
||||
api_base: Some(&api_base),
|
||||
custom_llm_provider: Some("bedrock"),
|
||||
extra_headers: None,
|
||||
optional_params,
|
||||
timeout: None,
|
||||
callbacks: Vec::new(),
|
||||
guardrails: Vec::new(),
|
||||
request_metadata: Default::default(),
|
||||
litellm_call_id: None,
|
||||
})
|
||||
.await
|
||||
.expect("transcription");
|
||||
assert_eq!(response, json!({"text": "hello"}));
|
||||
server.join().expect("server");
|
||||
}
|
||||
|
|
@ -0,0 +1,58 @@
|
|||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
|
||||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::integrations::custom_guardrail::CustomGuardrail;
|
||||
use crate::integrations::custom_logger::CustomLogger;
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
|
||||
pub struct AudioTranscriptionRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub audio: 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 PreparedAudioTranscriptionRequest {
|
||||
pub(crate) model: String,
|
||||
pub(crate) custom_llm_provider: String,
|
||||
pub(crate) litellm_call_id: String,
|
||||
pub(crate) audio: 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 PreparedAudioTranscriptionRequest {
|
||||
fn lifecycle_context(&self) -> CallLifecycleContext {
|
||||
CallLifecycleContext::new(
|
||||
"audio_transcription",
|
||||
self.model.clone(),
|
||||
self.custom_llm_provider.clone(),
|
||||
self.litellm_call_id.clone(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ProviderAudioTranscriptionRequest {
|
||||
pub(crate) model: String,
|
||||
pub(crate) config: &'static dyn AudioTranscriptionProviderConfig,
|
||||
pub(crate) url: String,
|
||||
pub(crate) body: Value,
|
||||
pub(crate) upstream_headers: Vec<(String, String)>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
|
|
@ -9,9 +9,9 @@
|
|||
//! runs during extraction, before the handler body. Routes never re-implement it.
|
||||
|
||||
use axum::extract::FromRequestParts;
|
||||
use axum::http::StatusCode;
|
||||
use axum::http::header::AUTHORIZATION;
|
||||
use axum::http::request::Parts;
|
||||
use axum::http::StatusCode;
|
||||
use sha2::{Digest, Sha256};
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
|
|
|
|||
|
|
@ -1,13 +1,13 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
const OCR_TIMEOUT_SECS: u64 = 600;
|
||||
const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
pub(super) fn http_client() -> &'static reqwest::Client {
|
||||
pub(crate) 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))
|
||||
.timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS))
|
||||
.build()
|
||||
.expect("failed to build reqwest client")
|
||||
})
|
||||
|
|
@ -28,3 +28,31 @@ 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";
|
||||
|
||||
/// Full-request timeout ceiling for Anthropic Messages provider calls, in
|
||||
/// seconds. Mirrors the Python Anthropic Messages default. The per-request
|
||||
/// timeout from `litellm_params` still overrides this on the request builder.
|
||||
pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Connect timeout for Anthropic Messages provider calls, in seconds.
|
||||
pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
/// Max characters of an upstream error body echoed across the host boundary
|
||||
/// before truncation, so provider bodies are bounded and data-minimized.
|
||||
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
pub(crate) const DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300;
|
||||
|
||||
/// HTTP path for the non-streaming Anthropic Messages route.
|
||||
#[cfg(feature = "server")]
|
||||
pub(crate) const MESSAGES_ROUTE_PATH: &str = "/v1/messages";
|
||||
|
||||
/// Provider name used by the Anthropic Messages route when a deployment's
|
||||
/// provider model does not carry an explicit provider prefix.
|
||||
pub(crate) const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
|
||||
|
||||
/// Request headers owned by the gateway and never forwarded upstream.
|
||||
#[cfg(feature = "server")]
|
||||
pub(crate) const MESSAGES_HEADERS_NOT_FORWARDED: &[&str] =
|
||||
&["authorization", "connection", "content-length", "host"];
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
pub use crate::audio_transcription::{AudioTranscriptionRequest, audio_transcription};
|
||||
1
litellm-rust/crates/ai-gateway/src/io/messages.rs
Normal file
1
litellm-rust/crates/ai-gateway/src/io/messages.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub use crate::messages::{MessagesRequest, messages};
|
||||
|
|
@ -1,3 +1,6 @@
|
|||
pub mod audio_transcription;
|
||||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod realtime;
|
||||
pub mod realtime_pool;
|
||||
pub mod responses_ws;
|
||||
|
|
|
|||
|
|
@ -1 +1 @@
|
|||
pub use crate::ocr::{ocr, OcrRequest};
|
||||
pub use crate::ocr::{OcrRequest, ocr};
|
||||
|
|
|
|||
|
|
@ -15,16 +15,16 @@ use std::time::Duration;
|
|||
|
||||
use futures_util::stream::{SplitSink, SplitStream};
|
||||
use futures_util::{Sink, SinkExt, Stream, StreamExt};
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::realtime::transformation::RealtimeProviderConfig;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::CoreResult;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
|
||||
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
|
||||
|
||||
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
|
||||
|
||||
|
|
@ -113,7 +113,7 @@ pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<Realt
|
|||
Message::Close(_) => {
|
||||
return Err(CoreError::Network(
|
||||
"upstream closed before first event".to_string(),
|
||||
))
|
||||
));
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,11 +28,11 @@ use std::sync::{Arc, Mutex};
|
|||
use std::time::{Duration, Instant};
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
|
||||
use crate::io::realtime::{
|
||||
dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs,
|
||||
UpstreamRx, UpstreamTx, UpstreamWs, dial_upstream, read_event, resolve_api_key,
|
||||
};
|
||||
|
||||
/// Default target warm sockets per key when pooling is enabled.
|
||||
|
|
@ -473,8 +473,8 @@ pub fn upstream_key(
|
|||
/// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an
|
||||
/// unexpected state. `Pending` (the healthy case) returns `false`.
|
||||
fn is_dead(rx: &mut UpstreamRx) -> bool {
|
||||
use futures_util::task::noop_waker_ref;
|
||||
use futures_util::Stream;
|
||||
use futures_util::task::noop_waker_ref;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
|
|
@ -523,15 +523,15 @@ mod tests {
|
|||
))
|
||||
.await;
|
||||
while let Some(Ok(msg)) = ws.next().await {
|
||||
if let Message::Text(text) = msg {
|
||||
if text.contains("response.create") {
|
||||
for frame in [
|
||||
r#"{"type":"response.created"}"#,
|
||||
r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#,
|
||||
r#"{"type":"response.done"}"#,
|
||||
] {
|
||||
let _ = ws.send(Message::Text(frame.to_string())).await;
|
||||
}
|
||||
if let Message::Text(text) = msg
|
||||
&& text.contains("response.create")
|
||||
{
|
||||
for frame in [
|
||||
r#"{"type":"response.created"}"#,
|
||||
r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#,
|
||||
r#"{"type":"response.done"}"#,
|
||||
] {
|
||||
let _ = ws.send(Message::Text(frame.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
549
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
Normal file
549
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
Normal file
|
|
@ -0,0 +1,549 @@
|
|||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::stream::{SplitSink, SplitStream};
|
||||
use futures_util::{Sink, SinkExt, Stream, StreamExt};
|
||||
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
|
||||
use litellm_core::responses::types::ResponsesWsEvent;
|
||||
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
|
||||
use litellm_core::{CoreError, CoreResult};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
|
||||
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
|
||||
|
||||
use crate::constants::{
|
||||
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
|
||||
};
|
||||
|
||||
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
|
||||
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
|
||||
|
||||
pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
|
||||
type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponsesWebSocketConnection {
|
||||
socket: Arc<Mutex<Option<ResponsesUpstreamWs>>>,
|
||||
}
|
||||
|
||||
impl ResponsesWebSocketConnection {
|
||||
pub async fn connect_url(
|
||||
url: &str,
|
||||
headers: &HashMap<String, String>,
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Self> {
|
||||
let mut request = url
|
||||
.into_client_request()
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
for (name, value) in headers {
|
||||
let header_name = name
|
||||
.parse::<HeaderName>()
|
||||
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
|
||||
let header_value = HeaderValue::from_str(value)
|
||||
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
|
||||
request.headers_mut().insert(header_name, header_value);
|
||||
}
|
||||
let connect = connect_async(request);
|
||||
let result = match timeout {
|
||||
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
|
||||
CoreError::Network("Responses WebSocket connection timed out".to_string())
|
||||
})?,
|
||||
None => connect.await,
|
||||
};
|
||||
let (socket, _) = result.map_err(|error| match error {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
|
||||
status: response.status().as_u16(),
|
||||
body: String::new(),
|
||||
},
|
||||
other => CoreError::Network(other.to_string()),
|
||||
})?;
|
||||
Ok(Self {
|
||||
socket: Arc::new(Mutex::new(Some(socket))),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn send_text(&self, text: String) -> CoreResult<()> {
|
||||
let mut socket = self.socket.lock().await;
|
||||
let Some(socket) = socket.as_mut() else {
|
||||
return Err(CoreError::Network(
|
||||
"Responses WebSocket is closed".to_string(),
|
||||
));
|
||||
};
|
||||
socket
|
||||
.send(Message::Text(text))
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))
|
||||
}
|
||||
|
||||
pub async fn recv_text(&self) -> CoreResult<Option<String>> {
|
||||
let mut socket_guard = self.socket.lock().await;
|
||||
let Some(socket) = socket_guard.as_mut() else {
|
||||
return Ok(None);
|
||||
};
|
||||
match socket.next().await {
|
||||
Some(Ok(Message::Text(text))) => Ok(Some(text)),
|
||||
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
|
||||
.map(Some)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string())),
|
||||
Some(Ok(Message::Close(_))) | None => Ok(None),
|
||||
Some(Ok(_)) => Ok(None),
|
||||
Some(Err(error)) => Err(CoreError::Network(error.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn close(&self) -> CoreResult<()> {
|
||||
let mut socket = self.socket.lock().await;
|
||||
if let Some(socket) = socket.as_mut() {
|
||||
socket
|
||||
.close(None)
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
}
|
||||
*socket = None;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| {
|
||||
std::env::var(OPENAI_API_KEY_ENV)
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
})
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
async fn dial_upstream(
|
||||
model: &str,
|
||||
api_key: &str,
|
||||
api_base: Option<&str>,
|
||||
) -> CoreResult<ResponsesUpstreamWs> {
|
||||
let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model);
|
||||
let mut request = url
|
||||
.as_str()
|
||||
.into_client_request()
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {api_key}"))
|
||||
.map_err(|error| CoreError::Auth(error.to_string()))?,
|
||||
);
|
||||
let result = tokio::time::timeout(
|
||||
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
|
||||
connect_async(request),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CoreError::Network("Responses WebSocket connection timed out".to_string()))?;
|
||||
result
|
||||
.map(|(socket, _)| socket)
|
||||
.map_err(|error| match error {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
|
||||
status: response.status().as_u16(),
|
||||
body: String::new(),
|
||||
},
|
||||
other => CoreError::Network(other.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
pub struct ResponsesWebSocketStreaming;
|
||||
|
||||
impl ResponsesWebSocketStreaming {
|
||||
pub async fn bidirectional_forward<In, Out>(
|
||||
model: &str,
|
||||
upstream_tx: UpstreamTx,
|
||||
upstream_rx: UpstreamRx,
|
||||
idle_timeout: Option<Duration>,
|
||||
observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
Out::Error: std::fmt::Display,
|
||||
{
|
||||
splice(
|
||||
model,
|
||||
upstream_tx,
|
||||
upstream_rx,
|
||||
idle_timeout,
|
||||
observe,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn splice<In, Out>(
|
||||
model: &str,
|
||||
mut upstream_tx: UpstreamTx,
|
||||
mut upstream_rx: UpstreamRx,
|
||||
idle_timeout: Option<Duration>,
|
||||
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
mut client_in: In,
|
||||
mut client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
Out::Error: std::fmt::Display,
|
||||
{
|
||||
let idle =
|
||||
idle_timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS));
|
||||
loop {
|
||||
tokio::select! {
|
||||
event = client_in.next() => {
|
||||
let Some(event) = event else { break };
|
||||
for outbound in OPENAI_RESPONSES_WS_CONFIG
|
||||
.transform_ws_request(&event, model)?
|
||||
.events
|
||||
{
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
|
||||
upstream_tx.send(Message::Text(payload))
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
}
|
||||
}
|
||||
message = upstream_rx.next() => {
|
||||
let Some(message) = message else { break };
|
||||
match message.map_err(|error| CoreError::Network(error.to_string()))? {
|
||||
Message::Text(text) => {
|
||||
let event = serde_json::from_str::<ResponsesWsEvent>(&text)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
|
||||
observe(&event);
|
||||
for outbound in OPENAI_RESPONSES_WS_CONFIG
|
||||
.transform_ws_response(&event, model)?
|
||||
.events
|
||||
{
|
||||
client_out.send(outbound)
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
_ = tokio::time::sleep(idle) => break,
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn async_responses_websocket<In, Out>(
|
||||
model: &str,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
first_frame: Option<ResponsesWsEvent>,
|
||||
idle_timeout: Option<Duration>,
|
||||
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
Out::Error: std::fmt::Display,
|
||||
{
|
||||
let key = resolve_api_key(api_key)?;
|
||||
let upstream = dial_upstream(model, &key, api_base).await?;
|
||||
let (mut upstream_tx, upstream_rx) = upstream.split();
|
||||
if let Some(first_frame) = first_frame {
|
||||
for outbound in OPENAI_RESPONSES_WS_CONFIG
|
||||
.transform_ws_request(&first_frame, model)?
|
||||
.events
|
||||
{
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
|
||||
upstream_tx
|
||||
.send(Message::Text(payload))
|
||||
.await
|
||||
.map_err(|error| CoreError::Network(error.to_string()))?;
|
||||
}
|
||||
}
|
||||
ResponsesWebSocketStreaming::bidirectional_forward(
|
||||
model,
|
||||
upstream_tx,
|
||||
upstream_rx,
|
||||
idle_timeout,
|
||||
&mut observe,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn responses_ws<In, Out>(
|
||||
model: &str,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
first_frame: Option<ResponsesWsEvent>,
|
||||
idle_timeout: Option<Duration>,
|
||||
observe: impl FnMut(&ResponsesWsEvent) + Send,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
Out::Error: std::fmt::Display,
|
||||
{
|
||||
async_responses_websocket(
|
||||
model,
|
||||
api_key,
|
||||
api_base,
|
||||
first_frame,
|
||||
idle_timeout,
|
||||
observe,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use futures_channel::mpsc;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_core::responses::types::ResponsesWsEventType;
|
||||
use serde_json::json;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_tungstenite::accept_async;
|
||||
|
||||
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
|
||||
let address = listener.local_addr().expect("local address");
|
||||
let task = tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.expect("accept");
|
||||
let mut socket = accept_async(stream).await.expect("websocket handshake");
|
||||
while let Some(Ok(Message::Text(text))) = socket.next().await {
|
||||
let request: serde_json::Value = serde_json::from_str(&text).expect("request json");
|
||||
let model = request
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.or_else(|| {
|
||||
request
|
||||
.get("response")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|response| {
|
||||
response.get("model").and_then(serde_json::Value::as_str)
|
||||
})
|
||||
})
|
||||
.expect("enforced model");
|
||||
socket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": format!("resp-{model}"),
|
||||
"model": model,
|
||||
"extra": "preserved"
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.await
|
||||
.expect("created event");
|
||||
socket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": format!("resp-{model}"),
|
||||
"model": model,
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.await
|
||||
.expect("completed event");
|
||||
}
|
||||
});
|
||||
(format!("http://{address}"), task)
|
||||
}
|
||||
|
||||
fn event(value: serde_json::Value) -> ResponsesWsEvent {
|
||||
serde_json::from_value(value).expect("event")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_nonblank_key_wins() {
|
||||
assert_eq!(
|
||||
resolve_api_key(Some(" explicit ")).expect("key"),
|
||||
"explicit"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blank_key_is_not_accepted_without_environment_key() {
|
||||
if std::env::var(OPENAI_API_KEY_ENV).is_err() {
|
||||
assert!(resolve_api_key(Some(" ")).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forwards_events_sequentially_and_enforces_model() {
|
||||
let (api_base, server) = websocket_base().await;
|
||||
let (client_tx, client_rx) = mpsc::unbounded();
|
||||
let (output_tx, mut output_rx) = mpsc::unbounded();
|
||||
let (observed_tx, observed_rx) = mpsc::unbounded();
|
||||
client_tx
|
||||
.unbounded_send(event(json!({
|
||||
"type": "response.create",
|
||||
"model": "wrong"
|
||||
})))
|
||||
.expect("first request");
|
||||
client_tx
|
||||
.unbounded_send(event(json!({
|
||||
"type": "response.create",
|
||||
"response": {"model": "also-wrong"}
|
||||
})))
|
||||
.expect("second request");
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
responses_ws(
|
||||
"authorized-model",
|
||||
Some("test-key"),
|
||||
Some(&api_base),
|
||||
None,
|
||||
Some(Duration::from_secs(1)),
|
||||
move |event| {
|
||||
observed_tx
|
||||
.unbounded_send(event.clone())
|
||||
.expect("observe event");
|
||||
},
|
||||
client_rx,
|
||||
output_tx,
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
let first = output_rx.next().await.expect("first output");
|
||||
let second = output_rx.next().await.expect("second output");
|
||||
let third = output_rx.next().await.expect("third output");
|
||||
let fourth = output_rx.next().await.expect("fourth output");
|
||||
drop(client_tx);
|
||||
task.await.expect("splice task").expect("successful splice");
|
||||
server.await.expect("server task");
|
||||
|
||||
assert_eq!(first.event_type, ResponsesWsEventType::ResponseCreated);
|
||||
assert_eq!(first.model(), Some("authorized-model"));
|
||||
assert_eq!(first.data["response"]["extra"], "preserved");
|
||||
assert_eq!(second.event_type, ResponsesWsEventType::ResponseCompleted);
|
||||
assert_eq!(third.event_type, ResponsesWsEventType::ResponseCreated);
|
||||
assert_eq!(fourth.event_type, ResponsesWsEventType::ResponseCompleted);
|
||||
let observed: Vec<_> = observed_rx.collect().await;
|
||||
assert_eq!(observed.len(), 4);
|
||||
assert!(
|
||||
observed
|
||||
.iter()
|
||||
.all(|event| event.event_type != ResponsesWsEventType::ResponseCreate)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn idle_timeout_ends_without_upstream_events() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
|
||||
let address = listener.local_addr().expect("address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.expect("accept");
|
||||
let _socket = accept_async(stream).await.expect("handshake");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
});
|
||||
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
|
||||
let (output_tx, mut output_rx) = mpsc::unbounded();
|
||||
let result = responses_ws(
|
||||
"model",
|
||||
Some("key"),
|
||||
Some(&format!("http://{address}")),
|
||||
None,
|
||||
Some(Duration::from_millis(20)),
|
||||
|_| {},
|
||||
client_rx,
|
||||
output_tx,
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_ok());
|
||||
assert!(output_rx.next().await.is_none());
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dial_http_status_is_preserved() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
|
||||
let address = listener.local_addr().expect("address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.expect("accept");
|
||||
stream
|
||||
.write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n")
|
||||
.await
|
||||
.expect("response");
|
||||
});
|
||||
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
|
||||
let (output_tx, _output_rx) = mpsc::unbounded();
|
||||
let error = responses_ws(
|
||||
"model",
|
||||
Some("key"),
|
||||
Some(&format!("http://{address}")),
|
||||
None,
|
||||
Some(Duration::from_millis(20)),
|
||||
|_| {},
|
||||
client_rx,
|
||||
output_tx,
|
||||
)
|
||||
.await
|
||||
.expect_err("status error");
|
||||
assert!(matches!(error, CoreError::Http { status: 401, .. }));
|
||||
server.await.expect("server task");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dial_http_500_status_is_preserved() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
|
||||
let address = listener.local_addr().expect("address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.expect("accept");
|
||||
stream
|
||||
.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n")
|
||||
.await
|
||||
.expect("response");
|
||||
});
|
||||
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
|
||||
let (output_tx, _output_rx) = mpsc::unbounded();
|
||||
let error = responses_ws(
|
||||
"model",
|
||||
Some("key"),
|
||||
Some(&format!("http://{address}")),
|
||||
None,
|
||||
Some(Duration::from_millis(20)),
|
||||
|_| {},
|
||||
client_rx,
|
||||
output_tx,
|
||||
)
|
||||
.await
|
||||
.expect_err("status error");
|
||||
assert!(matches!(error, CoreError::Http { status: 500, .. }));
|
||||
server.await.expect("server task");
|
||||
}
|
||||
}
|
||||
|
|
@ -11,7 +11,10 @@
|
|||
//! binary turns on. The `python-config` feature additionally pulls in [`python`]
|
||||
//! for the load-time config reader.
|
||||
|
||||
pub mod audio_transcription;
|
||||
mod client;
|
||||
pub mod io;
|
||||
pub mod messages;
|
||||
pub mod ocr;
|
||||
|
||||
/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
|
||||
|
|
@ -25,9 +28,6 @@ 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")]
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool};
|
||||
use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key};
|
||||
use litellm_ai_gateway::routes;
|
||||
use litellm_ai_gateway::state::AppState;
|
||||
use litellm_core::router::{Deployment, LiteLLMParams, Router};
|
||||
|
|
|
|||
15
litellm-rust/crates/ai-gateway/src/messages/client.rs
Normal file
15
litellm-rust/crates/ai-gateway/src/messages/client.rs
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS};
|
||||
|
||||
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(MESSAGES_TIMEOUT_SECS))
|
||||
.connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new())
|
||||
})
|
||||
}
|
||||
52
litellm-rust/crates/ai-gateway/src/messages/common_utils.rs
Normal file
52
litellm-rust/crates/ai-gateway/src/messages/common_utils.rs
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::{CoreError, json_type_name};
|
||||
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
|
||||
use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
|
||||
use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
|
||||
pub(super) fn messages_provider_config(
|
||||
provider: &str,
|
||||
) -> Option<&'static dyn AnthropicMessagesProviderConfig> {
|
||||
match provider {
|
||||
"anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG),
|
||||
"azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
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!(
|
||||
"messages extra_headers.{key} must be a string, got {}",
|
||||
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))
|
||||
}
|
||||
83
litellm-rust/crates/ai-gateway/src/messages/handler.rs
Normal file
83
litellm-rust/crates/ai-gateway/src/messages/handler.rs
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::common_utils::truncate_error_body;
|
||||
use super::types::ProviderMessagesRequest;
|
||||
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
|
||||
pub(super) async fn execute_messages_provider_call(
|
||||
request: ProviderMessagesRequest,
|
||||
) -> 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();
|
||||
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 = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid messages response JSON: {err}"))
|
||||
})?;
|
||||
let transformed = request
|
||||
.config
|
||||
.transform_response(&request.model, response)?;
|
||||
serde_json::to_value(transformed).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn execute_messages_provider_stream(
|
||||
request: ProviderMessagesRequest,
|
||||
) -> CoreResult<reqwest::Response> {
|
||||
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"streaming messages is not supported for this provider".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
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 !status.is_success() {
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
49
litellm-rust/crates/ai-gateway/src/messages/mod.rs
Normal file
49
litellm-rust/crates/ai-gateway/src/messages/mod.rs
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
use litellm_core::CoreResult;
|
||||
use serde_json::Value;
|
||||
|
||||
mod client;
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
mod types;
|
||||
|
||||
pub use types::MessagesRequest;
|
||||
|
||||
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
|
||||
use prepare::prepare_messages_call;
|
||||
|
||||
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<Value> {
|
||||
match execute_messages(request, false).await? {
|
||||
MessagesResponse::Json(body) => Ok(body),
|
||||
MessagesResponse::Stream(response) => {
|
||||
drop(response);
|
||||
Err(litellm_core::CoreError::InvalidResponse(
|
||||
"non-streaming messages execution returned a stream".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) enum MessagesResponse {
|
||||
Json(Value),
|
||||
Stream(reqwest::Response),
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_messages(
|
||||
request: MessagesRequest<'_>,
|
||||
stream: bool,
|
||||
) -> CoreResult<MessagesResponse> {
|
||||
let prepared = prepare_messages_call(request)?;
|
||||
if stream {
|
||||
execute_messages_provider_stream(prepared)
|
||||
.await
|
||||
.map(MessagesResponse::Stream)
|
||||
} else {
|
||||
execute_messages_provider_call(prepared)
|
||||
.await
|
||||
.map(MessagesResponse::Json)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
73
litellm-rust/crates/ai-gateway/src/messages/prepare.rs
Normal file
73
litellm-rust/crates/ai-gateway/src/messages/prepare.rs
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
use litellm_core::CoreError;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::messages::transformation::MessagesAuthStrategy;
|
||||
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::common_utils::{has_header, messages_provider_config, string_headers};
|
||||
use super::types::{MessagesRequest, ProviderMessagesRequest};
|
||||
|
||||
pub(super) fn prepare_messages_call(
|
||||
request: MessagesRequest<'_>,
|
||||
) -> CoreResult<ProviderMessagesRequest> {
|
||||
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
|
||||
.or_else(|| {
|
||||
request
|
||||
.custom_llm_provider
|
||||
.map(|provider| CustomLlmProvider {
|
||||
model: request.model,
|
||||
custom_llm_provider: provider,
|
||||
})
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidProvider(
|
||||
"unable to resolve custom_llm_provider for messages request".to_string(),
|
||||
)
|
||||
})?;
|
||||
let model = provider_info.model.to_string();
|
||||
let provider = provider_info.custom_llm_provider;
|
||||
|
||||
let config = messages_provider_config(provider)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
|
||||
let mut headers = string_headers(request.extra_headers)?;
|
||||
|
||||
let auth_strategy = config.auth_strategy();
|
||||
if !has_header(&headers, auth_strategy.header_name()) {
|
||||
let api_key = config.resolve_api_key(request.api_key, &env_lookup)?;
|
||||
let auth_header = match auth_strategy {
|
||||
MessagesAuthStrategy::Bearer => {
|
||||
("authorization".to_string(), format!("Bearer {api_key}"))
|
||||
}
|
||||
MessagesAuthStrategy::Header(name) => (name.to_string(), api_key),
|
||||
};
|
||||
headers.push(auth_header);
|
||||
}
|
||||
|
||||
for (name, value) in config.default_headers() {
|
||||
if !has_header(&headers, name) {
|
||||
headers.push((name.to_string(), value.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
let url = config.complete_url(request.api_base, &model, &env_lookup)?;
|
||||
let typed_request = serde_json::from_value(request.body).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
|
||||
})?;
|
||||
let transformed = config.transform_request(typed_request)?;
|
||||
let body = serde_json::to_value(transformed).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"failed to serialize Anthropic messages request: {err}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(ProviderMessagesRequest {
|
||||
provider: provider.to_string(),
|
||||
model,
|
||||
config,
|
||||
url,
|
||||
body,
|
||||
upstream_headers: headers,
|
||||
timeout: request.timeout,
|
||||
})
|
||||
}
|
||||
305
litellm-rust/crates/ai-gateway/src/messages/tests.rs
Normal file
305
litellm-rust/crates/ai-gateway/src/messages/tests.rs
Normal file
|
|
@ -0,0 +1,305 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use serde_json::{Map, Value, json};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use super::common_utils::{
|
||||
has_header, messages_provider_config, string_headers, truncate_error_body,
|
||||
};
|
||||
use super::{MessagesRequest, messages};
|
||||
|
||||
async fn read_http_request(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
let header_end = loop {
|
||||
let n = socket.read(&mut buffer).await.expect("reads request");
|
||||
if n == 0 {
|
||||
break request.len();
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
|
||||
break position + 4;
|
||||
}
|
||||
};
|
||||
let headers = String::from_utf8_lossy(&request[..header_end]);
|
||||
let content_length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().ok())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or(0);
|
||||
while request.len().saturating_sub(header_end) < content_length {
|
||||
let n = socket.read(&mut buffer).await.expect("reads body");
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
}
|
||||
String::from_utf8(request).expect("request is utf8")
|
||||
}
|
||||
|
||||
fn write_response(body: &str) -> String {
|
||||
format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_config_resolves_anthropic_and_azure_ai() {
|
||||
assert!(messages_provider_config("anthropic").is_some());
|
||||
assert!(messages_provider_config("azure_ai").is_some());
|
||||
assert!(messages_provider_config("openai").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_caps_long_payloads() {
|
||||
let body = "x".repeat(400);
|
||||
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 string_headers_rejects_non_string_values() {
|
||||
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
|
||||
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
|
||||
assert!(matches!(err, CoreError::InvalidRequest(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn has_header_is_case_insensitive() {
|
||||
let headers = vec![("X-Api-Key".to_string(), "secret".to_string())];
|
||||
assert!(has_header(&headers, "x-api-key"));
|
||||
assert!(!has_header(&headers, "authorization"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_round_trip_builds_azure_request_and_passes_response_through() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let addr = listener.local_addr().expect("addr");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let request = read_http_request(&mut socket).await;
|
||||
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
|
||||
socket
|
||||
.write_all(write_response(response_body).as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
request
|
||||
});
|
||||
|
||||
let response = messages(MessagesRequest {
|
||||
model: "claude-sonnet-4-5",
|
||||
body: json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 1024,
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "text",
|
||||
"text": "hi",
|
||||
"cache_control": {"type": "ephemeral", "scope": "global"}
|
||||
}]
|
||||
}]
|
||||
}),
|
||||
api_key: Some("sk-azure"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("azure_ai"),
|
||||
extra_headers: None,
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
})
|
||||
.await
|
||||
.expect("messages request succeeds");
|
||||
|
||||
assert_eq!(response["content"][0]["text"], "hi");
|
||||
assert_eq!(response["stop_reason"], "end_turn");
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
let (head, body) = request.split_once("\r\n\r\n").expect("has body");
|
||||
assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}");
|
||||
let head_lower = head.to_ascii_lowercase();
|
||||
assert!(head_lower.contains("x-api-key: sk-azure"), "{head}");
|
||||
assert!(
|
||||
head_lower.contains("anthropic-version: 2023-06-01"),
|
||||
"{head}"
|
||||
);
|
||||
assert!(
|
||||
head_lower.contains("content-type: application/json"),
|
||||
"{head}"
|
||||
);
|
||||
|
||||
let sent_body: Value = serde_json::from_str(body).expect("body is json");
|
||||
assert_eq!(
|
||||
sent_body["messages"][0]["content"][0]["cache_control"],
|
||||
json!({"type": "ephemeral"})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_round_trip_builds_native_anthropic_request() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let addr = listener.local_addr().expect("addr");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let request = read_http_request(&mut socket).await;
|
||||
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
|
||||
socket
|
||||
.write_all(write_response(response_body).as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
request
|
||||
});
|
||||
|
||||
let response = messages(MessagesRequest {
|
||||
model: "claude-sonnet-4-5",
|
||||
body: json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 1024,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
}),
|
||||
api_key: Some("sk-ant"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("anthropic"),
|
||||
extra_headers: None,
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
})
|
||||
.await
|
||||
.expect("messages request succeeds");
|
||||
|
||||
assert_eq!(response["content"][0]["text"], "hi");
|
||||
assert_eq!(response["stop_reason"], "end_turn");
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
let (head, _) = request.split_once("\r\n\r\n").expect("has body");
|
||||
assert!(head.starts_with("POST /v1/messages "), "{head}");
|
||||
let head_lower = head.to_ascii_lowercase();
|
||||
assert!(head_lower.contains("x-api-key: sk-ant"), "{head}");
|
||||
assert!(
|
||||
head_lower.contains("anthropic-version: 2023-06-01"),
|
||||
"{head}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let addr = listener.local_addr().expect("addr");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let request = read_http_request(&mut socket).await;
|
||||
let response_body =
|
||||
r#"{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"m"}"#;
|
||||
socket
|
||||
.write_all(write_response(response_body).as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
request
|
||||
});
|
||||
|
||||
let mut headers = Map::new();
|
||||
headers.insert(
|
||||
"x-api-key".to_string(),
|
||||
Value::String("from-python".to_string()),
|
||||
);
|
||||
headers.insert(
|
||||
"anthropic-beta".to_string(),
|
||||
Value::String("token-efficient-tools-2025-02-19".to_string()),
|
||||
);
|
||||
|
||||
messages(MessagesRequest {
|
||||
model: "claude-sonnet-4-5",
|
||||
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
|
||||
api_key: Some("rust-fallback-key"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("azure_ai"),
|
||||
extra_headers: Some(headers),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
})
|
||||
.await
|
||||
.expect("messages request succeeds");
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
let head = request
|
||||
.split_once("\r\n\r\n")
|
||||
.expect("has body")
|
||||
.0
|
||||
.to_ascii_lowercase();
|
||||
let api_key_count = head
|
||||
.lines()
|
||||
.filter(|line| line.starts_with("x-api-key:"))
|
||||
.count();
|
||||
assert_eq!(api_key_count, 1, "{head}");
|
||||
assert!(head.contains("x-api-key: from-python"), "{head}");
|
||||
assert!(
|
||||
head.contains("anthropic-beta: token-efficient-tools-2025-02-19"),
|
||||
"{head}"
|
||||
);
|
||||
assert!(!head.contains("rust-fallback-key"), "{head}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_maps_provider_error_status_to_http_error() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let addr = listener.local_addr().expect("addr");
|
||||
|
||||
tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let _ = read_http_request(&mut socket).await;
|
||||
let body = "unauthorized";
|
||||
let response = format!(
|
||||
"HTTP/1.1 401 Unauthorized\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
});
|
||||
|
||||
let err = messages(MessagesRequest {
|
||||
model: "claude-sonnet-4-5",
|
||||
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
|
||||
api_key: Some("sk-azure"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: Some("azure_ai"),
|
||||
extra_headers: None,
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
})
|
||||
.await
|
||||
.expect_err("provider error propagates");
|
||||
|
||||
assert!(matches!(err, CoreError::Http { status: 401, .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_rejects_unsupported_provider() {
|
||||
let err = messages(MessagesRequest {
|
||||
model: "claude-3-5-sonnet",
|
||||
body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}),
|
||||
api_key: Some("sk"),
|
||||
api_base: Some("http://127.0.0.1:1"),
|
||||
custom_llm_provider: Some("openai"),
|
||||
extra_headers: None,
|
||||
timeout: Some(Duration::from_millis(50)),
|
||||
})
|
||||
.await
|
||||
.expect_err("unsupported provider errors");
|
||||
|
||||
assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai"));
|
||||
}
|
||||
24
litellm-rust/crates/ai-gateway/src/messages/types.rs
Normal file
24
litellm-rust/crates/ai-gateway/src/messages/types.rs
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub struct MessagesRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub body: 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 timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub(crate) struct ProviderMessagesRequest {
|
||||
pub(crate) provider: String,
|
||||
pub(crate) model: String,
|
||||
pub(crate) config: &'static dyn AnthropicMessagesProviderConfig,
|
||||
pub(crate) url: String,
|
||||
pub(crate) body: Value,
|
||||
pub(crate) upstream_headers: Vec<(String, String)>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
|
|
@ -1,11 +1,11 @@
|
|||
use std::net::IpAddr;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::CoreResult;
|
||||
use reqwest::Url;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{
|
|||
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
|
||||
};
|
||||
|
||||
use super::client::http_client;
|
||||
use crate::client::http_client;
|
||||
|
||||
const ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
use litellm_core::CoreResult;
|
||||
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;
|
||||
use crate::client::http_client;
|
||||
|
||||
pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
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 serde_json::{Map, Value, json};
|
||||
|
||||
use super::common_utils::{
|
||||
convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers,
|
||||
|
|
@ -292,7 +292,7 @@ fn parse_ocr_pre_call_guardrail_request(
|
|||
Some(_) => {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"OCR pre_call guardrail optional_params must be an object".to_string(),
|
||||
))
|
||||
));
|
||||
}
|
||||
None => Map::new(),
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
use litellm_core::call_lifecycle::CallLifecycle;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::call_lifecycle::CallLifecycle;
|
||||
use serde_json::Value;
|
||||
|
||||
mod client;
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod hooks;
|
||||
|
|
@ -12,7 +11,7 @@ mod types;
|
|||
pub use types::OcrRequest;
|
||||
|
||||
use handler::execute_ocr_provider_call;
|
||||
use prepare::{prepare_ocr_call, PreparedOcrCall};
|
||||
use prepare::{PreparedOcrCall, prepare_ocr_call};
|
||||
|
||||
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
|
||||
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
|
||||
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::hooks::OcrLifecycleHooks;
|
||||
use super::types::{OcrRequest, PreparedOcrRequest};
|
||||
|
|
|
|||
|
|
@ -3,12 +3,12 @@ use std::time::Duration;
|
|||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrResponseHandling;
|
||||
use serde_json::{json, Map, Value};
|
||||
use serde_json::{Map, Value, json};
|
||||
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 super::{OcrRequest, ocr};
|
||||
use crate::integrations::custom_guardrail::{
|
||||
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
|
||||
GuardrailFuture, GuardrailRequest,
|
||||
|
|
@ -228,19 +228,23 @@ fn truncate_error_body_does_not_split_multibyte_chars() {
|
|||
#[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!(
|
||||
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("vertex_ai", "deepseek-ocr-maas")
|
||||
.expect("vertex deepseek config resolves")
|
||||
.supported_ocr_params()
|
||||
.contains(&"temperature")
|
||||
);
|
||||
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,9 +7,9 @@
|
|||
//!
|
||||
//! Compiled only under the `python-config` feature.
|
||||
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::router::{Deployment, Router};
|
||||
use litellm_core::CoreResult;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::gil;
|
||||
|
|
@ -17,7 +17,7 @@ use crate::gil;
|
|||
/// Load the router's `model_list` from `config_path` via the Python reader.
|
||||
pub fn load_router_from_config(config_path: &str) -> CoreResult<Router> {
|
||||
gil::record_acquisition();
|
||||
Python::with_gil(|py| {
|
||||
Python::attach(|py| {
|
||||
let model_list = py
|
||||
.import("litellm.proxy.read_model_list")
|
||||
.and_then(|module| module.getattr("read_model_list"))
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
//! Health probes. Simple-route template: a `router()` plus its handlers, in one file.
|
||||
|
||||
use axum::Router;
|
||||
use axum::http::StatusCode;
|
||||
use axum::routing::get;
|
||||
use axum::Router;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
|
|
|
|||
513
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
Normal file
513
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
Normal file
|
|
@ -0,0 +1,513 @@
|
|||
//! `POST /v1/messages`, the Anthropic Messages HTTP surface.
|
||||
|
||||
mod service;
|
||||
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
use axum::extract::{Json, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::post;
|
||||
use litellm_core::CoreError;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::auth::RequireMasterKey;
|
||||
use crate::constants::{MESSAGES_HEADERS_NOT_FORWARDED, MESSAGES_ROUTE_PATH};
|
||||
use crate::state::AppState;
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
_auth: RequireMasterKey,
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> Result<Response, MessagesRouteError> {
|
||||
let extra_headers = forwarded_headers(&headers)?;
|
||||
match service::run(&state.router, body, extra_headers)
|
||||
.await
|
||||
.map_err(MessagesRouteError::from)?
|
||||
{
|
||||
service::MessagesResponse::Json(body) => Ok(Json(body).into_response()),
|
||||
service::MessagesResponse::Stream(upstream) => stream_response(upstream),
|
||||
}
|
||||
}
|
||||
|
||||
fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRouteError> {
|
||||
let content_type = upstream
|
||||
.headers()
|
||||
.get(CONTENT_TYPE)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| HeaderValue::from_static("text/event-stream"));
|
||||
let mut response = Response::builder()
|
||||
.status(
|
||||
StatusCode::from_u16(upstream.status().as_u16()).map_err(|error| {
|
||||
MessagesRouteError(CoreError::InvalidResponse(format!(
|
||||
"invalid upstream response status: {error}"
|
||||
)))
|
||||
})?,
|
||||
)
|
||||
.header(CONTENT_TYPE, content_type);
|
||||
if let Some(value) = upstream.headers().get(CACHE_CONTROL) {
|
||||
response = response.header(CACHE_CONTROL, value);
|
||||
}
|
||||
response
|
||||
.body(Body::from_stream(upstream.bytes_stream()))
|
||||
.map_err(|error| {
|
||||
MessagesRouteError(CoreError::InvalidResponse(format!(
|
||||
"failed to build streaming response: {error}"
|
||||
)))
|
||||
})
|
||||
}
|
||||
|
||||
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, CoreError> {
|
||||
let forwarded = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
!MESSAGES_HEADERS_NOT_FORWARDED
|
||||
.iter()
|
||||
.any(|excluded| name.as_str().eq_ignore_ascii_case(excluded))
|
||||
})
|
||||
.map(|(name, value)| {
|
||||
let value = value.to_str().map_err(|_| {
|
||||
CoreError::InvalidRequest(format!("invalid value for header {}", name.as_str()))
|
||||
})?;
|
||||
Ok((name.to_string(), Value::String(value.to_string())))
|
||||
})
|
||||
.collect::<Result<Map<_, _>, CoreError>>()?;
|
||||
Ok((!forwarded.is_empty()).then_some(forwarded))
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct MessagesRouteError(CoreError);
|
||||
|
||||
impl From<CoreError> for MessagesRouteError {
|
||||
fn from(error: CoreError) -> Self {
|
||||
Self(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoResponse for MessagesRouteError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, message) = match self.0 {
|
||||
CoreError::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
|
||||
CoreError::InvalidProvider(_) | CoreError::Routing(_) => (
|
||||
StatusCode::NOT_FOUND,
|
||||
"no messages deployment is configured for this model".to_string(),
|
||||
),
|
||||
CoreError::Auth(_) => (
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"messages provider authentication failed".to_string(),
|
||||
),
|
||||
CoreError::Http { .. }
|
||||
| CoreError::Network(_)
|
||||
| CoreError::InvalidResponse(_)
|
||||
| CoreError::InvalidType { .. }
|
||||
| CoreError::MissingField(_) => (
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"messages provider request failed".to_string(),
|
||||
),
|
||||
};
|
||||
(
|
||||
status,
|
||||
Json(serde_json::json!({"error": {"message": message}})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::Request;
|
||||
use axum::http::StatusCode;
|
||||
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE};
|
||||
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use super::super::app;
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use crate::state::AppState;
|
||||
|
||||
fn state(model: &str, api_base: String, master_key: Option<&str>) -> AppState {
|
||||
state_with_provider(model, model, api_base, master_key)
|
||||
}
|
||||
|
||||
fn state_with_provider(
|
||||
model_alias: &str,
|
||||
provider_model: &str,
|
||||
api_base: String,
|
||||
master_key: Option<&str>,
|
||||
) -> AppState {
|
||||
AppState {
|
||||
router: Arc::new(ModelRouter::new(vec![Deployment {
|
||||
model_name: model_alias.to_string(),
|
||||
litellm_params: LiteLLMParams {
|
||||
model: format!("anthropic/{provider_model}"),
|
||||
api_key: Some("upstream-key".to_string()),
|
||||
api_base: Some(api_base),
|
||||
},
|
||||
}])),
|
||||
master_key: master_key.map(Arc::from),
|
||||
loggers: Arc::new(Vec::new()),
|
||||
realtime_pool: RealtimePool::disabled(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn upstream(listener: TcpListener) -> (String, tokio::task::JoinHandle<String>) {
|
||||
let address = listener.local_addr().expect("listener has address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 4096];
|
||||
loop {
|
||||
let read = socket.read(&mut buffer).await.expect("reads request");
|
||||
request.extend_from_slice(&buffer[..read]);
|
||||
if request.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let request = String::from_utf8(request).expect("request is utf8");
|
||||
let content_length = request
|
||||
.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);
|
||||
let header_end = request.find("\r\n\r\n").expect("request has headers") + 4;
|
||||
let mut full_request = request.into_bytes();
|
||||
while full_request.len().saturating_sub(header_end) < content_length {
|
||||
let read = socket.read(&mut buffer).await.expect("reads body");
|
||||
full_request.extend_from_slice(&buffer[..read]);
|
||||
}
|
||||
let request = String::from_utf8(full_request).expect("request is utf8");
|
||||
let body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-test"}"#;
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
request
|
||||
});
|
||||
(format!("http://{address}"), server)
|
||||
}
|
||||
|
||||
async fn streaming_upstream(
|
||||
listener: TcpListener,
|
||||
status: u16,
|
||||
content_type: &'static str,
|
||||
body: &'static str,
|
||||
) -> (String, tokio::task::JoinHandle<String>) {
|
||||
let address = listener.local_addr().expect("listener has address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts request");
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 4096];
|
||||
loop {
|
||||
let read = socket.read(&mut buffer).await.expect("reads request");
|
||||
request.extend_from_slice(&buffer[..read]);
|
||||
if request.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let request_text = String::from_utf8(request).expect("request is utf8");
|
||||
let content_length = request_text
|
||||
.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);
|
||||
let header_end = request_text.find("\r\n\r\n").expect("request has headers") + 4;
|
||||
let mut full_request = request_text.into_bytes();
|
||||
while full_request.len().saturating_sub(header_end) < content_length {
|
||||
let read = socket.read(&mut buffer).await.expect("reads body");
|
||||
full_request.extend_from_slice(&buffer[..read]);
|
||||
}
|
||||
let response = format!(
|
||||
"HTTP/1.1 {status} OK\r\ncontent-type: {content_type}\r\ncache-control: no-cache\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
String::from_utf8(full_request).expect("request is utf8")
|
||||
});
|
||||
(format!("http://{address}"), server)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn route_constructs_anthropic_upstream_request() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let (api_base, server) = upstream(listener).await;
|
||||
let app = app(state("claude-test", api_base, Some("master-key")));
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.header("authorization", "Bearer master-key")
|
||||
.header("x-api-key", "request-upstream-key")
|
||||
.header("anthropic-beta", "beta-feature")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hello"}]
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("request builds"),
|
||||
)
|
||||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("response body reads");
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&body).expect("json")["id"],
|
||||
"msg_1"
|
||||
);
|
||||
let upstream_request = server.await.expect("upstream task completes");
|
||||
let (head, body) = upstream_request
|
||||
.split_once("\r\n\r\n")
|
||||
.expect("upstream request has body");
|
||||
let head = head.to_ascii_lowercase();
|
||||
assert!(head.contains("x-api-key: request-upstream-key"));
|
||||
assert!(head.contains("anthropic-beta: beta-feature"));
|
||||
assert!(!head.contains("authorization: bearer master-key"));
|
||||
let body: serde_json::Value = serde_json::from_str(body).expect("upstream body is json");
|
||||
assert_eq!(body["model"], "claude-test");
|
||||
assert_eq!(body["messages"][0]["content"], "hello");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn route_substitutes_model_alias_with_provider_model_upstream() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let (api_base, server) = upstream(listener).await;
|
||||
let app = app(state_with_provider(
|
||||
"production",
|
||||
"claude-sonnet-4-5",
|
||||
api_base,
|
||||
Some("master-key"),
|
||||
));
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.header("authorization", "Bearer master-key")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "production",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hello"}]
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("request builds"),
|
||||
)
|
||||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let upstream_request = server.await.expect("upstream task completes");
|
||||
let (_, upstream_body) = upstream_request
|
||||
.split_once("\r\n\r\n")
|
||||
.expect("upstream request has body");
|
||||
let upstream_body: serde_json::Value =
|
||||
serde_json::from_str(upstream_body).expect("upstream body is json");
|
||||
assert_eq!(upstream_body["model"], "claude-sonnet-4-5");
|
||||
assert_ne!(upstream_body["model"], "production");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn route_streams_anthropic_events_without_buffering_or_reordering() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let events = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
|
||||
let (api_base, server) =
|
||||
streaming_upstream(listener, 200, "text/event-stream", events).await;
|
||||
let app = app(state("claude-test", api_base, Some("master-key")));
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.header("authorization", "Bearer master-key")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"max_tokens": 16,
|
||||
"stream": true,
|
||||
"messages": [{"role": "user", "content": "hello"}]
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("request builds"),
|
||||
)
|
||||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTENT_TYPE)
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap(),
|
||||
"text/event-stream"
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CACHE_CONTROL)
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap(),
|
||||
"no-cache"
|
||||
);
|
||||
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("response body reads");
|
||||
assert_eq!(response_body, events.as_bytes());
|
||||
let upstream_request = server.await.expect("upstream task completes");
|
||||
let (_, upstream_body) = upstream_request
|
||||
.split_once("\r\n\r\n")
|
||||
.expect("upstream request has body");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<serde_json::Value>(upstream_body)
|
||||
.expect("upstream body is json")["stream"],
|
||||
true
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn route_maps_streaming_upstream_errors_before_starting_response() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let (api_base, server) = streaming_upstream(
|
||||
listener,
|
||||
429,
|
||||
"application/json",
|
||||
r#"{"error":"rate limited"}"#,
|
||||
)
|
||||
.await;
|
||||
let app = app(state("claude-test", api_base, Some("master-key")));
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.header("authorization", "Bearer master-key")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"max_tokens": 16,
|
||||
"stream": true,
|
||||
"messages": [{"role": "user", "content": "hello"}]
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("request builds"),
|
||||
)
|
||||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
|
||||
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("response body reads");
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&response_body).expect("error is json")["error"]
|
||||
["message"],
|
||||
"messages provider request failed"
|
||||
);
|
||||
server.await.expect("upstream task completes");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn route_rejects_missing_master_key() {
|
||||
let app = app(state(
|
||||
"claude-test",
|
||||
"http://127.0.0.1:1".to_string(),
|
||||
Some("master-key"),
|
||||
));
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from("{}"))
|
||||
.expect("request builds"),
|
||||
)
|
||||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn route_rejects_invalid_master_key() {
|
||||
let app = app(state(
|
||||
"claude-test",
|
||||
"http://127.0.0.1:1".to_string(),
|
||||
Some("master-key"),
|
||||
));
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.header("authorization", "Bearer wrong-key")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from("{}"))
|
||||
.expect("request builds"),
|
||||
)
|
||||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn route_rejects_malformed_json_without_panicking() {
|
||||
let app = app(state(
|
||||
"claude-test",
|
||||
"http://127.0.0.1:1".to_string(),
|
||||
Some("master-key"),
|
||||
));
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/messages")
|
||||
.header("authorization", "Bearer master-key")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from("{not-json"))
|
||||
.expect("request builds"),
|
||||
)
|
||||
.await
|
||||
.expect("route responds");
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_core::router::Router;
|
||||
use litellm_core::{CoreError, CoreResult};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use crate::messages::{MessagesRequest, execute_messages};
|
||||
|
||||
pub(crate) enum MessagesResponse {
|
||||
Json(Value),
|
||||
Stream(reqwest::Response),
|
||||
}
|
||||
|
||||
pub async fn run(
|
||||
router: &Arc<Router>,
|
||||
body: Value,
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<MessagesResponse> {
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|model| !model.is_empty())
|
||||
.ok_or_else(|| CoreError::InvalidRequest("messages body requires a model".to_string()))?;
|
||||
let deployment = router.get_available_deployment(model).ok_or_else(|| {
|
||||
CoreError::Routing(format!("no deployment available for model '{model}'"))
|
||||
})?;
|
||||
let provider_model = deployment.litellm_params.model.as_str();
|
||||
let upstream_model = provider_model
|
||||
.split_once('/')
|
||||
.map_or(provider_model, |(_, model)| model);
|
||||
let custom_llm_provider = if provider_model.contains('/') {
|
||||
None
|
||||
} else {
|
||||
Some(ANTHROPIC_MESSAGES_PROVIDER)
|
||||
};
|
||||
let mut body = body;
|
||||
body.as_object_mut()
|
||||
.ok_or_else(|| CoreError::InvalidRequest("messages body must be an object".to_string()))?
|
||||
.insert(
|
||||
"model".to_string(),
|
||||
Value::String(upstream_model.to_string()),
|
||||
);
|
||||
|
||||
let request = MessagesRequest {
|
||||
model: provider_model,
|
||||
body,
|
||||
api_key: deployment.litellm_params.api_key.as_deref(),
|
||||
api_base: deployment.litellm_params.api_base.as_deref(),
|
||||
custom_llm_provider,
|
||||
extra_headers,
|
||||
timeout: None,
|
||||
};
|
||||
let stream = request.body.get("stream").and_then(Value::as_bool) == Some(true);
|
||||
execute_messages(request, stream)
|
||||
.await
|
||||
.map(|response| match response {
|
||||
crate::messages::MessagesResponse::Json(body) => MessagesResponse::Json(body),
|
||||
crate::messages::MessagesResponse::Stream(upstream) => {
|
||||
MessagesResponse::Stream(upstream)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -7,7 +7,9 @@
|
|||
|
||||
pub mod gil;
|
||||
pub mod health;
|
||||
pub mod messages;
|
||||
pub mod realtime;
|
||||
pub mod responses;
|
||||
|
||||
use axum::Router;
|
||||
|
||||
|
|
@ -18,6 +20,8 @@ pub fn app(state: AppState) -> Router {
|
|||
Router::new()
|
||||
.merge(health::router())
|
||||
.merge(gil::router())
|
||||
.merge(messages::router())
|
||||
.merge(realtime::router())
|
||||
.merge(responses::router())
|
||||
.with_state(state)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,17 +6,17 @@
|
|||
|
||||
mod service;
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use axum::Router;
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::Response;
|
||||
use axum::routing::get;
|
||||
use axum::Router;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router as ModelRouter;
|
||||
|
|
|
|||
|
|
@ -9,12 +9,12 @@
|
|||
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::io::realtime_pool::{upstream_key, RealtimePool};
|
||||
use crate::io::realtime_pool::{RealtimePool, upstream_key};
|
||||
use futures_util::{Sink, Stream};
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router;
|
||||
use litellm_core::CoreResult;
|
||||
|
||||
/// Select a deployment for `model` and splice the client stream to the provider.
|
||||
///
|
||||
|
|
|
|||
348
litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs
Normal file
348
litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs
Normal file
|
|
@ -0,0 +1,348 @@
|
|||
mod service;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use axum::Router;
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::Response;
|
||||
use axum::routing::get;
|
||||
use futures_util::{Sink, SinkExt, StreamExt};
|
||||
use litellm_core::responses::types::{ResponsesErrorFrame, ResponsesWsEvent, ResponsesWsEventType};
|
||||
use litellm_core::router::Router as ModelRouter;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::auth::RequireMasterKey;
|
||||
use crate::integrations::custom_logger::CustomLogger;
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
use crate::state::AppState;
|
||||
|
||||
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
fn new_call_id() -> String {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_nanos())
|
||||
.unwrap_or(0);
|
||||
let sequence = CALL_SEQ.fetch_add(1, Ordering::Relaxed);
|
||||
format!("respws-{nanos:x}-{sequence:x}")
|
||||
}
|
||||
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/v1/responses", get(handle))
|
||||
.route("/responses", get(handle))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ResponsesQuery {
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
_auth: RequireMasterKey,
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ResponsesQuery>,
|
||||
) -> Result<Response, (StatusCode, String)> {
|
||||
if let Some(model) = query.model.as_deref() {
|
||||
validate_model(&state.router, model)?;
|
||||
}
|
||||
let router = state.router.clone();
|
||||
let loggers = state.loggers.clone();
|
||||
let master_key = state.master_key.clone();
|
||||
Ok(ws.on_upgrade(move |socket| bridge(socket, router, loggers, master_key, query.model)))
|
||||
}
|
||||
|
||||
fn validate_model(router: &ModelRouter, model: &str) -> Result<(), (StatusCode, String)> {
|
||||
if model.trim().is_empty() {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing 'model' query param".to_string(),
|
||||
));
|
||||
}
|
||||
let Some(deployment) = router.get_available_deployment(model) else {
|
||||
return Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("no deployment for model '{model}'"),
|
||||
));
|
||||
};
|
||||
if deployment.litellm_params.model.contains('/')
|
||||
&& !deployment.litellm_params.model.starts_with("openai/")
|
||||
{
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"Responses WebSocket route supports OpenAI deployments only".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_error_and_close<S>(sink: &mut S, message: String)
|
||||
where
|
||||
S: futures_util::Sink<Message> + Unpin,
|
||||
S::Error: std::fmt::Display,
|
||||
{
|
||||
if let Ok(payload) = serde_json::to_string(&ResponsesErrorFrame::invalid_request(message)) {
|
||||
let _ = sink.send(Message::Text(payload)).await;
|
||||
}
|
||||
let _ = sink
|
||||
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
|
||||
code: 1008,
|
||||
reason: "Pre-call error".into(),
|
||||
})))
|
||||
.await;
|
||||
let _ = sink.close().await;
|
||||
}
|
||||
|
||||
struct ResponseClientSink {
|
||||
sink: futures_util::stream::SplitSink<WebSocket, Message>,
|
||||
}
|
||||
|
||||
impl Sink<ResponsesWsEvent> for ResponseClientSink {
|
||||
type Error = axum::Error;
|
||||
|
||||
fn poll_ready(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
context: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Result<(), Self::Error>> {
|
||||
std::pin::Pin::new(&mut self.sink).poll_ready(context)
|
||||
}
|
||||
|
||||
fn start_send(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
item: ResponsesWsEvent,
|
||||
) -> Result<(), Self::Error> {
|
||||
let payload = serde_json::to_string(&item).map_err(axum::Error::new)?;
|
||||
std::pin::Pin::new(&mut self.sink).start_send(Message::Text(payload))
|
||||
}
|
||||
|
||||
fn poll_flush(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
context: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Result<(), Self::Error>> {
|
||||
std::pin::Pin::new(&mut self.sink).poll_flush(context)
|
||||
}
|
||||
|
||||
fn poll_close(
|
||||
mut self: std::pin::Pin<&mut Self>,
|
||||
context: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Result<(), Self::Error>> {
|
||||
std::pin::Pin::new(&mut self.sink).poll_close(context)
|
||||
}
|
||||
}
|
||||
|
||||
impl ResponseClientSink {
|
||||
async fn close_with_code(&mut self, code: u16, reason: &'static str) {
|
||||
let _ = self
|
||||
.sink
|
||||
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
|
||||
code,
|
||||
reason: reason.into(),
|
||||
})))
|
||||
.await;
|
||||
let _ = self.sink.close().await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn bridge(
|
||||
socket: WebSocket,
|
||||
router: Arc<ModelRouter>,
|
||||
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
|
||||
master_key: Option<Arc<str>>,
|
||||
requested_model: Option<String>,
|
||||
) {
|
||||
let (mut ws_sink, ws_stream) = socket.split();
|
||||
let (model, first_frame, stream) = if let Some(model) = requested_model {
|
||||
(model, None, ws_stream)
|
||||
} else {
|
||||
let mut stream = ws_stream;
|
||||
let first = match stream.next().await {
|
||||
Some(Ok(Message::Text(text))) => {
|
||||
match serde_json::from_str::<ResponsesWsEvent>(&text) {
|
||||
Ok(event) => event,
|
||||
Err(_) => {
|
||||
send_error_and_close(
|
||||
&mut ws_sink,
|
||||
"Invalid JSON in response.create event".to_string(),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
send_error_and_close(&mut ws_sink, "Missing response.create event".to_string())
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let Some(model) = first.model().filter(|value| !value.trim().is_empty()) else {
|
||||
send_error_and_close(
|
||||
&mut ws_sink,
|
||||
"Missing model in response.create event".to_string(),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
};
|
||||
if first.event_type != ResponsesWsEventType::ResponseCreate {
|
||||
send_error_and_close(
|
||||
&mut ws_sink,
|
||||
"First frame must be a response.create event".to_string(),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
(model.to_string(), Some(first), stream)
|
||||
};
|
||||
if let Err((status, message)) = validate_model(&router, &model) {
|
||||
let _ = status;
|
||||
let _ = message;
|
||||
send_error_and_close(&mut ws_sink, "Unknown model deployment".to_string()).await;
|
||||
return;
|
||||
}
|
||||
|
||||
let call_id = new_call_id();
|
||||
let metadata = RequestMetadata {
|
||||
user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token),
|
||||
..RequestMetadata::default()
|
||||
};
|
||||
let client_in = Box::pin(stream.filter_map(|message| async move {
|
||||
match message {
|
||||
Ok(Message::Text(text)) => serde_json::from_str::<ResponsesWsEvent>(&text).ok(),
|
||||
_ => None,
|
||||
}
|
||||
}));
|
||||
let mut client_out = ResponseClientSink { sink: ws_sink };
|
||||
let result = service::run(
|
||||
&router,
|
||||
&model,
|
||||
first_frame,
|
||||
None,
|
||||
loggers,
|
||||
call_id,
|
||||
metadata,
|
||||
client_in,
|
||||
&mut client_out,
|
||||
)
|
||||
.await;
|
||||
if result.is_err() {
|
||||
client_out
|
||||
.close_with_code(1011, "Internal server error")
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use crate::state::AppState;
|
||||
use axum::body::Body;
|
||||
use axum::http::Request;
|
||||
use litellm_core::router::Router as ModelRouter;
|
||||
use serde_json::json;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
use tower::ServiceExt;
|
||||
|
||||
struct RecordingSink {
|
||||
messages: Vec<Message>,
|
||||
}
|
||||
|
||||
impl Sink<Message> for RecordingSink {
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
fn poll_ready(
|
||||
self: Pin<&mut Self>,
|
||||
_context: &mut Context<'_>,
|
||||
) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
|
||||
self.messages.push(item);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn poll_flush(
|
||||
self: Pin<&mut Self>,
|
||||
_context: &mut Context<'_>,
|
||||
) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_close(
|
||||
self: Pin<&mut Self>,
|
||||
_context: &mut Context<'_>,
|
||||
) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_call_error_matches_python_frame_and_close() {
|
||||
let mut sink = RecordingSink {
|
||||
messages: Vec::new(),
|
||||
};
|
||||
send_error_and_close(&mut sink, "missing model".to_string()).await;
|
||||
let Message::Text(payload) = &sink.messages[0] else {
|
||||
panic!("expected error text frame");
|
||||
};
|
||||
assert_eq!(
|
||||
serde_json::from_str::<serde_json::Value>(payload).expect("error json"),
|
||||
json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"message": "missing model"
|
||||
}
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
sink.messages[1],
|
||||
Message::Close(Some(axum::extract::ws::CloseFrame {
|
||||
code: 1008,
|
||||
reason: "Pre-call error".into(),
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
fn state() -> AppState {
|
||||
AppState {
|
||||
router: Arc::new(ModelRouter::default()),
|
||||
master_key: Some(Arc::from("master-key")),
|
||||
loggers: Arc::new(Vec::new()),
|
||||
realtime_pool: RealtimePool::disabled(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn auth_rejects_responses_upgrade_before_handler() {
|
||||
let request = Request::builder()
|
||||
.uri("/responses?model=known")
|
||||
.body(Body::empty())
|
||||
.expect("request");
|
||||
let response = router()
|
||||
.with_state(state())
|
||||
.oneshot(request)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_query_model_is_rejected_before_upgrade() {
|
||||
assert_eq!(
|
||||
validate_model(&ModelRouter::default(), "unknown").expect_err("unknown model"),
|
||||
(
|
||||
StatusCode::NOT_FOUND,
|
||||
"no deployment for model 'unknown'".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
156
litellm-rust/crates/ai-gateway/src/routes/responses/service.rs
Normal file
156
litellm-rust/crates/ai-gateway/src/routes/responses/service.rs
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::{Sink, Stream};
|
||||
use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext};
|
||||
use litellm_core::responses::instrumentation::{
|
||||
ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
|
||||
ResponsesWsMetadata,
|
||||
};
|
||||
use litellm_core::responses::types::ResponsesWsEvent;
|
||||
use litellm_core::{CoreError, CoreResult};
|
||||
|
||||
use crate::integrations::custom_logger::{
|
||||
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
|
||||
};
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn run<In, Out>(
|
||||
router: &litellm_core::router::Router,
|
||||
model: &str,
|
||||
first_frame: Option<ResponsesWsEvent>,
|
||||
idle_timeout: Option<Duration>,
|
||||
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
|
||||
call_id: String,
|
||||
metadata: RequestMetadata,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
|
||||
Out: Sink<ResponsesWsEvent> + Unpin + Send,
|
||||
Out::Error: std::fmt::Display,
|
||||
{
|
||||
let deployment = router.get_available_deployment(model).ok_or_else(|| {
|
||||
CoreError::Routing(format!("no deployment available for model '{model}'"))
|
||||
})?;
|
||||
let params = &deployment.litellm_params;
|
||||
let provider_model = params
|
||||
.model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(¶ms.model);
|
||||
if params.model.contains('/') && !params.model.starts_with("openai/") {
|
||||
return Err(CoreError::InvalidProvider(
|
||||
"Responses WebSocket route supports OpenAI deployments only".to_string(),
|
||||
));
|
||||
}
|
||||
let instrumentation = Arc::new(ResponsesWsInstrumentation::new(
|
||||
call_id.clone(),
|
||||
model,
|
||||
ResponsesWsMetadata {
|
||||
user_api_key_hash: metadata.user_api_key_hash,
|
||||
user_api_key_user_id: metadata.user_api_key_user_id,
|
||||
user_api_key_team_id: metadata.user_api_key_team_id,
|
||||
},
|
||||
));
|
||||
let observer_instrumentation = Arc::clone(&instrumentation);
|
||||
let context = CallLifecycleContext::new("responses_websocket", model, "openai", call_id);
|
||||
let result = CallLifecycle::default()
|
||||
.run(context, (), instrumentation.as_ref(), |_| async move {
|
||||
crate::io::responses_ws::async_responses_websocket(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
first_frame,
|
||||
idle_timeout,
|
||||
move |event| {
|
||||
observer_instrumentation.observe(event);
|
||||
},
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await;
|
||||
let outcome = instrumentation.take_or_build_outcome(result.is_ok());
|
||||
dispatch_outcome(loggers, outcome).await;
|
||||
result
|
||||
}
|
||||
|
||||
async fn dispatch_outcome(
|
||||
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
|
||||
outcome: ResponsesWsLogOutcome,
|
||||
) {
|
||||
let runner = CustomLoggerRunner::new(loggers.as_ref().clone());
|
||||
match outcome {
|
||||
ResponsesWsLogOutcome::Success { payload, callback } => {
|
||||
let (details, response, start_time, end_time) = logging_values(payload, callback, None);
|
||||
let _ = runner
|
||||
.async_log_success_event(
|
||||
&details,
|
||||
&response,
|
||||
CallbackTiming::new(start_time, end_time),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
ResponsesWsLogOutcome::Failure {
|
||||
payload,
|
||||
callback,
|
||||
error_message,
|
||||
error_kind,
|
||||
} => {
|
||||
let error = LoggingError {
|
||||
message: error_message,
|
||||
kind: error_kind,
|
||||
};
|
||||
let (details, response, start_time, end_time) =
|
||||
logging_values(payload, callback, Some(error));
|
||||
let _ = runner
|
||||
.async_log_failure_event(
|
||||
&details,
|
||||
Some(&response),
|
||||
CallbackTiming::new(start_time, end_time),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn logging_values(
|
||||
payload: litellm_core::responses::instrumentation::ResponsesWsLogPayload,
|
||||
callback: ResponsesWsCallbackPayload,
|
||||
error: Option<LoggingError>,
|
||||
) -> (ModelCallDetails, CallbackValue, f64, f64) {
|
||||
let start_time = payload.start_time;
|
||||
let end_time = payload.end_time;
|
||||
let callback = CallbackValue::new(callback.object, callback.value);
|
||||
let details = ModelCallDetails::from_standard_logging_payload(
|
||||
crate::integrations::types::StandardLoggingPayload {
|
||||
id: payload.id,
|
||||
litellm_call_id: payload.litellm_call_id,
|
||||
call_type: payload.call_type,
|
||||
model: payload.model,
|
||||
custom_llm_provider: payload.custom_llm_provider,
|
||||
response_cost: payload.response_cost,
|
||||
prompt_tokens: payload.usage.prompt_tokens,
|
||||
completion_tokens: payload.usage.completion_tokens,
|
||||
total_tokens: payload.usage.total_tokens,
|
||||
start_time: payload.start_time,
|
||||
end_time: payload.end_time,
|
||||
stream: payload.stream,
|
||||
metadata: crate::integrations::types::StandardLoggingMetadata {
|
||||
user_api_key_hash: payload.metadata.user_api_key_hash,
|
||||
user_api_key_user_id: payload.metadata.user_api_key_user_id,
|
||||
user_api_key_team_id: payload.metadata.user_api_key_team_id,
|
||||
..Default::default()
|
||||
},
|
||||
messages: None,
|
||||
},
|
||||
);
|
||||
let details = match error {
|
||||
Some(error) => details.with_failure_error(error),
|
||||
None => details,
|
||||
};
|
||||
(details, callback, start_time, end_time)
|
||||
}
|
||||
|
|
@ -10,6 +10,25 @@ rand.workspace = true
|
|||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
sha2.workspace = true
|
||||
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
|
||||
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
|
||||
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
|
||||
aws-sigv4 = { version = "1.5.1", optional = true }
|
||||
aws-types = { version = "1.4.0", optional = true }
|
||||
aws-smithy-runtime-api = { version = "1.13.0", optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
bedrock-auth = [
|
||||
"dep:aws-config",
|
||||
"dep:aws-credential-types",
|
||||
"dep:aws-sdk-sts",
|
||||
"dep:aws-sigv4",
|
||||
"dep:aws-types",
|
||||
"dep:aws-smithy-runtime-api",
|
||||
]
|
||||
|
||||
[dev-dependencies]
|
||||
reqwest.workspace = true
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
|
||||
|
|
|
|||
2
litellm-rust/crates/core/src/audio_transcription/mod.rs
Normal file
2
litellm-rust/crates/core/src/audio_transcription/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::CoreResult;
|
||||
|
||||
use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum AudioTranscriptionAuth {
|
||||
Bearer,
|
||||
AwsSigV4 {
|
||||
region: String,
|
||||
service: &'static str,
|
||||
},
|
||||
}
|
||||
|
||||
pub trait AudioTranscriptionProviderConfig: Sync {
|
||||
fn supported_transcription_params(&self) -> &'static [&'static str];
|
||||
|
||||
fn map_transcription_params(&self, params: &Map<String, Value>) -> Map<String, Value> {
|
||||
params
|
||||
.iter()
|
||||
.filter(|(key, _)| {
|
||||
self.supported_transcription_params()
|
||||
.contains(&key.as_str())
|
||||
})
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn transform_transcription_request(
|
||||
&self,
|
||||
model: &str,
|
||||
audio: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<AudioTranscriptionRequestData>;
|
||||
|
||||
fn transform_transcription_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<AudioTranscriptionResponseData>;
|
||||
|
||||
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 auth_strategy(
|
||||
&self,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<AudioTranscriptionAuth>;
|
||||
}
|
||||
20
litellm-rust/crates/core/src/audio_transcription/types.rs
Normal file
20
litellm-rust/crates/core/src/audio_transcription/types.rs
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AudioTranscriptionRequestData {
|
||||
pub body: Value,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AudioTranscriptionResponseData {
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
impl AudioTranscriptionResponseData {
|
||||
pub fn into_json(self) -> Value {
|
||||
serde_json::json!({
|
||||
"text": self.text,
|
||||
})
|
||||
}
|
||||
}
|
||||
258
litellm-rust/crates/core/src/caching/in_memory_cache.rs
Normal file
258
litellm-rust/crates/core/src/caching/in_memory_cache.rs
Normal file
|
|
@ -0,0 +1,258 @@
|
|||
use std::cmp::Reverse;
|
||||
use std::collections::{BinaryHeap, HashMap};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
pub struct InMemoryCache<V: Clone> {
|
||||
pub cache_dict: HashMap<String, V>,
|
||||
pub ttl_dict: HashMap<String, Duration>,
|
||||
pub expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
|
||||
pub max_size_in_memory: usize,
|
||||
pub default_ttl: Duration,
|
||||
now: Box<dyn Fn() -> Duration + Send + Sync>,
|
||||
}
|
||||
|
||||
impl<V: Clone> Default for InMemoryCache<V> {
|
||||
fn default() -> Self {
|
||||
Self::new(None, None)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone> InMemoryCache<V> {
|
||||
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
|
||||
Self::with_clock(max_size_in_memory, default_ttl, || {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_clock(
|
||||
max_size_in_memory: Option<usize>,
|
||||
default_ttl: Option<Duration>,
|
||||
now: impl Fn() -> Duration + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
cache_dict: HashMap::new(),
|
||||
ttl_dict: HashMap::new(),
|
||||
expiration_heap: BinaryHeap::new(),
|
||||
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
now: Box::new(now),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn evict_cache(&mut self) {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let current_time = (self.now)();
|
||||
while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() {
|
||||
if self.ttl_dict.get(&key).copied() != Some(expiration_time) {
|
||||
self.expiration_heap.pop();
|
||||
} else if expiration_time <= current_time {
|
||||
self.expiration_heap.pop();
|
||||
self.remove_key(&key);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
while self.cache_dict.len() >= self.max_size_in_memory {
|
||||
let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else {
|
||||
break;
|
||||
};
|
||||
if self.ttl_dict.get(&key).copied() == Some(expiration_time) {
|
||||
self.remove_key(&key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn allow_ttl_override(&self, key: &str) -> bool {
|
||||
match self.ttl_dict.get(key).copied() {
|
||||
None => true,
|
||||
Some(expiration_time) => expiration_time < (self.now)(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_cache(&mut self, key: impl Into<String>, value: V, ttl: Option<Duration>) {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
self.evict_cache();
|
||||
let key = key.into();
|
||||
self.cache_dict.insert(key.clone(), value);
|
||||
if self.allow_ttl_override(&key) {
|
||||
let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl);
|
||||
self.ttl_dict.insert(key.clone(), expiration_time);
|
||||
self.expiration_heap.push(Reverse((expiration_time, key)));
|
||||
}
|
||||
}
|
||||
|
||||
// Generic values intentionally omit Python's per-item size check.
|
||||
pub fn get_cache(&mut self, key: &str) -> Option<V> {
|
||||
if self.cache_dict.contains_key(key) {
|
||||
if self.is_key_expired(key) {
|
||||
self.remove_key(key);
|
||||
return None;
|
||||
}
|
||||
return self.cache_dict.get(key).cloned();
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub fn get_ttl(&self, key: &str) -> Option<Duration> {
|
||||
self.ttl_dict.get(key).copied()
|
||||
}
|
||||
|
||||
pub fn delete_cache(&mut self, key: &str) {
|
||||
self.remove_key(key);
|
||||
}
|
||||
|
||||
pub fn flush_cache(&mut self) {
|
||||
self.cache_dict.clear();
|
||||
self.ttl_dict.clear();
|
||||
self.expiration_heap.clear();
|
||||
}
|
||||
|
||||
fn is_key_expired(&self, key: &str) -> bool {
|
||||
self.ttl_dict
|
||||
.get(key)
|
||||
.is_some_and(|expiration_time| *expiration_time < (self.now)())
|
||||
}
|
||||
|
||||
fn remove_key(&mut self, key: &str) {
|
||||
self.cache_dict.remove(key);
|
||||
self.ttl_dict.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
};
|
||||
|
||||
use super::InMemoryCache;
|
||||
use std::time::Duration;
|
||||
|
||||
fn cache(now: Arc<AtomicU64>, max_size: usize, default_ttl: Duration) -> InMemoryCache<String> {
|
||||
InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || {
|
||||
Duration::from_secs(now.load(Ordering::Relaxed))
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ttl_expiry_is_deterministic() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
|
||||
cache.set_cache("key", "value".to_string(), None);
|
||||
assert_eq!(cache.get_cache("key"), Some("value".to_string()));
|
||||
now.store(161, Ordering::Relaxed);
|
||||
assert_eq!(cache.get_cache("key"), None);
|
||||
assert_eq!(cache.get_ttl("key"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_and_per_set_ttl_are_applied() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
|
||||
cache.set_cache("default", "value".to_string(), None);
|
||||
cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20)));
|
||||
assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160)));
|
||||
assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unexpired_entries_do_not_allow_ttl_override() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
|
||||
cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20)));
|
||||
cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80)));
|
||||
assert_eq!(cache.get_cache("key"), Some("second".to_string()));
|
||||
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120)));
|
||||
now.store(121, Ordering::Relaxed);
|
||||
cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80)));
|
||||
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_size_evicts_earliest_expiration() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 2, Duration::from_secs(60));
|
||||
cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10)));
|
||||
cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20)));
|
||||
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
|
||||
assert_eq!(cache.get_cache("early"), None);
|
||||
assert!(cache.get_cache("late").is_some());
|
||||
assert!(cache.get_cache("new").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_entries_are_evicted_before_live_entries() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 3, Duration::from_secs(60));
|
||||
cache.set_cache(
|
||||
"expired-one",
|
||||
"value".to_string(),
|
||||
Some(Duration::from_secs(10)),
|
||||
);
|
||||
cache.set_cache(
|
||||
"expired-two",
|
||||
"value".to_string(),
|
||||
Some(Duration::from_secs(20)),
|
||||
);
|
||||
cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100)));
|
||||
now.store(121, Ordering::Relaxed);
|
||||
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100)));
|
||||
assert_eq!(cache.get_cache("expired-one"), None);
|
||||
assert_eq!(cache.get_cache("expired-two"), None);
|
||||
assert!(cache.get_cache("live").is_some());
|
||||
assert!(cache.get_cache("new").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_heap_entries_are_skipped() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 1, Duration::from_secs(60));
|
||||
cache.set_cache(
|
||||
"removed",
|
||||
"value".to_string(),
|
||||
Some(Duration::from_secs(10)),
|
||||
);
|
||||
cache.delete_cache("removed");
|
||||
cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20)));
|
||||
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
|
||||
assert_eq!(cache.get_cache("removed"), None);
|
||||
assert_eq!(cache.get_cache("kept"), None);
|
||||
assert!(cache.get_cache("new").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn delete_and_flush_remove_values_and_ttls() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 10, Duration::from_secs(60));
|
||||
cache.set_cache("one", "value".to_string(), None);
|
||||
cache.set_cache("two", "value".to_string(), None);
|
||||
cache.delete_cache("one");
|
||||
assert_eq!(cache.get_cache("one"), None);
|
||||
cache.flush_cache();
|
||||
assert!(cache.cache_dict.is_empty());
|
||||
assert!(cache.ttl_dict.is_empty());
|
||||
assert!(cache.expiration_heap.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_max_size_does_not_cache() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 0, Duration::from_secs(60));
|
||||
cache.set_cache("key", "value".to_string(), None);
|
||||
assert_eq!(cache.get_cache("key"), None);
|
||||
assert!(cache.cache_dict.is_empty());
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/core/src/caching/mod.rs
Normal file
1
litellm-rust/crates/core/src/caching/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod in_memory_cache;
|
||||
3
litellm-rust/crates/core/src/constants.rs
Normal file
3
litellm-rust/crates/core/src/constants.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com";
|
||||
pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = "https://api.openai.com/v1";
|
||||
pub const OPENAI_RESPONSES_PATH: &str = "/responses";
|
||||
|
|
@ -19,9 +19,9 @@ pub enum CoreError {
|
|||
InvalidRequest(String),
|
||||
#[error("{0}")]
|
||||
Auth(String),
|
||||
#[error("OCR request failed with status {status}: {body}")]
|
||||
#[error("upstream request failed with status {status}: {body}")]
|
||||
Http { status: u16, body: String },
|
||||
#[error("OCR network error: {0}")]
|
||||
#[error("upstream network error: {0}")]
|
||||
Network(String),
|
||||
#[error("routing error: {0}")]
|
||||
Routing(String),
|
||||
|
|
|
|||
|
|
@ -1,8 +1,13 @@
|
|||
pub mod audio_transcription;
|
||||
pub mod caching;
|
||||
pub mod call_lifecycle;
|
||||
pub mod constants;
|
||||
pub mod error;
|
||||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod providers;
|
||||
pub mod realtime;
|
||||
pub mod responses;
|
||||
pub mod router;
|
||||
pub mod routing_utils;
|
||||
|
||||
|
|
|
|||
2
litellm-rust/crates/core/src/messages/mod.rs
Normal file
2
litellm-rust/crates/core/src/messages/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod transformation;
|
||||
pub mod types;
|
||||
59
litellm-rust/crates/core/src/messages/transformation.rs
Normal file
59
litellm-rust/crates/core/src/messages/transformation.rs
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
use crate::error::CoreResult;
|
||||
|
||||
use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum MessagesAuthStrategy {
|
||||
Bearer,
|
||||
Header(&'static str),
|
||||
}
|
||||
|
||||
impl MessagesAuthStrategy {
|
||||
pub fn header_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Bearer => "authorization",
|
||||
Self::Header(header_name) => header_name,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait AnthropicMessagesProviderConfig: Sync {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
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) -> MessagesAuthStrategy {
|
||||
MessagesAuthStrategy::Header("x-api-key")
|
||||
}
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[
|
||||
("anthropic-version", "2023-06-01"),
|
||||
("content-type", "application/json"),
|
||||
]
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
request: AnthropicMessagesRequest,
|
||||
) -> CoreResult<AnthropicMessagesRequest> {
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
_model: &str,
|
||||
response: AnthropicMessagesResponse,
|
||||
) -> CoreResult<AnthropicMessagesResponse> {
|
||||
Ok(response)
|
||||
}
|
||||
}
|
||||
110
litellm-rust/crates/core/src/messages/types.rs
Normal file
110
litellm-rust/crates/core/src/messages/types.rs
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum SystemPrompt {
|
||||
Text(String),
|
||||
Blocks(Vec<ContentBlock>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum MessageContent {
|
||||
Text(String),
|
||||
Blocks(Vec<ContentBlock>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ContentBlock {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_control: Option<CacheControl>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct CacheControl {
|
||||
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
||||
pub cache_type: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ttl: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub scope: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicMessage {
|
||||
pub role: String,
|
||||
pub content: MessageContent,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicMessagesRequest {
|
||||
pub model: String,
|
||||
pub messages: Vec<AnthropicMessage>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_tokens: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub system: Option<SystemPrompt>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub metadata: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stop_sequences: Option<Vec<String>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stream: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<Value>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub service_tier: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub container: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mcp_servers: Option<Vec<Value>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context_management: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_format: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_config: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub speed: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub inference_geo: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicMessagesResponse {
|
||||
pub id: String,
|
||||
#[serde(rename = "type")]
|
||||
pub message_type: String,
|
||||
pub role: String,
|
||||
pub model: String,
|
||||
pub content: Vec<Value>,
|
||||
// Anthropic always includes stop_reason / stop_sequence, null until the turn
|
||||
// ends; serialize them even when None so callers see the same shape as Python.
|
||||
pub stop_reason: Option<String>,
|
||||
pub stop_sequence: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub container: Option<Value>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,142 @@
|
|||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
|
||||
|
||||
const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
|
||||
const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE";
|
||||
const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com";
|
||||
const MESSAGES_PATH_SUFFIX: &str = "/v1/messages";
|
||||
|
||||
pub struct AnthropicMessagesConfig;
|
||||
|
||||
pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig;
|
||||
|
||||
pub fn non_empty(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
pub fn resolve_anthropic_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
non_empty(api_key)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \
|
||||
environment variable"
|
||||
.to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn complete_anthropic_url(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
let api_base = non_empty(api_base)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string());
|
||||
|
||||
let api_base = api_base.trim_end_matches('/');
|
||||
if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
|
||||
return api_base.to_string();
|
||||
}
|
||||
format!("{api_base}{MESSAGES_PATH_SUFFIX}")
|
||||
}
|
||||
|
||||
impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
Ok(complete_anthropic_url(api_base, env_lookup))
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_anthropic_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn auth_strategy(&self) -> MessagesAuthStrategy {
|
||||
MessagesAuthStrategy::Header("x-api-key")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn url_defaults_to_public_anthropic_endpoint() {
|
||||
assert_eq!(
|
||||
complete_anthropic_url(None, &|_| None),
|
||||
"https://api.anthropic.com/v1/messages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_appends_messages_suffix_to_custom_base() {
|
||||
assert_eq!(
|
||||
complete_anthropic_url(Some("https://proxy.internal"), &|_| None),
|
||||
"https://proxy.internal/v1/messages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_leaves_complete_messages_endpoint_untouched() {
|
||||
assert_eq!(
|
||||
complete_anthropic_url(Some("https://proxy.internal/v1/messages"), &|_| None),
|
||||
"https://proxy.internal/v1/messages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_falls_back_to_env_base() {
|
||||
let with_env = |key: &str| {
|
||||
(key == ANTHROPIC_API_BASE_ENV).then(|| "https://env.anthropic".to_string())
|
||||
};
|
||||
assert_eq!(
|
||||
complete_anthropic_url(Some(" "), &with_env),
|
||||
"https://env.anthropic/v1/messages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_key_prefers_param_then_env_then_errors() {
|
||||
assert_eq!(
|
||||
resolve_anthropic_api_key(Some("sk-param"), &|_| None).unwrap(),
|
||||
"sk-param"
|
||||
);
|
||||
let with_env = |key: &str| (key == ANTHROPIC_API_KEY_ENV).then(|| "sk-env".to_string());
|
||||
assert_eq!(
|
||||
resolve_anthropic_api_key(Some(" "), &with_env).unwrap(),
|
||||
"sk-env"
|
||||
);
|
||||
assert!(matches!(
|
||||
resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"),
|
||||
CoreError::Auth(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_strategy_and_default_headers_match_anthropic() {
|
||||
assert_eq!(
|
||||
ANTHROPIC_MESSAGES_CONFIG.auth_strategy().header_name(),
|
||||
"x-api-key"
|
||||
);
|
||||
assert_eq!(
|
||||
ANTHROPIC_MESSAGES_CONFIG.default_headers(),
|
||||
&[
|
||||
("anthropic-version", "2023-06-01"),
|
||||
("content-type", "application/json"),
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/core/src/providers/anthropic/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/anthropic/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod messages;
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,512 @@
|
|||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
|
||||
use crate::messages::types::{
|
||||
AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock,
|
||||
MessageContent, SystemPrompt,
|
||||
};
|
||||
use crate::providers::anthropic::messages::transformation::{
|
||||
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY";
|
||||
const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE";
|
||||
const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic";
|
||||
const MESSAGES_PATH_SUFFIX: &str = "/v1/messages";
|
||||
const SYSTEM_ROLE: &str = "system";
|
||||
const TEXT_BLOCK_TYPE: &str = "text";
|
||||
|
||||
pub struct AzureAnthropicMessagesConfig {
|
||||
anthropic: AnthropicMessagesConfig,
|
||||
}
|
||||
|
||||
pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig =
|
||||
AzureAnthropicMessagesConfig {
|
||||
anthropic: ANTHROPIC_MESSAGES_CONFIG,
|
||||
};
|
||||
|
||||
pub fn resolve_azure_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
non_empty(api_key)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
"Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable"
|
||||
.to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn complete_azure_anthropic_url(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let api_base = non_empty(api_base)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \
|
||||
Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
|
||||
.to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
let api_base = api_base.trim_end_matches('/');
|
||||
|
||||
if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
|
||||
return Ok(api_base.to_string());
|
||||
}
|
||||
|
||||
let with_anthropic = match api_base.split_once(ANTHROPIC_PATH_SEGMENT) {
|
||||
Some((prefix, _)) => format!("{prefix}{ANTHROPIC_PATH_SEGMENT}"),
|
||||
None => format!("{api_base}{ANTHROPIC_PATH_SEGMENT}"),
|
||||
};
|
||||
Ok(format!("{with_anthropic}{MESSAGES_PATH_SUFFIX}"))
|
||||
}
|
||||
|
||||
fn strip_scope_from_block(block: &mut ContentBlock) {
|
||||
if let Some(cache_control) = block.cache_control.as_mut() {
|
||||
cache_control.scope = None;
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_scope_from_system(system: &mut SystemPrompt) {
|
||||
if let SystemPrompt::Blocks(blocks) = system {
|
||||
blocks.iter_mut().for_each(strip_scope_from_block);
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_scope_from_message(message: &mut AnthropicMessage) {
|
||||
if let MessageContent::Blocks(blocks) = &mut message.content {
|
||||
blocks.iter_mut().for_each(strip_scope_from_block);
|
||||
}
|
||||
}
|
||||
|
||||
fn text_content_block(text: String) -> ContentBlock {
|
||||
let extra = Map::from_iter([
|
||||
(
|
||||
"type".to_string(),
|
||||
Value::String(TEXT_BLOCK_TYPE.to_string()),
|
||||
),
|
||||
("text".to_string(), Value::String(text)),
|
||||
]);
|
||||
ContentBlock {
|
||||
cache_control: None,
|
||||
extra,
|
||||
}
|
||||
}
|
||||
|
||||
fn content_into_blocks(content: MessageContent) -> Vec<ContentBlock> {
|
||||
match content {
|
||||
MessageContent::Text(text) => vec![text_content_block(text)],
|
||||
MessageContent::Blocks(blocks) => blocks,
|
||||
}
|
||||
}
|
||||
|
||||
fn system_into_blocks(system: Option<SystemPrompt>) -> Vec<ContentBlock> {
|
||||
match system {
|
||||
None => Vec::new(),
|
||||
Some(SystemPrompt::Text(text)) => vec![text_content_block(text)],
|
||||
Some(SystemPrompt::Blocks(blocks)) => blocks,
|
||||
}
|
||||
}
|
||||
|
||||
fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMessagesRequest {
|
||||
if !request.messages.iter().any(|msg| msg.role == SYSTEM_ROLE) {
|
||||
return request;
|
||||
}
|
||||
|
||||
let (system_messages, chat_messages): (Vec<AnthropicMessage>, Vec<AnthropicMessage>) = request
|
||||
.messages
|
||||
.into_iter()
|
||||
.partition(|msg| msg.role == SYSTEM_ROLE);
|
||||
|
||||
let folded_system: Vec<ContentBlock> = system_into_blocks(request.system)
|
||||
.into_iter()
|
||||
.chain(
|
||||
system_messages
|
||||
.into_iter()
|
||||
.flat_map(|msg| content_into_blocks(msg.content)),
|
||||
)
|
||||
.collect();
|
||||
|
||||
AnthropicMessagesRequest {
|
||||
messages: chat_messages,
|
||||
system: (!folded_system.is_empty()).then_some(SystemPrompt::Blocks(folded_system)),
|
||||
..request
|
||||
}
|
||||
}
|
||||
|
||||
impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_azure_anthropic_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_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn auth_strategy(&self) -> MessagesAuthStrategy {
|
||||
self.anthropic.auth_strategy()
|
||||
}
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
self.anthropic.default_headers()
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
request: AnthropicMessagesRequest,
|
||||
) -> CoreResult<AnthropicMessagesRequest> {
|
||||
let mut request = fold_system_role_messages(request);
|
||||
if let Some(system) = request.system.as_mut() {
|
||||
strip_scope_from_system(system);
|
||||
}
|
||||
request
|
||||
.messages
|
||||
.iter_mut()
|
||||
.for_each(strip_scope_from_message);
|
||||
self.anthropic.transform_request(request)
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: AnthropicMessagesResponse,
|
||||
) -> CoreResult<AnthropicMessagesResponse> {
|
||||
self.anthropic.transform_response(model, response)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn request_from(value: serde_json::Value) -> AnthropicMessagesRequest {
|
||||
serde_json::from_value(value).expect("valid request")
|
||||
}
|
||||
|
||||
fn to_value(request: AnthropicMessagesRequest) -> serde_json::Value {
|
||||
serde_json::to_value(request).expect("serializable request")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_appends_anthropic_and_messages_suffix() {
|
||||
let url =
|
||||
complete_azure_anthropic_url(Some("https://resource.services.ai.azure.com"), &|_| None)
|
||||
.expect("url builds");
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://resource.services.ai.azure.com/anthropic/v1/messages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_keeps_existing_anthropic_segment() {
|
||||
let url = complete_azure_anthropic_url(
|
||||
Some("https://resource.services.ai.azure.com/anthropic"),
|
||||
&|_| None,
|
||||
)
|
||||
.expect("url builds");
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://resource.services.ai.azure.com/anthropic/v1/messages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_leaves_complete_messages_endpoint_untouched() {
|
||||
for base in [
|
||||
"https://resource.services.ai.azure.com/anthropic/v1/messages",
|
||||
"https://resource.services.ai.azure.com/v1/messages",
|
||||
] {
|
||||
assert_eq!(
|
||||
complete_azure_anthropic_url(Some(base), &|_| None).expect("url builds"),
|
||||
base
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_trims_trailing_slash_and_truncates_after_anthropic() {
|
||||
let url = complete_azure_anthropic_url(
|
||||
Some("https://resource.services.ai.azure.com/anthropic/extra/"),
|
||||
&|_| None,
|
||||
)
|
||||
.expect("url builds");
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://resource.services.ai.azure.com/anthropic/v1/messages"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_falls_back_to_env_then_errors_when_absent() {
|
||||
let with_env = |key: &str| {
|
||||
(key == AZURE_API_BASE_ENV).then(|| "https://env.services.ai.azure.com".to_string())
|
||||
};
|
||||
assert_eq!(
|
||||
complete_azure_anthropic_url(None, &with_env).expect("url builds"),
|
||||
"https://env.services.ai.azure.com/anthropic/v1/messages"
|
||||
);
|
||||
let err = complete_azure_anthropic_url(Some(" "), &|_| None).expect_err("missing base");
|
||||
assert!(matches!(err, CoreError::Auth(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_api_key_prefers_param_then_env() {
|
||||
assert_eq!(
|
||||
resolve_azure_api_key(Some("sk-param"), &|_| None).unwrap(),
|
||||
"sk-param"
|
||||
);
|
||||
let with_env = |key: &str| (key == AZURE_API_KEY_ENV).then(|| "sk-env".to_string());
|
||||
assert_eq!(
|
||||
resolve_azure_api_key(Some(" "), &with_env).unwrap(),
|
||||
"sk-env"
|
||||
);
|
||||
assert!(matches!(
|
||||
resolve_azure_api_key(None, &|_| None).expect_err("missing key"),
|
||||
CoreError::Auth(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_strategy_is_x_api_key() {
|
||||
assert_eq!(
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.auth_strategy()
|
||||
.header_name(),
|
||||
"x-api-key"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_headers_match_python() {
|
||||
assert_eq!(
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG.default_headers(),
|
||||
&[
|
||||
("anthropic-version", "2023-06-01"),
|
||||
("content-type", "application/json"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_request_strips_scope_from_system_and_messages() {
|
||||
let request = request_from(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 1024,
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "sys",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"}
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "hi",
|
||||
"cache_control": {"type": "ephemeral", "scope": "global"}
|
||||
},
|
||||
{"type": "text", "text": "no cache control"}
|
||||
]
|
||||
}
|
||||
]
|
||||
}));
|
||||
|
||||
let transformed = to_value(
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request)
|
||||
.expect("request transforms"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
transformed["system"][0]["cache_control"],
|
||||
json!({"type": "ephemeral", "ttl": "1h"})
|
||||
);
|
||||
assert_eq!(
|
||||
transformed["messages"][0]["content"][0]["cache_control"],
|
||||
json!({"type": "ephemeral"})
|
||||
);
|
||||
assert_eq!(
|
||||
transformed["messages"][0]["content"][1],
|
||||
json!({"type": "text", "text": "no cache control"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_request_is_idempotent_and_preserves_string_system() {
|
||||
let request = request_from(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 16,
|
||||
"system": "plain string system",
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
}));
|
||||
let once = AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request)
|
||||
.expect("request transforms");
|
||||
let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(once.clone())
|
||||
.expect("request transforms");
|
||||
assert_eq!(once, twice);
|
||||
assert_eq!(to_value(once)["system"], json!("plain string system"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_request_preserves_all_supported_params() {
|
||||
let body = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 256,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"system": "be terse",
|
||||
"metadata": {"user_id": "u1"},
|
||||
"stop_sequences": ["STOP"],
|
||||
"stream": false,
|
||||
"temperature": 0.4,
|
||||
"top_p": 0.9,
|
||||
"top_k": 40,
|
||||
"tools": [{"name": "get_weather", "input_schema": {"type": "object"}}],
|
||||
"tool_choice": {"type": "auto"},
|
||||
"thinking": {"type": "enabled", "budget_tokens": 1024},
|
||||
"service_tier": "auto",
|
||||
"container": {"id": "c1"},
|
||||
"mcp_servers": [{"type": "url", "url": "https://mcp.example", "name": "x"}],
|
||||
"context_management": {"edits": []},
|
||||
"output_format": {"type": "json_schema"},
|
||||
"output_config": {"effort": "high"},
|
||||
"speed": "fast",
|
||||
"inference_geo": "us",
|
||||
"litellm_metadata": {"trace": "abc"}
|
||||
});
|
||||
let transformed = to_value(
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request_from(body.clone()))
|
||||
.expect("request transforms"),
|
||||
);
|
||||
assert_eq!(transformed, body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_request_folds_system_role_message_into_top_level_system() {
|
||||
let request = request_from(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 256,
|
||||
"system": [{"type": "text", "text": "base system"}],
|
||||
"messages": [
|
||||
{"role": "user", "content": "fix the bug"},
|
||||
{"role": "system", "content": "Available agent types: claude"}
|
||||
]
|
||||
}));
|
||||
|
||||
let transformed = to_value(
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request)
|
||||
.expect("request transforms"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
transformed["messages"],
|
||||
json!([{"role": "user", "content": "fix the bug"}])
|
||||
);
|
||||
assert_eq!(
|
||||
transformed["system"],
|
||||
json!([
|
||||
{"type": "text", "text": "base system"},
|
||||
{"type": "text", "text": "Available agent types: claude"}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_request_folds_system_role_when_no_top_level_system() {
|
||||
let request = request_from(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 256,
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
|
||||
{"role": "system", "content": [{"type": "text", "text": "sys block"}]}
|
||||
]
|
||||
}));
|
||||
|
||||
let transformed = to_value(
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request)
|
||||
.expect("request transforms"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
transformed["messages"],
|
||||
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
|
||||
);
|
||||
assert_eq!(
|
||||
transformed["system"],
|
||||
json!([{"type": "text", "text": "sys block"}])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_request_leaves_requests_without_system_role_untouched() {
|
||||
let body = json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 256,
|
||||
"system": "be terse",
|
||||
"messages": [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"}
|
||||
]
|
||||
});
|
||||
let transformed = to_value(
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request_from(body.clone()))
|
||||
.expect("request transforms"),
|
||||
);
|
||||
assert_eq!(transformed, body);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_request_rejects_non_object_body() {
|
||||
let err = serde_json::from_value::<AnthropicMessagesRequest>(json!("bad"))
|
||||
.expect_err("non-object body should error");
|
||||
assert!(err.is_data());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_response_passes_through() {
|
||||
let response: AnthropicMessagesResponse = serde_json::from_value(json!({
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "hello"}],
|
||||
"model": "claude-sonnet-4-5",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": null,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 2}
|
||||
}))
|
||||
.expect("valid response");
|
||||
let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_response("claude-sonnet-4-5", response)
|
||||
.expect("response transforms");
|
||||
let value = serde_json::to_value(transformed).expect("serializable");
|
||||
assert_eq!(value["stop_reason"], json!("end_turn"));
|
||||
assert_eq!(value["stop_sequence"], json!(null));
|
||||
assert_eq!(value["content"][0]["text"], json!("hello"));
|
||||
}
|
||||
}
|
||||
|
|
@ -1 +1,2 @@
|
|||
pub mod messages;
|
||||
pub mod ocr;
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::error::{json_type_name, CoreError, CoreResult};
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{json, Map, Value};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
|
|
@ -206,11 +206,11 @@ pub fn complete_document_intelligence_url(
|
|||
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);
|
||||
}
|
||||
if let Some(pages) = optional_params.get("pages")
|
||||
&& let Some(normalized) = normalize_pages_param(pages)?
|
||||
{
|
||||
url.push_str("&pages=");
|
||||
url.push_str(&normalized);
|
||||
}
|
||||
|
||||
Ok(url)
|
||||
|
|
@ -231,7 +231,7 @@ fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> {
|
|||
other => {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"Invalid document type: {other}. Must be 'document_url' or 'image_url'"
|
||||
)))
|
||||
)));
|
||||
}
|
||||
};
|
||||
object
|
||||
|
|
|
|||
|
|
@ -0,0 +1,310 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::audio_transcription::transformation::{
|
||||
AudioTranscriptionAuth, AudioTranscriptionProviderConfig,
|
||||
};
|
||||
use crate::audio_transcription::types::{
|
||||
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
|
||||
};
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
|
||||
use super::aws_base::AwsAuthConfig;
|
||||
use super::constants::{
|
||||
AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE,
|
||||
DEFAULT_BEDROCK_REGION,
|
||||
};
|
||||
|
||||
const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"];
|
||||
|
||||
pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig =
|
||||
BedrockAudioTranscriptionConfig;
|
||||
|
||||
pub struct BedrockAudioTranscriptionConfig;
|
||||
|
||||
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
||||
let mut stripped = model;
|
||||
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let mut region = None;
|
||||
if let Some((candidate, remainder)) = stripped.split_once('/')
|
||||
&& is_bedrock_region(candidate)
|
||||
{
|
||||
region = Some(candidate.to_string());
|
||||
stripped = remainder;
|
||||
}
|
||||
for prefix in ["nova-2/", "nova/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if region.is_none() {
|
||||
region = stripped
|
||||
.strip_prefix("arn:")
|
||||
.and_then(|value| value.split(':').nth(3))
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
}
|
||||
(stripped.to_string(), region)
|
||||
}
|
||||
|
||||
fn is_bedrock_region(value: &str) -> bool {
|
||||
value.len() > 3
|
||||
&& value.contains('-')
|
||||
&& value
|
||||
.chars()
|
||||
.all(|char| char.is_ascii_alphanumeric() || char == '-')
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_region(
|
||||
model_region: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
if let Some(region) = optional_params
|
||||
.get("aws_region_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return region.to_string();
|
||||
}
|
||||
if let Some(region) = model_region {
|
||||
return region.to_string();
|
||||
}
|
||||
env_lookup(AWS_REGION_NAME)
|
||||
.or_else(|| env_lookup(AWS_REGION))
|
||||
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
|
||||
}
|
||||
|
||||
fn audio_fields(audio: Value) -> CoreResult<(String, String)> {
|
||||
let object = audio.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&audio),
|
||||
})?;
|
||||
let data = object
|
||||
.get("data")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField("audio.data"))?;
|
||||
let format = object
|
||||
.get("format")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| matches!(*value, "wav" | "mp3" | "flac" | "ogg"))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string())
|
||||
})?;
|
||||
Ok((data.to_string(), format.to_string()))
|
||||
}
|
||||
|
||||
fn optional_string<'a>(params: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
|
||||
params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
||||
fn supported_transcription_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_PARAMS
|
||||
}
|
||||
|
||||
fn transform_transcription_request(
|
||||
&self,
|
||||
_model: &str,
|
||||
audio: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<AudioTranscriptionRequestData> {
|
||||
let (data, format) = audio_fields(audio)?;
|
||||
let mut instruction = "Transcribe the audio. Respond with only the transcript.".to_string();
|
||||
if let Some(language) = optional_string(&optional_params, "language") {
|
||||
instruction.push_str(&format!(" The audio language is {language}."));
|
||||
}
|
||||
if let Some(prompt) = optional_string(&optional_params, "prompt") {
|
||||
instruction.push_str(&format!(" Additional context: {prompt}"));
|
||||
}
|
||||
let mut inference_config = Map::from_iter([("maxTokens".to_string(), json!(4096))]);
|
||||
if let Some(temperature) = optional_params.get("temperature") {
|
||||
inference_config.insert("temperature".to_string(), temperature.clone());
|
||||
}
|
||||
Ok(AudioTranscriptionRequestData {
|
||||
body: json!({
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"audio": {"format": format, "source": {"bytes": data}}},
|
||||
{"text": instruction}
|
||||
]
|
||||
}],
|
||||
"system": [{"text": "You are a transcription assistant."}],
|
||||
"inferenceConfig": inference_config,
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_transcription_response(
|
||||
&self,
|
||||
_model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<AudioTranscriptionResponseData> {
|
||||
let content = response_json
|
||||
.get("output")
|
||||
.and_then(|value| value.get("message"))
|
||||
.and_then(|value| value.get("content"))
|
||||
.and_then(Value::as_array)
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("Bedrock response has no output content".to_string())
|
||||
})?;
|
||||
let mut text = String::new();
|
||||
for block in content {
|
||||
if let Some(value) = block.get("text").and_then(Value::as_str) {
|
||||
text.push_str(value);
|
||||
}
|
||||
}
|
||||
Ok(AudioTranscriptionResponseData { text })
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let (model_id, model_region) = bedrock_model_id_and_region(model);
|
||||
let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup);
|
||||
let endpoint = optional_params
|
||||
.get("aws_bedrock_runtime_endpoint")
|
||||
.and_then(Value::as_str)
|
||||
.or(api_base)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion));
|
||||
Ok(format!(
|
||||
"{}/model/{model_id}/converse",
|
||||
endpoint.trim_end_matches('/')
|
||||
))
|
||||
}
|
||||
|
||||
fn auth_strategy(
|
||||
&self,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<AudioTranscriptionAuth> {
|
||||
let (_, model_region) = bedrock_model_id_and_region(model);
|
||||
Ok(AudioTranscriptionAuth::AwsSigV4 {
|
||||
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
|
||||
service: BEDROCK_SERVICE,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn aws_auth_config(
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> AwsAuthConfig {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
};
|
||||
let env = |key: &str| env_lookup(key);
|
||||
AwsAuthConfig {
|
||||
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
|
||||
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
|
||||
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
|
||||
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
|
||||
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
|
||||
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
|
||||
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
|
||||
web_identity_token: value("aws_web_identity_token")
|
||||
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
|
||||
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
|
||||
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_matches_python_shape() {
|
||||
let params = Map::from_iter([
|
||||
("language".to_string(), json!("en")),
|
||||
("prompt".to_string(), json!("Speaker names")),
|
||||
("temperature".to_string(), json!(0)),
|
||||
("timestamp_granularities".to_string(), json!(["word"])),
|
||||
]);
|
||||
let params = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.map_transcription_params(¶ms);
|
||||
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
|
||||
.transform_transcription_request(
|
||||
"mistral.voxtral-mini-3b-2507",
|
||||
json!({"data": "AQI=", "format": "wav", "filename": "sample.wav"}),
|
||||
params,
|
||||
)
|
||||
.expect("request");
|
||||
assert_eq!(
|
||||
result.body,
|
||||
json!({
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"audio": {"format": "wav", "source": {"bytes": "AQI="}}},
|
||||
{"text": "Transcribe the audio. Respond with only the transcript. The audio language is en. Additional context: Speaker names"}
|
||||
]
|
||||
}],
|
||||
"system": [{"text": "You are a transcription assistant."}],
|
||||
"inferenceConfig": {"maxTokens": 4096, "temperature": 0}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_concatenates_content_blocks() {
|
||||
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
|
||||
.transform_transcription_response(
|
||||
"model",
|
||||
json!({"output": {"message": {"content": [{"text": "hello "}, {"text": "world"}]}}}),
|
||||
)
|
||||
.expect("response");
|
||||
assert_eq!(result.text, "hello world");
|
||||
assert_eq!(result.into_json(), json!({"text": "hello world"}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_audio_is_rejected() {
|
||||
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_transcription_request(
|
||||
"model",
|
||||
json!({"data": "AQI="}),
|
||||
Map::new(),
|
||||
);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn region_and_url_precedence_match_python() {
|
||||
let params = Map::from_iter([("aws_region_name".to_string(), json!("eu-west-1"))]);
|
||||
let url = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
|
||||
.complete_url(
|
||||
None,
|
||||
"bedrock/us-east-1/mistral.voxtral-mini-3b-2507",
|
||||
¶ms,
|
||||
&no_env,
|
||||
)
|
||||
.expect("url");
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://bedrock-runtime.eu-west-1.amazonaws.com/model/mistral.voxtral-mini-3b-2507/converse"
|
||||
);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue