Merge remote-tracking branch 'upstream/main' into vertex-location-generate-content
|
|
@ -3419,7 +3419,7 @@ workflows:
|
|||
name: integration-<< matrix.suite >>
|
||||
matrix:
|
||||
parameters:
|
||||
suite: [management, accounting, database, providers, mcp, sdk, cost, browser]
|
||||
suite: [management, accounting, database, providers, mcp, sdk, cost, security, browser]
|
||||
- integration_contracts:
|
||||
name: integration-extensions
|
||||
suite: extensions
|
||||
|
|
|
|||
|
|
@ -89,6 +89,7 @@ legacy_paths() {
|
|||
proxy-db-auth-checks)
|
||||
echo tests/unit/proxy/auth/test_auth_checks.py
|
||||
echo tests/unit/proxy/auth/test_user_api_key_auth.py
|
||||
echo tests/unit/proxy/test_credential_slot_registry.py
|
||||
echo tests/unit/proxy/test_deprecated_key_grace_period.py ;;
|
||||
proxy-db-budgets)
|
||||
echo tests/unit/proxy/auth/test_default_end_user_budget_simple.py
|
||||
|
|
@ -106,6 +107,7 @@ legacy_paths() {
|
|||
echo tests/unit/proxy/test_update_spend.py
|
||||
echo tests/unit/skills/test_skills_db.py ;;
|
||||
proxy-db-endpoints-and-responses)
|
||||
echo tests/unit/proxy/engine
|
||||
echo tests/unit/proxy/auth/test_models_fallback_endpoint.py
|
||||
echo tests/unit/proxy/common_utils/test_check_batch_cost.py
|
||||
echo tests/unit/proxy/common_utils/test_check_responses_cost.py
|
||||
|
|
@ -147,7 +149,10 @@ legacy_paths() {
|
|||
echo tests/unit/proxy/test_proxy_server.py ;;
|
||||
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;
|
||||
proxy-extras) echo tests/unit/litellm_proxy_extras ;;
|
||||
proxy-infra) echo tests/unit/gateway ;;
|
||||
proxy-infra)
|
||||
echo tests/unit/gateway
|
||||
echo tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py
|
||||
echo tests/unit/proxy/roi_calculator ;;
|
||||
responses-caching-types)
|
||||
find tests/unit/responses -name 'test_*.py' -not -path 'tests/unit/responses/mcp/*'
|
||||
echo tests/unit/types ;;
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ case "$subject" in
|
|||
;;
|
||||
esac
|
||||
|
||||
ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert"
|
||||
ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|security"
|
||||
# Description must not start with an uppercase letter — kept in sync with the
|
||||
# subjectPattern in .github/workflows/conventional-commits.yml so the local
|
||||
# hook is the strictly tighter of the two gates. (Without this guard, a commit
|
||||
|
|
@ -61,7 +61,7 @@ cat >&2 <<EOF
|
|||
Expected: <type>(<scope>)!: <description>
|
||||
(description must start with a lowercase letter)
|
||||
|
||||
Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert
|
||||
Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert, security
|
||||
Examples:
|
||||
feat(router): add weighted round-robin strategy
|
||||
fix(bedrock): decouple STS region from aws_region_name
|
||||
|
|
|
|||
8
.github/CODEOWNERS
vendored
|
|
@ -1,10 +1,2 @@
|
|||
/ui/ @yuneng-berri @ryan-crabbe-berri
|
||||
/litellm/proxy/_experimental/out/ @yuneng-berri @ryan-crabbe-berri
|
||||
/ui/Dockerfile
|
||||
/ui/nginx.conf
|
||||
/ui/litellm-dashboard/src/lib/http/schema.d.ts
|
||||
/ui/litellm-dashboard/tsconfig.tsbuildinfo
|
||||
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
|
||||
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
|
||||
/litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri
|
||||
/.github/CODEOWNERS @yuneng-berri
|
||||
|
|
|
|||
BIN
.github/assets/roi-calculator/00-original-setup.png
vendored
Normal file
|
After Width: | Height: | Size: 80 KiB |
BIN
.github/assets/roi-calculator/01-connect-github.png
vendored
Normal file
|
After Width: | Height: | Size: 58 KiB |
BIN
.github/assets/roi-calculator/02-repositories.png
vendored
Normal file
|
After Width: | Height: | Size: 63 KiB |
BIN
.github/assets/roi-calculator/03-estimator-schedule.png
vendored
Normal file
|
After Width: | Height: | Size: 70 KiB |
BIN
.github/assets/roi-calculator/04-backfill-progress.png
vendored
Normal file
|
After Width: | Height: | Size: 47 KiB |
BIN
.github/assets/roi-calculator/06-overview.png
vendored
Normal file
|
After Width: | Height: | Size: 76 KiB |
BIN
.github/assets/roi-calculator/07-people-unmatched.png
vendored
Normal file
|
After Width: | Height: | Size: 75 KiB |
BIN
.github/assets/roi-calculator/08-match-email.png
vendored
Normal file
|
After Width: | Height: | Size: 39 KiB |
BIN
.github/assets/roi-calculator/09-people-matched.png
vendored
Normal file
|
After Width: | Height: | Size: 70 KiB |
BIN
.github/assets/roi-calculator/10-pr-reasoning.png
vendored
Normal file
|
After Width: | Height: | Size: 93 KiB |
BIN
.github/assets/roi-calculator/11-settings.png
vendored
Normal file
|
After Width: | Height: | Size: 73 KiB |
BIN
.github/assets/roi-calculator/12-restart-setup.png
vendored
Normal file
|
After Width: | Height: | Size: 39 KiB |
BIN
.github/assets/roi-calculator/13-advanced-settings.png
vendored
Normal file
|
After Width: | Height: | Size: 81 KiB |
BIN
.github/assets/roi-calculator/14-overview-pulls.png
vendored
Normal file
|
After Width: | Height: | Size: 72 KiB |
BIN
.github/assets/roi-calculator/15-sample-preview.png
vendored
Normal file
|
After Width: | Height: | Size: 76 KiB |
BIN
.github/assets/roi-calculator/16-calculator-sidebar.png
vendored
Normal file
|
After Width: | Height: | Size: 50 KiB |
BIN
.github/assets/roi-calculator/19-matching-calculator-icons.png
vendored
Normal file
|
After Width: | Height: | Size: 59 KiB |
BIN
.github/assets/roi-calculator/20-partial-repository-report.png
vendored
Normal file
|
After Width: | Height: | Size: 57 KiB |
BIN
.github/assets/roi-calculator/21-empty-repository-preserved-report.png
vendored
Normal file
|
After Width: | Height: | Size: 56 KiB |
BIN
.github/assets/roi-calculator/22-partial-calculation-explanation.png
vendored
Normal file
|
After Width: | Height: | Size: 65 KiB |
BIN
.github/assets/roi-calculator/23-estimator-outage-preserved-report.png
vendored
Normal file
|
After Width: | Height: | Size: 55 KiB |
1
.github/workflows/conventional-commits.yml
vendored
|
|
@ -41,6 +41,7 @@ jobs:
|
|||
ci
|
||||
chore
|
||||
revert
|
||||
security
|
||||
requireScope: false
|
||||
subjectPattern: ^(?![A-Z]).+$
|
||||
subjectPatternError: |
|
||||
|
|
|
|||
13
.github/workflows/create-rc-branch.yml
vendored
|
|
@ -15,6 +15,8 @@ jobs:
|
|||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
outputs:
|
||||
version: ${{ steps.version.outputs.version }}
|
||||
steps:
|
||||
- name: Require main
|
||||
env:
|
||||
|
|
@ -64,3 +66,14 @@ jobs:
|
|||
sha: context.sha,
|
||||
});
|
||||
core.info(`Created branch ${branchName} at ${context.sha}`);
|
||||
|
||||
linear-release:
|
||||
name: Move the Linear release to rc
|
||||
needs: create-rc-branch
|
||||
permissions:
|
||||
contents: read
|
||||
uses: ./.github/workflows/linear-release.yml
|
||||
with:
|
||||
rc_version: ${{ needs.create-rc-branch.outputs.version }}
|
||||
secrets:
|
||||
LINEAR_API_KEY: ${{ secrets.LINEAR_API_KEY }}
|
||||
|
|
|
|||
60
.github/workflows/lens-worker.yml
vendored
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
name: Lens Worker Image
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches: [main, litellm_oss_branch, "litellm_**"]
|
||||
paths:
|
||||
- deploy/lens/**
|
||||
- litellm/proxy/engine/**
|
||||
- .github/workflows/lens-worker.yml
|
||||
push:
|
||||
branches: [main, litellm_agent_engine]
|
||||
paths:
|
||||
- deploy/lens/**
|
||||
- litellm/proxy/engine/**
|
||||
- .github/workflows/lens-worker.yml
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
lens-worker-image:
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
- name: Build Lens worker
|
||||
run: docker build -f deploy/lens/Dockerfile -t lens-worker:${{ github.sha }} .
|
||||
- name: Verify standalone imports with a read-only filesystem
|
||||
run: |
|
||||
docker run --rm --network none --read-only --cap-drop ALL \
|
||||
--security-opt no-new-privileges --entrypoint python \
|
||||
lens-worker:${{ github.sha }} -c '
|
||||
import os
|
||||
import engine.worker
|
||||
from engine.trace_store import trace_store
|
||||
assert os.getuid() == 65532
|
||||
with trace_store() as store:
|
||||
assert store.count() == 0
|
||||
'
|
||||
- name: Publish versioned Lens worker
|
||||
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
|
||||
env:
|
||||
REGISTRY_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
REGISTRY_USER: ${{ github.actor }}
|
||||
IMAGE: ghcr.io/berriai/litellm-lens-worker:sha-${{ github.sha }}
|
||||
run: |
|
||||
printf '%s' "$REGISTRY_TOKEN" | docker login ghcr.io -u "$REGISTRY_USER" --password-stdin
|
||||
docker tag lens-worker:${{ github.sha }} "$IMAGE"
|
||||
docker push "$IMAGE"
|
||||
printf 'Lens worker image: `%s`\n' "$IMAGE" >> "$GITHUB_STEP_SUMMARY"
|
||||
131
.github/workflows/linear-release.yml
vendored
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
name: Linear Release
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- "rc/**"
|
||||
release:
|
||||
types: [published]
|
||||
workflow_call:
|
||||
inputs:
|
||||
rc_version:
|
||||
description: "X.Y.0 release whose rc branch was just cut"
|
||||
required: true
|
||||
type: string
|
||||
secrets:
|
||||
LINEAR_API_KEY:
|
||||
required: true
|
||||
|
||||
permissions: {}
|
||||
|
||||
jobs:
|
||||
linear-release:
|
||||
name: Linear Release
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
|
||||
- name: Plan
|
||||
id: plan
|
||||
env:
|
||||
EVENT: ${{ github.event_name }}
|
||||
REF_NAME: ${{ github.ref_name }}
|
||||
BEFORE: ${{ github.event.before }}
|
||||
CREATED: ${{ github.event.created }}
|
||||
RC_VERSION: ${{ inputs.rc_version }}
|
||||
RELEASE_TAG: ${{ github.event.release.tag_name }}
|
||||
PRERELEASE: ${{ github.event.release.prerelease }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
sync_base="${BEFORE}"
|
||||
if [ "${CREATED}" = "true" ]; then
|
||||
sync_base=""
|
||||
fi
|
||||
if [ -n "${RC_VERSION}" ]; then
|
||||
echo "version=${RC_VERSION}" >> "$GITHUB_OUTPUT"
|
||||
echo "stage=rc" >> "$GITHUB_OUTPUT"
|
||||
elif [ "${EVENT}" = "release" ]; then
|
||||
if [ "${PRERELEASE}" = "true" ] || ! echo "${RELEASE_TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.0$'; then
|
||||
echo "::notice::${RELEASE_TAG} is not an X.Y.0 stable release; nothing to complete"
|
||||
exit 0
|
||||
fi
|
||||
echo "version=${RELEASE_TAG#v}" >> "$GITHUB_OUTPUT"
|
||||
echo "complete=true" >> "$GITHUB_OUTPUT"
|
||||
elif [ "${REF_NAME}" = "main" ]; then
|
||||
version="$(python3 .github/scripts/read_rc_version.py | cut -d= -f2)"
|
||||
status=0
|
||||
git ls-remote --exit-code --heads origin "rc/${version}" > /dev/null || status=$?
|
||||
case "${status}" in
|
||||
0)
|
||||
IFS=. read -r major minor _ <<< "${version}"
|
||||
version="${major}.$((minor + 1)).0"
|
||||
;;
|
||||
2) ;;
|
||||
*)
|
||||
echo "::error::could not check whether rc/${version} exists (git ls-remote exit ${status})"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
echo "version=${version}" >> "$GITHUB_OUTPUT"
|
||||
echo "sync_base=${sync_base}" >> "$GITHUB_OUTPUT"
|
||||
echo "main=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "version=${REF_NAME#rc/}" >> "$GITHUB_OUTPUT"
|
||||
echo "sync_base=${sync_base}" >> "$GITHUB_OUTPUT"
|
||||
echo "stage=rc" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Sync commits into the release
|
||||
if: steps.plan.outputs.sync_base != ''
|
||||
uses: linear/linear-release-action@d4af10092984f9bc6d5efa075b242bdf01333463 # v0.18.0
|
||||
with:
|
||||
access_key: ${{ secrets.LINEAR_API_KEY }}
|
||||
command: sync
|
||||
name: LiteLLM ${{ steps.plan.outputs.version }}
|
||||
version: ${{ steps.plan.outputs.version }}
|
||||
base_ref: ${{ steps.plan.outputs.sync_base }}
|
||||
cli_version: v0.18.0
|
||||
|
||||
- name: Keep the main stage unless the rc branch was cut during this run
|
||||
id: main_stage
|
||||
if: steps.plan.outputs.main == 'true'
|
||||
env:
|
||||
VERSION: ${{ steps.plan.outputs.version }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
status=0
|
||||
git ls-remote --exit-code --heads origin "rc/${VERSION}" > /dev/null || status=$?
|
||||
case "${status}" in
|
||||
0) echo "::notice::rc/${VERSION} was cut during this run; leaving the release in its rc stage" ;;
|
||||
2) echo "stage=main" >> "$GITHUB_OUTPUT" ;;
|
||||
*)
|
||||
echo "::error::could not check whether rc/${VERSION} exists (git ls-remote exit ${status})"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
- name: Move the release to its stage
|
||||
if: steps.plan.outputs.stage != '' || steps.main_stage.outputs.stage != ''
|
||||
uses: linear/linear-release-action@d4af10092984f9bc6d5efa075b242bdf01333463 # v0.18.0
|
||||
with:
|
||||
access_key: ${{ secrets.LINEAR_API_KEY }}
|
||||
command: update
|
||||
stage: ${{ steps.plan.outputs.stage || steps.main_stage.outputs.stage }}
|
||||
version: ${{ steps.plan.outputs.version }}
|
||||
cli_version: v0.18.0
|
||||
|
||||
- name: Complete the release
|
||||
if: steps.plan.outputs.complete == 'true'
|
||||
uses: linear/linear-release-action@d4af10092984f9bc6d5efa075b242bdf01333463 # v0.18.0
|
||||
with:
|
||||
access_key: ${{ secrets.LINEAR_API_KEY }}
|
||||
command: complete
|
||||
version: ${{ steps.plan.outputs.version }}
|
||||
cli_version: v0.18.0
|
||||
5
.github/workflows/test-code-quality.yml
vendored
|
|
@ -80,6 +80,11 @@ jobs:
|
|||
- name: test_e2e_changed_gate
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
|
||||
|
||||
- name: test_e2e_metadata
|
||||
env:
|
||||
PYTHONPATH: tests/e2e
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_metadata.py tests/code_coverage_tests/test_e2e_junit_report.py
|
||||
|
||||
- name: Check merge smoke harness
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_merge_smoke.py
|
||||
|
||||
|
|
|
|||
11
.github/workflows/test-postgres.yml
vendored
|
|
@ -24,6 +24,7 @@ jobs:
|
|||
timeout-minutes: ${{ matrix.job-timeout-minutes }}
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
|
||||
services:
|
||||
postgres:
|
||||
|
|
@ -134,9 +135,19 @@ jobs:
|
|||
env:
|
||||
TEST_PATH: ${{ matrix.test-path }}
|
||||
WORKERS: ${{ matrix.workers }}
|
||||
PYTEST_ADDOPTS: ${{ matrix.shard == 'proxy-behavior' && '--cov=litellm/proxy/engine --cov-report=xml:coverage-lens-postgres.xml' || '' }}
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10
|
||||
else
|
||||
uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10 -n "${WORKERS}"
|
||||
fi
|
||||
|
||||
- name: Upload Lens database coverage
|
||||
if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'proxy-behavior' && !cancelled()
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
with:
|
||||
use_oidc: true
|
||||
files: coverage-lens-postgres.xml
|
||||
flags: lens-postgres
|
||||
fail_ci_if_error: true
|
||||
|
|
|
|||
4
.github/workflows/test-unit.yml
vendored
|
|
@ -79,7 +79,9 @@ jobs:
|
|||
|
||||
- shard: integrations
|
||||
artifact-name: integrations
|
||||
test-path: ""
|
||||
test-path: >-
|
||||
tests/test_litellm/integrations
|
||||
tests/test_litellm/tracing
|
||||
unit-flag: integrations
|
||||
workers: 2
|
||||
reruns: 3
|
||||
|
|
|
|||
|
|
@ -81,6 +81,8 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
# Spend / analytics
|
||||
"/spend/",
|
||||
"/analytics/",
|
||||
"/engine/",
|
||||
"/v1/traces",
|
||||
"/global/",
|
||||
"/user_agent",
|
||||
"/usage/",
|
||||
|
|
@ -144,6 +146,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
|
|||
{
|
||||
"/",
|
||||
"/routes",
|
||||
"/engine",
|
||||
"/openapi.json",
|
||||
"/docs",
|
||||
"/docs/oauth2-redirect",
|
||||
|
|
|
|||
|
|
@ -99,7 +99,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44358
|
||||
"limit": 44802
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
|
|
|
|||
|
|
@ -1,37 +0,0 @@
|
|||
# Publish MCP servers in the AI Hub
|
||||
|
||||
Set `litellm_settings.public_mcp_servers` to the concrete IDs of the servers you want listed in the public AI Hub. Pin `server_id` in each configuration entry so the publication list stays stable across deployments
|
||||
|
||||
```yaml
|
||||
mcp_servers:
|
||||
documentation:
|
||||
server_id: documentation-mcp
|
||||
url: https://mcp.example.com/mcp
|
||||
transport: http
|
||||
available_on_public_internet: true
|
||||
|
||||
litellm_settings:
|
||||
public_mcp_hub_strict_whitelist: true
|
||||
public_mcp_servers:
|
||||
- documentation-mcp
|
||||
```
|
||||
|
||||
Use `documentation-mcp`, the `server_id`, in the publication list. The configuration key `documentation`, display names, and aliases are not publication IDs. Database-created servers use the ID returned by `/v1/mcp/server`
|
||||
|
||||
The dashboard's **AI Hub > MCP Hub > Manage MCP Hub Visibility** dialog edits this same list. Its YAML example includes the selected server IDs. With database-backed configuration (`store_model_in_db: true`), a value declared in YAML is owned by that file: edit the file and reload, or remove that key from YAML to let the dashboard manage it in the database. File-backed deployments can save the list directly to their configuration file
|
||||
|
||||
To remove all explicit entries, save an empty selection in the dialog or configure:
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
public_mcp_hub_strict_whitelist: true
|
||||
public_mcp_servers: []
|
||||
```
|
||||
|
||||
## Hub listing and network access
|
||||
|
||||
The **Hub listing** column in AI Hub identifies servers that appear in `/public/mcp_hub`. The dashboard derives this status from the current registry and publication settings. Setting `mcp_info.is_public` on a server does not publish it; that response field is derived metadata. `mcp_info.is_public_explicit` identifies registered servers included in the explicit publication list
|
||||
|
||||
Gateway cards and server details show **All Networks** when `available_on_public_internet` is enabled or the server is explicitly published in `public_mcp_servers`. They show **Internal Only** when both are false. The per-server flag defaults to `true`; explicit publication overrides a disabled flag for compatibility. Older proxies that omit the metadata needed to determine access show **Unknown**. These labels describe allowed client IPs; authentication and tool permissions still apply
|
||||
|
||||
The default `public_mcp_hub_strict_whitelist: true` lists only registered servers in `public_mcp_servers`. Legacy mode (`false`) additionally lists registered servers with `available_on_public_internet: true`. In legacy mode, clearing the explicit publication list leaves these automatically listed servers visible. Enable strict mode when the publication list should fully determine hub visibility
|
||||
|
|
@ -24,7 +24,7 @@ model_list:
|
|||
- model_name: sagemaker-completion-model
|
||||
litellm_params:
|
||||
model: sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4
|
||||
input_cost_per_second: 0.000420
|
||||
cost_per_second: 0.000420
|
||||
- model_name: text-embedding-ada-002
|
||||
litellm_params:
|
||||
model: azure/azure-embedding-model
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ def encode_image(image_path):
|
|||
|
||||
|
||||
# Path to your image
|
||||
image_path = "litellm/proxy/logo.jpg"
|
||||
image_path = "litellm/proxy/logo.png"
|
||||
|
||||
# Getting the Base64 string
|
||||
base64_image = encode_image(image_path)
|
||||
|
|
@ -27,7 +27,7 @@ response = client.responses.create(
|
|||
{"type": "input_text", "text": "what color is the image"},
|
||||
{
|
||||
"type": "input_image",
|
||||
"image_url": f"data:image/jpeg;base64,{base64_image}",
|
||||
"image_url": f"data:image/png;base64,{base64_image}",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
|
|
|||
7
deploy/lens/Dockerfile
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
FROM python:3.12-slim
|
||||
WORKDIR /app
|
||||
RUN pip install --no-cache-dir httpx==0.28.1 pydantic==2.11.7
|
||||
COPY litellm/proxy/engine/__init__.py litellm/proxy/engine/models.py litellm/proxy/engine/trace_store.py litellm/proxy/engine/analysis.py litellm/proxy/engine/worker.py /app/engine/
|
||||
VOLUME /tmp
|
||||
USER 65532:65532
|
||||
CMD ["python", "-m", "engine.worker"]
|
||||
8
deploy/lens/Dockerfile.dockerignore
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
**
|
||||
!litellm/
|
||||
!litellm/proxy/
|
||||
!litellm/proxy/engine/
|
||||
!litellm/proxy/engine/__init__.py
|
||||
!litellm/proxy/engine/models.py
|
||||
!litellm/proxy/engine/analysis.py
|
||||
!litellm/proxy/engine/worker.py
|
||||
105
deploy/lens/README.md
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
# Lens worker
|
||||
|
||||
Lens reviews recorded activity and saves evidence-linked findings in the LiteLLM dashboard under Observability, Lens (`/ui/lens/`)
|
||||
|
||||
## Start a worker
|
||||
|
||||
Upgrade your existing LiteLLM proxy to a release that includes Lens with PostgreSQL, agent tracing (`general_settings.tracing: {store: clickhouse}`), and ClickHouse configured through `CLICKHOUSE_URL` and a separate SELECT-only `CLICKHOUSE_READER_URL`. Enable the ClickHouse callback and request/response logging to analyze LLM requests. Lens can only inspect content you actually retain
|
||||
|
||||
In Lens, click **Connect worker**, then **Generate setup command**. The LiteLLM address is filled in for you; change it only if the server running Docker needs a different network address. Copy the command and run it on your server. The dialog changes to **Worker connected** when the container checks in
|
||||
|
||||
The command already contains the compatible worker image and one worker token. No separate API key, source checkout, environment file, or second LiteLLM deployment is needed. Keep the command private because it includes the token. The LiteLLM release provides the dashboard and APIs; the container only runs background analysis
|
||||
|
||||
The dashboard and Compose file pin a verified worker image by digest. The image uses Linux amd64, and the generated command selects that platform. Worker image releases are independent of proxy releases: update the pinned image when changing their API contract. CI also publishes immutable commit tags for reproducible builds
|
||||
|
||||
For deployments managed with Compose, download `compose.yaml` and provide `LITELLM_URL` and `LENS_WORKER_TOKEN` in an environment file. Its default image is already selected:
|
||||
|
||||
```bash
|
||||
docker compose --env-file /path/to/lens.env -f compose.yaml up -d
|
||||
```
|
||||
|
||||
Developers can build locally with `LENS_WORKER_IMAGE=litellm-lens-worker:local docker compose -f deploy/lens/compose.yaml -f deploy/lens/compose.build.yaml up -d --build`
|
||||
|
||||
The worker needs outbound HTTPS access to LiteLLM. It needs no inbound ports, provider keys, direct database access, or GPU. The proxy calls your selected model through its configured router; trace content reaches that model provider. Use a model with JSON output support and known token prices. One worker handles one scan at a time and can serve multiple lenses. For more throughput, start another worker with a separate credential
|
||||
|
||||
V1 setup, manual runs, feedback, and worker credentials are restricted to proxy administrators. Proxy-admin viewers can inspect results. Regular user and team keys cannot access the Lens API. Worker credentials can serve the administrator’s lenses. Revoke it in the connection dialog when retiring a worker. Redeploy the worker alongside proxy upgrades so their API versions match
|
||||
|
||||
## Configure a lens
|
||||
|
||||
Choose agent runs, individual LLM requests, or both. The matching-activity preview updates as you choose an application (the recorded OpenTelemetry service.name) or, for request activity, a LiteLLM model group and add metadata conditions. It shows run names, timestamps, and trace IDs; open a run to inspect its original steps before starting analysis. Suggestions come from up to 100 recent executions and may not include every recorded attribute. You can enter other exact keys and values. Leave service and filters blank for all activity your account can access. Filters are exact key/value matches, combined with AND. Trace filters match span or resource attributes on the same span. Request filters match logged metadata, including caller metadata stored under `requester_metadata`; `tag=value` matches request tags. `swarm=research` works only if your instrumentation records that attribute
|
||||
|
||||
Describe how the agent should behave and optionally add specific checks. Select the lookback window, team and metadata, then choose the percentage to review and an optional maximum. **100% with no maximum selects every matching run**. The preview pages through all matching activity and lets you select particular runs. Percentage sampling uses a stable hash order, rounds up, and applies the optional maximum after the percentage
|
||||
|
||||
Choose your analysis model, parallelism and monthly budget. Parallelism controls simultaneous model calls, not the number of runs selected. New lenses run once by default. Turn on monitoring to repeat the same setup at a custom interval. **Run now** uses the same saved settings immediately, including the same lookback window and sampling. Every scan recalculates the window, so overlapping windows can review the same activity again. Duplicate a lens when you want a separate investigation without changing an existing monitor
|
||||
|
||||
Pausing stops future scheduled scans; cancel the active scan separately if needed. The worker polls every 10 seconds; creating a lens or clicking Run now queues a scan, and due schedules are queued when the worker polls. Scans for the same lens never overlap, and its next interval starts after completion. Closing the browser does not stop the worker. Configuration edits apply to the next scan. A running scan retains its settings and selected execution IDs across retries
|
||||
|
||||
## Read the results
|
||||
|
||||
Needs attention shows issues, highest priority first. Patterns contains useful trends and successful behavior that may not need a fix. Each finding starts with a short explanation and a next step when useful. Expand the limitations for uncertainty and counterexamples. Evidence is grouped by run and collapsed until you need it; each quote opens the original step
|
||||
|
||||
Use the batch selector or Scans tab to reopen previous results. Each batch keeps its own findings, settings, selected runs, coverage and cost. Older batches created before snapshot support remain available through accumulated findings. The Runs tab lists the selected batch's sample and can filter per-run observations, including runs without an observed issue and runs with insufficient evidence. These observations precede the final evidence investigation. Linked-run counts on findings include cited counterexamples, so they are not failure counts
|
||||
|
||||
Choose **This is expected** and explain why to teach later scans about acceptable behavior. Feedback is kept with the lens and included in subsequent reviews. It does not alter historical evidence or exempt different problems
|
||||
|
||||
## What a scan does
|
||||
|
||||
The proxy selects executions received or updated within the configured lookback window, with a two-minute settling period. Older rows without receipt timestamps use execution end time. Overlapping scans do not increment a finding's occurrence count for the same execution ID
|
||||
|
||||
A trace is spans sharing a trace ID within one team, not an automatically reconstructed conversation session. Requests are individual LLM calls. When both sources are enabled, requests correlated to a recorded span by response ID are excluded to reduce double counting
|
||||
|
||||
The worker reviews the selected executions in parallel. It pages through their recorded spans and gives the first reviewer a catalog, task and outcome excerpts. The reviewer can read more original content to resolve uncertainties. Large catalogs and groups of observations are processed in bounded context windows, with every page available. Grouping retains supporting run IDs in code, so a pattern occurring thousands of times does not require a model to repeat thousands of IDs. Candidate investigators can page through supporting observations, other runs and original evidence
|
||||
|
||||
There is no fixed total run, span, candidate or investigation-turn cutoff. Repeated or empty evidence requests stop a stalled investigation. Context windows, the configured budget, available model capacity and recorded evidence still bound practical work. The dashboard reports completed work and gaps. The investigator has no shell, browsing, code-editing or production-action tools
|
||||
|
||||
Each model response must match a bounded JSON schema. A malformed response gets one repair attempt through the same budget controls; repeated invalid output fails the scan. Both the worker and proxy validate quoted evidence. Findings retain exact quotes and open the source trace or request. Resolve a finding after a fix, or dismiss it with a reason. A resolved finding reopens when new execution IDs support the same pattern; dismissed findings remain dismissed
|
||||
|
||||
Coverage distinguishes eligible, sampled, reviewed, partial, and unassessable executions. Findings describe observations in the sample, not population-wide success rates or proven causes. A root span does not prove that a trace contains every expected span. Long, missing, redacted, or expired content limits the conclusions
|
||||
|
||||
## Operations and limits
|
||||
|
||||
PostgreSQL stores configurations, findings and all scan history, returned in pages of 50 jobs. Workers claim jobs with optimistic concurrency and a five-minute lease, renewed every 30 seconds. A disconnected job can be reclaimed up to three times. Cancellation stops subsequent work; a model call already in flight may finish and incur cost
|
||||
|
||||
Before every model call, Lens reserves a conservative amount against the monthly lens budget. Successful calls reconcile to reported cost where pricing is available. Interrupted calls retain their reservation because the provider may have charged. A scan stops when the next reservation would exceed the limit, so it can stop with some budget remaining. Lens budgets are separate from virtual-key budgets; analysis calls use the proxy router directly
|
||||
|
||||
V1 requires ClickHouse for both sources. It does not reconstruct sessions from unrelated trace IDs, guarantee exhaustive reviews, cache all per-execution observations across scans, or automatically fix agent code. Trace contents can change as late spans arrive, even though a job's selected IDs are fixed. Findings should be reviewed by a person before acting on them
|
||||
|
||||
|
||||
## API access
|
||||
|
||||
The UI and API use the same scan lifecycle. Authenticate with a proxy administrator credential for writes, or a proxy-admin viewer credential for reads. Worker credentials are only for worker operations
|
||||
|
||||
```bash
|
||||
curl "$LITELLM_URL/engine" -H "Authorization: Bearer $LITELLM_API_KEY" \
|
||||
-H 'Content-Type: application/json' -d '{
|
||||
"name": "Research quality", "model": "your-model-alias",
|
||||
"context": "Answer the requested question using cited, retrieved evidence.",
|
||||
"source": "traces", "lookback_hours": 24,
|
||||
"sample_percent": 100, "sample_size": null, "concurrency": 8,
|
||||
"enabled": true, "interval_minutes": 1440, "monthly_budget": 50
|
||||
}'
|
||||
|
||||
curl "$LITELLM_URL/engine/$LENS_ID/runs" -X POST \
|
||||
-H "Authorization: Bearer $LITELLM_API_KEY" -H 'Content-Type: application/json' -d '{}'
|
||||
|
||||
curl "$LITELLM_URL/engine/$LENS_ID/runs?offset=0" -H "Authorization: Bearer $LITELLM_API_KEY"
|
||||
curl "$LITELLM_URL/engine/$LENS_ID/runs/$BATCH_ID" -H "Authorization: Bearer $LITELLM_API_KEY"
|
||||
```
|
||||
|
||||
Creation queues the first batch. Posting to `/engine/{id}/runs` queues another, or returns the existing active batch. The run response contains its ID under `jobs[0].id`. Poll the batch URL for status, findings and assessments. List responses omit large result payloads; request a batch to retrieve them. Supply an optional complete `settings` object on the runs POST for a one-off override; the saved lens stays unchanged. Selection accepts `team_id`, exact `filters`, and opaque `execution_ids` returned by `/engine/preview/sample`. Preview accepts `offset` and `as_of` to keep the time window fixed while paging. Feedback uses `PATCH /engine/{id}/findings/{finding_id}` with `status` and `reason`
|
||||
|
||||
## Quality evaluation
|
||||
|
||||
Run the checked-in cases against a configured real model. Expected labels are used only for scoring, never passed to the model. Dev and held-out cases include missing outcomes, failed tools, recovery, handoffs, unsupported claims, repeated work, long evidence and prompt injection. The background option adds clean arithmetic traces to test rare-issue discovery at scale; those repeated synthetic cases do not establish accuracy on every production workload
|
||||
|
||||
```bash
|
||||
python -m tests.proxy_behavior.lens.evaluate --api-base "$LITELLM_URL" \
|
||||
--model your-model-alias --split all --background 1000 --concurrency 16 \
|
||||
--output /tmp/lens-quality.json
|
||||
```
|
||||
|
||||
Set `LITELLM_API_KEY` privately. This makes paid model calls. Inspect missed and unexpected per-run labels, final findings and coverage; do not equate a passing dataset with guaranteed detection on arbitrary traces
|
||||
|
||||
The worker uses temporary disk space for trace content while reviewing it, and removes those files after each review. Its Docker image supplies a writable temporary volume while keeping the application filesystem read-only
|
||||
|
||||
To check that accepted behavior stays accepted without hiding new problems, run the evaluator with `--dataset tests/proxy_behavior/lens/feedback_cases.json`. Reports include elapsed time, model call count, reported cost when the proxy provides it, missed checks, unexpected checks, and inconclusive candidates
|
||||
6
deploy/lens/compose.build.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
services:
|
||||
lens-worker:
|
||||
build:
|
||||
context: ../..
|
||||
dockerfile: deploy/lens/Dockerfile
|
||||
image: litellm-lens-worker:local
|
||||
10
deploy/lens/compose.yaml
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
services:
|
||||
lens-worker:
|
||||
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:40fdb82113dd4474cb6e833cf28552487d87c8baf61693a1c3fc2863b7968c6a}
|
||||
environment:
|
||||
LITELLM_URL: ${LITELLM_URL:?Set the URL reachable from this container}
|
||||
LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:?Create a worker credential in the Lens UI}
|
||||
restart: unless-stopped
|
||||
read_only: true
|
||||
cap_drop: [ALL]
|
||||
security_opt: [no-new-privileges:true]
|
||||
BIN
deploy/lens/screenshots/after.png
Normal file
|
After Width: | Height: | Size: 95 KiB |
BIN
deploy/lens/screenshots/before.png
Normal file
|
After Width: | Height: | Size: 6.9 KiB |
BIN
deploy/lens/screenshots/finding.png
Normal file
|
After Width: | Height: | Size: 89 KiB |
BIN
deploy/lens/screenshots/progress.png
Normal file
|
After Width: | Height: | Size: 80 KiB |
BIN
deploy/lens/screenshots/setup.png
Normal file
|
After Width: | Height: | Size: 70 KiB |
BIN
deploy/lens/screenshots/trace.png
Normal file
|
After Width: | Height: | Size: 132 KiB |
|
|
@ -103,7 +103,7 @@ ENV LITELLM_NON_ROOT=true
|
|||
|
||||
RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
|
||||
cp -r /app/litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \
|
||||
cp /app/litellm/proxy/logo.jpg /var/lib/litellm/assets/logo.jpg && \
|
||||
cp /app/litellm/proxy/logo.png /var/lib/litellm/assets/logo.png && \
|
||||
touch /var/lib/litellm/ui/.litellm_ui_ready
|
||||
|
||||
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
||||
|
|
|
|||
62
docker/docker-compose.tracing.yml
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
name: litellm-tracing
|
||||
|
||||
services:
|
||||
litellm:
|
||||
build:
|
||||
context: ..
|
||||
target: runtime
|
||||
command: ["--config", "/app/tracing-config.yaml", "--port", "4000"]
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: local-tracing-master-key
|
||||
LITELLM_SALT_KEY: sk-local-tracing-salt-key
|
||||
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
|
||||
STORE_MODEL_IN_DB: "True"
|
||||
CLICKHOUSE_URL: http://default:local-tracing@clickhouse:8123
|
||||
CLICKHOUSE_READER_URL: http://default:local-tracing@clickhouse:8123
|
||||
CLICKHOUSE_DATABASE: litellm
|
||||
OPENAI_API_KEY: ${OPENAI_API_KEY:-}
|
||||
volumes:
|
||||
- ./tracing-config.yaml:/app/tracing-config.yaml:ro
|
||||
ports:
|
||||
- "127.0.0.1:4002:4000"
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
clickhouse:
|
||||
condition: service_healthy
|
||||
|
||||
db:
|
||||
image: postgres:16
|
||||
environment:
|
||||
POSTGRES_DB: litellm
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: litellm
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
ports:
|
||||
- "127.0.0.1:15432:5432"
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U litellm -d litellm"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 10
|
||||
|
||||
clickhouse:
|
||||
image: clickhouse/clickhouse-server:26.9.6.6
|
||||
environment:
|
||||
CLICKHOUSE_USER: default
|
||||
CLICKHOUSE_PASSWORD: local-tracing
|
||||
CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT: "1"
|
||||
volumes:
|
||||
- clickhouse_data:/var/lib/clickhouse
|
||||
ports:
|
||||
- "127.0.0.1:18123:8123"
|
||||
healthcheck:
|
||||
test: ["CMD", "clickhouse-client", "--user", "default", "--password", "local-tracing", "--query", "SELECT 1"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 20
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
clickhouse_data:
|
||||
10
docker/tracing-config.yaml
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
model_list:
|
||||
- model_name: gpt-6.1-sol
|
||||
litellm_params:
|
||||
model: openai/gpt-6.1-sol
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
tracing:
|
||||
store: clickhouse
|
||||
|
|
@ -22,10 +22,8 @@ from litellm._uuid import uuid
|
|||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_project_object
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership
|
||||
_set_object_metadata_field,
|
||||
)
|
||||
from litellm.proxy.management.teams.access import is_team_admin
|
||||
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
|
||||
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
management_endpoint_wrapper,
|
||||
|
|
@ -117,7 +115,7 @@ async def _check_user_permission_for_project(
|
|||
return False
|
||||
|
||||
team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump())
|
||||
return _is_user_team_admin(user_api_key_dict, team) or user_api_key_dict.user_id in (team.admins or [])
|
||||
return is_team_admin(user_api_key_dict, team) or user_api_key_dict.user_id in (team.admins or [])
|
||||
|
||||
|
||||
async def _validate_team_exists(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.71"
|
||||
version = "0.1.72"
|
||||
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.71"
|
||||
version = "0.1.72"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -73,6 +73,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/v1/containers",
|
||||
"/containers",
|
||||
"/v1/evals",
|
||||
"/v1/traces",
|
||||
"/v1/memory",
|
||||
"/queue/chat/",
|
||||
# Google data plane (v1beta is the Google AI Studio version)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,97 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "enabled" BOOLEAN NOT NULL DEFAULT true,
|
||||
ADD COLUMN IF NOT EXISTS "execution_mode" TEXT NOT NULL DEFAULT 'autonomous',
|
||||
ADD COLUMN IF NOT EXISTS "identity_managed" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "billing_agent_id" TEXT;
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_AgentIdentity" (
|
||||
"agent_id" TEXT NOT NULL,
|
||||
"active" BOOLEAN NOT NULL DEFAULT true,
|
||||
"provider" TEXT NOT NULL,
|
||||
"issuer" TEXT NOT NULL,
|
||||
"tenant_id" TEXT NOT NULL,
|
||||
"client_id" TEXT NOT NULL,
|
||||
"service_principal_id" TEXT,
|
||||
"required_roles" TEXT[] DEFAULT ARRAY[]::TEXT[],
|
||||
"required_scopes" TEXT[] DEFAULT ARRAY['user_impersonation']::TEXT[],
|
||||
"revision" TEXT NOT NULL,
|
||||
"last_authenticated_at" TIMESTAMP(3),
|
||||
|
||||
CONSTRAINT "LiteLLM_AgentIdentity_pkey" PRIMARY KEY ("agent_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgentIdentity" (
|
||||
"binding_id" TEXT NOT NULL,
|
||||
"agent_id" TEXT,
|
||||
"provider" TEXT NOT NULL,
|
||||
"issuer" TEXT NOT NULL,
|
||||
"tenant_id" TEXT NOT NULL,
|
||||
"client_id" TEXT NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_RetiredAgentIdentity_pkey" PRIMARY KEY ("binding_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgent" (
|
||||
"original_agent_id" TEXT NOT NULL,
|
||||
"retired_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_RetiredAgent_pkey" PRIMARY KEY ("original_agent_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_VerifiedSubject" (
|
||||
"subject_id" TEXT NOT NULL,
|
||||
"issuer" TEXT NOT NULL,
|
||||
"tenant_id" TEXT NOT NULL,
|
||||
"oid" TEXT NOT NULL,
|
||||
"kind" TEXT NOT NULL DEFAULT 'human',
|
||||
"user_id" TEXT,
|
||||
"verified_via" TEXT NOT NULL DEFAULT 'sso_interactive',
|
||||
"verified_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_VerifiedSubject_pkey" PRIMARY KEY ("subject_id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_AgentIdentity"("provider", "tenant_id", "client_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_issuer_service_principal_id_key" ON "LiteLLM_AgentIdentity"("issuer", "service_principal_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_RetiredAgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_RetiredAgentIdentity"("provider", "tenant_id", "client_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_user_id_idx" ON "LiteLLM_VerifiedSubject"("user_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_issuer_tenant_id_oid_key" ON "LiteLLM_VerifiedSubject"("issuer", "tenant_id", "oid");
|
||||
|
||||
-- AddForeignKey
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentIdentity_agent_id_fkey') THEN
|
||||
ALTER TABLE "LiteLLM_AgentIdentity" ADD CONSTRAINT "LiteLLM_AgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE CASCADE ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
-- AddForeignKey
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_RetiredAgentIdentity_agent_id_fkey') THEN
|
||||
ALTER TABLE "LiteLLM_RetiredAgentIdentity" ADD CONSTRAINT "LiteLLM_RetiredAgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE SET NULL ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
-- AddForeignKey
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerifiedSubject_user_id_fkey') THEN
|
||||
ALTER TABLE "LiteLLM_VerifiedSubject" ADD CONSTRAINT "LiteLLM_VerifiedSubject_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE CASCADE ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "pinned_tools" JSONB DEFAULT '{}';
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_DailyModelUsage" (
|
||||
"date" TEXT NOT NULL,
|
||||
"model_group" TEXT NOT NULL,
|
||||
"model" TEXT NOT NULL,
|
||||
"custom_llm_provider" TEXT NOT NULL,
|
||||
"task_type" TEXT NOT NULL,
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"prompt_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"completion_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"request_count" BIGINT NOT NULL DEFAULT 0,
|
||||
"successful_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"failed_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
CONSTRAINT "LiteLLM_DailyModelUsage_pkey" PRIMARY KEY ("date", "model_group", "model", "custom_llm_provider", "task_type")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyModelUsage_date_idx" ON "LiteLLM_DailyModelUsage"("date");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyModelUsage_model_group_idx" ON "LiteLLM_DailyModelUsage"("model_group");
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_Engine" (
|
||||
"id" TEXT NOT NULL PRIMARY KEY,
|
||||
"version" INTEGER NOT NULL DEFAULT 0,
|
||||
"data" JSONB NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_EngineWorker" (
|
||||
"id" TEXT NOT NULL PRIMARY KEY,
|
||||
"token_hash" TEXT NOT NULL UNIQUE,
|
||||
"data" JSONB NOT NULL
|
||||
);
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_EngineRun" (
|
||||
"id" TEXT NOT NULL PRIMARY KEY,
|
||||
"engine_id" TEXT NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL,
|
||||
"data" JSONB NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_EngineRun_engine_id_created_at_idx" ON "LiteLLM_EngineRun"("engine_id", "created_at");
|
||||
|
|
@ -78,6 +78,11 @@ model LiteLLM_AgentsTable {
|
|||
object_permission_id String?
|
||||
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
|
||||
spend Float @default(0.0)
|
||||
identity_managed Boolean @default(false)
|
||||
enabled Boolean @default(true)
|
||||
execution_mode String @default("autonomous")
|
||||
identity LiteLLM_AgentIdentity?
|
||||
retired_identities LiteLLM_RetiredAgentIdentity[]
|
||||
tpm_limit Int?
|
||||
rpm_limit Int?
|
||||
session_tpm_limit Int?
|
||||
|
|
@ -88,6 +93,56 @@ model LiteLLM_AgentsTable {
|
|||
updated_by String
|
||||
}
|
||||
|
||||
model LiteLLM_AgentIdentity {
|
||||
agent_id String @id
|
||||
active Boolean @default(true)
|
||||
agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade)
|
||||
provider String
|
||||
issuer String
|
||||
tenant_id String
|
||||
client_id String
|
||||
service_principal_id String?
|
||||
required_roles String[] @default([])
|
||||
required_scopes String[] @default(["user_impersonation"])
|
||||
revision String @default(uuid())
|
||||
last_authenticated_at DateTime?
|
||||
@@unique([provider, tenant_id, client_id])
|
||||
@@unique([issuer, service_principal_id])
|
||||
}
|
||||
|
||||
model LiteLLM_RetiredAgentIdentity {
|
||||
binding_id String @id @default(uuid())
|
||||
agent_id String?
|
||||
agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull)
|
||||
provider String
|
||||
issuer String
|
||||
tenant_id String
|
||||
client_id String
|
||||
@@unique([provider, tenant_id, client_id])
|
||||
}
|
||||
|
||||
model LiteLLM_RetiredAgent {
|
||||
original_agent_id String @id
|
||||
retired_at DateTime @default(now())
|
||||
}
|
||||
|
||||
model LiteLLM_VerifiedSubject {
|
||||
subject_id String @id @default(uuid())
|
||||
issuer String
|
||||
tenant_id String
|
||||
oid String
|
||||
kind String @default("human")
|
||||
user_id String?
|
||||
user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade)
|
||||
verified_via String @default("sso_interactive")
|
||||
verified_at DateTime @default(now())
|
||||
@@unique([issuer, tenant_id, oid])
|
||||
@@index([user_id])
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
model LiteLLM_OrganizationTable {
|
||||
organization_id String @id @default(uuid())
|
||||
organization_alias String
|
||||
|
|
@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable {
|
|||
|
||||
// Track spend, rate limit, budget Users
|
||||
model LiteLLM_UserTable {
|
||||
verified_subjects LiteLLM_VerifiedSubject[]
|
||||
user_id String @id
|
||||
user_alias String?
|
||||
team_id String?
|
||||
|
|
@ -322,6 +378,7 @@ model LiteLLM_MCPServerTable {
|
|||
allowed_tools String[] @default([])
|
||||
tool_name_to_display_name Json? @default("{}")
|
||||
tool_name_to_description Json? @default("{}")
|
||||
pinned_tools Json? @default("{}")
|
||||
extra_headers String[] @default([])
|
||||
static_headers Json? @default("{}")
|
||||
// Admin-configured environment variables interpolated into static_headers
|
||||
|
|
@ -674,6 +731,7 @@ model LiteLLM_SpendLogs {
|
|||
session_id String?
|
||||
status String?
|
||||
mcp_namespaced_tool_name String?
|
||||
billing_agent_id String?
|
||||
agent_id String?
|
||||
proxy_server_request Json? @default("{}")
|
||||
litellm_call_id String?
|
||||
|
|
@ -1259,6 +1317,26 @@ model LiteLLM_DailyToolSpend {
|
|||
@@id([date, tool_name])
|
||||
}
|
||||
|
||||
model LiteLLM_DailyModelUsage {
|
||||
date String
|
||||
model_group String
|
||||
model String
|
||||
custom_llm_provider String
|
||||
task_type String
|
||||
spend Float @default(0.0)
|
||||
prompt_tokens BigInt @default(0)
|
||||
completion_tokens BigInt @default(0)
|
||||
request_count BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@id([date, model_group, model, custom_llm_provider, task_type])
|
||||
@@index([date])
|
||||
@@index([model_group])
|
||||
}
|
||||
|
||||
// Gateway request counts recorded at the ASGI edge by
|
||||
// BillableRequestMetricsMiddleware. This is the source of truth for SGR
|
||||
// (successful gateway requests): it counts what the proxy actually answered,
|
||||
|
|
@ -1816,3 +1894,24 @@ model LiteLLM_WorkflowMessage {
|
|||
@@unique([run_id, sequence_number])
|
||||
@@index([run_id])
|
||||
}
|
||||
|
||||
model LiteLLM_Engine {
|
||||
id String @id
|
||||
version Int @default(0)
|
||||
data Json
|
||||
}
|
||||
|
||||
model LiteLLM_EngineRun {
|
||||
id String @id
|
||||
engine_id String
|
||||
created_at DateTime
|
||||
data Json
|
||||
|
||||
@@index([engine_id, created_at])
|
||||
}
|
||||
|
||||
model LiteLLM_EngineWorker {
|
||||
id String @id
|
||||
token_hash String @unique
|
||||
data Json
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.102"
|
||||
version = "0.4.103"
|
||||
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.102"
|
||||
version = "0.4.103"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
149
litellm-rust/Cargo.lock
generated
|
|
@ -1274,6 +1274,18 @@ version = "0.4.33"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e8ccc4ea9f6acc32d102c0f6d471d11d913ad15f20c04de743374861fa1d414"
|
||||
|
||||
[[package]]
|
||||
name = "const-hex"
|
||||
version = "1.19.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0e59eef12462b0f9b0a3620219be5d639afd79fe39dff0a42c3997061f9298b4"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.2.17",
|
||||
"proptest",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.9.6"
|
||||
|
|
@ -2372,9 +2384,9 @@ dependencies = [
|
|||
"http-body-util",
|
||||
"hyper 1.10.1",
|
||||
"lazy_static",
|
||||
"opentelemetry",
|
||||
"opentelemetry 0.32.0",
|
||||
"opentelemetry-semantic-conventions",
|
||||
"opentelemetry_sdk",
|
||||
"opentelemetry_sdk 0.32.1",
|
||||
"percent-encoding",
|
||||
"pin-project",
|
||||
"prost",
|
||||
|
|
@ -3552,8 +3564,10 @@ name = "litellm-cache-response"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-gcs",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-http",
|
||||
"py_literal",
|
||||
"redis",
|
||||
"redis-test",
|
||||
|
|
@ -3562,6 +3576,7 @@ dependencies = [
|
|||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3645,13 +3660,18 @@ dependencies = [
|
|||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-cache",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-framing",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-secrets",
|
||||
"litellm-tracing",
|
||||
"litellm-types",
|
||||
"mime_guess",
|
||||
"moka",
|
||||
"rand 0.8.7",
|
||||
|
|
@ -3667,6 +3687,7 @@ dependencies = [
|
|||
"time",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tokio-util",
|
||||
"tracing",
|
||||
"url",
|
||||
"veil",
|
||||
|
|
@ -3678,13 +3699,12 @@ name = "litellm-core-utils"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"fancy-regex 0.19.2",
|
||||
"litellm-llms-types",
|
||||
"litellm-tracing",
|
||||
"litellm-types",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
"serde_with",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"url",
|
||||
|
|
@ -3806,15 +3826,19 @@ dependencies = [
|
|||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-host",
|
||||
"litellm-host-http",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-router",
|
||||
"litellm-secrets",
|
||||
"litellm-types",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
|
|
@ -3907,12 +3931,24 @@ dependencies = [
|
|||
"futures-util",
|
||||
"http 1.4.2",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-host-native"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-host",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-host-python"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3973,9 +4009,9 @@ dependencies = [
|
|||
"litellm-framing",
|
||||
"litellm-host",
|
||||
"litellm-http",
|
||||
"litellm-llms-types",
|
||||
"litellm-python-compat",
|
||||
"litellm-secrets",
|
||||
"litellm-types",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -3989,13 +4025,26 @@ dependencies = [
|
|||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-llms-types"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"macro_rules_attribute",
|
||||
"rstest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"strum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-model-catalog"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"indexmap 2.14.0",
|
||||
"jsonschema",
|
||||
"litellm-types",
|
||||
"litellm-llms-types",
|
||||
"rstest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
|
|
@ -4033,12 +4082,13 @@ dependencies = [
|
|||
"litellm-host-python",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-secrets",
|
||||
"litellm-secrets-aws",
|
||||
"litellm-secrets-types",
|
||||
"litellm-token-counter",
|
||||
"litellm-traces",
|
||||
"litellm-tracing",
|
||||
"litellm-types",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"qdrant-client",
|
||||
|
|
@ -4314,6 +4364,26 @@ dependencies = [
|
|||
"tiktoken-rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-traces"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"flate2",
|
||||
"litellm-http",
|
||||
"opentelemetry-proto",
|
||||
"prost",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"testcontainers-modules",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-tracing"
|
||||
version = "0.1.0"
|
||||
|
|
@ -4328,17 +4398,6 @@ dependencies = [
|
|||
"tracing-subscriber",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-types"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"rstest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"strum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litemap"
|
||||
version = "0.8.2"
|
||||
|
|
@ -4734,6 +4793,33 @@ dependencies = [
|
|||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry"
|
||||
version = "0.33.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6cdb0b1b267eb9db3331b434ed9ddab10d50e280a9adf9d13e5233e2002b61b5"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"js-sys",
|
||||
"pin-project-lite",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry-proto"
|
||||
version = "0.33.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "25da1ac11a0aeccf38d7f77ee0348715adaf8340f65ad46c94a02c6b20e2f65d"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"const-hex",
|
||||
"opentelemetry 0.33.0",
|
||||
"opentelemetry_sdk 0.33.0",
|
||||
"prost",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry-semantic-conventions"
|
||||
version = "0.32.1"
|
||||
|
|
@ -4749,7 +4835,23 @@ dependencies = [
|
|||
"futures-channel",
|
||||
"futures-executor",
|
||||
"futures-util",
|
||||
"opentelemetry",
|
||||
"opentelemetry 0.32.0",
|
||||
"percent-encoding",
|
||||
"portable-atomic",
|
||||
"rand 0.9.5",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry_sdk"
|
||||
version = "0.33.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cb39533d9d1c912123efd7d41d7e0c29d16917b60ce15b4c8d87cb1af7f67520"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-executor",
|
||||
"futures-util",
|
||||
"opentelemetry 0.33.0",
|
||||
"percent-encoding",
|
||||
"portable-atomic",
|
||||
"rand 0.9.5",
|
||||
|
|
@ -5678,6 +5780,7 @@ checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029"
|
|||
dependencies = [
|
||||
"base64 0.23.1",
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2 0.4.15",
|
||||
|
|
@ -5689,6 +5792,7 @@ dependencies = [
|
|||
"hyper-util",
|
||||
"js-sys",
|
||||
"log",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"quinn",
|
||||
|
|
@ -6919,6 +7023,7 @@ dependencies = [
|
|||
"memchr",
|
||||
"parse-display",
|
||||
"pin-project-lite",
|
||||
"reqwest 0.13.5",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
|
|
@ -7479,7 +7584,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26"
|
||||
dependencies = [
|
||||
"js-sys",
|
||||
"opentelemetry",
|
||||
"opentelemetry 0.32.0",
|
||||
"tracing",
|
||||
"tracing-core",
|
||||
"tracing-subscriber",
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ repository = "https://github.com/BerriAI/litellm"
|
|||
litellm-config = { path = "crates/config" }
|
||||
litellm-router = { path = "crates/router" }
|
||||
litellm-tracing = { path = "crates/tracing" }
|
||||
litellm-traces = { path = "crates/traces" }
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
|
||||
litellm-gateway = { path = "crates/gateway" }
|
||||
|
|
@ -22,6 +23,7 @@ litellm-gateway-ui = { path = "crates/gateway-ui" }
|
|||
litellm-coroutine = { path = "crates/coroutine" }
|
||||
litellm-host = { path = "crates/host" }
|
||||
litellm-host-http = { path = "crates/host-http" }
|
||||
litellm-host-native = { path = "crates/host-native" }
|
||||
litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
|
||||
litellm-framing = { path = "crates/framer" }
|
||||
litellm-auth = { path = "crates/auth" }
|
||||
|
|
@ -38,7 +40,7 @@ litellm-secrets-azure = { path = "crates/secrets-azure" }
|
|||
litellm-secrets-cyberark = { path = "crates/secrets-cyberark" }
|
||||
litellm-http = { path = "crates/http" }
|
||||
litellm-llms = { path = "crates/llms" }
|
||||
litellm-types = { path = "crates/types" }
|
||||
litellm-llms-types = { path = "crates/llms-types" }
|
||||
litellm-core-utils = { path = "crates/core-utils" }
|
||||
litellm-db = { path = "crates/db" }
|
||||
litellm-db-testing = { path = "crates/db-testing" }
|
||||
|
|
@ -73,6 +75,7 @@ proptest = "1.7.0"
|
|||
pyo3 = "0.29.2"
|
||||
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
||||
rand = "0.8"
|
||||
macro_rules_attribute = "0.2.3"
|
||||
schemars = "1"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
|
||||
qdrant-client = { version = "1.19.0", default-features = false }
|
||||
|
|
|
|||
|
|
@ -39,7 +39,33 @@ impl TokenProviderHandle {
|
|||
Self(caller)
|
||||
}
|
||||
|
||||
pub fn from_callback<F, Fut>(acquire: F) -> Self
|
||||
where
|
||||
F: Fn() -> Fut + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<ResolvedCredential, Error>> + Send + 'static,
|
||||
{
|
||||
Self::new(Arc::new(CallbackTokenProvider(acquire)))
|
||||
}
|
||||
|
||||
pub async fn acquire(&self) -> Result<ResolvedCredential, Error> {
|
||||
self.0.acquire().await
|
||||
}
|
||||
}
|
||||
|
||||
struct CallbackTokenProvider<F>(F);
|
||||
|
||||
impl<F> std::fmt::Debug for CallbackTokenProvider<F> {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("CallbackTokenProvider")
|
||||
}
|
||||
}
|
||||
|
||||
impl<F, Fut> TokenProvider for CallbackTokenProvider<F>
|
||||
where
|
||||
F: Fn() -> Fut + Send + Sync,
|
||||
Fut: Future<Output = Result<ResolvedCredential, Error>> + Send + 'static,
|
||||
{
|
||||
fn acquire(&self) -> TokenFuture<'_> {
|
||||
Box::pin((self.0)())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
104
litellm-rust/crates/auth-types/tests/token.rs
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
use std::{
|
||||
error::Error as StdError,
|
||||
future::{Future, poll_fn},
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
},
|
||||
task::Poll,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use litellm_auth_types::{
|
||||
Error, ErrorDetail, ResolvedCredential, SecretValue, TokenProviderHandle,
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
fn credential(index: usize, access_token: bool) -> ResolvedCredential {
|
||||
let token = SecretValue::new(format!("credential-{index}"));
|
||||
if access_token {
|
||||
return ResolvedCredential::AccessToken {
|
||||
token,
|
||||
expires_on: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(index as u64)),
|
||||
};
|
||||
}
|
||||
ResolvedCredential::Static(token)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::static_secret(false)]
|
||||
#[case::access_token(true)]
|
||||
#[tokio::test]
|
||||
async fn callbacks_acquire_fresh_credentials_on_demand(#[case] access_token: bool) {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let callback_calls = calls.clone();
|
||||
let provider = TokenProviderHandle::from_callback(move || {
|
||||
let index = callback_calls.fetch_add(1, Ordering::SeqCst);
|
||||
async move {
|
||||
tokio::task::yield_now().await;
|
||||
Ok(credential(index, access_token))
|
||||
}
|
||||
});
|
||||
let cloned = provider.clone();
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(
|
||||
provider.acquire().await.unwrap(),
|
||||
credential(0, access_token)
|
||||
);
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(cloned.acquire().await.unwrap(), credential(1, access_token));
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn callback_errors_preserve_the_original_source() {
|
||||
let provider = TokenProviderHandle::from_callback(|| async {
|
||||
Err(Error::CredentialAcquisition(ErrorDetail::failed(
|
||||
"caller credential",
|
||||
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
|
||||
)))
|
||||
});
|
||||
|
||||
let error = provider.acquire().await.unwrap_err();
|
||||
assert!(matches!(error, Error::CredentialAcquisition(_)));
|
||||
let source = std::iter::successors(Some(&error as &(dyn StdError + 'static)), |error| {
|
||||
(*error).source()
|
||||
})
|
||||
.find_map(|error| error.downcast_ref::<std::io::Error>())
|
||||
.unwrap();
|
||||
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
|
||||
}
|
||||
|
||||
struct Release(Arc<AtomicBool>);
|
||||
|
||||
impl Drop for Release {
|
||||
fn drop(&mut self) {
|
||||
self.0.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn cancelling_acquisition_drops_the_callback_future() {
|
||||
let released = Arc::new(AtomicBool::new(false));
|
||||
let callback_released = released.clone();
|
||||
let provider = TokenProviderHandle::from_callback(move || {
|
||||
let released = callback_released.clone();
|
||||
async move {
|
||||
let _release = Release(released);
|
||||
std::future::pending().await
|
||||
}
|
||||
});
|
||||
|
||||
let mut acquisition = Box::pin(provider.acquire());
|
||||
poll_fn(|context| {
|
||||
assert!(acquisition.as_mut().poll(context).is_pending());
|
||||
assert!(!released.load(Ordering::SeqCst));
|
||||
Poll::Ready(())
|
||||
})
|
||||
.await;
|
||||
drop(acquisition);
|
||||
assert!(released.load(Ordering::SeqCst));
|
||||
}
|
||||
|
|
@ -56,6 +56,7 @@ async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer
|
|||
)]
|
||||
#[case::missing("missing", ResponseTemplate::new(404), Ok(None))]
|
||||
#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))]
|
||||
#[case::unauthorized("unauthorized", ResponseTemplate::new(401), Err(Error::Unavailable))]
|
||||
#[case::invalid(
|
||||
"invalid",
|
||||
ResponseTemplate::new(200).set_body_string("not json"),
|
||||
|
|
|
|||
29
litellm-rust/crates/cache-response/AGENTS.md
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
# Response caching
|
||||
|
||||
Design this crate for shared Rust execution used by the Python SDK and the Rust gateway. The Python SDK will remain, with more core execution moving to Rust and Python callbacks staying in Python. The Rust gateway is still evolving and is intended to replace the Python proxy. Keep response-cache policy independent of Python, HTTP serving, and either proxy's configuration format
|
||||
|
||||
Separate what is cached, how a hit is matched, and where entries are stored. Chat Completions, Messages, Responses, and embeddings are API workloads. Exact and semantic matching are lookup behaviors. Memory, Redis, disk, and object stores are storage choices. Embeddings are inference too, so do not use an inference-cache name to imply a category that excludes embeddings. Consult the existing Python cache and caching handler for behavior and compatibility contracts without copying their class structure
|
||||
|
||||
Storage traits, codecs, and backend capabilities belong in `litellm-cache` and the storage crates. Keep storage reusable for value types beyond LLM responses. This crate owns response entries, matching and freshness semantics, the Python-compatible response codec, and deferred-write policy. Core owns route-specific request identity, response encoding and reconstruction, embedding partial-hit orchestration, and stream capture and replay. Boundaries own configuration translation, resource construction, and caller identity
|
||||
|
||||
Construct and inject the response-cache service at the Python bridge or gateway boundary, as with the HTTP client. Reuse it across calls. Core and provider code must not discover cache configuration through Python globals, process configuration, or backend-specific factories
|
||||
|
||||
Keep `ResponseCache<B>` generic over its storage backend. Preserve typed backend contexts and capability bounds internally. Inject an object-safe service into core for runtime backend selection, so storage types do not spread through route and host types. Keep API request and response types statically typed. Add a generic parameter only where it preserves a useful type relationship or capability
|
||||
|
||||
Keep the core service contract narrow. Lookup and store must not require connection testing, ping, flush, deletion, counters, queues, or scripts. Require batch operations where a consumer needs partial hits, and keep management capabilities on their own interfaces. An exact-only adapter must remain explicit about its matching restriction. Supporting semantic matching requires a defined lookup-context and embedding execution contract, not just a renamed trait
|
||||
|
||||
Separate reusable resources from per-call policy. Backend configuration, namespace, default expiry, and entry limits belong to the configured service or backend. Read/write controls, expiry and freshness overrides, and authenticated caller scope belong to the call. Passing call options must not replace or mutate the route's configured service
|
||||
|
||||
Keep cache misses and storage failures distinguishable in return values. Core owns the decision to continue with provider execution after a cache failure. A read can reject an entry for freshness while the backend still retains it. Preserve the timestamp at which a response was produced when writing it later
|
||||
|
||||
Define lookup placement explicitly relative to authorization, deployment and credential resolution, and request-transforming callbacks. Cache identity must account for every input that affects reuse, including API surface and caller scope, while preserving intentional Python caching groups. Preserve existing keys and response formats unless changing them is an explicit migration decision
|
||||
|
||||
Cache normalized provider results before caller-specific response transformations. Hits must still run the applicable response processing, success callbacks, and cache-hit accounting. Keep callback execution in the host. Python cache implementations and semantic embedders that require the caller's task must use the existing host-operation mechanism rather than Python calls from a Rust worker. Preserve legacy fallback until that contract is supported
|
||||
|
||||
Keep unary caching independent of stream-only methods. Store streams only after successful exhaustion and protocol completion. Errors, incomplete streams, cancellation, and oversized entries must not populate the cache. Embedding batches need ordered partial results and reconstruction around the uncached inputs
|
||||
|
||||
Test each contract in its owner: storage capabilities in backend tests, envelopes and freshness here, reuse and replay in core, Python callback and fallback behavior at the bridge, and HTTP behavior at the gateway. Run backend contract checks and Python response-codec fixtures before exposing a new backend
|
||||
|
||||
`ScopedCache` requires an explicit shared or isolated scope at construction. Per-call `CachePolicy` controls reads, writes, expiry, and freshness without replacing the attached scope or service. `CacheOptions` binds that policy to an explicit scope for storage requests and has no default sharing policy. Versioned native envelopes reject incompatible API surfaces and versions as misses; this envelope is distinct from the legacy Python response codec
|
||||
|
||||
Response storage is not the source of budget or rate-limit coordination dependencies. Keep counters, reservations, and atomic admission operations out of `ResponseCacheService`, including when both services happen to use Redis
|
||||
|
|
@ -13,9 +13,12 @@ serde_json.workspace = true
|
|||
sha2.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-gcs.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = "1.7.0"
|
||||
redis-test = "1.0.4"
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock = "0.6.5"
|
||||
|
|
|
|||
|
|
@ -1,51 +0,0 @@
|
|||
# Response cache
|
||||
|
||||
`ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
|
||||
|
||||
## Ownership
|
||||
|
||||
`litellm-cache` defines typed storage, codec, and capability traits. `BaseCache` is only get, set, TTL, and pipeline writes. Everything else is an optional capability a backend implements only where its Python class defines the method: `DisconnectCache`, `ConnectionCache` (`test_connection`), `PingCache`, `BatchCache`, `DeleteCache`, `FlushCache`, counters, queues, TTL, scan, and scripts. Memory, Redis, disk, S3, GCS, and Azure Blob implement those traits without depending on response policy, so other consumers can store their own value types in the same backends
|
||||
|
||||
Semantic backends (Redis, Valkey, Qdrant) are generic over their embedder and codec, and share one prompt and embedding contract from `litellm_cache::semantic`. They take a `SemanticCacheContext`, so `ResponseCache` drives them the same way it drives exact backends
|
||||
|
||||
`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python
|
||||
|
||||
`ExactResponseCache` is the object-safe view of a `ResponseCache` over an exact backend. `ConnectionProbe` is the object-safe `test_connection`, implemented only when the backend implements `ConnectionCache`, so a host holds one next to its `ExactResponseCache` and reports the operation as unsupported otherwise, as Python's `BaseCache` does. Lookup, store, batch, and flush never require it
|
||||
|
||||
## Native Rust use
|
||||
|
||||
```rust
|
||||
use std::{sync::Arc, time::Duration};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest};
|
||||
use serde_json::json;
|
||||
|
||||
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
|
||||
let request = ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("example:key".into()),
|
||||
..Default::default()
|
||||
});
|
||||
let now = Duration::from_secs(100);
|
||||
cache.store(&request, json!({"answer": 7}), now)?;
|
||||
assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7})));
|
||||
```
|
||||
|
||||
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved
|
||||
|
||||
Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it
|
||||
|
||||
## Python integration
|
||||
|
||||
The bridge activates backends through the Rust catalog in `litellm/rust_bridge/catalog.py`. Every cache rule ships as `PYTHON_ONLY`, so SDK, Router, and proxy calls stay on Python and construct no native cache resources until a rule is changed
|
||||
|
||||
When a rule selects a backend, the Python `Cache` facade builds the native runtime from its own configuration and routes its storage calls (sync and async lookup and store, and pipelined batch store) to it. Stream replay, embedding partial-hit merging, response reconstruction, and callbacks stay in Python on top of that native store. The Python backend object remains for its direct API
|
||||
|
||||
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
|
||||
|
||||
Native cache handles must be recreated after fork. Native errors propagate to the host, which owns the existing fail-open and logging policy
|
||||
|
||||
## Adding another backend
|
||||
|
||||
Implement `BaseCache` for the backend with its associated value type and the capability traits its Python class supports, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation
|
||||
|
||||
Run the `litellm-cache-testing` contract checks the backend's capabilities allow, and run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before adding a catalog rule
|
||||
|
|
@ -4,6 +4,7 @@ mod codec;
|
|||
mod embedding;
|
||||
mod exact;
|
||||
mod response;
|
||||
mod service;
|
||||
|
||||
pub use buffer::WriteBuffer;
|
||||
pub use caching::{
|
||||
|
|
@ -14,3 +15,8 @@ pub use codec::ResponseCacheCodec;
|
|||
pub use embedding::PartialHits;
|
||||
pub use exact::{ConnectionProbe, ExactResponseCache};
|
||||
pub use response::{ResponseCache, ResponseCacheRequest};
|
||||
|
||||
pub use service::{
|
||||
CacheOptions, CachePolicy, CacheScope, ResponseCacheConfig, ResponseCacheService,
|
||||
ResponseEnvelope, ScopedCache,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -7,7 +7,9 @@ use litellm_cache::{
|
|||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
|
||||
use crate::{
|
||||
CacheControls, CacheEntry, CacheKeyInput, PartialHits, ResponseCacheConfig, cache_key,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
|
||||
|
|
@ -50,6 +52,7 @@ where
|
|||
B::Context: Default + PartialEq,
|
||||
{
|
||||
backend: Arc<B>,
|
||||
config: ResponseCacheConfig,
|
||||
}
|
||||
|
||||
impl<B> ResponseCache<B>
|
||||
|
|
@ -58,7 +61,18 @@ where
|
|||
B::Context: Default + PartialEq,
|
||||
{
|
||||
pub fn new(backend: Arc<B>) -> Self {
|
||||
Self { backend }
|
||||
Self {
|
||||
backend,
|
||||
config: ResponseCacheConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_config(self, config: ResponseCacheConfig) -> Self {
|
||||
Self { config, ..self }
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &ResponseCacheConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn backend(&self) -> &B {
|
||||
|
|
@ -221,7 +235,7 @@ where
|
|||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
if !request.controls.writes() {
|
||||
if !request.controls.writes() || !self.fits(&response) {
|
||||
return Ok(());
|
||||
}
|
||||
self.backend.set_cache(
|
||||
|
|
@ -240,7 +254,7 @@ where
|
|||
response: Value,
|
||||
now: Duration,
|
||||
) -> Result<(), Error> {
|
||||
if !request.controls.writes() {
|
||||
if !request.controls.writes() || !self.fits(&response) {
|
||||
return Ok(());
|
||||
}
|
||||
self.backend
|
||||
|
|
@ -277,7 +291,7 @@ where
|
|||
) -> Result<(), Error> {
|
||||
let writable = entries
|
||||
.into_iter()
|
||||
.filter(|(request, _, _)| request.controls.writes())
|
||||
.filter(|(request, response, _)| request.controls.writes() && self.fits(response))
|
||||
.map(|(request, response, now)| {
|
||||
(
|
||||
cache_key(&request.key),
|
||||
|
|
@ -312,6 +326,11 @@ where
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn fits(&self, response: &Value) -> bool {
|
||||
self.config.max_entry_bytes == usize::MAX
|
||||
|| response.to_string().len() <= self.config.max_entry_bytes
|
||||
}
|
||||
|
||||
fn partial_hits(
|
||||
requests: &[ResponseCacheRequest<B::Context>],
|
||||
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
|
||||
|
|
|
|||
185
litellm-rust/crates/cache-response/src/service.rs
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
use std::{future::Future, pin::Pin, time::Duration};
|
||||
|
||||
use litellm_cache::{BaseCache, Error, ExactCacheContext};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
|
||||
};
|
||||
|
||||
type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponseCacheConfig {
|
||||
pub namespace: String,
|
||||
pub max_entry_bytes: usize,
|
||||
}
|
||||
|
||||
impl Default for ResponseCacheConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
namespace: String::new(),
|
||||
max_entry_bytes: usize::MAX,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ResponseCacheService: Send + Sync {
|
||||
fn config(&self) -> &ResponseCacheConfig;
|
||||
|
||||
fn lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, Option<Value>>;
|
||||
|
||||
fn store<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, ()>;
|
||||
}
|
||||
|
||||
impl<B> ResponseCacheService for ResponseCache<B>
|
||||
where
|
||||
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>,
|
||||
{
|
||||
fn config(&self) -> &ResponseCacheConfig {
|
||||
self.config()
|
||||
}
|
||||
|
||||
fn lookup<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, Option<Value>> {
|
||||
Box::pin(self.async_lookup(request, now))
|
||||
}
|
||||
|
||||
fn store<'a>(
|
||||
&'a self,
|
||||
request: &'a ResponseCacheRequest,
|
||||
response: Value,
|
||||
now: Duration,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
Box::pin(self.async_store(request, response, now))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum CacheScope {
|
||||
Shared,
|
||||
Isolated(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Default)]
|
||||
pub struct CachePolicy {
|
||||
pub caching: Option<bool>,
|
||||
pub no_cache: bool,
|
||||
pub no_store: bool,
|
||||
pub ttl: Option<Duration>,
|
||||
pub max_age: Option<Duration>,
|
||||
}
|
||||
|
||||
impl CachePolicy {
|
||||
pub fn enabled(&self) -> bool {
|
||||
self.caching != Some(false) && !(self.no_cache && self.no_store)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CacheOptions {
|
||||
pub policy: CachePolicy,
|
||||
pub scope: CacheScope,
|
||||
}
|
||||
|
||||
impl CacheOptions {
|
||||
pub fn new(scope: CacheScope) -> Self {
|
||||
Self {
|
||||
policy: CachePolicy::default(),
|
||||
scope,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn request(self, namespace: &str, surface: &str, mut input: Value) -> ResponseCacheRequest {
|
||||
input.sort_all_objects();
|
||||
let scope = match self.scope {
|
||||
CacheScope::Shared => String::new(),
|
||||
CacheScope::Isolated(scope) => serde_json::json!(["isolated", scope]).to_string(),
|
||||
};
|
||||
ResponseCacheRequest {
|
||||
key: CacheKeyInput {
|
||||
namespace: Some(format!("{namespace}:inference-v2")),
|
||||
fields: [
|
||||
("surface", surface.to_owned()),
|
||||
("scope", scope),
|
||||
("request", input.to_string()),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(name, value)| CacheKeyField {
|
||||
name: name.into(),
|
||||
value: Some(value),
|
||||
api_parameter: true,
|
||||
internal_parameter: false,
|
||||
})
|
||||
.collect(),
|
||||
..Default::default()
|
||||
},
|
||||
controls: CacheControls {
|
||||
configured: true,
|
||||
supported_call_type: true,
|
||||
native_backend: true,
|
||||
default_on: true,
|
||||
caching: self.policy.caching,
|
||||
no_cache: self.policy.no_cache,
|
||||
no_store: self.policy.no_store,
|
||||
..Default::default()
|
||||
},
|
||||
context: ExactCacheContext {
|
||||
ttl: self.policy.ttl,
|
||||
},
|
||||
max_age: self.policy.max_age,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
pub struct ResponseEnvelope<T> {
|
||||
version: u32,
|
||||
surface: String,
|
||||
output: T,
|
||||
}
|
||||
|
||||
impl<T> ResponseEnvelope<T> {
|
||||
pub fn new(surface: &str, output: T) -> Self {
|
||||
Self {
|
||||
version: 1,
|
||||
surface: surface.into(),
|
||||
output,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decode(self, surface: &str) -> Option<T> {
|
||||
(self.version == 1 && self.surface == surface).then_some(self.output)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ScopedCache {
|
||||
pub service: std::sync::Arc<dyn ResponseCacheService>,
|
||||
pub scope: CacheScope,
|
||||
}
|
||||
|
||||
impl ScopedCache {
|
||||
pub fn new(service: std::sync::Arc<dyn ResponseCacheService>, scope: CacheScope) -> Self {
|
||||
Self { service, scope }
|
||||
}
|
||||
|
||||
pub fn options(&self, policy: Option<CachePolicy>) -> CacheOptions {
|
||||
CacheOptions {
|
||||
policy: policy.unwrap_or_default(),
|
||||
scope: self.scope.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -18,7 +18,7 @@ use litellm_cache_response::{
|
|||
WriteBuffer, cache_key,
|
||||
};
|
||||
use redis_test::MockCmd;
|
||||
use rstest::rstest;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use support::{keyed, memory, redis, request};
|
||||
|
||||
|
|
@ -648,3 +648,129 @@ async fn write_buffer_clear_drops_pending_entries(memory: Memory, request: Respo
|
|||
assert_eq!(memory.lookup(&request, now).unwrap(), None);
|
||||
assert_eq!(memory.lookup(&other, now).unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::python_sync("{'timestamp': 100.0, 'response': '{\"answer\": 7}'}")]
|
||||
#[case::python_async(r#"{"timestamp":100.0,"response":{"answer":7}}"#)]
|
||||
#[case::bare_response(r#"{"answer":7}"#)]
|
||||
#[tokio::test]
|
||||
async fn gcs_reads_python_entries_and_writes_python_compatible_envelopes(
|
||||
#[case] encoded: &str,
|
||||
#[values(false, true)] asynchronous: bool,
|
||||
#[future(awt)] gcs: (wiremock::MockServer, Gcs),
|
||||
) {
|
||||
use wiremock::{
|
||||
Mock, ResponseTemplate,
|
||||
matchers::{body_json, header, method, path, query_param},
|
||||
};
|
||||
|
||||
let (server, cache) = gcs;
|
||||
let response = json!({"answer": 7});
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/cache%2Fpython"))
|
||||
.and(query_param("alt", "media"))
|
||||
.and(header("authorization", "Bearer token"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string(encoded))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/upload/storage/v1/b/bucket/o"))
|
||||
.and(query_param("uploadType", "media"))
|
||||
.and(query_param("name", "cache/native"))
|
||||
.and(header("authorization", "Bearer token"))
|
||||
.and(header("content-type", "application/json"))
|
||||
.and(body_json(json!({"timestamp": 102.0, "response": response})))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let lookup = if asynchronous {
|
||||
cache
|
||||
.async_lookup(&keyed("python"), Duration::from_secs(102))
|
||||
.await
|
||||
} else {
|
||||
cache.lookup(&keyed("python"), Duration::from_secs(102))
|
||||
};
|
||||
assert_eq!(lookup.unwrap(), Some(response.clone()));
|
||||
let request = ResponseCacheRequest {
|
||||
context: litellm_cache::ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(12)),
|
||||
},
|
||||
..keyed("native")
|
||||
};
|
||||
let stored = if asynchronous {
|
||||
cache
|
||||
.async_store(&request, response, Duration::from_secs(102))
|
||||
.await
|
||||
} else {
|
||||
cache.store(&request, response, Duration::from_secs(102))
|
||||
};
|
||||
assert_eq!(stored, Ok(()));
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let upload = requests
|
||||
.iter()
|
||||
.find(|request| request.method.as_str() == "POST")
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
upload.url.query(),
|
||||
Some("uploadType=media&name=cache%2Fnative")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn gcs_batch_reads_preserve_order_and_treat_invalid_entries_as_misses(
|
||||
#[future(awt)] gcs: (wiremock::MockServer, Gcs),
|
||||
) {
|
||||
use wiremock::{
|
||||
Mock, ResponseTemplate,
|
||||
matchers::{method, path},
|
||||
};
|
||||
|
||||
let (server, cache) = gcs;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/cache%2Fhit"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(json!({"timestamp": 100.0, "response": {"answer":7}})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/cache%2Finvalid"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string("not an entry"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/storage/v1/b/bucket/o/cache%2Fmissing"))
|
||||
.respond_with(ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let requests = [keyed("hit"), keyed("missing"), keyed("invalid")];
|
||||
let partial = cache
|
||||
.async_lookup_batch(&requests, Duration::from_secs(102))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(partial.values, vec![Some(json!({"answer":7})), None, None]);
|
||||
assert_eq!(partial.missing_indices, vec![1, 2]);
|
||||
}
|
||||
|
||||
type Gcs = ResponseCache<litellm_cache_gcs::GcsCache<litellm_cache_response::ResponseCacheCodec>>;
|
||||
|
||||
#[fixture]
|
||||
async fn gcs() -> (wiremock::MockServer, Gcs) {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
let cache = ResponseCache::new(Arc::new(litellm_cache_gcs::GcsCache::with_token_source(
|
||||
litellm_cache_gcs::GcsConfig {
|
||||
bucket_name: "bucket".into(),
|
||||
gcs_path: Some("cache".into()),
|
||||
path_service_account: None,
|
||||
endpoint: server.uri(),
|
||||
},
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_cache_response::ResponseCacheCodec,
|
||||
Arc::new(litellm_cache_gcs::StaticTokenSource("token".into())),
|
||||
)));
|
||||
(server, cache)
|
||||
}
|
||||
|
|
|
|||
185
litellm-rust/crates/cache-response/tests/service.rs
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::ExactCacheContext;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheConfig, ResponseCacheRequest,
|
||||
ResponseCacheService,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn service_honors_per_call_expiry_and_freshness() {
|
||||
let clock = Arc::new(AtomicU64::new(0));
|
||||
let cache_clock = clock.clone();
|
||||
let cache: Arc<dyn ResponseCacheService> = Arc::new(ResponseCache::new(Arc::new(
|
||||
InMemoryCache::with_clock(Some(100), Some(Duration::from_secs(60)), move || {
|
||||
Duration::from_secs(cache_clock.load(Ordering::SeqCst))
|
||||
}),
|
||||
)));
|
||||
let request = ResponseCacheRequest {
|
||||
context: ExactCacheContext {
|
||||
ttl: Some(Duration::from_secs(5)),
|
||||
},
|
||||
..ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some("entry".into()),
|
||||
..Default::default()
|
||||
})
|
||||
};
|
||||
cache
|
||||
.store(&request, json!({"answer":7}), Duration::ZERO)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::ZERO).await.unwrap(),
|
||||
Some(json!({"answer":7}))
|
||||
);
|
||||
let stale_request = ResponseCacheRequest {
|
||||
max_age: Some(Duration::from_secs(1)),
|
||||
..request.clone()
|
||||
};
|
||||
clock.store(2, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&stale_request, Duration::from_secs(2))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.lookup(&request, Duration::from_secs(2))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
clock.store(6, Ordering::SeqCst);
|
||||
assert_eq!(
|
||||
cache
|
||||
.lookup(&request, Duration::from_secs(6))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn entry_limit_applies_to_sync_async_and_batch_writes() {
|
||||
let storage = Arc::new(InMemoryCache::<CacheEntry>::default());
|
||||
let cache = ResponseCache::new(storage.clone()).with_config(ResponseCacheConfig {
|
||||
namespace: "service-test".into(),
|
||||
max_entry_bytes: json!({"answer":7}).to_string().len(),
|
||||
});
|
||||
let small = json!({"answer":7});
|
||||
let large = json!({"answer":"too large"});
|
||||
let request = |key: &str| {
|
||||
ResponseCacheRequest::new(CacheKeyInput {
|
||||
preset: Some(key.into()),
|
||||
..Default::default()
|
||||
})
|
||||
};
|
||||
cache
|
||||
.store(&request("sync"), large.clone(), Duration::ZERO)
|
||||
.unwrap();
|
||||
cache
|
||||
.async_store(&request("async"), large.clone(), Duration::ZERO)
|
||||
.await
|
||||
.unwrap();
|
||||
cache
|
||||
.async_store_batch(
|
||||
vec![
|
||||
(request("batch-large"), large),
|
||||
(request("batch-small"), small.clone()),
|
||||
],
|
||||
Duration::ZERO,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let service: Arc<dyn ResponseCacheService> = Arc::new(cache);
|
||||
service
|
||||
.store(&request("service"), small.clone(), Duration::ZERO)
|
||||
.await
|
||||
.unwrap();
|
||||
for key in ["sync", "async", "batch-large"] {
|
||||
assert!(storage.get_cache(key).unwrap().is_none());
|
||||
}
|
||||
for key in ["batch-small", "service"] {
|
||||
assert_eq!(
|
||||
service.lookup(&request(key), Duration::ZERO).await.unwrap(),
|
||||
Some(small.clone())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::same_scope("tenant-a", "tenant-a", true)]
|
||||
#[case::different_scope("tenant-a", "tenant-b", false)]
|
||||
#[case::empty_isolated_scope("", "", true)]
|
||||
#[tokio::test]
|
||||
async fn isolated_policy_controls_actual_entry_reuse(
|
||||
#[case] first: &str,
|
||||
#[case] second: &str,
|
||||
#[case] hit: bool,
|
||||
#[values(false, true)] override_policy: bool,
|
||||
) {
|
||||
use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
|
||||
let service = Arc::new(ResponseCache::new(Arc::new(
|
||||
InMemoryCache::<CacheEntry>::default(),
|
||||
)));
|
||||
let request = |scope| {
|
||||
ScopedCache::new(service.clone(), scope)
|
||||
.options(override_policy.then_some(CachePolicy {
|
||||
ttl: Some(Duration::from_secs(30)),
|
||||
..CachePolicy::default()
|
||||
}))
|
||||
.request("test", "messages", json!({"prompt":"hello"}))
|
||||
};
|
||||
service
|
||||
.async_store(
|
||||
&request(CacheScope::Isolated(first.into())),
|
||||
json!({"answer":7}),
|
||||
Duration::ZERO,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
service
|
||||
.async_lookup(
|
||||
&request(CacheScope::Isolated(second.into())),
|
||||
Duration::ZERO
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
hit.then(|| json!({"answer":7}))
|
||||
);
|
||||
assert_eq!(
|
||||
service
|
||||
.async_lookup(&request(CacheScope::Shared), Duration::ZERO)
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::valid(1, "messages", Some(7))]
|
||||
#[case::unknown_version(2, "messages", None)]
|
||||
#[case::another_surface(1, "responses", None)]
|
||||
fn envelopes_require_a_matching_surface_and_version(
|
||||
#[case] version: u32,
|
||||
#[case] surface: &str,
|
||||
#[case] expected: Option<u32>,
|
||||
) {
|
||||
let envelope: litellm_cache_response::ResponseEnvelope<u32> =
|
||||
serde_json::from_value(json!({"version":version,"surface":surface,"output":7})).unwrap();
|
||||
assert_eq!(envelope.decode("messages"), expected);
|
||||
}
|
||||
|
|
@ -1,12 +1,13 @@
|
|||
- Target invariants, not completion claims
|
||||
- This crate is the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on)
|
||||
- This crate owns compatibility for all existing Python callbacks and loggers, including `CustomLogger`. `mapping.rs` owns the executable call bindings and the inventory of Python-owned hooks. A Python-owned entry records an existing path, never permission to invoke it a second time. The native call adapter preserves the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on)
|
||||
- Smell test: if a future callback host (`callbacks-v1-python`, WASM, in-process Rust) could share a piece of this crate, it does not belong here
|
||||
- SDK request policy (credential inheritance, the budget and retry-count limits) is the driver's preflight, supplied by `python-bridge`; this crate only adopts the keyword view it produces
|
||||
- The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks`; they never learn which Python objects consume a call
|
||||
- SDK request policy (credential inheritance, the budget and retry-count limits) is a separate hook supplied by `python-bridge`; compose it after this adapter so logging adopts the final keyword view before policy mutates or rejects it
|
||||
- The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks` using the shared `CallEvent`; they never learn which Python objects consume a call
|
||||
- Every litellm Python internal Rust still borrows is a variant of `LegacyPython`, grouped by subsystem, with its signature pinned in `python_contract.json`
|
||||
- The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it
|
||||
- Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython`
|
||||
- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords; routes hand it over through `run_legacy_call` and keep no copy
|
||||
- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords; shared bridge composition hands it to `LegacyLogging`; routes use the neutral call boundary
|
||||
- `LoggingOperation` selects legacy logging entrypoints and response handling. It belongs here rather than in shared inference data contracts
|
||||
- `setup` reuses a `Logging` passed as `litellm_logging_obj` (the proxy and Router) and otherwise builds one through `function_setup`; which callbacks run is `Logging`'s decision, never this crate's
|
||||
- Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation
|
||||
- Retain complete boundary arguments, opaque values, aliases, omitted/default distinctions and deliberate copies; preserve the deployment-hook kwargs view
|
||||
|
|
|
|||
|
|
@ -2,12 +2,14 @@
|
|||
//! raises is answered with the same `Logging` calls, in the same order, as the Python
|
||||
//! `@client` path makes them.
|
||||
|
||||
use crate::LoggingOperation;
|
||||
use litellm_host_python::PythonOwned;
|
||||
|
||||
use litellm_host::event::{
|
||||
FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds,
|
||||
use litellm_host::{
|
||||
interceptors::{RawResponse, RequestContext, WireRequest},
|
||||
lifecycle::{FailureOrigin, Timing, epoch_seconds},
|
||||
};
|
||||
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks, from_py, missing_state, to_py};
|
||||
use litellm_host_python::{HookStep, from_py, missing_state, to_py};
|
||||
use pyo3::{
|
||||
exceptions::{PyBaseException, PyException},
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
|
|
@ -24,22 +26,10 @@ use crate::{
|
|||
setup,
|
||||
};
|
||||
|
||||
/// What the legacy contract needs to know about the route it is logging.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct LegacySurface {
|
||||
pub call_type: &'static str,
|
||||
/// What `Logging.pre_call` is told the input was.
|
||||
pub input_description: &'static str,
|
||||
/// How a streamed response is billed; `None` for a route that never streams.
|
||||
pub stream: Option<PassThroughStream>,
|
||||
}
|
||||
|
||||
/// The pass-through billing a streamed response goes through once its chunks are in.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct PassThroughStream {
|
||||
pub url_route: &'static str,
|
||||
/// A value of Python's `EndpointType`.
|
||||
pub endpoint_type: &'static str,
|
||||
struct PassThroughStream {
|
||||
url_route: &'static str,
|
||||
endpoint_type: &'static str,
|
||||
}
|
||||
|
||||
/// What the Messages stream iterator keeps for its end-of-stream billing.
|
||||
|
|
@ -55,7 +45,7 @@ struct LoggedRequest {
|
|||
}
|
||||
|
||||
pub struct LegacyLogging {
|
||||
surface: LegacySurface,
|
||||
operation: LoggingOperation,
|
||||
call: PublicCall,
|
||||
logger: Option<PythonLogger>,
|
||||
start: Py<PyAny>,
|
||||
|
|
@ -66,6 +56,7 @@ pub struct LegacyLogging {
|
|||
stream: Option<DeliveredStream>,
|
||||
asynchronous: bool,
|
||||
internal: bool,
|
||||
cache_key: Option<String>,
|
||||
}
|
||||
|
||||
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
|
||||
|
|
@ -79,12 +70,12 @@ fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
|
|||
impl LegacyLogging {
|
||||
pub fn new(
|
||||
py: Python<'_>,
|
||||
surface: LegacySurface,
|
||||
operation: LoggingOperation,
|
||||
call: PublicCall,
|
||||
asynchronous: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
surface,
|
||||
operation,
|
||||
call,
|
||||
logger: None,
|
||||
start: py.None(),
|
||||
|
|
@ -95,9 +86,47 @@ impl LegacyLogging {
|
|||
stream: None,
|
||||
asynchronous,
|
||||
internal: false,
|
||||
cache_key: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn call_type(&self) -> &'static str {
|
||||
match (self.operation, self.asynchronous) {
|
||||
(LoggingOperation::Completion, false) => "completion",
|
||||
(LoggingOperation::Completion, true) => "acompletion",
|
||||
(LoggingOperation::Responses, false) => "responses",
|
||||
(LoggingOperation::Responses, true) => "aresponses",
|
||||
(LoggingOperation::Messages, _) => "anthropic_messages",
|
||||
(LoggingOperation::Ocr, false) => "ocr",
|
||||
(LoggingOperation::Ocr, true) => "aocr",
|
||||
}
|
||||
}
|
||||
|
||||
fn input_description(&self) -> &'static str {
|
||||
match self.operation {
|
||||
LoggingOperation::Completion => "Chat completions",
|
||||
LoggingOperation::Responses => "Responses",
|
||||
LoggingOperation::Messages => "Messages",
|
||||
LoggingOperation::Ocr => "OCR document processing",
|
||||
}
|
||||
}
|
||||
|
||||
fn stream_billing(&self) -> Option<PassThroughStream> {
|
||||
match self.operation {
|
||||
LoggingOperation::Messages => Some(PassThroughStream {
|
||||
url_route: "/v1/messages",
|
||||
endpoint_type: "anthropic",
|
||||
}),
|
||||
LoggingOperation::Completion | LoggingOperation::Responses | LoggingOperation::Ocr => {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn adopt_arguments(&mut self, py: Python<'_>, arguments: &Py<PyDict>) {
|
||||
self.call.set_kwargs(arguments.clone_ref(py));
|
||||
}
|
||||
|
||||
/// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never
|
||||
/// runs them.
|
||||
fn runs_deployment_hooks(&self) -> bool {
|
||||
|
|
@ -112,7 +141,6 @@ impl LegacyLogging {
|
|||
|
||||
/// The keyword view the rest of the call reads: a copy, so the deployment hook's own
|
||||
/// dict is left as the hook returned it, carrying the logger as `@client` injects it.
|
||||
/// The driver's preflight rewrites this same dict before the host projects from it.
|
||||
fn prepare(&mut self, py: Python<'_>) -> PyResult<HookStep<Self, Py<PyDict>>> {
|
||||
let prepared = self.call.kwargs().bind(py).copy()?;
|
||||
prepared.set_item("litellm_logging_obj", self.logger()?.object(py))?;
|
||||
|
|
@ -181,17 +209,17 @@ impl LegacyLogging {
|
|||
|
||||
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
|
||||
let logger = self.logger()?;
|
||||
let billing = self.surface.stream.ok_or_else(missing_state)?;
|
||||
let billing = self.stream_billing().ok_or_else(missing_state)?;
|
||||
let billed = Streaming::Success.call(
|
||||
py,
|
||||
(
|
||||
logger.object(py),
|
||||
billing.url_route,
|
||||
billing.endpoint_type,
|
||||
&self
|
||||
.request
|
||||
.as_ref()
|
||||
.map(|request| request.body.clone_ref(py)),
|
||||
&self.request.as_ref().map_or_else(
|
||||
|| self.call.kwargs().clone_ref(py),
|
||||
|request| request.body.clone_ref(py),
|
||||
),
|
||||
&stream.chunks,
|
||||
&self.start,
|
||||
&self.end,
|
||||
|
|
@ -211,9 +239,12 @@ impl LegacyLogging {
|
|||
/// partial usage. The sync path has no loop to schedule that on, so it falls back to
|
||||
/// the plain failure handler.
|
||||
fn stream_failure(&mut self, py: Python<'_>) -> PyResult<HookStep<Self, ()>> {
|
||||
let (Some(logger), Some(error), Some(stream), Some(billing)) =
|
||||
(&self.logger, &self.error, &self.stream, self.surface.stream)
|
||||
else {
|
||||
let (Some(logger), Some(error), Some(stream), Some(billing)) = (
|
||||
&self.logger,
|
||||
&self.error,
|
||||
&self.stream,
|
||||
self.stream_billing(),
|
||||
) else {
|
||||
return Ok(HookStep::Ready(()));
|
||||
};
|
||||
if !self.asynchronous {
|
||||
|
|
@ -224,10 +255,10 @@ impl LegacyLogging {
|
|||
(
|
||||
logger.object(py),
|
||||
billing.endpoint_type,
|
||||
&self
|
||||
.request
|
||||
.as_ref()
|
||||
.map(|request| request.body.clone_ref(py)),
|
||||
&self.request.as_ref().map_or_else(
|
||||
|| self.call.kwargs().clone_ref(py),
|
||||
|request| request.body.clone_ref(py),
|
||||
),
|
||||
&stream.chunks,
|
||||
error,
|
||||
),
|
||||
|
|
@ -309,8 +340,8 @@ impl LegacyLogging {
|
|||
}
|
||||
}
|
||||
|
||||
impl PythonCallHooks for LegacyLogging {
|
||||
fn prepare_arguments(
|
||||
impl LegacyLogging {
|
||||
pub(crate) fn prepare_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
arguments: Py<PyDict>,
|
||||
|
|
@ -321,7 +352,7 @@ impl PythonCallHooks for LegacyLogging {
|
|||
self.internal = is_internal_call(py)?;
|
||||
let result = setup(
|
||||
py,
|
||||
self.surface.call_type,
|
||||
self.call_type(),
|
||||
self.call.args(),
|
||||
self.call.kwargs(),
|
||||
&self.start,
|
||||
|
|
@ -331,14 +362,14 @@ impl PythonCallHooks for LegacyLogging {
|
|||
self.call.set_kwargs(result.kwargs()?);
|
||||
if self.runs_deployment_hooks() {
|
||||
return Ok(HookStep::Await(
|
||||
DeploymentHooks::before_call(py, self.call.kwargs(), self.surface.call_type)?,
|
||||
DeploymentHooks::before_call(py, self.call.kwargs(), self.call_type())?,
|
||||
Self::resume_begin,
|
||||
));
|
||||
}
|
||||
self.prepare(py)
|
||||
}
|
||||
|
||||
fn before_provider_request(
|
||||
pub(crate) fn pre_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
wire: Box<WireRequest>,
|
||||
|
|
@ -367,7 +398,7 @@ impl PythonCallHooks for LegacyLogging {
|
|||
});
|
||||
self.logger()?.pre_call(
|
||||
py,
|
||||
self.surface.input_description,
|
||||
self.input_description(),
|
||||
context.api_key.as_ref().map(|api_key| api_key.expose()),
|
||||
&body,
|
||||
&headers,
|
||||
|
|
@ -384,7 +415,7 @@ impl PythonCallHooks for LegacyLogging {
|
|||
})))
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
pub(crate) fn transform_public_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
|
|
@ -398,7 +429,7 @@ impl PythonCallHooks for LegacyLogging {
|
|||
py,
|
||||
self.call.kwargs(),
|
||||
&self.response,
|
||||
self.surface.call_type,
|
||||
self.call_type(),
|
||||
)?,
|
||||
Self::resume_after_success,
|
||||
));
|
||||
|
|
@ -406,67 +437,105 @@ impl PythonCallHooks for LegacyLogging {
|
|||
self.finalize(py)
|
||||
}
|
||||
|
||||
fn on_event(&mut self, py: Python<'_>, event: HookEvent<'_>) -> PyResult<HookStep<Self, ()>> {
|
||||
match event {
|
||||
HookEvent::Started { .. } => Ok(HookStep::Ready(())),
|
||||
HookEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
|
||||
let api_key = self
|
||||
.request
|
||||
.as_ref()
|
||||
.and_then(|request| request.context.api_key.as_ref())
|
||||
.map(|api_key| api_key.expose());
|
||||
self.logger()?.post_call(
|
||||
py,
|
||||
&raw.body,
|
||||
api_key,
|
||||
self.request.as_ref().map(|request| &request.body),
|
||||
self.request.as_ref().map(|request| &request.headers),
|
||||
)?;
|
||||
Ok(HookStep::Ready(()))
|
||||
}
|
||||
HookEvent::Succeeded { timing, response } => {
|
||||
self.end = Some(datetime(py, timing.end_time)?);
|
||||
self.response = Some(response.clone_ref(py));
|
||||
match &self.stream {
|
||||
Some(stream) => self.stream_success(py, stream)?,
|
||||
None => self.dispatch_success(py)?,
|
||||
}
|
||||
Ok(HookStep::Ready(()))
|
||||
}
|
||||
HookEvent::Failed {
|
||||
timing,
|
||||
origin,
|
||||
error,
|
||||
} => {
|
||||
self.end = Some(datetime(py, timing.end_time)?);
|
||||
self.error = Some(error.clone_ref(py).into_value(py));
|
||||
if self.stream.is_some() {
|
||||
return self.stream_failure(py);
|
||||
}
|
||||
if origin == FailureOrigin::Call
|
||||
&& self.logger.is_some()
|
||||
&& self.runs_deployment_hooks()
|
||||
{
|
||||
let error = self.error.as_ref().ok_or_else(missing_state)?;
|
||||
return Ok(HookStep::Await(
|
||||
DeploymentHooks::after_failure(
|
||||
py,
|
||||
self.call.kwargs(),
|
||||
error,
|
||||
self.surface.call_type,
|
||||
)?,
|
||||
Self::resume_deployment_failure,
|
||||
));
|
||||
}
|
||||
self.dispatch_failure(py)
|
||||
}
|
||||
}
|
||||
pub(crate) fn result_ready(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
facts: &litellm_host::interceptors::ExecutionFacts,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
use litellm_host::interceptors::ResultSource;
|
||||
|
||||
let logger = self.logger()?.object(py);
|
||||
let params = logger
|
||||
.getattr("litellm_params")?
|
||||
.cast_into::<PyDict>()?
|
||||
.copy()?;
|
||||
params.set_item("custom_llm_provider", &facts.provider.provider)?;
|
||||
crate::python::Logging::Update.call(
|
||||
py,
|
||||
(
|
||||
&logger,
|
||||
self.call.kwargs(),
|
||||
&facts.provider.model,
|
||||
logger.getattr("optional_params")?,
|
||||
params,
|
||||
&facts.provider.provider,
|
||||
),
|
||||
)?;
|
||||
let details = logger.getattr("model_call_details")?;
|
||||
self.cache_key = match &facts.source {
|
||||
ResultSource::Provider => None,
|
||||
ResultSource::Cache { key } => Some(key.clone()),
|
||||
};
|
||||
details.set_item("cache_hit", self.cache_key.is_some())?;
|
||||
details.set_item("cache_key", self.cache_key.as_deref())?;
|
||||
Ok(HookStep::Ready(()))
|
||||
}
|
||||
|
||||
fn on_stream_open(&mut self, py: Python<'_>) -> PyResult<()> {
|
||||
if self.surface.stream.is_none() {
|
||||
pub(crate) fn post_call(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
raw: &RawResponse,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
let api_key = self
|
||||
.request
|
||||
.as_ref()
|
||||
.and_then(|request| request.context.api_key.as_ref())
|
||||
.map(|api_key| api_key.expose());
|
||||
self.logger()?.post_call(
|
||||
py,
|
||||
&raw.body,
|
||||
api_key,
|
||||
self.request.as_ref().map(|request| &request.body),
|
||||
self.request.as_ref().map(|request| &request.headers),
|
||||
)?;
|
||||
Ok(HookStep::Ready(()))
|
||||
}
|
||||
|
||||
pub(crate) fn succeeded(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
timing: Timing,
|
||||
response: &Py<PyAny>,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
self.end = Some(datetime(py, timing.end_time)?);
|
||||
self.response = Some(response.clone_ref(py));
|
||||
match &self.stream {
|
||||
Some(stream) => self.stream_success(py, stream)?,
|
||||
None => self.dispatch_success(py)?,
|
||||
}
|
||||
Ok(HookStep::Ready(()))
|
||||
}
|
||||
|
||||
pub(crate) fn failed(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
timing: Timing,
|
||||
origin: FailureOrigin,
|
||||
error: &PyErr,
|
||||
) -> PyResult<HookStep<Self, ()>> {
|
||||
self.end = Some(datetime(py, timing.end_time)?);
|
||||
self.error = Some(error.clone_ref(py).into_value(py));
|
||||
if self.stream.is_some() {
|
||||
return self.stream_failure(py);
|
||||
}
|
||||
if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() {
|
||||
let error = self.error.as_ref().ok_or_else(missing_state)?;
|
||||
return Ok(HookStep::Await(
|
||||
DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type())?,
|
||||
Self::resume_deployment_failure,
|
||||
));
|
||||
}
|
||||
self.dispatch_failure(py)
|
||||
}
|
||||
|
||||
pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
|
||||
if self.stream_billing().is_none() {
|
||||
return Err(missing_state());
|
||||
}
|
||||
if let Some(key) = &self.cache_key {
|
||||
head.bind(py).set_item("cache_key", key)?;
|
||||
head.bind(py).set_item("cache_hit", true)?;
|
||||
}
|
||||
Streaming::Opened.call(py, (self.logger()?.object(py),))?;
|
||||
self.stream = Some(DeliveredStream {
|
||||
chunks: PyList::empty(py).unbind(),
|
||||
|
|
@ -475,7 +544,7 @@ impl PythonCallHooks for LegacyLogging {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
|
||||
pub(crate) fn stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
|
||||
let stream = self.stream.as_mut().ok_or_else(missing_state)?;
|
||||
if stream.first_chunk.is_none() {
|
||||
stream.first_chunk = Some(datetime(py, epoch_seconds())?);
|
||||
|
|
@ -519,8 +588,9 @@ impl PythonOwned for LegacyLogging {
|
|||
mod deployment_hooks_tests {
|
||||
use std::ffi::CStr;
|
||||
|
||||
use litellm_host::event::{FailureOrigin, Timing};
|
||||
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks};
|
||||
use litellm_host::hooks::CallHooks;
|
||||
use litellm_host::lifecycle::{FailureOrigin, Timing};
|
||||
use litellm_host_python::{HookStep, PythonCallEvent};
|
||||
use pyo3::exceptions::asyncio::CancelledError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyDict;
|
||||
|
|
@ -579,6 +649,47 @@ kwargs = {'logger': logger, 'document': document}
|
|||
matches!(step, HookStep::Await(_, _))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::sync_completion(crate::LoggingOperation::Completion, false, "completion")]
|
||||
#[case::async_completion(crate::LoggingOperation::Completion, true, "acompletion")]
|
||||
#[case::sync_responses(crate::LoggingOperation::Responses, false, "responses")]
|
||||
#[case::async_responses(crate::LoggingOperation::Responses, true, "aresponses")]
|
||||
#[case::sync_messages(crate::LoggingOperation::Messages, false, "anthropic_messages")]
|
||||
#[case::async_messages(crate::LoggingOperation::Messages, true, "anthropic_messages")]
|
||||
#[case::sync_ocr(crate::LoggingOperation::Ocr, false, "ocr")]
|
||||
#[case::async_ocr(crate::LoggingOperation::Ocr, true, "aocr")]
|
||||
fn operation_selects_the_legacy_setup_and_deployment_hook_contract(
|
||||
#[case] operation: crate::LoggingOperation,
|
||||
#[case] asynchronous: bool,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = namespace(py, CALL);
|
||||
let mut logging = LegacyLogging {
|
||||
operation,
|
||||
..legacy_call(py, &locals, asynchronous)
|
||||
};
|
||||
let kwargs = local(&locals, "kwargs")
|
||||
.cast_into::<PyDict>()
|
||||
.unwrap()
|
||||
.unbind();
|
||||
let step = logging.prepare_arguments(py, kwargs, 0.0).unwrap();
|
||||
assert_eq!(awaits_deployment_hook(&step), asynchronous);
|
||||
locals.set_item("expected", expected).unwrap();
|
||||
locals.set_item("asynchronous", asynchronous).unwrap();
|
||||
run(
|
||||
py,
|
||||
&locals,
|
||||
c"
|
||||
assert logger.setup_call_type == expected
|
||||
if asynchronous:
|
||||
assert logger.calls == [('pre_hook', expected)]
|
||||
",
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::synchronous(false)]
|
||||
#[case::asynchronous(true)]
|
||||
|
|
@ -775,7 +886,7 @@ assert finalized is replacement
|
|||
)
|
||||
.unwrap();
|
||||
let failure = PyErr::from_value(local(&locals, "failure"));
|
||||
let failed = HookEvent::Failed {
|
||||
let failed = PythonCallEvent::Failed {
|
||||
timing: TIMING,
|
||||
origin: FailureOrigin::Call,
|
||||
error: &failure,
|
||||
|
|
@ -808,8 +919,10 @@ mod payload_tests {
|
|||
use std::ffi::CStr;
|
||||
|
||||
use litellm_auth::SecretValue;
|
||||
use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest};
|
||||
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks, PythonOwned, to_py};
|
||||
use litellm_host::hooks::CallHooks;
|
||||
use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest};
|
||||
use litellm_host::lifecycle::ExecutionEvent;
|
||||
use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py};
|
||||
use proptest::prelude::*;
|
||||
use pyo3::gc::{PyTraverseError, PyVisit};
|
||||
use pyo3::prelude::*;
|
||||
|
|
@ -833,6 +946,7 @@ class PayloadLogger(StubLogger):
|
|||
def pre_call(self, input, api_key, additional_args):
|
||||
self.record('pre_call', None)
|
||||
self.pre = additional_args
|
||||
self.pre_input = input
|
||||
self.pre_api_key = api_key
|
||||
on_pre_call(additional_args)
|
||||
|
||||
|
|
@ -924,13 +1038,18 @@ check = lambda: None
|
|||
let step = logging
|
||||
.before_provider_request(py, Box::new(wire), context)
|
||||
.unwrap();
|
||||
let raw = MachineEvent::ResponseReceived {
|
||||
raw: RawResponse {
|
||||
body: "raw response".into(),
|
||||
},
|
||||
let raw = RawResponse {
|
||||
body: "raw response".into(),
|
||||
};
|
||||
assert!(matches!(
|
||||
logging.on_event(py, HookEvent::Machine(&raw)).unwrap(),
|
||||
logging
|
||||
.on_event(
|
||||
py,
|
||||
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived {
|
||||
raw: &raw
|
||||
})
|
||||
)
|
||||
.unwrap(),
|
||||
HookStep::Ready(())
|
||||
));
|
||||
(logging, step)
|
||||
|
|
@ -975,6 +1094,57 @@ check = lambda: None
|
|||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::completion(crate::LoggingOperation::Completion, "Chat completions")]
|
||||
#[case::responses(crate::LoggingOperation::Responses, "Responses")]
|
||||
#[case::messages(crate::LoggingOperation::Messages, "Messages")]
|
||||
#[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")]
|
||||
fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases(
|
||||
#[case] operation: crate::LoggingOperation,
|
||||
#[case] description: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = namespace(py, PAYLOAD_LOGGER);
|
||||
run(
|
||||
py,
|
||||
&locals,
|
||||
c"
|
||||
original = [0]
|
||||
replacement = [1]
|
||||
kwargs['pages'] = original
|
||||
prepared = {'pages': replacement}
|
||||
",
|
||||
);
|
||||
let mut logging = LegacyLogging {
|
||||
operation,
|
||||
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
|
||||
..legacy_call(py, &locals, false)
|
||||
};
|
||||
let prepared = local(&locals, "prepared")
|
||||
.cast_into::<pyo3::types::PyDict>()
|
||||
.unwrap()
|
||||
.unbind();
|
||||
logging.arguments_prepared(py, &prepared).unwrap();
|
||||
let wire = WireRequest {
|
||||
body: json!({"pages": [1]}),
|
||||
..route_wire()
|
||||
};
|
||||
let (_, step) = send_and_receive(py, &mut logging, wire, &route_context());
|
||||
assert!(matches!(step, HookStep::Ready(_)));
|
||||
locals.set_item("description", description).unwrap();
|
||||
run(
|
||||
py,
|
||||
&locals,
|
||||
c"
|
||||
assert logger.pre['complete_input_dict']['pages'] is replacement
|
||||
assert logger.pre_input == description
|
||||
assert original == [0]
|
||||
",
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn a_cycle_through_the_retained_headers_is_collected() {
|
||||
Python::initialize();
|
||||
|
|
@ -1459,8 +1629,9 @@ def check():
|
|||
mod terminal_tests {
|
||||
use std::ffi::CStr;
|
||||
|
||||
use litellm_host::event::{FailureOrigin, Timing};
|
||||
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks, PythonOwned};
|
||||
use litellm_host::hooks::CallHooks;
|
||||
use litellm_host::lifecycle::{FailureOrigin, Timing};
|
||||
use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned};
|
||||
use pyo3::exceptions::PyRuntimeError;
|
||||
use pyo3::exceptions::asyncio::CancelledError;
|
||||
use pyo3::prelude::*;
|
||||
|
|
@ -1492,7 +1663,7 @@ mod terminal_tests {
|
|||
logging
|
||||
.on_event(
|
||||
py,
|
||||
HookEvent::Succeeded {
|
||||
PythonCallEvent::Succeeded {
|
||||
timing: TIMING,
|
||||
response: &response,
|
||||
},
|
||||
|
|
@ -1509,7 +1680,7 @@ mod terminal_tests {
|
|||
logging
|
||||
.on_event(
|
||||
py,
|
||||
HookEvent::Failed {
|
||||
PythonCallEvent::Failed {
|
||||
timing: TIMING,
|
||||
origin: FailureOrigin::Host,
|
||||
error: &failure,
|
||||
|
|
@ -1558,6 +1729,77 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
|
|||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn dropped_observations_preserve_deferred_success_and_response_identity() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = namespace(
|
||||
py,
|
||||
c"response = object()\nlogger._defer_async_logging = True",
|
||||
);
|
||||
let mut logging = logged(py, &locals, true);
|
||||
let response = local(&locals, "response").unbind();
|
||||
let event = PythonCallEvent::Succeeded {
|
||||
timing: TIMING,
|
||||
response: &response,
|
||||
};
|
||||
let (sender, receiver) = litellm_host::observation::observation_channel(
|
||||
std::num::NonZeroUsize::new(1).unwrap(),
|
||||
);
|
||||
drop(receiver);
|
||||
sender.emit(event.snapshot());
|
||||
assert!(matches!(
|
||||
logging.on_event(py, event).unwrap(),
|
||||
HookStep::Ready(())
|
||||
));
|
||||
assert_eq!(sender.dropped_events(), 1);
|
||||
run(py, &locals, c"
|
||||
assert logger.names() == ['sync_success_for_async_call'], logger.calls
|
||||
logger._native_pending_logging.release(True)
|
||||
logger._native_pending_logging.release(True)
|
||||
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls
|
||||
assert logger.calls[0][1] is response
|
||||
assert logger.calls[1][1] is response
|
||||
");
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn stream_bindings_deliver_collected_chunks_in_order_without_success_fan_out() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = namespace(py, c"first = b'first'\nlast = b'last'\nresponse = None");
|
||||
let mut logging = LegacyLogging {
|
||||
operation: crate::LoggingOperation::Messages,
|
||||
..logged(py, &locals, true)
|
||||
};
|
||||
logging
|
||||
.on_stream_open(py, &pyo3::types::PyDict::new(py).into_any().unbind())
|
||||
.unwrap();
|
||||
logging
|
||||
.on_stream_chunk(py, &local(&locals, "first").unbind())
|
||||
.unwrap();
|
||||
logging
|
||||
.on_stream_chunk(py, &local(&locals, "last").unbind())
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
succeed(py, &locals, &mut logging),
|
||||
HookStep::Ready(())
|
||||
));
|
||||
run(
|
||||
py,
|
||||
&locals,
|
||||
c"
|
||||
assert logger.names() == ['stream_opened', 'stream_success'], logger.calls
|
||||
chunks = logger.calls[1][1]
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0] is first
|
||||
assert chunks[1] is last
|
||||
",
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::synchronous(false, &["failure_handler"])]
|
||||
#[case::asynchronous(true, &[])]
|
||||
|
|
|
|||
|
|
@ -3,16 +3,13 @@
|
|||
//! lifetime. No other callback host has that obligation, which is why nothing outside
|
||||
//! this crate holds them.
|
||||
|
||||
use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
|
||||
use litellm_host_python::{Preflight, PythonBinding, PythonHostCalls, lookup, run_call};
|
||||
use litellm_host_python::lookup;
|
||||
use pyo3::{
|
||||
gc::{PyTraverseError, PyVisit},
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use crate::{LegacyLogging, LegacySurface};
|
||||
|
||||
pub struct PublicCall {
|
||||
args: Py<PyTuple>,
|
||||
kwargs: Py<PyDict>,
|
||||
|
|
@ -34,6 +31,10 @@ impl PublicCall {
|
|||
})
|
||||
}
|
||||
|
||||
pub fn arguments(&self, py: Python<'_>) -> Py<PyDict> {
|
||||
self.kwargs.clone_ref(py)
|
||||
}
|
||||
|
||||
pub(crate) fn args(&self) -> &Py<PyTuple> {
|
||||
&self.args
|
||||
}
|
||||
|
|
@ -64,35 +65,6 @@ impl PublicCall {
|
|||
}
|
||||
}
|
||||
|
||||
/// Runs one native call under the legacy `Logging` contract: the protocol host projects from
|
||||
/// the keyword view the contract prepares and `preflight` rewrites, and the contract
|
||||
/// observes the call.
|
||||
pub fn run_legacy_call<H, M>(
|
||||
py: Python<'_>,
|
||||
surface: LegacySurface,
|
||||
call: PublicCall,
|
||||
start: impl FnOnce(<H::Protocol as Protocol>::Request) -> M + Send + Sync + 'static,
|
||||
host: H,
|
||||
preflight: Preflight,
|
||||
asynchronous: bool,
|
||||
) -> PyResult<Py<PyAny>>
|
||||
where
|
||||
H: PythonBinding + PythonHostCalls<H::Protocol> + 'static,
|
||||
M: Machine<Protocol = H::Protocol> + 'static,
|
||||
M::Complete: Into<HostedCompletion<<H::Protocol as Protocol>::Response>>,
|
||||
{
|
||||
let arguments = call.kwargs.clone_ref(py);
|
||||
run_call(
|
||||
py,
|
||||
start,
|
||||
host,
|
||||
LegacyLogging::new(py, surface, call, asynchronous),
|
||||
preflight,
|
||||
arguments,
|
||||
asynchronous,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
//! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls
|
||||
//! duplication. All of it expires with the legacy callback contract.
|
||||
|
||||
use litellm_host::event::{RequestContext, WireRequest};
|
||||
use litellm_host::interceptors::{RequestContext, WireRequest};
|
||||
use litellm_host_python::to_py;
|
||||
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
|
||||
|
||||
|
|
|
|||
|
|
@ -2,25 +2,31 @@
|
|||
//! sync and async callback registries it fans out to, the deployment hooks and the deferred
|
||||
//! proxy release. All of it sits behind one
|
||||
//! [`PythonCallHooks`](litellm_host_python::PythonCallHooks), so the driver, the routes and
|
||||
//! core never learn which Python object is on the other end. The SDK's own request policy
|
||||
//! (credential inheritance, the budget and retry limits) is the driver's preflight, not this
|
||||
//! crate's.
|
||||
//! core never learn which Python object is on the other end.
|
||||
//!
|
||||
//! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`]
|
||||
//! is where those objects live, and [`run_legacy_call`] is how a route hands them over
|
||||
//! without keeping a copy.
|
||||
//! is where those objects live.
|
||||
|
||||
mod adapter;
|
||||
mod call;
|
||||
mod callbacks;
|
||||
mod deferred;
|
||||
mod logger;
|
||||
mod mapping;
|
||||
mod python;
|
||||
pub(crate) use adapter::LegacyLogging;
|
||||
pub use adapter::{LegacySurface, PassThroughStream};
|
||||
pub use call::{PublicCall, run_legacy_call};
|
||||
pub use adapter::LegacyLogging;
|
||||
pub use call::PublicCall;
|
||||
pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
|
||||
pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
|
||||
pub use mapping::{CallBoundary, CallbackMapping, Dispatch, callback_mappings};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum LoggingOperation {
|
||||
Completion,
|
||||
Responses,
|
||||
Messages,
|
||||
Ocr,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_support;
|
||||
|
|
|
|||
288
litellm-rust/crates/callbacks-legacy-python/src/mapping.rs
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
use litellm_host::{
|
||||
hooks::CallHooks,
|
||||
interceptors::{RawResponse, RequestContext, WireRequest},
|
||||
lifecycle::{ExecutionEvent, FailureOrigin, Timing},
|
||||
};
|
||||
use litellm_host_python::{HookStep, PythonCallEvent, PythonRuntime};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use crate::LegacyLogging;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum CallBoundary {
|
||||
PrepareArguments,
|
||||
BeforeProviderRequest,
|
||||
AfterProviderResponse,
|
||||
TransformResponse,
|
||||
Succeeded,
|
||||
Failed,
|
||||
StreamOpened,
|
||||
StreamChunk,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Dispatch {
|
||||
Call(CallBoundary),
|
||||
Python(&'static str),
|
||||
DeclarationOnly,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct CallbackMapping {
|
||||
pub callback: &'static str,
|
||||
pub dispatch: Dispatch,
|
||||
}
|
||||
|
||||
struct Binding<H> {
|
||||
boundary: CallBoundary,
|
||||
invoke: H,
|
||||
callbacks: &'static [&'static str],
|
||||
}
|
||||
|
||||
impl<H> Binding<H> {
|
||||
fn mappings(&self) -> impl Iterator<Item = CallbackMapping> {
|
||||
self.callbacks.iter().map(|callback| CallbackMapping {
|
||||
callback,
|
||||
dispatch: Dispatch::Call(self.boundary),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type Step<T> = PyResult<HookStep<LegacyLogging, T>>;
|
||||
type Prepare = fn(&mut LegacyLogging, Python<'_>, Py<PyDict>, f64) -> Step<Py<PyDict>>;
|
||||
type Before =
|
||||
fn(&mut LegacyLogging, Python<'_>, Box<WireRequest>, &RequestContext) -> Step<Box<WireRequest>>;
|
||||
type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>;
|
||||
type Transform = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>, Timing) -> Step<Py<PyAny>>;
|
||||
type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py<PyAny>) -> Step<()>;
|
||||
type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>;
|
||||
type Open = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
|
||||
type Chunk = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
|
||||
|
||||
const PREPARE: Binding<Prepare> = Binding {
|
||||
boundary: CallBoundary::PrepareArguments,
|
||||
invoke: LegacyLogging::prepare_call,
|
||||
callbacks: &["async_pre_call_deployment_hook"],
|
||||
};
|
||||
|
||||
const BEFORE: Binding<Before> = Binding {
|
||||
boundary: CallBoundary::BeforeProviderRequest,
|
||||
invoke: LegacyLogging::pre_call,
|
||||
callbacks: &["log_pre_api_call", "log_input_event"],
|
||||
};
|
||||
|
||||
const AFTER: Binding<After> = Binding {
|
||||
boundary: CallBoundary::AfterProviderResponse,
|
||||
invoke: LegacyLogging::post_call,
|
||||
callbacks: &["log_post_api_call"],
|
||||
};
|
||||
|
||||
const TRANSFORM: Binding<Transform> = Binding {
|
||||
boundary: CallBoundary::TransformResponse,
|
||||
invoke: LegacyLogging::transform_public_response,
|
||||
callbacks: &["async_post_call_success_deployment_hook"],
|
||||
};
|
||||
|
||||
const SUCCESS: Binding<Success> = Binding {
|
||||
boundary: CallBoundary::Succeeded,
|
||||
invoke: LegacyLogging::succeeded,
|
||||
callbacks: &[
|
||||
"log_success_event",
|
||||
"async_log_success_event",
|
||||
"logging_hook",
|
||||
"async_logging_hook",
|
||||
"redact_standard_logging_payload_from_model_call_details",
|
||||
"log_event",
|
||||
"async_log_event",
|
||||
],
|
||||
};
|
||||
|
||||
const FAILURE: Binding<Failure> = Binding {
|
||||
boundary: CallBoundary::Failed,
|
||||
invoke: LegacyLogging::failed,
|
||||
callbacks: &[
|
||||
"async_post_call_failure_deployment_hook",
|
||||
"log_failure_event",
|
||||
"async_log_failure_event",
|
||||
"log_model_group_rate_limit_error",
|
||||
"log_event",
|
||||
"async_log_event",
|
||||
],
|
||||
};
|
||||
|
||||
const OPEN: Binding<Open> = Binding {
|
||||
boundary: CallBoundary::StreamOpened,
|
||||
invoke: LegacyLogging::stream_opened,
|
||||
callbacks: &[],
|
||||
};
|
||||
|
||||
const CHUNK: Binding<Chunk> = Binding {
|
||||
boundary: CallBoundary::StreamChunk,
|
||||
invoke: LegacyLogging::stream_chunk,
|
||||
callbacks: &[],
|
||||
};
|
||||
|
||||
pub fn callback_mappings() -> impl Iterator<Item = CallbackMapping> {
|
||||
PREPARE
|
||||
.mappings()
|
||||
.chain(BEFORE.mappings())
|
||||
.chain(AFTER.mappings())
|
||||
.chain(TRANSFORM.mappings())
|
||||
.chain(SUCCESS.mappings())
|
||||
.chain(FAILURE.mappings())
|
||||
.chain(OPEN.mappings())
|
||||
.chain(CHUNK.mappings())
|
||||
.chain(PYTHON_CALLBACKS.iter().copied())
|
||||
}
|
||||
|
||||
impl CallHooks<PythonRuntime> for LegacyLogging {
|
||||
fn prepare_arguments(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
arguments: Py<PyDict>,
|
||||
started_at: f64,
|
||||
) -> Step<Py<PyDict>> {
|
||||
(PREPARE.invoke)(self, py, arguments, started_at)
|
||||
}
|
||||
|
||||
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
|
||||
self.adopt_arguments(py, arguments);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn before_provider_request(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
wire: Box<WireRequest>,
|
||||
context: &RequestContext,
|
||||
) -> Step<Box<WireRequest>> {
|
||||
(BEFORE.invoke)(self, py, wire, context)
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
response: Py<PyAny>,
|
||||
timing: Timing,
|
||||
) -> Step<Py<PyAny>> {
|
||||
(TRANSFORM.invoke)(self, py, response, timing)
|
||||
}
|
||||
|
||||
fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> Step<()> {
|
||||
match event {
|
||||
PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => {
|
||||
Ok(HookStep::Ready(()))
|
||||
}
|
||||
PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts }) => {
|
||||
self.result_ready(py, &facts)
|
||||
}
|
||||
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
|
||||
(AFTER.invoke)(self, py, raw)
|
||||
}
|
||||
PythonCallEvent::Succeeded { timing, response } => {
|
||||
(SUCCESS.invoke)(self, py, timing, response)
|
||||
}
|
||||
PythonCallEvent::Failed {
|
||||
timing,
|
||||
origin,
|
||||
error,
|
||||
} => (FAILURE.invoke)(self, py, timing, origin, error),
|
||||
}
|
||||
}
|
||||
|
||||
fn on_stream_open(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
|
||||
(OPEN.invoke)(self, py, head)
|
||||
}
|
||||
|
||||
fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
|
||||
(CHUNK.invoke)(self, py, chunk)
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! python_callbacks {
|
||||
($($dispatch:expr => [$($callback:literal),* $(,)?]),* $(,)?) => {
|
||||
const PYTHON_CALLBACKS: &[CallbackMapping] = &[
|
||||
$($(CallbackMapping { callback: $callback, dispatch: $dispatch },)*)*
|
||||
];
|
||||
};
|
||||
}
|
||||
|
||||
python_callbacks! {
|
||||
Dispatch::Python("litellm.router") => [
|
||||
"async_pre_routing_hook",
|
||||
"async_filter_deployments",
|
||||
"pre_call_check",
|
||||
"async_pre_call_check",
|
||||
],
|
||||
Dispatch::Python("litellm.router_utils.fallback_event_handlers") => [
|
||||
"log_success_fallback_event",
|
||||
"log_failure_fallback_event",
|
||||
],
|
||||
Dispatch::Python("litellm.proxy.utils") => [
|
||||
"async_pre_call_hook",
|
||||
"async_post_call_response_headers_hook",
|
||||
"async_post_call_failure_hook",
|
||||
"async_post_call_success_hook",
|
||||
"async_moderation_hook",
|
||||
"async_post_call_streaming_hook",
|
||||
"async_post_call_streaming_iterator_hook",
|
||||
"async_filter_listed_models",
|
||||
],
|
||||
Dispatch::Python("litellm.litellm_core_utils.litellm_logging") => [
|
||||
"async_get_chat_completion_prompt",
|
||||
"get_chat_completion_prompt",
|
||||
"log_stream_event",
|
||||
"async_log_stream_event",
|
||||
"async_post_mcp_tool_call_hook",
|
||||
],
|
||||
Dispatch::Python("litellm.llms.anthropic.pass_through.messages.handler") => [
|
||||
"async_pre_request_hook",
|
||||
],
|
||||
Dispatch::Python("litellm.litellm_core_utils.streaming_handler") => [
|
||||
"async_post_call_streaming_deployment_hook",
|
||||
],
|
||||
Dispatch::Python("litellm.responses.streaming_iterator") => [
|
||||
"async_post_call_streaming_deployment_hook",
|
||||
],
|
||||
Dispatch::Python("litellm.main") => [
|
||||
"translate_completion_input_params",
|
||||
"translate_completion_output_params",
|
||||
"translate_completion_output_params_streaming",
|
||||
],
|
||||
Dispatch::Python("litellm.integrations.argilla") => ["async_dataset_hook"],
|
||||
Dispatch::Python("litellm.proxy.management_helpers.audit_logs") => ["async_log_audit_log_event"],
|
||||
Dispatch::Python("litellm.llms.custom_httpx.llm_http_handler") => [
|
||||
"async_should_run_agentic_loop",
|
||||
"async_run_agentic_loop",
|
||||
"async_build_agentic_loop_plan",
|
||||
"async_post_agentic_loop_response_hook",
|
||||
"async_agentic_loop_cleanup_hook",
|
||||
"async_should_run_chat_completion_agentic_loop",
|
||||
"async_run_chat_completion_agentic_loop",
|
||||
"async_build_chat_completion_agentic_loop_plan",
|
||||
],
|
||||
Dispatch::Python("litellm.litellm_core_utils.chat_completion_agentic_loop") => [
|
||||
"async_should_run_agentic_loop",
|
||||
"async_run_agentic_loop",
|
||||
"async_build_agentic_loop_plan",
|
||||
"async_post_agentic_loop_response_hook",
|
||||
"async_agentic_loop_cleanup_hook",
|
||||
],
|
||||
Dispatch::Python("litellm.llms.openai.openai") => [
|
||||
"async_should_run_chat_completion_agentic_loop",
|
||||
"async_run_chat_completion_agentic_loop",
|
||||
],
|
||||
Dispatch::Python("litellm.proxy.spend_tracking.cold_storage_handler") => [
|
||||
"get_proxy_server_request_from_cold_storage_with_object_key",
|
||||
],
|
||||
Dispatch::Python("litellm.integrations.custom_logger") => [
|
||||
"truncate_standard_logging_payload_content",
|
||||
"redacts_messages_itself",
|
||||
"handle_callback_failure",
|
||||
"get_callback_env_vars",
|
||||
],
|
||||
Dispatch::DeclarationOnly => [
|
||||
"async_log_pre_api_call",
|
||||
"async_log_input_event",
|
||||
],
|
||||
}
|
||||
|
|
@ -3,7 +3,7 @@ use std::ffi::CStr;
|
|||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyDict, PyTuple};
|
||||
|
||||
use crate::{LegacyLogging, LegacySurface, PublicCall};
|
||||
use crate::{LegacyLogging, PublicCall};
|
||||
|
||||
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
|
||||
/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
|
||||
|
|
@ -45,11 +45,14 @@ def contracted(name, fake):
|
|||
if not hasattr(legacy, 'is_internal'):
|
||||
legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False)
|
||||
|
||||
def setup(call_type, args, kwargs, start, asynchronous):
|
||||
logger = kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger']
|
||||
logger.setup_call_type = call_type
|
||||
return types.SimpleNamespace(logger=logger, kwargs=kwargs)
|
||||
|
||||
|
||||
FAKES = {
|
||||
'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace(
|
||||
logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'],
|
||||
kwargs=kwargs,
|
||||
),
|
||||
'setup': setup,
|
||||
'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response),
|
||||
'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
|
|
@ -82,10 +85,10 @@ FAKES = {
|
|||
),
|
||||
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
|
||||
'stream_opened': lambda logger: logger.record('stream_opened', None),
|
||||
'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record(
|
||||
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record(
|
||||
'stream_success', list(chunks)
|
||||
),
|
||||
'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error),
|
||||
'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error),
|
||||
}
|
||||
assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys())
|
||||
for name, fake in FAKES.items():
|
||||
|
|
@ -186,14 +189,5 @@ pub(crate) fn legacy_call(
|
|||
.map(|kwargs| kwargs.cast_into::<PyDict>().unwrap())
|
||||
.unwrap_or_else(|| PyDict::new(py));
|
||||
let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
|
||||
LegacyLogging::new(
|
||||
py,
|
||||
LegacySurface {
|
||||
call_type: "test",
|
||||
input_description: "test input",
|
||||
stream: None,
|
||||
},
|
||||
call,
|
||||
asynchronous,
|
||||
)
|
||||
LegacyLogging::new(py, crate::LoggingOperation::Ocr, call, asynchronous)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,11 +8,10 @@ repository.workspace = true
|
|||
[dependencies]
|
||||
fancy-regex.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-llms-types.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_path_to_error = "0.1"
|
||||
serde_with.workspace = true
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
url.workspace = true
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use litellm_types::utils::{ChatCompletionsUsage, PromptTokensDetails};
|
||||
use litellm_llms_types::formats::chat_completions::{ChatCompletionsUsage, PromptTokensDetails};
|
||||
|
||||
/// OpenAI finish reasons, mirroring Python's `_FINISH_REASON_MAP` for the
|
||||
/// reasons the providers on this route can emit. Python warns and falls back to
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders};
|
||||
use litellm_llms_types::headers::{ProviderSpecificHeader, ProviderSpecificHeaders};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub fn get_provider_specific_headers(
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@
|
|||
//! `_bedrock_converse_messages_pt` for the text-only surface this route
|
||||
//! accepts; anything richer is declined upstream by the capability gate.
|
||||
|
||||
use litellm_types::llms::openai::{ChatMessage, ChatMessageContent};
|
||||
use litellm_llms_types::formats::chat_completions::{ChatMessage, ChatMessageContent};
|
||||
use strum::IntoStaticStr;
|
||||
|
||||
pub const EMPTY_TEXT_PLACEHOLDER: &str =
|
||||
|
|
|
|||
|
|
@ -1,12 +1,3 @@
|
|||
use serde::{
|
||||
Deserializer,
|
||||
de::{Error, Visitor},
|
||||
};
|
||||
use serde_with::DeserializeAs;
|
||||
|
||||
pub struct LaxI64;
|
||||
pub struct FiniteF64;
|
||||
|
||||
pub fn parse_str_bool(value: &str) -> Option<bool> {
|
||||
let token = value.trim_matches(|character: char| {
|
||||
character.is_whitespace() || matches!(character, '\u{1c}'..='\u{1f}')
|
||||
|
|
@ -22,129 +13,12 @@ pub fn parse_redis_bool(value: &str) -> bool {
|
|||
value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes")
|
||||
}
|
||||
|
||||
impl<'de> DeserializeAs<'de, i64> for LaxI64 {
|
||||
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
|
||||
deserializer.deserialize_any(Self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Visitor<'de> for LaxI64 {
|
||||
type Value = i64;
|
||||
|
||||
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("an integer in the i64 range")
|
||||
}
|
||||
|
||||
fn visit_i64<E: Error>(self, value: i64) -> Result<i64, E> {
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn visit_u64<E: Error>(self, value: u64) -> Result<i64, E> {
|
||||
i64::try_from(value).map_err(E::custom)
|
||||
}
|
||||
|
||||
fn visit_f64<E: Error>(self, value: f64) -> Result<i64, E> {
|
||||
integral_float(value).ok_or_else(|| E::custom("expected an integer in the i64 range"))
|
||||
}
|
||||
|
||||
fn visit_str<E: Error>(self, value: &str) -> Result<i64, E> {
|
||||
integer_string(value.trim())
|
||||
.ok_or_else(|| E::custom("expected an integer in the i64 range"))
|
||||
}
|
||||
|
||||
fn visit_bool<E: Error>(self, value: bool) -> Result<i64, E> {
|
||||
Ok(i64::from(value))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> DeserializeAs<'de, f64> for FiniteF64 {
|
||||
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<f64, D::Error> {
|
||||
deserializer.deserialize_any(Self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Visitor<'de> for FiniteF64 {
|
||||
type Value = f64;
|
||||
|
||||
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str("a finite number")
|
||||
}
|
||||
|
||||
fn visit_i64<E: Error>(self, value: i64) -> Result<f64, E> {
|
||||
Ok(value as f64)
|
||||
}
|
||||
|
||||
fn visit_u64<E: Error>(self, value: u64) -> Result<f64, E> {
|
||||
Ok(value as f64)
|
||||
}
|
||||
|
||||
fn visit_f64<E: Error>(self, value: f64) -> Result<f64, E> {
|
||||
value
|
||||
.is_finite()
|
||||
.then_some(value)
|
||||
.ok_or_else(|| E::custom("expected a finite number"))
|
||||
}
|
||||
|
||||
fn visit_str<E: Error>(self, value: &str) -> Result<f64, E> {
|
||||
self.visit_f64(value.trim().parse::<f64>().map_err(E::custom)?)
|
||||
}
|
||||
|
||||
fn visit_bool<E: Error>(self, value: bool) -> Result<f64, E> {
|
||||
Ok(f64::from(value))
|
||||
}
|
||||
}
|
||||
|
||||
fn integer_string(value: &str) -> Option<i64> {
|
||||
let integer = match value.split_once('.') {
|
||||
Some((integer, fraction)) => {
|
||||
if fraction.is_empty() || !fraction.bytes().all(|byte| byte == b'0') {
|
||||
return None;
|
||||
}
|
||||
integer
|
||||
}
|
||||
None => value,
|
||||
};
|
||||
if integer.starts_with('_') || integer.ends_with('_') || integer.contains("__") {
|
||||
return None;
|
||||
}
|
||||
let digits = integer.strip_prefix(['+', '-']).unwrap_or(integer);
|
||||
if digits.is_empty()
|
||||
|| digits.starts_with('_')
|
||||
|| !digits
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit() || byte == b'_')
|
||||
{
|
||||
return None;
|
||||
}
|
||||
integer.replace('_', "").parse().ok()
|
||||
}
|
||||
|
||||
fn integral_float(value: f64) -> Option<i64> {
|
||||
(value.is_finite()
|
||||
&& value.fract() == 0.0
|
||||
&& value >= i64::MIN as f64
|
||||
&& value < -(i64::MIN as f64))
|
||||
.then_some(value as i64)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use serde_with::serde_as;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[serde_as]
|
||||
#[derive(Debug, Deserialize, Serialize, PartialEq)]
|
||||
struct Numbers {
|
||||
#[serde_as(deserialize_as = "Option<Vec<LaxI64>>")]
|
||||
integers: Option<Vec<i64>>,
|
||||
#[serde_as(deserialize_as = "Option<FiniteF64>")]
|
||||
float: Option<f64>,
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::trimmed_true(" True ", Some(true))]
|
||||
#[case::control_whitespace_true("\u{1c}TRUE\u{1f}", Some(true))]
|
||||
|
|
@ -160,73 +34,4 @@ mod tests {
|
|||
) {
|
||||
assert_eq!(parse_str_bool(input), expected, "{input:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn adapters_compose_and_serialize_as_numbers() {
|
||||
let numbers: Numbers = serde_json::from_value(json!({
|
||||
"integers": ["9007199254740993.0", "1_000", " +2.000 ", 3.0, true],
|
||||
"float": " 1.5 "
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(numbers).unwrap(),
|
||||
json!({
|
||||
"integers": [9_007_199_254_740_993_i64, 1000, 2, 3, 1], "float": 1.5
|
||||
})
|
||||
);
|
||||
for input in [json!({}), json!({"integers": null, "float": null})] {
|
||||
assert_eq!(
|
||||
serde_json::from_value::<Numbers>(input).unwrap(),
|
||||
Numbers {
|
||||
integers: None,
|
||||
float: None,
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn integer_bounds_and_invalid_values_are_checked() {
|
||||
for input in [
|
||||
json!(i64::MIN),
|
||||
json!(i64::MAX),
|
||||
json!(i64::MAX.to_string()),
|
||||
] {
|
||||
assert!(serde_json::from_value::<Numbers>(json!({"integers": [input]})).is_ok());
|
||||
}
|
||||
for input in [
|
||||
json!(u64::MAX),
|
||||
json!(9_223_372_036_854_775_808_u64),
|
||||
json!(9_223_372_036_854_775_808.0),
|
||||
json!("-9223372036854775809"),
|
||||
json!("1.0000000000000001"),
|
||||
json!("1e3"),
|
||||
json!("2."),
|
||||
json!(".0"),
|
||||
json!("_2"),
|
||||
json!("2__0"),
|
||||
json!(2.5),
|
||||
json!(null),
|
||||
json!({}),
|
||||
] {
|
||||
assert!(serde_json::from_value::<Numbers>(json!({"integers": [input]})).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn floats_reject_nonfinite_and_invalid_values() {
|
||||
for input in [
|
||||
json!("NaN"),
|
||||
json!("inf"),
|
||||
json!("-inf"),
|
||||
json!("1e999"),
|
||||
json!([]),
|
||||
] {
|
||||
assert!(serde_json::from_value::<Numbers>(json!({"float": input})).is_err());
|
||||
}
|
||||
for (input, expected) in [(json!(2), 2.0), (json!(2.5), 2.5), (json!(true), 1.0)] {
|
||||
let numbers: Numbers = serde_json::from_value(json!({"float": input})).unwrap();
|
||||
assert_eq!(numbers.float, Some(expected));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ litellm-core owns route orchestration. Messages and HTTP Responses return `litel
|
|||
|
||||
Hosts assemble route objects from shared `CoreResources`, HTTP settings, and secret sources. Each route owns its provider client and authentication dependencies. Gateway routes live for the gateway lifetime; Python assembles routes per call from its settings snapshot
|
||||
|
||||
Chat Completions, Messages, and OCR execute through their route objects. Calls pass `RouteHooks` directly; use `&()` when no hooks are needed. Construction does no work; preparation and lifecycle observation begin when the future is polled. Handlers accept `RouteHooks`, never a concrete `ChannelHooks`. Native observers receive start and terminal events through the shared call runner; a stream retains its lifecycle until exhaustion, error, or drop. A host channel has no native observer because its driver owns terminal dispatch
|
||||
Chat Completions, Messages, Responses, and OCR execute through their route objects. Calls pass `Interceptors` and an optional `ObservationSender` separately; use `&()` for no hooks and `None` for no observer. Construction does no work; preparation and lifecycle observation begin when the future is polled. Handlers accept `Interceptors`, never a concrete `ChannelInterceptors`. Native observers receive start and terminal events through the shared call runner; a stream retains its lifecycle until exhaustion, error, or drop. Hosted routes leave terminal observation to their driver
|
||||
|
||||
`route.rs` declares the concrete `Protocol` and implements a route method that accepts a typed request and constructs a `litellm_host::call::HostedMachine` with `hosted_call`. The shared call plumbing owns stream opening, delivery, backpressure, and detachment. Request decoding belongs to the boundary before the machine starts. Route closures only supply execution dependencies and route-specific host capabilities such as an OCR token provider. Use `run_hosted` for a native host so detachment is reported as cancellation. Python uses its own shared driver and preserves caller-task callback execution
|
||||
|
||||
|
|
@ -10,17 +10,17 @@ Responses WebSocket sessions remain separate from the HTTP call driver because a
|
|||
|
||||
## Crate layering
|
||||
|
||||
For Messages, Responses, Chat Completions, OCR, and other API formats, `core/src/<format>/` owns orchestration. Shared API data contracts belong in `litellm-types`, adapter contracts and shared transformation machinery in `llms/src/base_llm/<format>/`, and provider policy in `llms/src/<provider>/<format>/`. A repeated format directory name does not imply interchangeable responsibilities. Select concrete adapters here, then invoke their contracts instead of applying one provider's policy to every call. Route types describe call envelopes and execution state, not duplicate public payload schemas
|
||||
For Messages, Responses, Chat Completions, OCR, and other API formats, `core/src/<format>/` owns orchestration. Shared API data contracts belong in `litellm-llms-types`, adapter contracts and shared transformation machinery in `llms/src/base_llm/<format>/`, and provider policy in `llms/src/<provider>/<format>/`. A repeated format directory name does not imply interchangeable responsibilities. Select concrete adapters here, then invoke their contracts instead of applying one provider's policy to every call. Route types describe call envelopes and execution state, not duplicate public payload schemas
|
||||
|
||||
Each crate mirrors one top-level Python package, so a Rust path reads as its Python path with the crate name in place of the package directory. Dependencies only point down:
|
||||
Crates separate API data, transformations, transport, and orchestration. Python package names identify counterparts, not ownership. Dependencies only point down:
|
||||
|
||||
- `litellm-types` mirrors `litellm/types/`: pure serde data, no I/O
|
||||
- `litellm-llms-types` owns shared inference API contracts, grouped by format: pure serde data and shape validation, no I/O
|
||||
- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments, settings lookup and layer merge), no network I/O
|
||||
- `litellm-http` is Rust-only and route-neutral: settings resolution, the pooled `reqwest` clients, TLS, proxies, the SSRF-safe media fetcher, request and header helpers, and transport errors. Python's `litellm/llms/custom_httpx/` is split by responsibility instead of mirrored: its transport half lives here, its OCR handler in `litellm-llms`
|
||||
- `litellm-llms` mirrors `litellm/llms/`: `base_llm/<api>/transformation.rs`, `<provider>/<api>/transformation.rs`, and `base_llm/ocr/handler.rs` (the OCR request handler)
|
||||
- `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks
|
||||
|
||||
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, the provider-level hooks OCR implements over its host until it folds into `litellm_host::hooks::RouteHooks`. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
|
||||
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, the provider-level hooks OCR implements over its host until it folds into `litellm_host::interceptors::Interceptors`. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
|
||||
|
||||
## Error placement
|
||||
|
||||
|
|
@ -33,3 +33,17 @@ Scope follows the concept, not the first caller. An error type under `litellm-ll
|
|||
`litellm_llms::Error` (`crates/llms/src/error.rs`) is the one transformation error for every provider and API. `base_llm/ocr/error.rs` is the recorded exception until OCR folds into it
|
||||
|
||||
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business.
|
||||
|
||||
## Response caching and accounting boundary
|
||||
|
||||
Attach a `litellm_cache_response::ScopedCache` with `route.with_cache(cache)`. Cached and uncached routes use the same `execute` and `machine` methods. `CallOptions` carries a scope-free `CachePolicy` and observation; per-call policy never replaces the attached scope or service
|
||||
|
||||
Messages groups per-call dependencies in `CallContext` and explicitly sequences cache lookup, provider execution, result acceptance, and cache storage. Provider transport does not own cache orchestration. Stream capture remains in the shared cache implementation
|
||||
|
||||
Core owns request identity, typed response reconstruction and stream capture/replay. `cache-response` owns cache policy, namespacing, scope encoding, versioned envelopes and freshness. The SDK explicitly chooses shared scope. The gateway derives isolated scope from authenticated identity before attaching its service
|
||||
|
||||
Core delivers `ExecutionFacts` through the awaited `ResultReady` host operation for both provider and cached results, before public response processing or stream opening. Facts carry resolved model/provider and result source, including the hit key. Usage remains in the typed response or delivered stream, where completion and cancellation determine what was actually reported. Passive observation is not an accounting delivery mechanism
|
||||
|
||||
Core does not calculate prices, charge budgets, or update rate-limit counters. The legacy Python callback adapter translates execution facts into the existing Python logging contract; Python remains the accounting owner on that path. Native gateway accounting belongs to gateway dependencies, independently of `host-python`. Response-cache services expose no coordination counters or reservation APIs. A shared Redis deployment does not make response storage and accounting coordination the same dependency
|
||||
|
||||
Cache lookup follows provider preparation, credential resolution and the request interceptor. Keys describe the effective provider URL, authenticated headers and rewritten body. Signed requests bypass caching until the signing identity has a stable cache representation
|
||||
|
|
|
|||
|
|
@ -6,8 +6,12 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
litellm-framing.workspace = true
|
||||
tokio-util = { version = "0.7", features = ["codec"] }
|
||||
litellm-secrets.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-llms-types.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
litellm-host.workspace = true
|
||||
bytes.workspace = true
|
||||
|
|
@ -36,8 +40,10 @@ url.workspace = true
|
|||
veil.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-auth-gcp.workspace = true
|
||||
litellm-host-native.workspace = true
|
||||
litellm-llms = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
rstest_reuse.workspace = true
|
||||
|
|
|
|||
376
litellm-rust/crates/core/src/caching.rs
Normal file
|
|
@ -0,0 +1,376 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
marker::PhantomData,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use litellm_cache_response::{
|
||||
CacheOptions, CachePolicy, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
|
||||
ScopedCache, cache_key,
|
||||
};
|
||||
use litellm_host::{
|
||||
call::{CallOutput, OutputOf},
|
||||
interceptors::{ExecutionFacts, Interceptors, ProviderIdentity, ResultSource, WireRequest},
|
||||
lifecycle::{CallEvent, ExecutionEvent},
|
||||
observation::ObservationSender,
|
||||
protocol::Protocol,
|
||||
};
|
||||
use serde::{Deserialize, Serialize, de::DeserializeOwned};
|
||||
use serde_json::Value;
|
||||
use tokio_util::codec::Decoder;
|
||||
|
||||
use crate::RouteError;
|
||||
|
||||
pub trait Cachable: Protocol<Error = RouteError> {
|
||||
const SURFACE: &'static str;
|
||||
|
||||
fn reusable(_response: &Self::Response) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CacheRequest {
|
||||
pub identity: ProviderIdentity,
|
||||
pub input: Value,
|
||||
}
|
||||
|
||||
impl CacheRequest {
|
||||
pub fn from_wire(identity: ProviderIdentity, wire: Option<&WireRequest>) -> Self {
|
||||
Self {
|
||||
input: wire.map_or(Value::Null, |wire| {
|
||||
serde_json::json!({
|
||||
"provider": identity.provider,
|
||||
"model": identity.model,
|
||||
"url": wire.url,
|
||||
"headers": wire.headers,
|
||||
"body": wire.body,
|
||||
})
|
||||
}),
|
||||
identity,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait StreamCachable: Cachable {
|
||||
const TERMINAL_EVENT: &'static str;
|
||||
|
||||
fn replay(data: Bytes) -> Option<OutputOf<Self>>;
|
||||
fn bytes(chunk: &Self::Chunk) -> &[u8];
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", content = "value")]
|
||||
pub enum CachedOutput<R> {
|
||||
Response(R),
|
||||
Stream(String),
|
||||
}
|
||||
|
||||
struct CacheSession {
|
||||
service: Arc<dyn ResponseCacheService>,
|
||||
request: ResponseCacheRequest,
|
||||
}
|
||||
|
||||
impl CacheSession {
|
||||
fn prepare<P: Cachable>(
|
||||
service: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
request: &CacheRequest,
|
||||
) -> Option<Self> {
|
||||
let options = options.filter(|options| options.policy.enabled())?;
|
||||
let service = service?;
|
||||
let input = request.input.clone();
|
||||
let request = options.request(&service.config().namespace, P::SURFACE, input);
|
||||
Some(Self { service, request })
|
||||
}
|
||||
|
||||
async fn lookup<P: Cachable>(&self) -> Option<CachedOutput<P::Response>>
|
||||
where
|
||||
P::Response: DeserializeOwned,
|
||||
{
|
||||
if !self.request.controls.reads() {
|
||||
return None;
|
||||
}
|
||||
match self.service.lookup(&self.request, now()).await {
|
||||
Ok(Some(value)) => {
|
||||
serde_json::from_value::<ResponseEnvelope<CachedOutput<P::Response>>>(value)
|
||||
.ok()
|
||||
.and_then(|entry| entry.decode(P::SURFACE))
|
||||
}
|
||||
Ok(None) => None,
|
||||
Err(_) => {
|
||||
tracing::warn!("response cache lookup failed");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn store(&self, entry: Value) {
|
||||
if !self.request.controls.writes() {
|
||||
return;
|
||||
}
|
||||
if self
|
||||
.service
|
||||
.store(&self.request, entry, now())
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
tracing::warn!("response cache write failed");
|
||||
}
|
||||
}
|
||||
|
||||
async fn store_response<P: Cachable>(&self, response: &P::Response)
|
||||
where
|
||||
P::Response: Serialize,
|
||||
{
|
||||
if !self.request.controls.writes() || !P::reusable(response) {
|
||||
return;
|
||||
}
|
||||
if let Ok(value) = serde_json::to_value(response)
|
||||
&& let Ok(entry) = serde_json::to_value(ResponseEnvelope::new(
|
||||
P::SURFACE,
|
||||
CachedOutput::Response(value),
|
||||
))
|
||||
{
|
||||
self.store(entry).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute_unary<P, F, Fut>(
|
||||
request: CacheRequest,
|
||||
cache: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
interceptors: &impl Interceptors<RouteError>,
|
||||
observers: Option<&ObservationSender>,
|
||||
provider: F,
|
||||
) -> Result<P::Response, RouteError>
|
||||
where
|
||||
P: Cachable,
|
||||
P::Response: Serialize + DeserializeOwned,
|
||||
F: FnOnce() -> Fut,
|
||||
Fut: Future<Output = Result<P::Response, RouteError>>,
|
||||
{
|
||||
let identity = request.identity.clone();
|
||||
crate::diagnostic::provider(&identity.model, &identity.provider);
|
||||
let session = CacheSession::prepare::<P>(cache, options, &request);
|
||||
let hit = match &session {
|
||||
Some(session) => session.lookup::<P>().await.and_then(|entry| match entry {
|
||||
CachedOutput::Response(response) => Some((response, cache_key(&session.request.key))),
|
||||
CachedOutput::Stream(_) => None,
|
||||
}),
|
||||
None => None,
|
||||
};
|
||||
let (response, source) = match hit {
|
||||
Some((response, key)) => (response, ResultSource::Cache { key }),
|
||||
None => (provider().await?, ResultSource::Provider),
|
||||
};
|
||||
let from_provider = source == ResultSource::Provider;
|
||||
publish(
|
||||
ExecutionFacts {
|
||||
provider: identity,
|
||||
source,
|
||||
},
|
||||
interceptors,
|
||||
observers,
|
||||
)
|
||||
.await?;
|
||||
if from_provider && let Some(session) = session {
|
||||
session.store_response::<P>(&response).await;
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
pub async fn execute_streaming<P, F, Fut>(
|
||||
request: CacheRequest,
|
||||
cache: Option<Arc<dyn ResponseCacheService>>,
|
||||
options: Option<CacheOptions>,
|
||||
interceptors: &impl Interceptors<RouteError>,
|
||||
observers: Option<&ObservationSender>,
|
||||
provider: F,
|
||||
) -> Result<OutputOf<P>, RouteError>
|
||||
where
|
||||
P: StreamCachable,
|
||||
P::Response: Serialize + DeserializeOwned,
|
||||
F: FnOnce() -> Fut,
|
||||
Fut: Future<Output = Result<OutputOf<P>, RouteError>>,
|
||||
{
|
||||
let identity = request.identity.clone();
|
||||
crate::diagnostic::provider(&identity.model, &identity.provider);
|
||||
let session = CacheSession::prepare::<P>(cache, options, &request);
|
||||
let cache = CallCache::<P> {
|
||||
session,
|
||||
protocol: PhantomData,
|
||||
};
|
||||
let hit = cache.lookup().await;
|
||||
let (output, source) = match hit {
|
||||
Some(hit) => hit,
|
||||
None => (provider().await?, ResultSource::Provider),
|
||||
};
|
||||
publish(
|
||||
ExecutionFacts {
|
||||
provider: identity,
|
||||
source: source.clone(),
|
||||
},
|
||||
interceptors,
|
||||
observers,
|
||||
)
|
||||
.await?;
|
||||
Ok(cache.finish(output, &source).await)
|
||||
}
|
||||
|
||||
pub(crate) struct CallCache<P> {
|
||||
session: Option<CacheSession>,
|
||||
protocol: PhantomData<P>,
|
||||
}
|
||||
|
||||
impl<P: StreamCachable> CallCache<P> {
|
||||
pub(crate) fn from_wire(
|
||||
cache: Option<&ScopedCache>,
|
||||
policy: CachePolicy,
|
||||
identity: &ProviderIdentity,
|
||||
wire: &WireRequest,
|
||||
) -> Self {
|
||||
let session = cache.and_then(|cache| {
|
||||
if !policy.enabled() {
|
||||
return None;
|
||||
}
|
||||
let options = cache.options(Some(policy));
|
||||
let request = CacheRequest::from_wire(identity.clone(), Some(wire));
|
||||
Some(CacheSession {
|
||||
request: options.request(
|
||||
&cache.service.config().namespace,
|
||||
P::SURFACE,
|
||||
request.input,
|
||||
),
|
||||
service: cache.service.clone(),
|
||||
})
|
||||
});
|
||||
Self {
|
||||
session,
|
||||
protocol: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn lookup(&self) -> Option<(OutputOf<P>, ResultSource)>
|
||||
where
|
||||
P::Response: DeserializeOwned,
|
||||
{
|
||||
let session = self.session.as_ref()?;
|
||||
let output = match session.lookup::<P>().await? {
|
||||
CachedOutput::Response(response) => CallOutput::Complete(response),
|
||||
CachedOutput::Stream(data) => P::replay(Bytes::from(data))?,
|
||||
};
|
||||
Some((
|
||||
output,
|
||||
ResultSource::Cache {
|
||||
key: cache_key(&session.request.key),
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn finish(self, output: OutputOf<P>, source: &ResultSource) -> OutputOf<P>
|
||||
where
|
||||
P::Response: Serialize,
|
||||
{
|
||||
let Some(session) = self.session.filter(|session| {
|
||||
*source == ResultSource::Provider && session.request.controls.writes()
|
||||
}) else {
|
||||
return output;
|
||||
};
|
||||
match output {
|
||||
CallOutput::Complete(response) => {
|
||||
session.store_response::<P>(&response).await;
|
||||
CallOutput::Complete(response)
|
||||
}
|
||||
CallOutput::Stream { head, chunks } => CallOutput::Stream {
|
||||
head,
|
||||
chunks: capture_stream::<P>(chunks, session),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn capture_stream<P: StreamCachable>(
|
||||
chunks: futures_util::stream::BoxStream<'static, Result<P::Chunk, RouteError>>,
|
||||
session: CacheSession,
|
||||
) -> futures_util::stream::BoxStream<'static, Result<P::Chunk, RouteError>> {
|
||||
stream::try_unfold(
|
||||
(chunks, Some(Vec::<u8>::new()), session),
|
||||
|(mut chunks, captured, session)| async move {
|
||||
match chunks.try_next().await? {
|
||||
Some(chunk) => {
|
||||
let captured = captured.and_then(|mut data| {
|
||||
let bytes = P::bytes(&chunk);
|
||||
if data.len().saturating_add(bytes.len())
|
||||
> session.service.config().max_entry_bytes
|
||||
{
|
||||
return None;
|
||||
}
|
||||
data.extend_from_slice(bytes);
|
||||
Some(data)
|
||||
});
|
||||
Ok(Some((chunk, (chunks, captured, session))))
|
||||
}
|
||||
None => {
|
||||
if let Some(data) = captured
|
||||
&& let Ok(text) = String::from_utf8(data)
|
||||
&& successful_stream(&text, P::TERMINAL_EVENT)
|
||||
&& let Ok(entry) = serde_json::to_value(ResponseEnvelope::new(
|
||||
P::SURFACE,
|
||||
CachedOutput::<Value>::Stream(text),
|
||||
))
|
||||
{
|
||||
session.store(entry).await;
|
||||
}
|
||||
Ok::<_, RouteError>(None)
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn now() -> Duration {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn successful_stream(text: &str, terminal: &str) -> bool {
|
||||
let mut pending = BytesMut::from(text.as_bytes());
|
||||
let mut codec = litellm_framing::sse::SseCodec::default();
|
||||
let mut complete = false;
|
||||
loop {
|
||||
let event = match codec.decode(&mut pending) {
|
||||
Ok(Some(event)) => event,
|
||||
Ok(None) => return complete && pending.is_empty(),
|
||||
Err(_) => return false,
|
||||
};
|
||||
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
|
||||
return false;
|
||||
};
|
||||
let Some(kind) = value.get("type").and_then(Value::as_str) else {
|
||||
return false;
|
||||
};
|
||||
if matches!(kind, "error" | "response.failed" | "response.incomplete") {
|
||||
return false;
|
||||
}
|
||||
complete |= kind == terminal;
|
||||
}
|
||||
}
|
||||
|
||||
async fn publish(
|
||||
facts: ExecutionFacts,
|
||||
interceptors: &impl Interceptors<RouteError>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<(), RouteError> {
|
||||
if let Some(observers) = observers {
|
||||
observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady {
|
||||
facts: facts.clone(),
|
||||
}));
|
||||
}
|
||||
interceptors.result_ready(facts).await
|
||||
}
|
||||
|
|
@ -1,16 +1,14 @@
|
|||
use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender};
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_host::{
|
||||
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
|
||||
hooks::RouteHooks,
|
||||
};
|
||||
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
|
||||
use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body};
|
||||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, resolve_auth},
|
||||
chat::transformation::ProviderChatResponseData,
|
||||
};
|
||||
use litellm_types::utils::ChatCompletionsResponse;
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::Error;
|
||||
|
|
@ -23,7 +21,10 @@ pub(super) async fn execute(
|
|||
http: &Client,
|
||||
auth: &AuthServices,
|
||||
request: ProviderChatCompletionsRequest,
|
||||
hooks: &impl RouteHooks<Error>,
|
||||
cache: Option<litellm_cache_response::ScopedCache>,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let ProviderChatCompletionsRequest {
|
||||
model,
|
||||
|
|
@ -45,7 +46,11 @@ pub(super) async fn execute(
|
|||
api_key,
|
||||
};
|
||||
let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?;
|
||||
let wire = hooks
|
||||
let identity = litellm_host::interceptors::ProviderIdentity {
|
||||
model: context.model.clone(),
|
||||
provider: context.custom_llm_provider.clone(),
|
||||
};
|
||||
let wire = interceptors
|
||||
.before_provider_request(
|
||||
WireRequest {
|
||||
url,
|
||||
|
|
@ -55,55 +60,72 @@ pub(super) async fn execute(
|
|||
context,
|
||||
)
|
||||
.await?;
|
||||
let outbound = outbound_request(
|
||||
Authenticated {
|
||||
headers: wire.headers,
|
||||
signer: authenticated.signer,
|
||||
let cache = cache.filter(|_| authenticated.signer.is_none());
|
||||
let cache_request =
|
||||
crate::caching::CacheRequest::from_wire(identity, cache.as_ref().map(|_| &wire));
|
||||
crate::caching::execute_unary::<super::route::ChatCompletions, _, _>(
|
||||
cache_request,
|
||||
cache.as_ref().map(|cache| cache.service.clone()),
|
||||
cache.as_ref().map(|cache| cache.options(cache_options)),
|
||||
interceptors,
|
||||
observers,
|
||||
|| async move {
|
||||
let outbound = outbound_request(
|
||||
Authenticated {
|
||||
headers: wire.headers,
|
||||
signer: authenticated.signer,
|
||||
},
|
||||
wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
)?;
|
||||
|
||||
let response = crate::outbound::send(outbound, http).await.map_err(|err| {
|
||||
// Failing to establish the connection means the request never went out,
|
||||
// so the host can still serve it. Everything else here, a timeout
|
||||
// above all, may have reached the provider and been answered.
|
||||
if err.is_connect() || err.is_builder() {
|
||||
Error::Transport(litellm_http::transport::Error::Connect(err.to_string()))
|
||||
} else {
|
||||
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
|
||||
}
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response.text().await.map_err(|err| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
|
||||
})?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(Error::Transport(litellm_http::transport::Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
}));
|
||||
}
|
||||
let raw = RawResponse { body: text.clone() };
|
||||
if let Some(observers) = observers {
|
||||
observers.emit(litellm_host::lifecycle::CallEvent::Execution(
|
||||
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
|
||||
));
|
||||
}
|
||||
interceptors
|
||||
.after_provider_response(raw)
|
||||
.await
|
||||
.map_err(Error::post_call)?;
|
||||
|
||||
let body: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
|
||||
"chat completions response JSON",
|
||||
err,
|
||||
))
|
||||
})?;
|
||||
config
|
||||
.transform_response(&model, ProviderChatResponseData { body })
|
||||
.map_err(Error::from)
|
||||
.map_err(as_response_error)
|
||||
},
|
||||
wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
)?;
|
||||
|
||||
let response = crate::outbound::send(outbound, http).await.map_err(|err| {
|
||||
// Failing to establish the connection means the request never went out,
|
||||
// so the host can still serve it. Everything else here, a timeout
|
||||
// above all, may have reached the provider and been answered.
|
||||
if err.is_connect() || err.is_builder() {
|
||||
Error::Transport(litellm_http::transport::Error::Connect(err.to_string()))
|
||||
} else {
|
||||
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
|
||||
}
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response.text().await.map_err(|err| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
|
||||
})?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(Error::Transport(litellm_http::transport::Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
}));
|
||||
}
|
||||
hooks
|
||||
.on_event(MachineEvent::ResponseReceived {
|
||||
raw: RawResponse { body: text.clone() },
|
||||
})
|
||||
.await
|
||||
.map_err(Error::post_call)?;
|
||||
|
||||
let body: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
|
||||
"chat completions response JSON",
|
||||
err,
|
||||
))
|
||||
})?;
|
||||
config
|
||||
.transform_response(&model, ProviderChatResponseData { body })
|
||||
.map_err(Error::from)
|
||||
.map_err(as_response_error)
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Re-tag an error raised while normalizing a response the provider already
|
||||
|
|
@ -168,7 +190,7 @@ mod tests {
|
|||
raw: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl RouteHooks<Error> for RecordingHooks {
|
||||
impl Interceptors<Error> for RecordingHooks {
|
||||
async fn before_provider_request(
|
||||
&self,
|
||||
wire: WireRequest,
|
||||
|
|
@ -188,8 +210,7 @@ mod tests {
|
|||
})
|
||||
}
|
||||
|
||||
async fn on_event(&self, event: MachineEvent) -> Result<(), Error> {
|
||||
let MachineEvent::ResponseReceived { raw } = event;
|
||||
async fn after_provider_response(&self, raw: RawResponse) -> Result<(), Error> {
|
||||
self.raw.lock().unwrap().push(raw.body);
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -223,13 +244,16 @@ mod tests {
|
|||
)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let hooks = RecordingHooks::default();
|
||||
let interceptors = RecordingHooks::default();
|
||||
|
||||
execute(
|
||||
&Client::plain_for_test(),
|
||||
&AuthServices::default(),
|
||||
prepared(&upstream.uri()),
|
||||
&hooks,
|
||||
None,
|
||||
None,
|
||||
&interceptors,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("chat completions call succeeds");
|
||||
|
|
@ -240,14 +264,17 @@ mod tests {
|
|||
assert_eq!(sent["system"], "added by the host");
|
||||
assert_eq!(request.headers["x-host"], "seen");
|
||||
assert_eq!(request.headers["x-api-key"], "sk-test");
|
||||
let [context] = <[RequestContext; 1]>::try_from(hooks.contexts.into_inner().unwrap())
|
||||
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
|
||||
let [context] =
|
||||
<[RequestContext; 1]>::try_from(interceptors.contexts.into_inner().unwrap())
|
||||
.unwrap_or_else(|seen| {
|
||||
panic!("before_provider_request runs once, saw {}", seen.len())
|
||||
});
|
||||
assert_eq!(
|
||||
(context.model.as_str(), context.custom_llm_provider.as_str()),
|
||||
("claude-sonnet-4-5", "anthropic")
|
||||
);
|
||||
assert_eq!(context.optional_params, json!({"max_tokens": 16}));
|
||||
assert_eq!(hooks.raw.into_inner().unwrap(), [ANTHROPIC_MESSAGE]);
|
||||
assert_eq!(interceptors.raw.into_inner().unwrap(), [ANTHROPIC_MESSAGE]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -258,13 +285,16 @@ mod tests {
|
|||
.respond_with(ResponseTemplate::new(500).set_body_string("boom"))
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let hooks = RecordingHooks::default();
|
||||
let interceptors = RecordingHooks::default();
|
||||
|
||||
let error = execute(
|
||||
&Client::plain_for_test(),
|
||||
&AuthServices::default(),
|
||||
prepared(&upstream.uri()),
|
||||
&hooks,
|
||||
None,
|
||||
None,
|
||||
&interceptors,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect_err("the upstream failure fails the call");
|
||||
|
|
@ -273,7 +303,7 @@ mod tests {
|
|||
error,
|
||||
Error::Transport(litellm_http::transport::Error::Http { status: 500, .. })
|
||||
));
|
||||
assert!(hooks.raw.into_inner().unwrap().is_empty());
|
||||
assert!(interceptors.raw.into_inner().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
use litellm_host::observation::ObservationSender;
|
||||
pub mod route;
|
||||
pub mod types;
|
||||
pub use crate::error::RouteError as Error;
|
||||
mod common_utils;
|
||||
pub(crate) mod handler;
|
||||
mod prepare;
|
||||
use litellm_types::utils::ChatCompletionsResponse;
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
|
||||
use prepare::{prepare_provider_request, resolve_request};
|
||||
|
||||
use crate::chat_completions::types::ChatCompletionsRequest;
|
||||
|
|
@ -17,6 +18,7 @@ pub struct ChatCompletionsRoute {
|
|||
http: litellm_http::Client,
|
||||
auth: Arc<AuthServices>,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
cache: Option<litellm_cache_response::ScopedCache>,
|
||||
}
|
||||
|
||||
impl ChatCompletionsRoute {
|
||||
|
|
@ -29,44 +31,63 @@ impl ChatCompletionsRoute {
|
|||
http,
|
||||
auth,
|
||||
secrets,
|
||||
cache: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_cache(self, cache: litellm_cache_response::ScopedCache) -> Self {
|
||||
Self {
|
||||
cache: Some(cache),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(
|
||||
&self,
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
options: impl Into<crate::CallOptions>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
litellm_host::lifecycle::observe_unary(hooks.observer(), self.run(request, hooks)).await
|
||||
let crate::CallOptions {
|
||||
cache: cache_options,
|
||||
observers,
|
||||
} = options.into();
|
||||
litellm_host::lifecycle::observe_unary(
|
||||
observers.clone(),
|
||||
self.run_call(
|
||||
request.into(),
|
||||
cache_options,
|
||||
interceptors,
|
||||
observers.as_ref(),
|
||||
),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "chat_completions",
|
||||
model = %request.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
async fn run(
|
||||
&self,
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
crate::diagnostic::unary(async {
|
||||
let resolved = resolve_request(request)?;
|
||||
let snapshot = self
|
||||
.secrets
|
||||
.resolve(&resolved.config.secret_names())
|
||||
.await?;
|
||||
let prepared = prepare_provider_request(resolved, snapshot)?;
|
||||
crate::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
|
||||
let execute: futures_util::future::BoxFuture<
|
||||
'_,
|
||||
Result<ChatCompletionsResponse, Error>,
|
||||
> = Box::pin(handler::execute(&self.http, &self.auth, prepared, hooks));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
let resolved = resolve_request(request)?;
|
||||
let snapshot = self
|
||||
.secrets
|
||||
.resolve(&resolved.config.secret_names())
|
||||
.await?;
|
||||
let prepared = prepare_provider_request(resolved, snapshot)?;
|
||||
crate::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<ChatCompletionsResponse, Error>> =
|
||||
Box::pin(handler::execute(
|
||||
&self.http,
|
||||
&self.auth,
|
||||
prepared,
|
||||
self.cache.clone(),
|
||||
cache_options,
|
||||
interceptors,
|
||||
observers,
|
||||
));
|
||||
execute.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@ use litellm_auth::SecretValue;
|
|||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_http::request::with_default_headers;
|
||||
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
|
||||
use litellm_llms_types::formats::chat_completions::ChatMessage;
|
||||
use litellm_secrets::source::Secrets;
|
||||
use litellm_types::llms::openai::ChatMessage;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
use litellm_host::observation::ObservationSender;
|
||||
use std::convert::Infallible;
|
||||
|
||||
use litellm_host::{
|
||||
call::{CallOutput, HostedMachine, hosted_call},
|
||||
protocol::Protocol,
|
||||
};
|
||||
use litellm_types::utils::ChatCompletionsResponse;
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
|
||||
|
||||
use super::{
|
||||
ChatCompletionsRoute, Error,
|
||||
|
|
@ -23,22 +24,59 @@ impl Protocol for ChatCompletions {
|
|||
}
|
||||
|
||||
impl ChatCompletionsRoute {
|
||||
pub fn machine(self, call: ChatCompletionsCall) -> HostedMachine<ChatCompletions> {
|
||||
pub fn machine(
|
||||
self,
|
||||
call: ChatCompletionsCall,
|
||||
options: impl Into<crate::CallOptions>,
|
||||
) -> HostedMachine<ChatCompletions> {
|
||||
let crate::CallOptions {
|
||||
cache: cache_options,
|
||||
observers,
|
||||
} = options.into();
|
||||
hosted_call(
|
||||
call,
|
||||
move |call: ChatCompletionsCall, _, hooks| async move {
|
||||
let request = ChatCompletionsRequest {
|
||||
model: &call.model,
|
||||
messages: call.messages,
|
||||
optional_params: call.optional_params,
|
||||
api_key: call.api_key.as_deref(),
|
||||
api_base: call.api_base.as_deref(),
|
||||
custom_llm_provider: call.custom_llm_provider.as_deref(),
|
||||
extra_headers: call.extra_headers,
|
||||
timeout: call.timeout,
|
||||
};
|
||||
self.run(request, &hooks).await.map(CallOutput::Complete)
|
||||
observers,
|
||||
move |call, _, interceptors, observers| async move {
|
||||
self.run_call(call, cache_options, &interceptors, observers.as_ref())
|
||||
.await
|
||||
.map(CallOutput::Complete)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "chat_completions",
|
||||
model = %call.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
pub(super) async fn run_call(
|
||||
&self,
|
||||
call: ChatCompletionsCall,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
crate::diagnostic::unary(async {
|
||||
let request = ChatCompletionsRequest {
|
||||
model: &call.model,
|
||||
messages: call.messages,
|
||||
optional_params: call.optional_params,
|
||||
api_key: call.api_key.as_deref(),
|
||||
api_base: call.api_base.as_deref(),
|
||||
custom_llm_provider: call.custom_llm_provider.as_deref(),
|
||||
extra_headers: call.extra_headers,
|
||||
timeout: call.timeout,
|
||||
};
|
||||
self.run(request, cache_options, interceptors, observers)
|
||||
.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::caching::Cachable for ChatCompletions {
|
||||
const SURFACE: &'static str = "chat_completions";
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ use std::time::Duration;
|
|||
|
||||
use litellm_auth::SecretValue;
|
||||
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
|
||||
use litellm_types::llms::openai::ChatMessage;
|
||||
use litellm_llms_types::formats::chat_completions::ChatMessage;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// A `/chat/completions` call as it crosses into the core.
|
||||
|
|
|
|||
48
litellm-rust/crates/core/src/context.rs
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
use litellm_cache_response::CachePolicy;
|
||||
use litellm_host::{
|
||||
interceptors::{ExecutionFacts, Interceptors, RawResponse},
|
||||
lifecycle::{CallEvent, ExecutionEvent},
|
||||
observation::ObservationSender,
|
||||
};
|
||||
|
||||
use crate::{CallOptions, RouteError};
|
||||
|
||||
pub(crate) struct CallContext<'a, I> {
|
||||
pub interceptors: &'a I,
|
||||
pub observers: Option<ObservationSender>,
|
||||
pub cache: CachePolicy,
|
||||
}
|
||||
|
||||
impl<'a, I: Interceptors<RouteError>> CallContext<'a, I> {
|
||||
pub fn new(interceptors: &'a I, options: CallOptions) -> Self {
|
||||
Self {
|
||||
interceptors,
|
||||
observers: options.observers,
|
||||
cache: options.cache.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> {
|
||||
if let Some(observers) = &self.observers {
|
||||
observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady {
|
||||
facts: facts.clone(),
|
||||
}));
|
||||
}
|
||||
self.interceptors.result_ready(facts).await
|
||||
}
|
||||
|
||||
pub async fn response_received(&self, body: &str) -> Result<(), RouteError> {
|
||||
let raw = RawResponse {
|
||||
body: body.to_owned(),
|
||||
};
|
||||
if let Some(observers) = &self.observers {
|
||||
observers.emit(CallEvent::Execution(
|
||||
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
|
||||
));
|
||||
}
|
||||
self.interceptors
|
||||
.after_provider_response(raw)
|
||||
.await
|
||||
.map_err(RouteError::post_call)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
mod context;
|
||||
mod diagnostic;
|
||||
|
||||
pub mod audio_transcription;
|
||||
pub mod caching;
|
||||
pub mod chat_completions;
|
||||
pub mod constants;
|
||||
pub mod error;
|
||||
|
|
@ -12,3 +14,27 @@ pub mod resources;
|
|||
pub mod responses;
|
||||
|
||||
pub use error::RouteError;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct CallOptions {
|
||||
pub cache: Option<litellm_cache_response::CachePolicy>,
|
||||
pub observers: Option<litellm_host::observation::ObservationSender>,
|
||||
}
|
||||
|
||||
impl From<Option<litellm_host::observation::ObservationSender>> for CallOptions {
|
||||
fn from(observers: Option<litellm_host::observation::ObservationSender>) -> Self {
|
||||
Self {
|
||||
cache: None,
|
||||
observers,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<litellm_cache_response::CachePolicy> for CallOptions {
|
||||
fn from(cache: litellm_cache_response::CachePolicy) -> Self {
|
||||
Self {
|
||||
cache: Some(cache),
|
||||
observers: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
This directory owns provider-independent Messages call orchestration: the entrypoint, call envelopes, provider selection, credential resolution, transport coordination, hooks, and stream lifecycle. Shared API data contracts belong in `litellm-types::messages`, adapter contracts and execution inputs in `llms/src/base_llm/messages`, and provider implementations in `llms/src/<provider>/messages`
|
||||
This directory owns provider-independent Messages call orchestration: the entrypoint, call envelopes, provider selection, credential resolution, transport coordination, hooks, and stream lifecycle. Shared API data contracts belong in `litellm-llms-types::formats::messages`, adapter contracts and execution inputs in `llms/src/base_llm/messages`, and provider implementations in `llms/src/<provider>/messages`
|
||||
|
||||
Select concrete provider adapters and invoke their contracts. Delegate authentication policy, beta selection, payload rewriting, and response interpretation to those adapters. Keep provider policy out of request preparation and transport handlers. Calling a concrete provider helper for every provider is still a policy dependency
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ pub(super) use litellm_http::request::truncate_error_body;
|
|||
use litellm_llms::{
|
||||
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
|
||||
azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
|
||||
base_llm::messages::transformation::BaseAnthropicMessagesConfig,
|
||||
base_llm::messages::transformation::BaseMessagesConfig,
|
||||
bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
|
@ -30,7 +30,7 @@ impl MessagesProvider {
|
|||
.into()
|
||||
}
|
||||
|
||||
pub(crate) fn config(self) -> &'static dyn BaseAnthropicMessagesConfig {
|
||||
pub(crate) fn config(self) -> &'static dyn BaseMessagesConfig {
|
||||
match self {
|
||||
Self::Anthropic => &ANTHROPIC_MESSAGES_CONFIG,
|
||||
Self::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG,
|
||||
|
|
|
|||
|
|
@ -2,95 +2,144 @@ use std::time::Duration;
|
|||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_host::{
|
||||
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
|
||||
hooks::RouteHooks,
|
||||
};
|
||||
use litellm_host::interceptors::{Interceptors, ProviderIdentity, RequestContext, WireRequest};
|
||||
use litellm_http::transport::Error as TransportError;
|
||||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, resolve_auth},
|
||||
messages::{
|
||||
streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
|
||||
transformation::BaseAnthropicMessagesConfig,
|
||||
transformation::BaseMessagesConfig,
|
||||
},
|
||||
};
|
||||
use litellm_llms_types::formats::messages::MessagesResponse;
|
||||
use litellm_tracing::ByteChunk;
|
||||
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
Error, MessagesResponse, common_utils::truncate_error_body, prepare::ProviderMessagesRequest,
|
||||
Error, MessagesCallResponse, MessagesRoute, common_utils::truncate_error_body,
|
||||
prepare::ProviderMessagesRequest,
|
||||
};
|
||||
use crate::{constants::MESSAGES_TIMEOUT_SECS, outbound::outbound_request};
|
||||
use crate::{constants::MESSAGES_TIMEOUT_SECS, context::CallContext, outbound::outbound_request};
|
||||
|
||||
pub(super) async fn execute(
|
||||
http: &litellm_http::Client,
|
||||
auth: &AuthServices,
|
||||
request: ProviderMessagesRequest,
|
||||
hooks: &impl RouteHooks<Error>,
|
||||
) -> Result<MessagesResponse, Error> {
|
||||
let ProviderMessagesRequest {
|
||||
provider,
|
||||
url,
|
||||
body,
|
||||
environment,
|
||||
timeout,
|
||||
api_key,
|
||||
} = request;
|
||||
let stream = body.params.stream == Some(true);
|
||||
let context = RequestContext {
|
||||
model: body.model.clone(),
|
||||
custom_llm_provider: provider.as_str().to_string(),
|
||||
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
|
||||
secret_fields: Vec::new(),
|
||||
api_key,
|
||||
};
|
||||
let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?;
|
||||
let wire = hooks
|
||||
.before_provider_request(
|
||||
WireRequest {
|
||||
url,
|
||||
headers: authenticated.headers,
|
||||
body: serde_json::to_value(&body).map_err(serialize_failure)?,
|
||||
pub(super) struct ProviderCall {
|
||||
pub identity: ProviderIdentity,
|
||||
pub wire: WireRequest,
|
||||
provider: super::common_utils::MessagesProvider,
|
||||
signer: Option<litellm_auth_aws::SigV4Signer>,
|
||||
timeout: Option<Duration>,
|
||||
stream: bool,
|
||||
}
|
||||
|
||||
impl ProviderCall {
|
||||
pub fn cacheable(&self) -> bool {
|
||||
self.signer.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
impl MessagesRoute {
|
||||
pub(super) async fn prepare_outbound(
|
||||
&self,
|
||||
request: ProviderMessagesRequest,
|
||||
context: &CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<ProviderCall, Error> {
|
||||
let ProviderMessagesRequest {
|
||||
provider,
|
||||
url,
|
||||
body,
|
||||
environment,
|
||||
timeout,
|
||||
api_key,
|
||||
} = request;
|
||||
let request_context = RequestContext {
|
||||
model: body.model.clone(),
|
||||
custom_llm_provider: provider.as_str().to_string(),
|
||||
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
|
||||
secret_fields: Vec::new(),
|
||||
api_key,
|
||||
};
|
||||
let authenticated =
|
||||
resolve_auth(&self.auth, environment, &|key| std::env::var(key).ok()).await?;
|
||||
let identity = ProviderIdentity {
|
||||
model: request_context.model.clone(),
|
||||
provider: request_context.custom_llm_provider.clone(),
|
||||
};
|
||||
let wire = context
|
||||
.interceptors
|
||||
.before_provider_request(
|
||||
WireRequest {
|
||||
url,
|
||||
headers: authenticated.headers,
|
||||
body: serde_json::to_value(&body).map_err(serialize_failure)?,
|
||||
},
|
||||
request_context,
|
||||
)
|
||||
.await?;
|
||||
let stream = match wire.body.get("stream") {
|
||||
None | Some(Value::Null) => false,
|
||||
Some(Value::Bool(stream)) => *stream,
|
||||
Some(value) => {
|
||||
return Err(Error::InvalidRequest(
|
||||
litellm_llms::ErrorDetail::InvalidValue {
|
||||
field: "stream",
|
||||
expected: "a boolean",
|
||||
actual: value.clone(),
|
||||
},
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(ProviderCall {
|
||||
identity,
|
||||
wire,
|
||||
provider,
|
||||
signer: authenticated.signer,
|
||||
timeout,
|
||||
stream,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn call_provider(
|
||||
&self,
|
||||
request: ProviderCall,
|
||||
context: &CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
let ProviderCall {
|
||||
identity,
|
||||
wire,
|
||||
provider,
|
||||
signer,
|
||||
timeout,
|
||||
stream,
|
||||
} = request;
|
||||
let provider_name = provider.as_str();
|
||||
log_request_body(provider_name, stream, &wire.body);
|
||||
let response = send(
|
||||
&self.http,
|
||||
Authenticated {
|
||||
headers: wire.headers,
|
||||
signer,
|
||||
},
|
||||
context,
|
||||
&wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
)
|
||||
.await?;
|
||||
let provider_name = provider.as_str();
|
||||
log_request_body(provider_name, stream, &wire.body);
|
||||
let response = send(
|
||||
http,
|
||||
Authenticated {
|
||||
headers: wire.headers,
|
||||
signer: authenticated.signer,
|
||||
},
|
||||
&wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
)
|
||||
.await?;
|
||||
if !response.status().is_success() {
|
||||
return Err(provider_error(response).await);
|
||||
if !response.status().is_success() {
|
||||
return Err(provider_error(response).await);
|
||||
}
|
||||
let config = provider.config();
|
||||
if stream {
|
||||
return Ok(streaming_response(
|
||||
response,
|
||||
config.stream_decoder(),
|
||||
provider_name,
|
||||
));
|
||||
}
|
||||
let text = response.text().await.map_err(network)?;
|
||||
log_response_body(&text);
|
||||
context.response_received(&text).await?;
|
||||
decode_response(config, &identity.model, &text)
|
||||
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
|
||||
}
|
||||
let config = provider.config();
|
||||
if stream {
|
||||
return Ok(streaming_response(
|
||||
response,
|
||||
config.stream_decoder(),
|
||||
provider_name,
|
||||
));
|
||||
}
|
||||
let text = response.text().await.map_err(network)?;
|
||||
log_response_body(&text);
|
||||
hooks
|
||||
.on_event(MachineEvent::ResponseReceived {
|
||||
raw: RawResponse { body: text.clone() },
|
||||
})
|
||||
.await
|
||||
.map_err(Error::post_call)?;
|
||||
decode_response(config, &body.model, &text)
|
||||
.map(|message| MessagesResponse::Complete(Box::new(message)))
|
||||
}
|
||||
|
||||
fn serialize_failure(err: serde_json::Error) -> Error {
|
||||
|
|
@ -135,10 +184,10 @@ async fn provider_error(response: reqwest::Response) -> Error {
|
|||
}
|
||||
|
||||
fn decode_response(
|
||||
config: &dyn BaseAnthropicMessagesConfig,
|
||||
config: &dyn BaseMessagesConfig,
|
||||
model: &str,
|
||||
text: &str,
|
||||
) -> Result<AnthropicMessagesResponse, Error> {
|
||||
) -> Result<MessagesResponse, Error> {
|
||||
let response = serde_json::from_str(text).map_err(|err| {
|
||||
Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
|
||||
"messages response JSON",
|
||||
|
|
@ -154,7 +203,7 @@ fn streaming_response(
|
|||
response: reqwest::Response,
|
||||
decoder: Option<StreamDecoder>,
|
||||
provider: &'static str,
|
||||
) -> MessagesResponse {
|
||||
) -> MessagesCallResponse {
|
||||
let headers = response
|
||||
.headers()
|
||||
.iter()
|
||||
|
|
@ -171,7 +220,7 @@ fn streaming_response(
|
|||
.boxed(),
|
||||
Some(decode) => decoded_chunks(response, decode, provider),
|
||||
};
|
||||
MessagesResponse::Stream {
|
||||
MessagesCallResponse::Stream {
|
||||
head: super::route::MessagesStreamHead { headers },
|
||||
chunks,
|
||||
}
|
||||
|
|
@ -245,7 +294,7 @@ mod tests {
|
|||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
let MessagesResponse::Stream { mut chunks, .. } =
|
||||
let MessagesCallResponse::Stream { mut chunks, .. } =
|
||||
streaming_response(response, Some(anthropic_sse_event_stream), "test")
|
||||
else {
|
||||
panic!("a streaming response returns chunks");
|
||||
|
|
|
|||
|
|
@ -4,18 +4,23 @@ mod prepare;
|
|||
pub mod route;
|
||||
mod types;
|
||||
|
||||
use futures_util::FutureExt;
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_host::interceptors::{ExecutionFacts, Interceptors, ResultSource};
|
||||
|
||||
use crate::{caching::CallCache, context::CallContext};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use crate::error::RouteError as Error;
|
||||
pub use types::{MessagesCall, MessagesResponse, MessagesShaping, messages_body};
|
||||
pub use types::{MessagesCall, MessagesCallResponse, MessagesShaping, messages_body};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MessagesRoute {
|
||||
http: litellm_http::Client,
|
||||
auth: Arc<AuthServices>,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
cache: Option<litellm_cache_response::ScopedCache>,
|
||||
}
|
||||
|
||||
impl MessagesRoute {
|
||||
|
|
@ -28,15 +33,27 @@ impl MessagesRoute {
|
|||
http,
|
||||
auth,
|
||||
secrets,
|
||||
cache: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_cache(self, cache: litellm_cache_response::ScopedCache) -> Self {
|
||||
Self {
|
||||
cache: Some(cache),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(
|
||||
&self,
|
||||
call: MessagesCall,
|
||||
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
|
||||
) -> Result<MessagesResponse, Error> {
|
||||
litellm_host::lifecycle::observe_call(hooks.observer(), self.run(call, hooks)).await
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
options: impl Into<crate::CallOptions>,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
let context = CallContext::new(interceptors, options.into());
|
||||
litellm_host::lifecycle::observe_call(context.observers.clone(), self.run(call, context))
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
|
|
@ -50,14 +67,33 @@ impl MessagesRoute {
|
|||
async fn run(
|
||||
&self,
|
||||
call: MessagesCall,
|
||||
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
|
||||
) -> Result<MessagesResponse, Error> {
|
||||
context: CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
crate::diagnostic::call(async {
|
||||
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
crate::diagnostic::provider(&request.body.model, request.provider.as_str());
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<MessagesResponse, Error>> =
|
||||
Box::pin(handler::execute(&self.http, &self.auth, request, hooks));
|
||||
execute.await
|
||||
let prepared = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
crate::diagnostic::provider(&prepared.body.model, prepared.provider.as_str());
|
||||
let request = self.prepare_outbound(prepared, &context).boxed().await?;
|
||||
let cache = CallCache::<route::Messages>::from_wire(
|
||||
self.cache.as_ref().filter(|_| request.cacheable()),
|
||||
context.cache,
|
||||
&request.identity,
|
||||
&request.wire,
|
||||
);
|
||||
let identity = request.identity.clone();
|
||||
let (output, source) = match cache.lookup().await {
|
||||
Some(hit) => hit,
|
||||
None => (
|
||||
self.call_provider(request, &context).await?,
|
||||
ResultSource::Provider,
|
||||
),
|
||||
};
|
||||
context
|
||||
.result_ready(ExecutionFacts {
|
||||
provider: identity,
|
||||
source: source.clone(),
|
||||
})
|
||||
.await?;
|
||||
Ok(cache.finish(output, &source).await)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
|
|
|||