Merge remote-tracking branch 'origin/main' into litellm_batch_jsonl_line_item_callbacks
|
|
@ -486,6 +486,7 @@ jobs:
|
|||
- install_rust
|
||||
- run:
|
||||
name: Build the wheel
|
||||
no_output_timeout: 30m
|
||||
environment:
|
||||
UV_HTTP_TIMEOUT: "300"
|
||||
command: |
|
||||
|
|
|
|||
|
|
@ -168,6 +168,7 @@ start_proxy() {
|
|||
"${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
|
||||
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
|
||||
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
|
||||
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
|
||||
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
|
||||
|
|
@ -228,6 +229,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
|
|||
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
|
||||
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
|
||||
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
INTEGRATION_SEED="$INTEGRATION_SEED" \
|
||||
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
|
||||
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
|
||||
|
|
|
|||
|
|
@ -107,7 +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/lens
|
||||
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
|
||||
|
|
@ -116,7 +116,7 @@ legacy_paths() {
|
|||
echo tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py
|
||||
echo tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py
|
||||
echo tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py
|
||||
echo tests/unit/proxy/response_polling/test_response_polling_handler.py
|
||||
echo tests/unit/proxy/response_polling
|
||||
echo tests/unit/proxy/test_custom_tokenizer_bug.py
|
||||
echo tests/unit/proxy/test_get_favicon.py
|
||||
echo tests/unit/proxy/test_get_image.py
|
||||
|
|
@ -151,6 +151,7 @@ legacy_paths() {
|
|||
proxy-extras) echo tests/unit/litellm_proxy_extras ;;
|
||||
proxy-infra)
|
||||
echo tests/unit/gateway
|
||||
echo tests/unit/proxy/management
|
||||
echo tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py
|
||||
echo tests/unit/proxy/roi_calculator ;;
|
||||
responses-caching-types)
|
||||
|
|
|
|||
13
.github/actions/cache-cargo-build/action.yml
vendored
|
|
@ -25,6 +25,7 @@ runs:
|
|||
using: composite
|
||||
steps:
|
||||
- name: Restore the Cargo registry and target directory
|
||||
if: github.ref == 'refs/heads/main'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
@ -34,3 +35,15 @@ runs:
|
|||
key: ${{ runner.os }}-maturin-${{ inputs.profile }}-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-maturin-${{ inputs.profile }}-
|
||||
|
||||
- name: Restore the Cargo registry and target directory
|
||||
if: github.ref != 'refs/heads/main'
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
litellm-rust/target
|
||||
key: ${{ runner.os }}-maturin-${{ inputs.profile }}-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-maturin-${{ inputs.profile }}-
|
||||
|
|
|
|||
10
.github/actions/cache-prisma-binaries/action.yml
vendored
|
|
@ -30,6 +30,7 @@ runs:
|
|||
echo "version=${version}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Restore Prisma binaries
|
||||
if: github.ref == 'refs/heads/main'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
# ~/.cache/prisma-python holds the npm install tree prisma-client-py
|
||||
|
|
@ -38,3 +39,12 @@ runs:
|
|||
~/.cache/prisma-python
|
||||
~/.cache/prisma
|
||||
key: ${{ runner.os }}-prisma-binaries-${{ steps.version.outputs.version }}
|
||||
|
||||
- name: Restore Prisma binaries
|
||||
if: github.ref != 'refs/heads/main'
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/prisma-python
|
||||
~/.cache/prisma
|
||||
key: ${{ runner.os }}-prisma-binaries-${{ steps.version.outputs.version }}
|
||||
|
|
|
|||
25
.github/actions/cache-uv-downloads/action.yml
vendored
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
name: "Cache uv downloads"
|
||||
description: >-
|
||||
Restore the uv download cache on every run and save it only from main, so pull
|
||||
requests reuse main's cache instead of evicting it with their own copies.
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- name: Restore and save the uv download cache
|
||||
if: github.ref == 'refs/heads/main'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: ${{ env.UV_CACHE_DIR }}
|
||||
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-
|
||||
|
||||
- name: Restore the uv download cache
|
||||
if: github.ref != 'refs/heads/main'
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: ${{ env.UV_CACHE_DIR }}
|
||||
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-
|
||||
|
|
@ -17,6 +17,7 @@ runs:
|
|||
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
|
||||
with:
|
||||
version: ${{ inputs.version }}
|
||||
save-cache: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
||||
- name: Wait before attempt 2
|
||||
if: steps.attempt-1.outcome == 'failure'
|
||||
|
|
@ -30,6 +31,7 @@ runs:
|
|||
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
|
||||
with:
|
||||
version: ${{ inputs.version }}
|
||||
save-cache: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
||||
- name: Wait before attempt 3
|
||||
if: steps.attempt-2.outcome == 'failure'
|
||||
|
|
@ -41,3 +43,4 @@ runs:
|
|||
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
|
||||
with:
|
||||
version: ${{ inputs.version }}
|
||||
save-cache: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
|
|
|||
1
.github/e2e-stack/select_tests.py
vendored
|
|
@ -11,6 +11,7 @@ UNSUPPORTED: Final = re.compile(
|
|||
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
|
||||
r"|^tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e\.py$"
|
||||
r"|^tests/e2e/logging/test_langsmith_batch_serialization_e2e\.py$"
|
||||
r"|^tests/e2e/logging/test_s3_log_e2e\.py$"
|
||||
r"|^tests/e2e/secret_manager/"
|
||||
)
|
||||
HARNESS: Final = re.compile(
|
||||
|
|
|
|||
48
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -9,6 +9,7 @@ import sys
|
|||
import warnings
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import yaml
|
||||
|
|
@ -35,7 +36,7 @@ GLOB_CHARS = frozenset("*?")
|
|||
# itself decomposed one level deeper and is checked through its own entry.
|
||||
SHARDED_ROOTS: tuple[str, ...] = (
|
||||
"tests/test_litellm",
|
||||
"tests/test_litellm/proxy",
|
||||
"tests/unit/proxy",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -119,11 +120,48 @@ def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
|
|||
)
|
||||
|
||||
|
||||
def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]:
|
||||
SELECTION_ARM_RE = re.compile(r"(?ms)^\s*([A-Za-z0-9_|*-]+)\)\s*(.*?);;")
|
||||
|
||||
|
||||
def _unit_selection_arms(repo_root: pathlib.Path = REPO_ROOT) -> Mapping[str, frozenset[str]]:
|
||||
script: Final = repo_root / ".circleci/scripts/unit_selection.sh"
|
||||
if not script.is_file():
|
||||
return frozenset()
|
||||
return frozenset(match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(_uncommented(script.read_text())))
|
||||
return MappingProxyType({})
|
||||
text: Final = _uncommented(script.read_text())
|
||||
return MappingProxyType(
|
||||
{
|
||||
label: frozenset(
|
||||
match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(body)
|
||||
)
|
||||
for label, body in SELECTION_ARM_RE.findall(text)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]:
|
||||
return frozenset(
|
||||
token for tokens in _unit_selection_arms(repo_root).values() for token in tokens
|
||||
)
|
||||
|
||||
|
||||
def _wired_unit_flags(scalars: Iterable[Scalar]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
scalar.value
|
||||
for scalar in scalars
|
||||
if scalar.key == "unit-flag" and "${{" not in scalar.value
|
||||
)
|
||||
|
||||
|
||||
def _shard_tokens(
|
||||
scalars: Iterable[Scalar], arms: Mapping[str, frozenset[str]]
|
||||
) -> frozenset[str]:
|
||||
wired: Final = _wired_unit_flags(scalars)
|
||||
return _invoked_test_tokens(scalars) | frozenset(
|
||||
token
|
||||
for label, tokens in arms.items()
|
||||
if label in wired
|
||||
for token in tokens
|
||||
)
|
||||
|
||||
|
||||
def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
|
||||
|
|
@ -480,7 +518,7 @@ def _check_slices() -> int:
|
|||
|
||||
|
||||
def _check_shards() -> int:
|
||||
findings = _unassigned_shard_children(_invoked_test_tokens(_all_scalars()))
|
||||
findings = _unassigned_shard_children(_shard_tokens(_all_scalars(), _unit_selection_arms()))
|
||||
if findings:
|
||||
_report(
|
||||
"test directories and files that no shard claims",
|
||||
|
|
|
|||
11
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -132,12 +132,7 @@ jobs:
|
|||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: ${{ env.UV_CACHE_DIR }}
|
||||
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-
|
||||
uses: ./.github/actions/cache-uv-downloads
|
||||
|
||||
- name: Cache the Rust build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
|
|
@ -274,7 +269,7 @@ jobs:
|
|||
- name: Upload to Codecov
|
||||
id: codecov-upload
|
||||
continue-on-error: true
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
|
||||
with:
|
||||
use_oidc: true
|
||||
directory: coverage-reports
|
||||
|
|
@ -285,7 +280,7 @@ jobs:
|
|||
- name: Upload to Codecov (retry)
|
||||
if: steps.codecov-upload.outcome == 'failure'
|
||||
continue-on-error: true
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
|
||||
with:
|
||||
use_oidc: true
|
||||
directory: coverage-reports
|
||||
|
|
|
|||
10
.github/workflows/lens-worker.yml
vendored
|
|
@ -5,13 +5,13 @@ on:
|
|||
branches: [main, litellm_oss_branch, "litellm_**"]
|
||||
paths:
|
||||
- deploy/lens/**
|
||||
- litellm/proxy/engine/**
|
||||
- litellm/proxy/lens/**
|
||||
- .github/workflows/lens-worker.yml
|
||||
push:
|
||||
branches: [main, litellm_agent_engine]
|
||||
branches: [main]
|
||||
paths:
|
||||
- deploy/lens/**
|
||||
- litellm/proxy/engine/**
|
||||
- litellm/proxy/lens/**
|
||||
- .github/workflows/lens-worker.yml
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
@ -41,8 +41,8 @@ jobs:
|
|||
--security-opt no-new-privileges --entrypoint python \
|
||||
lens-worker:${{ github.sha }} -c '
|
||||
import os
|
||||
import engine.worker
|
||||
from engine.trace_store import trace_store
|
||||
import lens.worker
|
||||
from lens.trace_store import trace_store
|
||||
assert os.getuid() == 65532
|
||||
with trace_store() as store:
|
||||
assert store.count() == 0
|
||||
|
|
|
|||
12
.github/workflows/mutation-test.yml
vendored
|
|
@ -44,6 +44,7 @@ jobs:
|
|||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: github.ref == 'refs/heads/main'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
@ -53,6 +54,17 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: github.ref != 'refs/heads/main'
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
|
|
|
|||
12
.github/workflows/test-code-quality.yml
vendored
|
|
@ -44,6 +44,7 @@ jobs:
|
|||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: github.ref == 'refs/heads/main'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
@ -53,6 +54,17 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: github.ref != 'refs/heads/main'
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
|
|
|
|||
14
.github/workflows/test-postgres.yml
vendored
|
|
@ -95,7 +95,7 @@ jobs:
|
|||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
if: steps.changes.outputs.decision != 'skip' && github.ref == 'refs/heads/main'
|
||||
timeout-minutes: 5
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
|
|
@ -106,6 +106,18 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-postgres-
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip' && github.ref != 'refs/heads/main'
|
||||
timeout-minutes: 5
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-postgres-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-postgres-
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 12
|
||||
|
|
|
|||
2
.github/workflows/test-redis-compat.yml
vendored
|
|
@ -98,7 +98,7 @@ jobs:
|
|||
|
||||
- name: Upload Redis coverage
|
||||
if: matrix.redis-version == '5.3.1'
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
|
||||
with:
|
||||
use_oidc: true
|
||||
files: coverage-redis.xml
|
||||
|
|
|
|||
3
.github/workflows/test-rust.yml
vendored
|
|
@ -83,6 +83,7 @@ jobs:
|
|||
with:
|
||||
workspaces: litellm-rust
|
||||
cache-on-failure: true
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
||||
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
|
||||
|
||||
|
|
@ -121,6 +122,7 @@ jobs:
|
|||
with:
|
||||
workspaces: litellm-rust
|
||||
cache-on-failure: true
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
||||
- run: cargo nextest run --workspace --locked
|
||||
|
||||
|
|
@ -162,6 +164,7 @@ jobs:
|
|||
with:
|
||||
workspaces: litellm-rust
|
||||
cache-on-failure: true
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
||||
- run: uv build --wheel --out-dir dist
|
||||
|
||||
|
|
|
|||
12
.github/workflows/test-terraform-provider.yml
vendored
|
|
@ -77,6 +77,7 @@ jobs:
|
|||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: github.ref == 'refs/heads/main'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
@ -86,6 +87,17 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: github.ref != 'refs/heads/main'
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
|
|
|
|||
13
.github/workflows/test-unit-documentation.yml
vendored
|
|
@ -54,7 +54,7 @@ jobs:
|
|||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
if: steps.changes.outputs.decision != 'skip' && github.ref == 'refs/heads/main'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
@ -64,6 +64,17 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip' && github.ref != 'refs/heads/main'
|
||||
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
|
|
|||
141
.github/workflows/test-unit.yml
vendored
|
|
@ -140,38 +140,46 @@ jobs:
|
|||
- shard: proxy-endpoints
|
||||
artifact-name: proxy-endpoints
|
||||
test-path: >-
|
||||
tests/test_litellm/proxy/analytics_endpoints
|
||||
tests/test_litellm/proxy/management_endpoints
|
||||
tests/test_litellm/proxy/list_api
|
||||
tests/test_litellm/proxy/memory
|
||||
tests/test_litellm/proxy/guardrails
|
||||
tests/test_litellm/proxy/management_helpers
|
||||
tests/test_litellm/proxy/anthropic_endpoints
|
||||
tests/test_litellm/proxy/google_endpoints
|
||||
tests/test_litellm/proxy/openai_files_endpoint
|
||||
tests/test_litellm/proxy/batches_endpoints
|
||||
tests/test_litellm/proxy/container_endpoints
|
||||
tests/test_litellm/proxy/fine_tuning_endpoints
|
||||
tests/test_litellm/proxy/vector_store_files_endpoints
|
||||
tests/test_litellm/proxy/video_endpoints
|
||||
tests/test_litellm/proxy/response_api_endpoints
|
||||
tests/test_litellm/proxy/image_endpoints
|
||||
tests/test_litellm/proxy/ocr_endpoints
|
||||
tests/test_litellm/proxy/vector_store_endpoints
|
||||
tests/test_litellm/proxy/agent_endpoints
|
||||
tests/test_litellm/proxy/a2a
|
||||
tests/test_litellm/proxy/credential_endpoints
|
||||
tests/test_litellm/proxy/discovery_endpoints
|
||||
tests/test_litellm/proxy/health_endpoints
|
||||
tests/test_litellm/proxy/shutdown
|
||||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/rerank_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
tests/test_litellm/proxy/config_resolvers
|
||||
tests/test_litellm/proxy/utils
|
||||
tests/unit/proxy/analytics_endpoints
|
||||
tests/unit/proxy/management_endpoints
|
||||
tests/unit/proxy/list_api
|
||||
tests/unit/proxy/memory
|
||||
tests/unit/proxy/guardrails
|
||||
tests/unit/proxy/management_helpers
|
||||
--ignore=tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py
|
||||
--ignore=tests/unit/proxy/management_endpoints/test_key_generate_prisma.py
|
||||
--ignore=tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py
|
||||
--ignore=tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
|
||||
--ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py
|
||||
--ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py
|
||||
--ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py
|
||||
--ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py
|
||||
tests/unit/proxy/anthropic_endpoints
|
||||
tests/unit/proxy/google_endpoints
|
||||
tests/unit/proxy/openai_files_endpoint
|
||||
tests/unit/proxy/batches_endpoints
|
||||
tests/unit/proxy/container_endpoints
|
||||
tests/unit/proxy/fine_tuning_endpoints
|
||||
tests/unit/proxy/vector_store_files_endpoints
|
||||
tests/unit/proxy/video_endpoints
|
||||
tests/unit/proxy/response_api_endpoints
|
||||
tests/unit/proxy/image_endpoints
|
||||
tests/unit/proxy/ocr_endpoints
|
||||
tests/unit/proxy/vector_store_endpoints
|
||||
tests/unit/proxy/agent_endpoints
|
||||
tests/unit/proxy/a2a
|
||||
tests/unit/proxy/credential_endpoints
|
||||
tests/unit/proxy/discovery_endpoints
|
||||
tests/unit/proxy/health_endpoints
|
||||
tests/unit/proxy/shutdown
|
||||
tests/unit/proxy/public_endpoints
|
||||
tests/unit/proxy/prompts
|
||||
tests/unit/proxy/rag_endpoints
|
||||
tests/unit/proxy/rerank_endpoints
|
||||
tests/unit/proxy/realtime_endpoints
|
||||
tests/unit/proxy/ui_crud_endpoints
|
||||
tests/unit/proxy/config_resolvers
|
||||
tests/unit/proxy/utils
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
|
|
@ -179,7 +187,7 @@ jobs:
|
|||
|
||||
- shard: proxy-server
|
||||
artifact-name: proxy-server
|
||||
test-path: "tests/test_litellm/proxy/proxy_server"
|
||||
test-path: "tests/unit/proxy/proxy_server"
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 60
|
||||
|
|
@ -188,25 +196,66 @@ jobs:
|
|||
- shard: proxy-infra
|
||||
artifact-name: proxy-infra
|
||||
test-path: >-
|
||||
tests/test_litellm/proxy/db
|
||||
tests/test_litellm/proxy/middleware
|
||||
tests/test_litellm/proxy/spend_tracking
|
||||
tests/test_litellm/proxy/pass_through_endpoints
|
||||
tests/test_litellm/proxy/_experimental
|
||||
tests/test_litellm/proxy/experimental
|
||||
tests/test_litellm/proxy/common_utils
|
||||
tests/test_litellm/proxy/enterprise_billing
|
||||
tests/test_litellm/proxy/types_utils
|
||||
tests/test_litellm/proxy/logging_endpoints
|
||||
tests/test_litellm/proxy/test_*.py
|
||||
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
|
||||
tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py
|
||||
tests/unit/proxy/db
|
||||
--ignore=tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py
|
||||
--ignore=tests/unit/proxy/db/test_update_daily_tag_spend.py
|
||||
tests/unit/proxy/middleware
|
||||
--ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py
|
||||
tests/unit/proxy/spend_tracking
|
||||
--ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py
|
||||
tests/unit/proxy/pass_through_endpoints
|
||||
tests/unit/proxy/_experimental
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server
|
||||
tests/unit/proxy/experimental
|
||||
tests/unit/proxy/common_utils
|
||||
--ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py
|
||||
--ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py
|
||||
--ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py
|
||||
--ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py
|
||||
--ignore=tests/unit/proxy/common_utils/test_realtime_cache.py
|
||||
tests/unit/proxy/enterprise_billing
|
||||
tests/unit/proxy/types_utils
|
||||
tests/unit/proxy/logging_endpoints
|
||||
unit-flag: proxy-infra
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: proxy-infra-root
|
||||
artifact-name: proxy-infra-root
|
||||
test-path: >-
|
||||
tests/unit/proxy/test_*.py
|
||||
--ignore=tests/unit/proxy/test_aproxy_startup.py
|
||||
--ignore=tests/unit/proxy/test_credential_slot_registry.py
|
||||
--ignore=tests/unit/proxy/test_custom_callback_input.py
|
||||
--ignore=tests/unit/proxy/test_custom_logger_s3_gcs.py
|
||||
--ignore=tests/unit/proxy/test_custom_tokenizer_bug.py
|
||||
--ignore=tests/unit/proxy/test_db_schema_changes.py
|
||||
--ignore=tests/unit/proxy/test_deprecated_key_grace_period.py
|
||||
--ignore=tests/unit/proxy/test_get_favicon.py
|
||||
--ignore=tests/unit/proxy/test_get_image.py
|
||||
--ignore=tests/unit/proxy/test_prisma_client_backoff_retry.py
|
||||
--ignore=tests/unit/proxy/test_prompt_test_endpoint.py
|
||||
--ignore=tests/unit/proxy/test_proxy_config_unit_test.py
|
||||
--ignore=tests/unit/proxy/test_proxy_custom_auth.py
|
||||
--ignore=tests/unit/proxy/test_proxy_reject_logging.py
|
||||
--ignore=tests/unit/proxy/test_proxy_server.py
|
||||
--ignore=tests/unit/proxy/test_proxy_setting_guardrails.py
|
||||
--ignore=tests/unit/proxy/test_proxy_token_counter.py
|
||||
--ignore=tests/unit/proxy/test_proxy_utils.py
|
||||
--ignore=tests/unit/proxy/test_reducto_ocr_route.py
|
||||
--ignore=tests/unit/proxy/test_response_polling_pre_call_checks.py
|
||||
--ignore=tests/unit/proxy/test_server_root_path.py
|
||||
--ignore=tests/unit/proxy/test_ui_path_detection.py
|
||||
--ignore=tests/unit/proxy/test_unit_test_proxy_hooks.py
|
||||
--ignore=tests/unit/proxy/test_update_spend.py
|
||||
--ignore=tests/unit/proxy/test_zero_cost_model_budget_bypass.py
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: caching-local
|
||||
artifact-name: caching-local
|
||||
test-path: ""
|
||||
|
|
|
|||
2
.gitignore
vendored
|
|
@ -104,7 +104,7 @@ litellm_config.yaml
|
|||
.cursor
|
||||
litellm/proxy/to_delete_loadtest_work/*
|
||||
update_model_cost_map.py
|
||||
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
|
||||
scripts/test_vertex_ai_search.py
|
||||
LAZY_LOADING_IMPROVEMENTS.md
|
||||
STABILIZATION_TODO.md
|
||||
|
|
|
|||
|
|
@ -98,7 +98,7 @@ Add your tests to the [`tests/unit/` directory](https://github.com/BerriAI/litel
|
|||
|
||||
The `tests/unit/` directory follows the same structure as `litellm/`:
|
||||
|
||||
- `litellm/proxy/caching_routes.py` → `tests/test_litellm/proxy/test_caching_routes.py`
|
||||
- `litellm/proxy/caching_routes.py` → `tests/unit/proxy/test_caching_routes.py`
|
||||
- `litellm/utils.py` → `tests/unit/test_utils.py`
|
||||
|
||||
### Example Test
|
||||
|
|
|
|||
12
Makefile
|
|
@ -1,7 +1,7 @@
|
|||
# LiteLLM Makefile
|
||||
# Simple Makefile for running tests and basic development tasks
|
||||
|
||||
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
|
||||
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc test-unit-proxy-root \
|
||||
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
|
||||
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
|
||||
test-rust-extension rust-sqlx-prepare \
|
||||
|
|
@ -47,6 +47,7 @@ help:
|
|||
@echo " make test-unit-proxy-guardrails - Run proxy guardrails+mgmt tests (~51 files)"
|
||||
@echo " make test-unit-proxy-core - Run proxy auth+client+db+hooks tests (~52 files)"
|
||||
@echo " make test-unit-proxy-misc - Run proxy misc tests (~77 files)"
|
||||
@echo " make test-unit-proxy-root - Run proxy root-file tests (tests/unit/proxy/test_*.py)"
|
||||
@echo " make test-unit-integrations - Run integration tests (~60 files)"
|
||||
@echo " make test-unit-core-utils - Run core utils tests (~32 files)"
|
||||
@echo " make test-unit-other - Run other tests (caching, responses, etc., ~69 files)"
|
||||
|
|
@ -321,13 +322,16 @@ test-unit-llms: install-test-deps
|
|||
$(UV_RUN) pytest tests/unit/llms --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-proxy-guardrails: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/proxy/guardrails tests/test_litellm/proxy/management_endpoints tests/test_litellm/proxy/management_helpers --tb=short -vv -n 4 --durations=20
|
||||
$(UV_RUN) pytest tests/unit/proxy/guardrails tests/unit/proxy/management_endpoints tests/unit/proxy/management_helpers --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-proxy-core: install-test-deps
|
||||
$(UV_RUN) pytest tests/unit/proxy/auth tests/unit/proxy/client tests/test_litellm/proxy/db tests/unit/proxy/hooks tests/unit/proxy/policy_engine --tb=short -vv -n 4 --durations=20
|
||||
$(UV_RUN) pytest tests/unit/proxy/auth tests/unit/proxy/client tests/unit/proxy/db tests/unit/proxy/hooks tests/unit/proxy/policy_engine --ignore=tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py --ignore=tests/unit/proxy/db/test_update_daily_tag_spend.py --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-proxy-misc: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py tests/unit/proxy/test_proxy_server_endpoints_and_startup.py tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py --tb=short -vv -n 4 --durations=20
|
||||
$(UV_RUN) pytest tests/unit/proxy/agent_endpoints tests/unit/proxy/anthropic_endpoints tests/unit/proxy/common_utils --ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py --ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py --ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py --ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py --ignore=tests/unit/proxy/common_utils/test_realtime_cache.py tests/unit/proxy/discovery_endpoints tests/unit/proxy/experimental tests/unit/proxy/google_endpoints tests/unit/proxy/health_endpoints tests/unit/proxy/image_endpoints tests/unit/proxy/middleware --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py tests/unit/proxy/openai_files_endpoint tests/unit/proxy/pass_through_endpoints tests/unit/proxy/prompts tests/unit/proxy/public_endpoints tests/unit/proxy/response_api_endpoints tests/unit/proxy/shutdown tests/unit/proxy/spend_tracking --ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py tests/unit/proxy/ui_crud_endpoints tests/unit/proxy/vector_store_endpoints tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py --ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py --ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py --ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py --ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-proxy-root: install-test-deps
|
||||
$(UV_RUN) pytest tests/unit/proxy/test_*.py --ignore=tests/unit/proxy/test_aproxy_startup.py --ignore=tests/unit/proxy/test_credential_slot_registry.py --ignore=tests/unit/proxy/test_custom_callback_input.py --ignore=tests/unit/proxy/test_custom_logger_s3_gcs.py --ignore=tests/unit/proxy/test_custom_tokenizer_bug.py --ignore=tests/unit/proxy/test_db_schema_changes.py --ignore=tests/unit/proxy/test_deprecated_key_grace_period.py --ignore=tests/unit/proxy/test_get_favicon.py --ignore=tests/unit/proxy/test_get_image.py --ignore=tests/unit/proxy/test_prisma_client_backoff_retry.py --ignore=tests/unit/proxy/test_prompt_test_endpoint.py --ignore=tests/unit/proxy/test_proxy_config_unit_test.py --ignore=tests/unit/proxy/test_proxy_custom_auth.py --ignore=tests/unit/proxy/test_proxy_reject_logging.py --ignore=tests/unit/proxy/test_proxy_server.py --ignore=tests/unit/proxy/test_proxy_setting_guardrails.py --ignore=tests/unit/proxy/test_proxy_token_counter.py --ignore=tests/unit/proxy/test_proxy_utils.py --ignore=tests/unit/proxy/test_reducto_ocr_route.py --ignore=tests/unit/proxy/test_response_polling_pre_call_checks.py --ignore=tests/unit/proxy/test_server_root_path.py --ignore=tests/unit/proxy/test_ui_path_detection.py --ignore=tests/unit/proxy/test_unit_test_proxy_hooks.py --ignore=tests/unit/proxy/test_update_spend.py --ignore=tests/unit/proxy/test_zero_cost_model_budget_bypass.py --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-integrations: install-test-deps
|
||||
$(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20
|
||||
|
|
|
|||
|
|
@ -81,7 +81,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
# Spend / analytics
|
||||
"/spend/",
|
||||
"/analytics/",
|
||||
"/engine/",
|
||||
"/lens/",
|
||||
"/v1/traces",
|
||||
"/global/",
|
||||
"/user_agent",
|
||||
|
|
@ -146,7 +146,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
|
|||
{
|
||||
"/",
|
||||
"/routes",
|
||||
"/engine",
|
||||
"/lens",
|
||||
"/openapi.json",
|
||||
"/docs",
|
||||
"/docs/oauth2-redirect",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
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/
|
||||
COPY litellm/proxy/lens/__init__.py litellm/proxy/lens/models.py litellm/proxy/lens/trace_store.py litellm/proxy/lens/analysis.py litellm/proxy/lens/worker.py /app/lens/
|
||||
USER 65532:65532
|
||||
CMD ["python", "-m", "engine.worker"]
|
||||
CMD ["python", "-m", "lens.worker"]
|
||||
|
|
|
|||
|
|
@ -1,8 +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
|
||||
!litellm/proxy/lens/
|
||||
!litellm/proxy/lens/__init__.py
|
||||
!litellm/proxy/lens/models.py
|
||||
!litellm/proxy/lens/analysis.py
|
||||
!litellm/proxy/lens/worker.py
|
||||
|
|
|
|||
|
|
@ -74,7 +74,7 @@ V1 requires ClickHouse for both sources. It does not reconstruct sessions from u
|
|||
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" \
|
||||
curl "$LITELLM_URL/lens" -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.",
|
||||
|
|
@ -83,14 +83,14 @@ curl "$LITELLM_URL/engine" -H "Authorization: Bearer $LITELLM_API_KEY" \
|
|||
"enabled": true, "interval_minutes": 1440, "monthly_budget": 50
|
||||
}'
|
||||
|
||||
curl "$LITELLM_URL/engine/$LENS_ID/runs" -X POST \
|
||||
curl "$LITELLM_URL/lens/$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"
|
||||
curl "$LITELLM_URL/lens/$LENS_ID/runs?offset=0" -H "Authorization: Bearer $LITELLM_API_KEY"
|
||||
curl "$LITELLM_URL/lens/$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`
|
||||
Creation queues the first batch. Posting to `/lens/{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 `/lens/preview/sample`. Preview accepts `offset` and `as_of` to keep the time window fixed while paging. Feedback uses `PATCH /lens/{id}/findings/{finding_id}` with `status` and `reason`
|
||||
|
||||
## Quality evaluation
|
||||
|
||||
|
|
@ -107,3 +107,11 @@ Set `LITELLM_API_KEY` privately. This makes paid model calls. Inspect missed and
|
|||
The worker uses temporary disk space for trace content while reviewing it, and removes those files after each review. The Docker command supplies a writable temporary mount 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
|
||||
|
||||
## Upgrading from the original Lens API
|
||||
|
||||
The Lens API now uses `/lens` instead of `/engine`, list responses use `lenses`, and worker claims use `lens_id`. Upgrade the proxy and recreate every worker with the image shown by the upgraded dashboard before starting new scans. Update API clients to the new paths and response fields. Old worker images cannot poll the renamed API
|
||||
|
||||
Stop workers and let active scans finish before upgrading. Deploy proxy instances together: older proxies cannot use the renamed database tables. The schema migration renames the three Lens tables and the run-history identifier column in place, preserving saved investigations, findings, history, worker credentials, and billing assignments. Existing migration files retain their original names and checksums
|
||||
|
||||
Upgrades using `--use_prisma_db_push` stop before schema changes if any legacy Lens table exists, preventing Prisma from dropping saved data. Apply `litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql` to the configured database schema before retrying. Deployments already using migration history can instead start without `--use_prisma_db_push` to apply the shipped migration normally. Fresh databases and databases already using the renamed tables can continue using database push
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
services:
|
||||
lens-worker:
|
||||
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:c41e932eaf3e4efbcaf8cc5027c7e93021e5b2823f21cb8785cd107e37b91c9a}
|
||||
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:a8e8731d954916594eea462969946b9292fb771681ff515a9fd296b53f856c77}
|
||||
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}
|
||||
|
|
|
|||
|
Before Width: | Height: | Size: 95 KiB |
|
Before Width: | Height: | Size: 6.9 KiB |
|
Before Width: | Height: | Size: 89 KiB |
|
Before Width: | Height: | Size: 80 KiB |
|
Before Width: | Height: | Size: 70 KiB |
|
Before Width: | Height: | Size: 132 KiB |
|
Before Width: | Height: | Size: 59 KiB |
|
Before Width: | Height: | Size: 54 KiB |
|
|
@ -0,0 +1,18 @@
|
|||
DO $$
|
||||
BEGIN
|
||||
ALTER TABLE IF EXISTS "LiteLLM_Engine" RENAME TO "LiteLLM_Lens";
|
||||
ALTER TABLE IF EXISTS "LiteLLM_EngineRun" RENAME TO "LiteLLM_LensRun";
|
||||
ALTER TABLE IF EXISTS "LiteLLM_EngineWorker" RENAME TO "LiteLLM_LensWorker";
|
||||
IF EXISTS (
|
||||
SELECT 1 FROM pg_attribute
|
||||
WHERE attrelid = to_regclass('"LiteLLM_LensRun"')
|
||||
AND attname = 'engine_id' AND NOT attisdropped
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_LensRun" RENAME COLUMN "engine_id" TO "lens_id";
|
||||
END IF;
|
||||
ALTER INDEX IF EXISTS "LiteLLM_Engine_pkey" RENAME TO "LiteLLM_Lens_pkey";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_EngineRun_pkey" RENAME TO "LiteLLM_LensRun_pkey";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_EngineWorker_pkey" RENAME TO "LiteLLM_LensWorker_pkey";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_EngineWorker_token_hash_key" RENAME TO "LiteLLM_LensWorker_token_hash_key";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_EngineRun_engine_id_created_at_idx" RENAME TO "LiteLLM_LensRun_lens_id_created_at_idx";
|
||||
END $$;
|
||||
|
|
@ -1895,22 +1895,22 @@ model LiteLLM_WorkflowMessage {
|
|||
@@index([run_id])
|
||||
}
|
||||
|
||||
model LiteLLM_Engine {
|
||||
model LiteLLM_Lens {
|
||||
id String @id
|
||||
version Int @default(0)
|
||||
data Json
|
||||
}
|
||||
|
||||
model LiteLLM_EngineRun {
|
||||
model LiteLLM_LensRun {
|
||||
id String @id
|
||||
engine_id String
|
||||
lens_id String
|
||||
created_at DateTime
|
||||
data Json
|
||||
|
||||
@@index([engine_id, created_at])
|
||||
@@index([lens_id, created_at])
|
||||
}
|
||||
|
||||
model LiteLLM_EngineWorker {
|
||||
model LiteLLM_LensWorker {
|
||||
id String @id
|
||||
token_hash String @unique
|
||||
data Json
|
||||
|
|
|
|||
|
|
@ -590,6 +590,36 @@ class ProxyExtrasDBManager:
|
|||
f"Failed to resolve migration {migration_name}: {e.stderr}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def raise_if_lens_rename_pending() -> None:
|
||||
database_url: Final = os.environ.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
return
|
||||
try:
|
||||
import psycopg
|
||||
except ImportError as exc:
|
||||
raise RuntimeError("Install psycopg to verify Lens data safety before prisma db push.") from exc
|
||||
try:
|
||||
with psycopg.connect(
|
||||
ProxyExtrasDBManager._strip_prisma_query_params(database_url), connect_timeout=10, autocommit=True
|
||||
) as connection:
|
||||
legacy: Final = connection.execute(
|
||||
"SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid=c.relnamespace "
|
||||
"WHERE n.nspname=%s AND c.relname IN ('LiteLLM_Engine', 'LiteLLM_EngineRun', 'LiteLLM_EngineWorker') "
|
||||
"LIMIT 1",
|
||||
(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",),
|
||||
).fetchone()
|
||||
except psycopg.Error as exc:
|
||||
raise RuntimeError(
|
||||
"Cannot verify Lens data safety; refusing prisma db push. Check database connectivity and psycopg installation."
|
||||
) from exc
|
||||
if legacy is not None:
|
||||
raise RuntimeError(
|
||||
"Legacy Lens tables exist. prisma db push would drop saved Lens data. "
|
||||
"Apply the shipped 20261001100000_rename_lens migration to this database schema before retrying. "
|
||||
"Deployments using migration history can upgrade without --use_prisma_db_push instead."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def spend_logs_is_partitioned() -> bool:
|
||||
"""True when the connected database's LiteLLM_SpendLogs is a
|
||||
|
|
@ -895,6 +925,7 @@ class ProxyExtrasDBManager:
|
|||
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
|
||||
|
||||
if not use_migrate:
|
||||
ProxyExtrasDBManager.raise_if_lens_rename_pending()
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
original_dir = os.getcwd()
|
||||
|
|
@ -1398,6 +1429,7 @@ class ProxyExtrasDBManager:
|
|||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
# Use prisma db push with increased timeout
|
||||
ProxyExtrasDBManager.raise_if_lens_rename_pending()
|
||||
prisma_toolchain.run_prisma(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=prisma_command_timeout(),
|
||||
|
|
|
|||
|
|
@ -4,6 +4,10 @@ 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"
|
||||
dependencies = [
|
||||
"psycopg>=3.2,<4.0",
|
||||
"psycopg-binary>=3.2,<4.0",
|
||||
]
|
||||
license = "MIT"
|
||||
license-files = ["LICENSE"]
|
||||
authors = [
|
||||
|
|
|
|||
|
|
@ -69,6 +69,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
|
|||
# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
|
||||
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
|
||||
S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY"
|
||||
S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY"
|
||||
MAX_FILE_LIST_LIMIT: Final = 10000
|
||||
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
|
||||
|
|
|
|||
|
|
@ -498,6 +498,13 @@
|
|||
"ui_name": "Log Prompts Only",
|
||||
"description": "Log request messages to S3 but drop the model response from each logged object",
|
||||
"required": false
|
||||
},
|
||||
"s3_partition_granularity": {
|
||||
"type": "select",
|
||||
"ui_name": "Folder Partitioning",
|
||||
"description": "day writes one folder per date, hour adds an hour folder below each date (s3_v2 only)",
|
||||
"options": ["day", "hour"],
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "S3 Bucket (AWS) Logging Integration"
|
||||
|
|
|
|||
|
|
@ -16,8 +16,10 @@ from litellm.constants import (
|
|||
MAX_S3_OBJECT_KEY_BYTES,
|
||||
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES,
|
||||
S3_LOG_PROMPTS_ONLY_ENV_VAR,
|
||||
S3_PARTITION_GRANULARITY_ENV_VAR,
|
||||
S3_PREFIX_DIGEST_CHARS,
|
||||
)
|
||||
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
_S3_BOOL: Final = TypeAdapter(bool)
|
||||
|
|
@ -36,6 +38,18 @@ def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] |
|
|||
return True
|
||||
|
||||
|
||||
def resolve_s3_partition_granularity(
|
||||
configured: object, environ: Mapping[str, str] | None = None
|
||||
) -> S3PartitionGranularity:
|
||||
env: Final = os.environ if environ is None else environ
|
||||
raw: Final = env.get(S3_PARTITION_GRANULARITY_ENV_VAR) if configured is None else configured
|
||||
if raw == "hour":
|
||||
return "hour"
|
||||
if raw is not None and raw not in ("", "day"):
|
||||
verbose_logger.warning("s3 logging: s3_partition_granularity=%r is not one of day, hour, using day", raw)
|
||||
return "day"
|
||||
|
||||
|
||||
def _resolve_positive_int(setting: str, configured: object, fallback: int, *, reject_bool: bool) -> int:
|
||||
if configured is None or configured == "":
|
||||
return fallback
|
||||
|
|
@ -371,10 +385,11 @@ def get_s3_object_key(
|
|||
prefix: str,
|
||||
start_time: datetime,
|
||||
s3_file_name: str,
|
||||
partition_granularity: S3PartitionGranularity = "day",
|
||||
) -> str:
|
||||
sanitized_s3_file_name: Final = s3_file_name.replace("/", "_").replace(":", "_")
|
||||
configured_prefix: Final = (s3_path.rstrip("/") + "/" if s3_path else "") + prefix
|
||||
date_segment: Final = start_time.strftime("%Y-%m-%d") + "/"
|
||||
date_segment: Final = start_time.strftime("%Y-%m-%d/%H/" if partition_granularity == "hour" else "%Y-%m-%d/")
|
||||
# we need the s3 key to include the time, so we log cache hits too
|
||||
s3_object_key: Final = configured_prefix + date_segment + sanitized_s3_file_name + ".json"
|
||||
if len(s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES:
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ NOTE 1: S3 does not provide a BATCH PUT API endpoint; by default each element is
|
|||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
|
|
@ -28,6 +29,7 @@ from litellm.constants import (
|
|||
DEFAULT_S3_FLUSH_INTERVAL_SECONDS,
|
||||
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY,
|
||||
DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
|
||||
S3_PARTITION_GRANULARITY_ENV_VAR,
|
||||
)
|
||||
from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample
|
||||
from litellm.integrations.s3 import (
|
||||
|
|
@ -42,6 +44,7 @@ from litellm.integrations.s3 import (
|
|||
resolve_s3_max_concurrent_uploads,
|
||||
resolve_s3_max_queue_size,
|
||||
resolve_s3_max_retry_age_seconds,
|
||||
resolve_s3_partition_granularity,
|
||||
resolve_sse_params,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
|
|
@ -53,7 +56,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||||
from litellm.types.integrations.s3_v2 import S3PartitionGranularity, s3BatchLoggingElement
|
||||
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
|
||||
|
||||
from .custom_batch_logger import CustomBatchLogger
|
||||
|
|
@ -119,6 +122,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
_upload_limiter: asyncio.Semaphore | AdaptiveConcurrencyLimiter | None = None
|
||||
s3_drop_on_terminal_error: bool = True
|
||||
s3_max_retry_age_seconds: int | None = 3600
|
||||
s3_partition_granularity: object = None
|
||||
_partition_granularity_cache: tuple[object, S3PartitionGranularity] | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -147,6 +152,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
s3_server_side_encryption: str | None = None,
|
||||
s3_sse_kms_key_id: str | None = None,
|
||||
s3_log_prompts_only: bool | None = None,
|
||||
s3_partition_granularity: str | None = None,
|
||||
s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
|
||||
s3_max_queue_size: int | None = None,
|
||||
s3_max_retry_age_seconds: int | None = 3600,
|
||||
|
|
@ -195,6 +201,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
s3_server_side_encryption=s3_server_side_encryption,
|
||||
s3_sse_kms_key_id=s3_sse_kms_key_id,
|
||||
s3_log_prompts_only=s3_log_prompts_only,
|
||||
s3_partition_granularity=s3_partition_granularity,
|
||||
s3_max_concurrent_uploads=s3_max_concurrent_uploads,
|
||||
s3_max_queue_size=s3_max_queue_size,
|
||||
s3_max_retry_age_seconds=s3_max_retry_age_seconds,
|
||||
|
|
@ -271,6 +278,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
s3_server_side_encryption: str | None = None,
|
||||
s3_sse_kms_key_id: str | None = None,
|
||||
s3_log_prompts_only: bool | None = None,
|
||||
s3_partition_granularity: str | None = None,
|
||||
s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
|
||||
s3_max_queue_size: int | None = None,
|
||||
s3_max_retry_age_seconds: int | None = 3600,
|
||||
|
|
@ -331,6 +339,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only
|
||||
)
|
||||
|
||||
self.s3_partition_granularity = (
|
||||
params.get("s3_partition_granularity") if s3_partition_granularity is None else s3_partition_granularity
|
||||
)
|
||||
self._partition_granularity_cache = None
|
||||
|
||||
self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
|
||||
params.get("s3_server_side_encryption") or s3_server_side_encryption,
|
||||
params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id,
|
||||
|
|
@ -482,6 +495,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
"audit_logs/",
|
||||
now,
|
||||
f"{now.strftime('%H-%M-%S')}_{audit_log_id}",
|
||||
partition_granularity=self.resolve_partition_granularity(),
|
||||
)
|
||||
|
||||
element: Final = s3BatchLoggingElement(
|
||||
|
|
@ -758,6 +772,19 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
),
|
||||
)
|
||||
|
||||
def resolve_partition_granularity(self) -> S3PartitionGranularity:
|
||||
raw: Final = (
|
||||
os.environ.get(S3_PARTITION_GRANULARITY_ENV_VAR)
|
||||
if self.s3_partition_granularity is None
|
||||
else self.s3_partition_granularity
|
||||
)
|
||||
cached: Final = self._partition_granularity_cache
|
||||
if cached is not None and cached[0] == raw:
|
||||
return cached[1]
|
||||
resolved: Final = resolve_s3_partition_granularity(raw)
|
||||
self._partition_granularity_cache = (raw, resolved)
|
||||
return resolved
|
||||
|
||||
def create_s3_batch_logging_element(
|
||||
self,
|
||||
start_time: datetime,
|
||||
|
|
@ -803,11 +830,27 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
prefix_path,
|
||||
s3_file_name,
|
||||
)
|
||||
s3_object_key: Final = get_s3_object_key(
|
||||
s3_path=cast(str | None, self.s3_path) or "",
|
||||
prefix=prefix_path,
|
||||
start_time=start_time,
|
||||
s3_file_name=s3_file_name,
|
||||
|
||||
def object_key(partition_granularity: S3PartitionGranularity) -> str:
|
||||
return get_s3_object_key(
|
||||
s3_path=cast(str | None, self.s3_path) or "",
|
||||
prefix=prefix_path,
|
||||
start_time=start_time,
|
||||
s3_file_name=s3_file_name,
|
||||
partition_granularity=partition_granularity,
|
||||
)
|
||||
|
||||
metadata: Final = standard_logging_payload.get("metadata")
|
||||
cold_storage_object_key: Final = (
|
||||
metadata.get("cold_storage_object_key")
|
||||
if metadata is not None and litellm.cold_storage_custom_logger == "s3_v2"
|
||||
else None
|
||||
)
|
||||
s3_object_key: Final = (
|
||||
cold_storage_object_key
|
||||
if cold_storage_object_key is not None
|
||||
and cold_storage_object_key in (object_key("day"), object_key("hour"))
|
||||
else object_key(self.resolve_partition_granularity())
|
||||
)
|
||||
verbose_logger.debug("s3_object_key=%s", s3_object_key)
|
||||
|
||||
|
|
|
|||
|
|
@ -121,6 +121,7 @@ from litellm.llms.base_llm.search.transformation import SearchResponse
|
|||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
from litellm.types.containers.main import ContainerObject
|
||||
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
|
||||
from litellm.types.interactions import (
|
||||
InteractionsAPIResponse,
|
||||
InteractionsAPIStreamingResponse,
|
||||
|
|
@ -6086,6 +6087,7 @@ class StandardLoggingPayloadSetup:
|
|||
|
||||
# Get the actual s3_path from the configured cold storage logger instance
|
||||
s3_path = "" # default value
|
||||
partition_granularity: S3PartitionGranularity = "day"
|
||||
|
||||
# Try to get the actual logger instance from the logger name
|
||||
try:
|
||||
|
|
@ -6094,6 +6096,8 @@ class StandardLoggingPayloadSetup:
|
|||
)
|
||||
if custom_logger and hasattr(custom_logger, "s3_path") and getattr(custom_logger, "s3_path"):
|
||||
s3_path = getattr(custom_logger, "s3_path")
|
||||
if isinstance(custom_logger, S3V2Logger):
|
||||
partition_granularity = custom_logger.resolve_partition_granularity()
|
||||
except Exception:
|
||||
# If any error occurs in getting the logger instance, use default empty s3_path
|
||||
pass
|
||||
|
|
@ -6103,6 +6107,7 @@ class StandardLoggingPayloadSetup:
|
|||
prefix="", # Don't split by team alias for cold storage
|
||||
start_time=start_time,
|
||||
s3_file_name=s3_file_name,
|
||||
partition_granularity=partition_granularity,
|
||||
)
|
||||
|
||||
return s3_object_key
|
||||
|
|
|
|||
|
|
@ -34666,52 +34666,6 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/engine/workers/register": {
|
||||
"post": {
|
||||
"operationId": "register_worker_engine_workers_register_post",
|
||||
"requestBody": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/WorkerName"
|
||||
}
|
||||
}
|
||||
},
|
||||
"required": true
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/WorkerCreated"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Register Worker",
|
||||
"tags": [
|
||||
"mcp_discoverable"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/guardrails/register": {
|
||||
"post": {
|
||||
"description": "Register a guardrail for onboarding (team submission).\n\nAccepts a guardrail config in the\n[Generic Guardrail API](https://docs.litellm.ai/docs/adding_provider/generic_guardrail_api) format.\nThe submission is stored with status `pending_review` until an admin approves it.",
|
||||
|
|
@ -34804,6 +34758,52 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/lens/workers/register": {
|
||||
"post": {
|
||||
"operationId": "register_worker_lens_workers_register_post",
|
||||
"requestBody": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/WorkerName"
|
||||
}
|
||||
}
|
||||
},
|
||||
"required": true
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/WorkerCreated"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Register Worker",
|
||||
"tags": [
|
||||
"mcp_discoverable"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/register": {
|
||||
"post": {
|
||||
"operationId": "register_client_register_post",
|
||||
|
|
|
|||
|
|
@ -520,16 +520,16 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/v1/rag/ingest",
|
||||
"/rag/query",
|
||||
"/v1/rag/query",
|
||||
"/engine",
|
||||
"/engine/{engine_id}",
|
||||
"/engine/{engine_id}/runs",
|
||||
"/engine/{engine_id}/runs/{job_id}",
|
||||
"/engine/{engine_id}/executions/{execution_id}",
|
||||
"/engine/{engine_id}/cancel",
|
||||
"/engine/{engine_id}/findings/{finding_id}",
|
||||
"/engine/preview/sample",
|
||||
"/engine/workers/register",
|
||||
"/engine/workers/{worker_id}",
|
||||
"/lens",
|
||||
"/lens/{lens_id}",
|
||||
"/lens/{lens_id}/runs",
|
||||
"/lens/{lens_id}/runs/{job_id}",
|
||||
"/lens/{lens_id}/executions/{execution_id}",
|
||||
"/lens/{lens_id}/cancel",
|
||||
"/lens/{lens_id}/findings/{finding_id}",
|
||||
"/lens/preview/sample",
|
||||
"/lens/workers/register",
|
||||
"/lens/workers/{worker_id}",
|
||||
"/v1/traces",
|
||||
"/v1/traces/{trace_id}",
|
||||
"/v1/traces/{trace_id}/spans/{span_id}",
|
||||
|
|
@ -3961,6 +3961,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
|
|||
"AWS_SECRET_ACCESS_KEY",
|
||||
"AWS_REGION_NAME",
|
||||
"S3_LOG_PROMPTS_ONLY",
|
||||
"S3_PARTITION_GRANULARITY",
|
||||
],
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -953,6 +953,9 @@ class PrismaManager:
|
|||
verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
|
||||
return False
|
||||
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
|
||||
ProxyExtrasDBManager.raise_if_lens_rename_pending()
|
||||
PrismaManager._raise_if_partitioned_spend_logs()
|
||||
run_prisma(
|
||||
[
|
||||
|
|
|
|||
|
|
@ -13,17 +13,17 @@ from pydantic import AwareDatetime, BaseModel, Field, TypeAdapter
|
|||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
|
||||
from litellm.proxy.engine.billing import validate_key
|
||||
from litellm.proxy.engine.models import (
|
||||
from litellm.proxy.lens.billing import validate_key
|
||||
from litellm.proxy.lens.models import (
|
||||
Claim,
|
||||
Engine,
|
||||
EngineList,
|
||||
EngineSettings,
|
||||
Execution,
|
||||
ExecutionContent,
|
||||
FindingDraft,
|
||||
FindingUpdate,
|
||||
Job,
|
||||
Lens,
|
||||
LensList,
|
||||
LensSettings,
|
||||
ModelRequest,
|
||||
ModelResult,
|
||||
Progress,
|
||||
|
|
@ -34,9 +34,9 @@ from litellm.proxy.engine.models import (
|
|||
Worker,
|
||||
WorkerCreated,
|
||||
)
|
||||
from litellm.proxy.engine.repository import EngineRepository, WriterDatabase
|
||||
from litellm.proxy.engine.sources import SourceReader, parse_execution
|
||||
from litellm.proxy.engine.state import (
|
||||
from litellm.proxy.lens.repository import LensRepository, WriterDatabase
|
||||
from litellm.proxy.lens.sources import SourceReader, parse_execution
|
||||
from litellm.proxy.lens.state import (
|
||||
can_access,
|
||||
claim_job,
|
||||
current_job,
|
||||
|
|
@ -46,17 +46,17 @@ from litellm.proxy.engine.state import (
|
|||
snapshot_finding,
|
||||
)
|
||||
|
||||
router: Final = APIRouter(prefix="/engine", tags=["Lens"]) # mutable-ok: FastAPI requires list
|
||||
router: Final = APIRouter(prefix="/lens", tags=["Lens"]) # mutable-ok: FastAPI requires list
|
||||
_bearer: Final = HTTPBearer()
|
||||
Auth: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
|
||||
|
||||
|
||||
def repository() -> EngineRepository:
|
||||
def repository() -> LensRepository:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(503, "Lens needs a connected Postgres database")
|
||||
return EngineRepository(WriterDatabase(writer_wrapper(prisma_client.db)))
|
||||
return LensRepository(WriterDatabase(writer_wrapper(prisma_client.db)))
|
||||
|
||||
|
||||
def source_reader() -> SourceReader:
|
||||
|
|
@ -73,11 +73,11 @@ def user_scope(auth: UserAPIKeyAuth, write: bool = False) -> Scope:
|
|||
raise HTTPException(403, "Lens requires proxy administrator access")
|
||||
|
||||
|
||||
async def get_engine(engine_id: str, scope: Scope) -> Engine:
|
||||
engine: Final = await repository().get(engine_id)
|
||||
if engine is None or not can_access(scope, engine.scope):
|
||||
async def get_lens(lens_id: str, scope: Scope) -> Lens:
|
||||
lens: Final = await repository().get(lens_id)
|
||||
if lens is None or not can_access(scope, lens.scope):
|
||||
raise HTTPException(404, "Lens not found")
|
||||
return engine
|
||||
return lens
|
||||
|
||||
|
||||
async def worker_auth(credentials: Annotated[HTTPAuthorizationCredentials, Depends(_bearer)]) -> Worker:
|
||||
|
|
@ -90,9 +90,9 @@ async def worker_auth(credentials: Annotated[HTTPAuthorizationCredentials, Depen
|
|||
WorkerAuth: TypeAlias = Annotated[Worker, Depends(worker_auth)]
|
||||
|
||||
|
||||
async def assigned(engine_id: str, job_id: str, worker: Worker) -> tuple[Engine, Job]:
|
||||
engine: Final = await get_engine(engine_id, worker.scope)
|
||||
job: Final = current_job(engine)
|
||||
async def assigned(lens_id: str, job_id: str, worker: Worker) -> tuple[Lens, Job]:
|
||||
lens: Final = await get_lens(lens_id, worker.scope)
|
||||
job: Final = current_job(lens)
|
||||
if (
|
||||
job is None
|
||||
or job.id != job_id
|
||||
|
|
@ -102,16 +102,16 @@ async def assigned(engine_id: str, job_id: str, worker: Worker) -> tuple[Engine,
|
|||
or job.lease_until <= datetime.now(timezone.utc)
|
||||
):
|
||||
raise HTTPException(409, "This worker no longer owns the job")
|
||||
return engine, job
|
||||
return lens, job
|
||||
|
||||
|
||||
def required(engine: Engine | None) -> Engine:
|
||||
if engine is None:
|
||||
def required(lens: Lens | None) -> Lens:
|
||||
if lens is None:
|
||||
raise HTTPException(409, "Lens changed concurrently; retry the operation")
|
||||
return engine
|
||||
return lens
|
||||
|
||||
|
||||
def validate_selection(settings: EngineSettings) -> None:
|
||||
def validate_selection(settings: LensSettings) -> None:
|
||||
for identity in settings.execution_ids:
|
||||
try:
|
||||
source, _, _, _ = parse_execution(identity)
|
||||
|
|
@ -121,7 +121,7 @@ def validate_selection(settings: EngineSettings) -> None:
|
|||
raise HTTPException(422, "Choose execution IDs returned by the activity preview")
|
||||
|
||||
|
||||
def validate_model(settings: EngineSettings, auth: UserAPIKeyAuth) -> None:
|
||||
def validate_model(settings: LensSettings, auth: UserAPIKeyAuth) -> None:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
validate_selection(settings)
|
||||
|
|
@ -137,24 +137,24 @@ def validate_model(settings: EngineSettings, auth: UserAPIKeyAuth) -> None:
|
|||
raise HTTPException(403, "This key does not have access to the analysis model")
|
||||
|
||||
|
||||
@router.get("", response_model=EngineList)
|
||||
async def list_engines(auth: Auth) -> EngineList:
|
||||
@router.get("", response_model=LensList)
|
||||
async def list_lenses(auth: Auth) -> LensList:
|
||||
from litellm.proxy import tracing_endpoints
|
||||
|
||||
scope: Final = user_scope(auth)
|
||||
return EngineList(
|
||||
engines=tuple(e for e in await repository().engines() if can_access(scope, e.scope)),
|
||||
return LensList(
|
||||
lenses=tuple(e for e in await repository().lenses() if can_access(scope, e.scope)),
|
||||
workers=tuple(w for w in await repository().workers() if can_access(scope, w.scope)),
|
||||
tracing_enabled=tracing_endpoints.receiver is not None,
|
||||
)
|
||||
|
||||
|
||||
@router.post("", response_model=Engine)
|
||||
async def create_engine(settings: EngineSettings, auth: Auth) -> Engine:
|
||||
@router.post("", response_model=Lens)
|
||||
async def create_lens(settings: LensSettings, auth: Auth) -> Lens:
|
||||
scope: Final = user_scope(auth, write=True)
|
||||
validate_model(settings, auth)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
engine: Final = Engine(
|
||||
lens: Final = Lens(
|
||||
id=str(uuid4()),
|
||||
scope=scope,
|
||||
settings=settings,
|
||||
|
|
@ -162,16 +162,16 @@ async def create_engine(settings: EngineSettings, auth: Auth) -> Engine:
|
|||
next_run_at=now,
|
||||
budget_month=now.strftime("%Y-%m"),
|
||||
)
|
||||
return await repository().create(queue_job(engine, now, str(uuid4())))
|
||||
return await repository().create(queue_job(lens, now, str(uuid4())))
|
||||
|
||||
|
||||
@router.put("/{engine_id}", response_model=Engine)
|
||||
async def update_engine(engine_id: str, settings: EngineSettings, auth: Auth) -> Engine:
|
||||
await get_engine(engine_id, user_scope(auth, write=True))
|
||||
@router.put("/{lens_id}", response_model=Lens)
|
||||
async def update_lens(lens_id: str, settings: LensSettings, auth: Auth) -> Lens:
|
||||
await get_lens(lens_id, user_scope(auth, write=True))
|
||||
validate_model(settings, auth)
|
||||
return required(
|
||||
await repository().update(
|
||||
engine_id,
|
||||
lens_id,
|
||||
lambda e: e.model_copy(
|
||||
update=MappingProxyType(
|
||||
{
|
||||
|
|
@ -184,47 +184,47 @@ async def update_engine(engine_id: str, settings: EngineSettings, auth: Auth) ->
|
|||
)
|
||||
|
||||
|
||||
@router.post("/{engine_id}/runs", response_model=Engine)
|
||||
async def run_engine(engine_id: str, body: RunRequest, auth: Auth) -> Engine:
|
||||
await get_engine(engine_id, user_scope(auth, write=True))
|
||||
@router.post("/{lens_id}/runs", response_model=Lens)
|
||||
async def run_lens(lens_id: str, body: RunRequest, auth: Auth) -> Lens:
|
||||
await get_lens(lens_id, user_scope(auth, write=True))
|
||||
if body.settings is not None:
|
||||
validate_model(body.settings, auth)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
job_id: Final = str(uuid4())
|
||||
return required(
|
||||
await repository().update(engine_id, lambda e: queue_job(e, now, job_id, body.lookback_hours, body.settings))
|
||||
await repository().update(lens_id, lambda e: queue_job(e, now, job_id, body.lookback_hours, body.settings))
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{engine_id}", response_model=Engine)
|
||||
async def read_engine(engine_id: str, auth: Auth) -> Engine:
|
||||
return await get_engine(engine_id, user_scope(auth))
|
||||
@router.get("/{lens_id}", response_model=Lens)
|
||||
async def read_lens(lens_id: str, auth: Auth) -> Lens:
|
||||
return await get_lens(lens_id, user_scope(auth))
|
||||
|
||||
|
||||
@router.get("/{engine_id}/runs", response_model=tuple[Job, ...])
|
||||
async def list_runs(engine_id: str, auth: Auth, offset: int = Query(default=0, ge=0)) -> tuple[Job, ...]:
|
||||
await get_engine(engine_id, user_scope(auth))
|
||||
@router.get("/{lens_id}/runs", response_model=tuple[Job, ...])
|
||||
async def list_runs(lens_id: str, auth: Auth, offset: int = Query(default=0, ge=0)) -> tuple[Job, ...]:
|
||||
await get_lens(lens_id, user_scope(auth))
|
||||
return tuple(
|
||||
j.model_copy(update=MappingProxyType({"sample": None, "findings": None, "assessments": ()}))
|
||||
for j in await repository().jobs(engine_id, offset)
|
||||
for j in await repository().jobs(lens_id, offset)
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{engine_id}/runs/{job_id}", response_model=Job)
|
||||
async def read_run(engine_id: str, job_id: str, auth: Auth) -> Job:
|
||||
await get_engine(engine_id, user_scope(auth))
|
||||
job: Final = await repository().job(engine_id, job_id)
|
||||
@router.get("/{lens_id}/runs/{job_id}", response_model=Job)
|
||||
async def read_run(lens_id: str, job_id: str, auth: Auth) -> Job:
|
||||
await get_lens(lens_id, user_scope(auth))
|
||||
job: Final = await repository().job(lens_id, job_id)
|
||||
if job is None:
|
||||
raise HTTPException(404, "Investigation not found")
|
||||
return job
|
||||
|
||||
|
||||
@router.post("/{engine_id}/cancel", response_model=Engine)
|
||||
async def cancel_engine(engine_id: str, auth: Auth) -> Engine:
|
||||
await get_engine(engine_id, user_scope(auth, write=True))
|
||||
@router.post("/{lens_id}/cancel", response_model=Lens)
|
||||
async def cancel_lens(lens_id: str, auth: Auth) -> Lens:
|
||||
await get_lens(lens_id, user_scope(auth, write=True))
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
|
||||
def cancel(e: Engine) -> Engine:
|
||||
def cancel(e: Lens) -> Lens:
|
||||
job: Final = current_job(e)
|
||||
if job is None:
|
||||
return e
|
||||
|
|
@ -235,15 +235,15 @@ async def cancel_engine(engine_id: str, auth: Auth) -> Engine:
|
|||
update=MappingProxyType({"next_run_at": now + timedelta(minutes=e.settings.interval_minutes)})
|
||||
)
|
||||
|
||||
return required(await repository().update(engine_id, cancel))
|
||||
return required(await repository().update(lens_id, cancel))
|
||||
|
||||
|
||||
@router.patch("/{engine_id}/findings/{finding_id}", response_model=Engine)
|
||||
async def update_finding(engine_id: str, finding_id: str, body: FindingUpdate, auth: Auth) -> Engine:
|
||||
await get_engine(engine_id, user_scope(auth, write=True))
|
||||
@router.patch("/{lens_id}/findings/{finding_id}", response_model=Lens)
|
||||
async def update_finding(lens_id: str, finding_id: str, body: FindingUpdate, auth: Auth) -> Lens:
|
||||
await get_lens(lens_id, user_scope(auth, write=True))
|
||||
return required(
|
||||
await repository().update(
|
||||
engine_id,
|
||||
lens_id,
|
||||
lambda e: e.model_copy(
|
||||
update=MappingProxyType(
|
||||
{
|
||||
|
|
@ -260,7 +260,7 @@ async def update_finding(engine_id: str, finding_id: str, body: FindingUpdate, a
|
|||
class Preview(BaseModel):
|
||||
as_of: AwareDatetime | None = None
|
||||
offset: int = Field(default=0, ge=0)
|
||||
settings: EngineSettings
|
||||
settings: LensSettings
|
||||
lookback_hours: int = Field(default=24, ge=1, le=720)
|
||||
|
||||
|
||||
|
|
@ -335,7 +335,7 @@ async def claim(worker: WorkerAuth, protocol_version: int = 1) -> Claim | None:
|
|||
raise HTTPException(409, "Assign an analysis key to this worker in Lens setup")
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
await repository().heartbeat(worker.id, now.isoformat())
|
||||
for candidate in await repository().engines():
|
||||
for candidate in await repository().lenses():
|
||||
if not can_access(worker.scope, candidate.scope):
|
||||
continue
|
||||
if claimed := await claim_candidate(candidate, worker, now):
|
||||
|
|
@ -343,12 +343,12 @@ async def claim(worker: WorkerAuth, protocol_version: int = 1) -> Claim | None:
|
|||
return None
|
||||
|
||||
|
||||
@router.post("/worker/{engine_id}/{job_id}/progress", response_model=bool)
|
||||
async def progress(engine_id: str, job_id: str, body: Progress, worker: WorkerAuth) -> bool:
|
||||
await assigned(engine_id, job_id, worker)
|
||||
@router.post("/worker/{lens_id}/{job_id}/progress", response_model=bool)
|
||||
async def progress(lens_id: str, job_id: str, body: Progress, worker: WorkerAuth) -> bool:
|
||||
await assigned(lens_id, job_id, worker)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
|
||||
def renew(e: Engine) -> Engine:
|
||||
def renew(e: Lens) -> Lens:
|
||||
job: Final = current_job(e)
|
||||
if job is None or job.id != job_id or job.worker_id != worker.id:
|
||||
return e
|
||||
|
|
@ -361,21 +361,21 @@ async def progress(engine_id: str, job_id: str, body: Progress, worker: WorkerAu
|
|||
),
|
||||
)
|
||||
|
||||
required(await repository().update(engine_id, renew))
|
||||
required(await repository().update(lens_id, renew))
|
||||
await repository().heartbeat(worker.id, now.isoformat())
|
||||
return True
|
||||
|
||||
|
||||
@router.get("/worker/{engine_id}/{job_id}/sample", response_model=Sample)
|
||||
async def sample(engine_id: str, job_id: str, worker: WorkerAuth) -> Sample:
|
||||
engine, job = await assigned(engine_id, job_id, worker)
|
||||
@router.get("/worker/{lens_id}/{job_id}/sample", response_model=Sample)
|
||||
async def sample(lens_id: str, job_id: str, worker: WorkerAuth) -> Sample:
|
||||
lens, job = await assigned(lens_id, job_id, worker)
|
||||
if job.sample is not None:
|
||||
return job.sample
|
||||
pages: list[Sample] = [] # mutable-ok: freeze selection after stable cursor traversal
|
||||
cursor = "" # rebind-ok: advance by immutable identity, never by shifting row positions
|
||||
while True:
|
||||
page = await source_reader().sample(
|
||||
engine.scope,
|
||||
lens.scope,
|
||||
job.settings,
|
||||
int(job.start.timestamp() * 1000),
|
||||
int(job.end.timestamp() * 1000),
|
||||
|
|
@ -390,7 +390,7 @@ async def sample(engine_id: str, job_id: str, worker: WorkerAuth) -> Sample:
|
|||
) # comprehension-ok: flatten query pages
|
||||
selected: Final = Sample(executions=executions, eligible=pages[0].eligible, selected=len(executions))
|
||||
|
||||
def freeze(e: Engine) -> Engine:
|
||||
def freeze(e: Lens) -> Lens:
|
||||
active: Final = current_job(e)
|
||||
if active is None or active.id != job_id or active.worker_id != worker.id:
|
||||
raise HTTPException(409, "Job was cancelled or reassigned")
|
||||
|
|
@ -400,45 +400,45 @@ async def sample(engine_id: str, job_id: str, worker: WorkerAuth) -> Sample:
|
|||
else e
|
||||
)
|
||||
|
||||
updated: Final = required(await repository().update(engine_id, freeze))
|
||||
updated: Final = required(await repository().update(lens_id, freeze))
|
||||
frozen: Final = next(j for j in updated.jobs if j.id == job_id).sample
|
||||
if frozen is None:
|
||||
raise HTTPException(409, "Could not freeze the sample")
|
||||
return frozen
|
||||
|
||||
|
||||
@router.get("/worker/{engine_id}/{job_id}/content", response_model=ExecutionContent)
|
||||
@router.get("/worker/{lens_id}/{job_id}/content", response_model=ExecutionContent)
|
||||
async def content(
|
||||
engine_id: str,
|
||||
lens_id: str,
|
||||
job_id: str,
|
||||
execution_id: str,
|
||||
worker: WorkerAuth,
|
||||
cursor: str = "",
|
||||
offset: int = Query(default=0, ge=0),
|
||||
) -> ExecutionContent:
|
||||
engine, job = await assigned(engine_id, job_id, worker)
|
||||
lens, job = await assigned(lens_id, job_id, worker)
|
||||
selected: Final = job.sample or Sample(executions=(), eligible=0)
|
||||
execution: Final = next((e for e in selected.executions if e.id == execution_id), None)
|
||||
if execution is None:
|
||||
raise HTTPException(404, "Execution is outside this job's sample")
|
||||
return await source_reader().content(engine.scope, execution, cursor, offset)
|
||||
return await source_reader().content(lens.scope, execution, cursor, offset)
|
||||
|
||||
|
||||
@router.post("/worker/{engine_id}/{job_id}/model", response_model=ModelResult)
|
||||
async def model(engine_id: str, job_id: str, body: ModelRequest, worker: WorkerAuth, request: Request) -> ModelResult:
|
||||
from litellm.proxy.engine.inference import analyze
|
||||
@router.post("/worker/{lens_id}/{job_id}/model", response_model=ModelResult)
|
||||
async def model(lens_id: str, job_id: str, body: ModelRequest, worker: WorkerAuth, request: Request) -> ModelResult:
|
||||
from litellm.proxy.lens.inference import analyze
|
||||
|
||||
engine, job = await assigned(engine_id, job_id, worker)
|
||||
return await analyze(repository(), engine, job, worker, body, request)
|
||||
lens, job = await assigned(lens_id, job_id, worker)
|
||||
return await analyze(repository(), lens, job, worker, body, request)
|
||||
|
||||
|
||||
@router.post("/worker/{engine_id}/{job_id}/result", response_model=Engine)
|
||||
async def result(engine_id: str, job_id: str, body: Result, worker: WorkerAuth) -> Engine:
|
||||
engine: Final = await get_engine(engine_id, worker.scope)
|
||||
old: Final = next((j for j in engine.jobs if j.id == job_id), None)
|
||||
@router.post("/worker/{lens_id}/{job_id}/result", response_model=Lens)
|
||||
async def result(lens_id: str, job_id: str, body: Result, worker: WorkerAuth) -> Lens:
|
||||
lens: Final = await get_lens(lens_id, worker.scope)
|
||||
old: Final = next((j for j in lens.jobs if j.id == job_id), None)
|
||||
if old and old.status in ("completed", "failed") and old.worker_id == worker.id:
|
||||
return engine
|
||||
_, job = await assigned(engine_id, job_id, worker)
|
||||
return lens
|
||||
_, job = await assigned(lens_id, job_id, worker)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
selected: Final = job.sample or Sample(executions=(), eligible=0)
|
||||
allowed: Final = frozenset(e.id for e in selected.executions)
|
||||
|
|
@ -455,9 +455,9 @@ async def result(engine_id: str, job_id: str, body: Result, worker: WorkerAuth)
|
|||
raise HTTPException(422, "Finding references evidence outside the job")
|
||||
|
||||
for finding in body.findings:
|
||||
await validate_finding(engine, selected, finding)
|
||||
await validate_finding(lens, selected, finding)
|
||||
|
||||
def finish(e: Engine) -> Engine:
|
||||
def finish(e: Lens) -> Lens:
|
||||
active: Final = current_job(e)
|
||||
if active is None or active.id != job_id or active.worker_id != worker.id:
|
||||
return e
|
||||
|
|
@ -488,29 +488,29 @@ async def result(engine_id: str, job_id: str, body: Result, worker: WorkerAuth)
|
|||
)
|
||||
)
|
||||
|
||||
return required(await repository().update(engine_id, finish))
|
||||
return required(await repository().update(lens_id, finish))
|
||||
|
||||
|
||||
def merge_results(engine: Engine, result: Result, revision: int, now: datetime) -> Engine:
|
||||
def merge_one(current: Engine, draft: FindingDraft) -> Engine:
|
||||
def merge_results(lens: Lens, result: Result, revision: int, now: datetime) -> Lens:
|
||||
def merge_one(current: Lens, draft: FindingDraft) -> Lens:
|
||||
finding: Final = merge_finding(current, draft, revision, now)
|
||||
return current.model_copy(
|
||||
update=MappingProxyType({"findings": (finding, *(f for f in current.findings if f.id != finding.id))})
|
||||
)
|
||||
|
||||
return reduce(merge_one, result.findings, engine)
|
||||
return reduce(merge_one, result.findings, lens)
|
||||
|
||||
|
||||
@router.post("/worker/{engine_id}/{job_id}/heartbeat", response_model=bool)
|
||||
async def heartbeat(engine_id: str, job_id: str, worker: WorkerAuth) -> bool:
|
||||
_, job = await assigned(engine_id, job_id, worker)
|
||||
return await progress(engine_id, job_id, Progress(stage=job.stage, coverage=job.coverage), worker)
|
||||
@router.post("/worker/{lens_id}/{job_id}/heartbeat", response_model=bool)
|
||||
async def heartbeat(lens_id: str, job_id: str, worker: WorkerAuth) -> bool:
|
||||
_, job = await assigned(lens_id, job_id, worker)
|
||||
return await progress(lens_id, job_id, Progress(stage=job.stage, coverage=job.coverage), worker)
|
||||
|
||||
|
||||
async def claim_candidate(candidate: Engine, worker: Worker, now: datetime) -> Claim | None:
|
||||
async def claim_candidate(candidate: Lens, worker: Worker, now: datetime) -> Claim | None:
|
||||
job_id: Final = str(uuid4())
|
||||
|
||||
def schedule(e: Engine) -> Engine:
|
||||
def schedule(e: Lens) -> Lens:
|
||||
scheduled: Final = queue_job(e, now, job_id) if e.settings.enabled and e.next_run_at <= now else e
|
||||
return claim_job(scheduled, worker, now)
|
||||
|
||||
|
|
@ -519,31 +519,31 @@ async def claim_candidate(candidate: Engine, worker: Worker, now: datetime) -> C
|
|||
return None
|
||||
job: Final = current_job(updated)
|
||||
if job and job.worker_id == worker.id and job.status == "running" and job != current_job(candidate):
|
||||
return Claim(engine_id=updated.id, job=job, findings=updated.findings)
|
||||
return Claim(lens_id=updated.id, job=job, findings=updated.findings)
|
||||
return None
|
||||
|
||||
|
||||
async def validate_finding(engine: Engine, selected: Sample, finding: FindingDraft) -> None:
|
||||
previous: Final = next((f for f in engine.findings if f.id == finding.existing_finding_id), None)
|
||||
async def validate_finding(lens: Lens, selected: Sample, finding: FindingDraft) -> None:
|
||||
previous: Final = next((f for f in lens.findings if f.id == finding.existing_finding_id), None)
|
||||
if finding.existing_finding_id and (previous is None or previous.check_id != finding.check_id):
|
||||
raise HTTPException(422, "Existing finding must belong to the same check")
|
||||
for evidence in finding.evidence:
|
||||
if not await source_reader().verify_evidence(
|
||||
engine.scope, next(e for e in selected.executions if e.id == evidence.execution_id), evidence
|
||||
lens.scope, next(e for e in selected.executions if e.id == evidence.execution_id), evidence
|
||||
):
|
||||
raise HTTPException(422, "Evidence quote does not match stored content")
|
||||
|
||||
|
||||
@router.get("/{engine_id}/executions/{execution_id}", response_model=ExecutionContent)
|
||||
@router.get("/{lens_id}/executions/{execution_id}", response_model=ExecutionContent)
|
||||
async def evidence_content(
|
||||
engine_id: str, execution_id: str, auth: Auth, cursor: str = "", offset: int = Query(default=0, ge=0)
|
||||
lens_id: str, execution_id: str, auth: Auth, cursor: str = "", offset: int = Query(default=0, ge=0)
|
||||
) -> ExecutionContent:
|
||||
engine: Final = await get_engine(engine_id, user_scope(auth))
|
||||
lens: Final = await get_lens(lens_id, user_scope(auth))
|
||||
try:
|
||||
source, team, trace_id, trace_ref = parse_execution(execution_id)
|
||||
except ValueError:
|
||||
raise HTTPException(404, "Execution not found")
|
||||
if source not in ("traces", "requests") or (not engine.scope.all_teams and team != engine.scope.team_id):
|
||||
if source not in ("traces", "requests") or (not lens.scope.all_teams and team != lens.scope.team_id):
|
||||
raise HTTPException(404, "Execution not found")
|
||||
execution: Final = Execution(
|
||||
id=execution_id,
|
||||
|
|
@ -556,4 +556,4 @@ async def evidence_content(
|
|||
span_count=1,
|
||||
root_seen=source == "requests",
|
||||
)
|
||||
return await source_reader().content(engine.scope, execution, cursor, offset)
|
||||
return await source_reader().content(lens.scope, execution, cursor, offset)
|
||||
|
|
@ -8,10 +8,10 @@ from pydantic import BaseModel, ConfigDict, Field
|
|||
import litellm
|
||||
from litellm.integrations.clickhouse.context import lens_analysis
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import inherit_message_logging_privacy
|
||||
from litellm.proxy.engine.billing import complete, validate_key
|
||||
from litellm.proxy.engine.models import Engine, Job, ModelRequest, ModelResult, Worker
|
||||
from litellm.proxy.engine.repository import EngineRepository
|
||||
from litellm.proxy.engine.state import current_job, renew_budget, replace_job
|
||||
from litellm.proxy.lens.billing import complete, validate_key
|
||||
from litellm.proxy.lens.models import Job, Lens, ModelRequest, ModelResult, Worker
|
||||
from litellm.proxy.lens.repository import LensRepository
|
||||
from litellm.proxy.lens.state import current_job, renew_budget, replace_job
|
||||
from litellm.types.utils import CostPerToken, ModelResponse
|
||||
|
||||
|
||||
|
|
@ -87,7 +87,7 @@ def quote(deployments: tuple[Deployment, ...], prompt: str) -> float:
|
|||
|
||||
|
||||
async def analyze(
|
||||
repo: EngineRepository, engine: Engine, job: Job, worker: Worker, body: ModelRequest, request: Request
|
||||
repo: LensRepository, lens: Lens, job: Job, worker: Worker, body: ModelRequest, request: Request
|
||||
) -> ModelResult:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
|
|
@ -106,7 +106,7 @@ async def analyze(
|
|||
estimate: Final = quote(deployments, body.prompt)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
|
||||
def reserve(e: Engine) -> Engine:
|
||||
def reserve(e: Lens) -> Lens:
|
||||
current: Final = renew_budget(e, now)
|
||||
active: Final = current_job(current)
|
||||
if (
|
||||
|
|
@ -124,7 +124,7 @@ async def analyze(
|
|||
).model_copy(update=MappingProxyType({"spent": current.spent + estimate}))
|
||||
|
||||
async def reserve_budget() -> None:
|
||||
if await repo.update(engine.id, reserve) is None:
|
||||
if await repo.update(lens.id, reserve) is None:
|
||||
raise HTTPException(409, "Could not reserve analysis budget")
|
||||
|
||||
data: Final[dict[str, object]] = { # mutable-ok: proxy processing enriches request data
|
||||
|
|
@ -140,8 +140,8 @@ async def analyze(
|
|||
"disable_fallbacks": True,
|
||||
"response_format": {"type": "json_object"}, # mutable-ok: provider response-format JSON
|
||||
"metadata": { # mutable-ok: request processing enriches metadata
|
||||
"tags": ["litellm-engine"], # mutable-ok: logging callbacks require a list
|
||||
"lens_id": engine.id,
|
||||
"tags": ["litellm-lens"], # mutable-ok: logging callbacks require a list
|
||||
"lens_id": lens.id,
|
||||
"lens_run_id": job.id,
|
||||
"lens_worker_id": worker.id,
|
||||
"user_api_key_team_id": team_id,
|
||||
|
|
@ -153,7 +153,7 @@ async def analyze(
|
|||
parsed: Final = Completion.model_validate_json(response.model_dump_json())
|
||||
cost: Final = billed_cost if billed_cost is not None else completion_charge(deployments, response, estimate)
|
||||
|
||||
def settle(e: Engine) -> Engine:
|
||||
def settle(e: Lens) -> Lens:
|
||||
charged: Final = next((j for j in e.jobs if j.id == job.id), None)
|
||||
adjusted: Final = (
|
||||
e.model_copy(update=MappingProxyType({"spent": max(0, e.spent - estimate + cost)}))
|
||||
|
|
@ -168,7 +168,7 @@ async def analyze(
|
|||
else adjusted
|
||||
)
|
||||
|
||||
await repo.update(engine.id, settle)
|
||||
await repo.update(lens.id, settle)
|
||||
return ModelResult(content=parsed.choices[0].message.content or "{}", cost=cost)
|
||||
|
||||
|
||||
|
|
@ -25,7 +25,7 @@ class Check(Record):
|
|||
enabled: bool = True
|
||||
|
||||
|
||||
class EngineSettings(Record):
|
||||
class LensSettings(Record):
|
||||
name: str = Field(min_length=1, max_length=100)
|
||||
context: str = Field(default="", max_length=6000)
|
||||
source: Literal["traces", "requests", "both"] = "traces"
|
||||
|
|
@ -44,7 +44,7 @@ class EngineSettings(Record):
|
|||
monthly_budget: float = Field(default=20, gt=0, le=100000, allow_inf_nan=False)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def unique_checks(self) -> "EngineSettings":
|
||||
def unique_checks(self) -> "LensSettings":
|
||||
if len(frozenset(c.id for c in self.checks)) != len(self.checks):
|
||||
raise ValueError("Each check must have a unique ID")
|
||||
if not self.context.strip() and not any(c.enabled for c in self.checks):
|
||||
|
|
@ -163,7 +163,7 @@ class Job(Record):
|
|||
created_at: datetime
|
||||
start: datetime
|
||||
end: datetime
|
||||
settings: EngineSettings
|
||||
settings: LensSettings
|
||||
revision: int
|
||||
worker_id: str | None = None
|
||||
lease_until: datetime | None = None
|
||||
|
|
@ -177,10 +177,10 @@ class Job(Record):
|
|||
assessments: tuple[RunAssessment, ...] = ()
|
||||
|
||||
|
||||
class Engine(Record):
|
||||
class Lens(Record):
|
||||
id: str
|
||||
scope: Scope
|
||||
settings: EngineSettings
|
||||
settings: LensSettings
|
||||
revision: int = 1
|
||||
version: int = 0
|
||||
created_at: datetime
|
||||
|
|
@ -206,14 +206,14 @@ class WorkerCreated(Record):
|
|||
token: str
|
||||
|
||||
|
||||
class EngineList(Record):
|
||||
engines: tuple[Engine, ...]
|
||||
class LensList(Record):
|
||||
lenses: tuple[Lens, ...]
|
||||
workers: tuple[Worker, ...]
|
||||
tracing_enabled: bool
|
||||
|
||||
|
||||
class RunRequest(Record):
|
||||
settings: EngineSettings | None = None
|
||||
settings: LensSettings | None = None
|
||||
lookback_hours: int | None = Field(default=None, ge=1, le=720)
|
||||
|
||||
|
||||
|
|
@ -223,7 +223,7 @@ class FindingUpdate(Record):
|
|||
|
||||
|
||||
class Claim(Record):
|
||||
engine_id: str
|
||||
lens_id: str
|
||||
job: Job
|
||||
findings: tuple[Finding, ...]
|
||||
|
||||
|
|
@ -5,7 +5,7 @@ from typing import Final, Protocol
|
|||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.engine.models import Engine, Job, Worker
|
||||
from litellm.proxy.lens.models import Job, Lens, Worker
|
||||
|
||||
|
||||
class Database(Protocol):
|
||||
|
|
@ -20,44 +20,44 @@ class Row(BaseModel):
|
|||
_ROWS: Final = TypeAdapter(tuple[Row, ...])
|
||||
|
||||
|
||||
class EngineRepository:
|
||||
class LensRepository:
|
||||
def __init__(self, db: Database) -> None:
|
||||
self.db: Final = db
|
||||
|
||||
async def engines(self) -> tuple[Engine, ...]:
|
||||
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_Engine" ORDER BY id'))
|
||||
return tuple(Engine.model_validate(row.data) for row in rows)
|
||||
async def lenses(self) -> tuple[Lens, ...]:
|
||||
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_Lens" ORDER BY id'))
|
||||
return tuple(Lens.model_validate(row.data) for row in rows)
|
||||
|
||||
async def get(self, engine_id: str) -> Engine | None:
|
||||
async def get(self, lens_id: str) -> Lens | None:
|
||||
rows: Final = _ROWS.validate_python(
|
||||
await self.db.query_raw(
|
||||
'SELECT data FROM "LiteLLM_Engine" WHERE id=$1',
|
||||
engine_id,
|
||||
'SELECT data FROM "LiteLLM_Lens" WHERE id=$1',
|
||||
lens_id,
|
||||
)
|
||||
)
|
||||
return Engine.model_validate(rows[0].data) if rows else None
|
||||
return Lens.model_validate(rows[0].data) if rows else None
|
||||
|
||||
async def create(self, engine: Engine) -> Engine:
|
||||
async def create(self, lens: Lens) -> Lens:
|
||||
await self.db.execute_raw(
|
||||
'INSERT INTO "LiteLLM_Engine" (id, version, data) VALUES ($1,0,$2::jsonb)',
|
||||
engine.id,
|
||||
engine.model_dump_json(),
|
||||
'INSERT INTO "LiteLLM_Lens" (id, version, data) VALUES ($1,0,$2::jsonb)',
|
||||
lens.id,
|
||||
lens.model_dump_json(),
|
||||
)
|
||||
return engine
|
||||
return lens
|
||||
|
||||
async def update(
|
||||
self, engine_id: str, transform: Callable[[Engine], Engine], attempts: int = 8, *, changed_only: bool = False
|
||||
) -> Engine | None:
|
||||
self, lens_id: str, transform: Callable[[Lens], Lens], attempts: int = 8, *, changed_only: bool = False
|
||||
) -> Lens | None:
|
||||
for _ in range(attempts):
|
||||
completed, updated = await self._try_update(engine_id, transform, changed_only)
|
||||
completed, updated = await self._try_update(lens_id, transform, changed_only)
|
||||
if completed:
|
||||
return updated
|
||||
return None
|
||||
|
||||
async def _try_update(
|
||||
self, engine_id: str, transform: Callable[[Engine], Engine], changed_only: bool
|
||||
) -> tuple[bool, Engine | None]:
|
||||
previous: Final = await self.get(engine_id)
|
||||
self, lens_id: str, transform: Callable[[Lens], Lens], changed_only: bool
|
||||
) -> tuple[bool, Lens | None]:
|
||||
previous: Final = await self.get(lens_id)
|
||||
if previous is None:
|
||||
return True, None
|
||||
candidate: Final = transform(previous)
|
||||
|
|
@ -67,12 +67,12 @@ class EngineRepository:
|
|||
rows: Final = _ROWS.validate_python(
|
||||
await self.db.query_raw(
|
||||
"""WITH previous AS MATERIALIZED (
|
||||
SELECT data FROM "LiteLLM_Engine" WHERE id=$2 AND version=$3 FOR UPDATE
|
||||
SELECT data FROM "LiteLLM_Lens" WHERE id=$2 AND version=$3 FOR UPDATE
|
||||
), updated AS (
|
||||
UPDATE "LiteLLM_Engine" SET data=$1::jsonb, version=version+1
|
||||
UPDATE "LiteLLM_Lens" SET data=$1::jsonb, version=version+1
|
||||
WHERE id=$2 AND version=$3 AND EXISTS (SELECT 1 FROM previous) RETURNING id
|
||||
)
|
||||
, archived AS (INSERT INTO "LiteLLM_EngineRun" (id, engine_id, created_at, data)
|
||||
, archived AS (INSERT INTO "LiteLLM_LensRun" (id, lens_id, created_at, data)
|
||||
SELECT job->>'id', $2, (job->>'created_at')::timestamp, job
|
||||
FROM previous, jsonb_array_elements(previous.data->'jobs') AS job
|
||||
WHERE EXISTS (SELECT 1 FROM updated)
|
||||
|
|
@ -81,46 +81,46 @@ class EngineRepository:
|
|||
ON CONFLICT (id) DO NOTHING)
|
||||
SELECT to_jsonb(count(*)) AS data FROM updated""",
|
||||
updated.model_dump_json(),
|
||||
engine_id,
|
||||
lens_id,
|
||||
previous.version,
|
||||
)
|
||||
)
|
||||
return bool(rows and rows[0].data == 1), updated
|
||||
|
||||
async def jobs(self, engine_id: str, offset: int = 0) -> tuple[Job, ...]:
|
||||
async def jobs(self, lens_id: str, offset: int = 0) -> tuple[Job, ...]:
|
||||
rows: Final = _ROWS.validate_python(
|
||||
await self.db.query_raw(
|
||||
"""SELECT data FROM (
|
||||
SELECT data FROM "LiteLLM_EngineRun" WHERE engine_id=$1
|
||||
SELECT data FROM "LiteLLM_LensRun" WHERE lens_id=$1
|
||||
UNION ALL
|
||||
SELECT jsonb_array_elements(data->'jobs') AS data FROM "LiteLLM_Engine" WHERE id=$1
|
||||
SELECT jsonb_array_elements(data->'jobs') AS data FROM "LiteLLM_Lens" WHERE id=$1
|
||||
) AS jobs ORDER BY data->>'created_at' DESC, data->>'id' DESC LIMIT 50 OFFSET $2""",
|
||||
engine_id,
|
||||
lens_id,
|
||||
offset,
|
||||
)
|
||||
)
|
||||
return tuple(Job.model_validate(row.data) for row in rows)
|
||||
|
||||
async def job(self, engine_id: str, job_id: str) -> Job | None:
|
||||
async def job(self, lens_id: str, job_id: str) -> Job | None:
|
||||
rows: Final = _ROWS.validate_python(
|
||||
await self.db.query_raw(
|
||||
"""SELECT data FROM "LiteLLM_EngineRun" WHERE engine_id=$1 AND id=$2
|
||||
UNION ALL SELECT job AS data FROM "LiteLLM_Engine", jsonb_array_elements(data->'jobs') AS job
|
||||
"""SELECT data FROM "LiteLLM_LensRun" WHERE lens_id=$1 AND id=$2
|
||||
UNION ALL SELECT job AS data FROM "LiteLLM_Lens", jsonb_array_elements(data->'jobs') AS job
|
||||
WHERE id=$1 AND job->>'id'=$2 LIMIT 1""",
|
||||
engine_id,
|
||||
lens_id,
|
||||
job_id,
|
||||
)
|
||||
)
|
||||
return Job.model_validate(rows[0].data) if rows else None
|
||||
|
||||
async def workers(self) -> tuple[Worker, ...]:
|
||||
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_EngineWorker"'))
|
||||
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_LensWorker"'))
|
||||
return tuple(Worker.model_validate(row.data) for row in rows)
|
||||
|
||||
async def worker(self, token_hash: str) -> Worker | None:
|
||||
rows: Final = _ROWS.validate_python(
|
||||
await self.db.query_raw(
|
||||
'SELECT data FROM "LiteLLM_EngineWorker" WHERE token_hash=$1',
|
||||
'SELECT data FROM "LiteLLM_LensWorker" WHERE token_hash=$1',
|
||||
token_hash,
|
||||
)
|
||||
)
|
||||
|
|
@ -129,20 +129,20 @@ class EngineRepository:
|
|||
async def save_worker(self, worker: Worker, token_hash: str | None = None) -> None:
|
||||
if token_hash is not None:
|
||||
await self.db.execute_raw(
|
||||
'INSERT INTO "LiteLLM_EngineWorker" (id,token_hash,data) VALUES ($1,$2,$3::jsonb)',
|
||||
'INSERT INTO "LiteLLM_LensWorker" (id,token_hash,data) VALUES ($1,$2,$3::jsonb)',
|
||||
worker.id,
|
||||
token_hash,
|
||||
worker.model_dump_json(),
|
||||
)
|
||||
return
|
||||
await self.db.execute_raw(
|
||||
'UPDATE "LiteLLM_EngineWorker" SET data=$1::jsonb WHERE id=$2', worker.model_dump_json(), worker.id
|
||||
'UPDATE "LiteLLM_LensWorker" SET data=$1::jsonb WHERE id=$2', worker.model_dump_json(), worker.id
|
||||
)
|
||||
|
||||
async def set_worker_billing(self, worker_id: str, key_id: str) -> Worker | None:
|
||||
rows: Final = _ROWS.validate_python(
|
||||
await self.db.query_raw(
|
||||
"""UPDATE "LiteLLM_EngineWorker"
|
||||
"""UPDATE "LiteLLM_LensWorker"
|
||||
SET data=jsonb_set(data, '{analysis_key_id}', to_jsonb($1::text))
|
||||
WHERE id=$2 AND COALESCE((data->>'revoked')::boolean, false)=false RETURNING data""",
|
||||
key_id,
|
||||
|
|
@ -153,13 +153,13 @@ class EngineRepository:
|
|||
|
||||
async def revoke_worker(self, worker_id: str) -> None:
|
||||
await self.db.execute_raw(
|
||||
"""UPDATE "LiteLLM_EngineWorker" SET data=jsonb_set(data, '{revoked}', 'true') WHERE id=$1""",
|
||||
"""UPDATE "LiteLLM_LensWorker" SET data=jsonb_set(data, '{revoked}', 'true') WHERE id=$1""",
|
||||
worker_id,
|
||||
)
|
||||
|
||||
async def heartbeat(self, worker_id: str, now: str) -> None:
|
||||
await self.db.execute_raw(
|
||||
"""UPDATE "LiteLLM_EngineWorker" SET data=jsonb_set(data, '{last_seen}', to_jsonb($1::text)) WHERE id=$2""",
|
||||
"""UPDATE "LiteLLM_LensWorker" SET data=jsonb_set(data, '{last_seen}', to_jsonb($1::text)) WHERE id=$2""",
|
||||
now,
|
||||
worker_id,
|
||||
)
|
||||
|
|
@ -6,11 +6,11 @@ from typing import Final, Literal, Protocol
|
|||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm.proxy.engine.models import (
|
||||
EngineSettings,
|
||||
from litellm.proxy.lens.models import (
|
||||
Evidence,
|
||||
Execution,
|
||||
ExecutionContent,
|
||||
LensSettings,
|
||||
MetadataFilter,
|
||||
Sample,
|
||||
Scope,
|
||||
|
|
@ -93,7 +93,7 @@ class SourceReader:
|
|||
async def sample(
|
||||
self,
|
||||
scope: Scope,
|
||||
settings: EngineSettings,
|
||||
settings: LensSettings,
|
||||
start: int,
|
||||
end: int,
|
||||
offset: int = 0,
|
||||
|
|
@ -3,7 +3,7 @@ from datetime import datetime, timedelta
|
|||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.proxy.engine.models import Engine, EngineSettings, Finding, FindingDraft, Job, Scope, Worker
|
||||
from litellm.proxy.lens.models import Finding, FindingDraft, Job, Lens, LensSettings, Scope, Worker
|
||||
|
||||
|
||||
def can_access(viewer: Scope, target: Scope) -> bool:
|
||||
|
|
@ -14,46 +14,46 @@ def can_access(viewer: Scope, target: Scope) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def current_job(engine: Engine) -> Job | None:
|
||||
return next((job for job in engine.jobs if job.status in ("queued", "running")), None)
|
||||
def current_job(lens: Lens) -> Job | None:
|
||||
return next((job for job in lens.jobs if job.status in ("queued", "running")), None)
|
||||
|
||||
|
||||
def replace_job(engine: Engine, job: Job) -> Engine:
|
||||
return engine.model_copy(
|
||||
update=MappingProxyType({"jobs": tuple(job if old.id == job.id else old for old in engine.jobs)})
|
||||
def replace_job(lens: Lens, job: Job) -> Lens:
|
||||
return lens.model_copy(
|
||||
update=MappingProxyType({"jobs": tuple(job if old.id == job.id else old for old in lens.jobs)})
|
||||
)
|
||||
|
||||
|
||||
def queue_job(
|
||||
engine: Engine,
|
||||
lens: Lens,
|
||||
now: datetime,
|
||||
job_id: str,
|
||||
lookback_hours: int | None = None,
|
||||
settings: EngineSettings | None = None,
|
||||
) -> Engine:
|
||||
if current_job(engine):
|
||||
return engine
|
||||
selected: Final = settings or engine.settings
|
||||
settings: LensSettings | None = None,
|
||||
) -> Lens:
|
||||
if current_job(lens):
|
||||
return lens
|
||||
selected: Final = settings or lens.settings
|
||||
job: Final = Job(
|
||||
id=job_id,
|
||||
created_at=now,
|
||||
start=now - timedelta(hours=lookback_hours if lookback_hours is not None else selected.lookback_hours),
|
||||
end=now - timedelta(minutes=2),
|
||||
settings=selected,
|
||||
revision=engine.revision,
|
||||
revision=lens.revision,
|
||||
)
|
||||
return engine.model_copy(update=MappingProxyType({"jobs": (job,)}))
|
||||
return lens.model_copy(update=MappingProxyType({"jobs": (job,)}))
|
||||
|
||||
|
||||
def claim_job(engine: Engine, worker: Worker, now: datetime) -> Engine:
|
||||
job: Final = current_job(engine)
|
||||
if job is None or not can_access(worker.scope, engine.scope):
|
||||
return engine
|
||||
def claim_job(lens: Lens, worker: Worker, now: datetime) -> Lens:
|
||||
job: Final = current_job(lens)
|
||||
if job is None or not can_access(worker.scope, lens.scope):
|
||||
return lens
|
||||
if job.status == "running" and job.lease_until is not None and job.lease_until > now:
|
||||
return engine
|
||||
return lens
|
||||
if job.attempts >= 3:
|
||||
return replace_job(
|
||||
engine,
|
||||
lens,
|
||||
job.model_copy(
|
||||
update=MappingProxyType(
|
||||
{
|
||||
|
|
@ -64,11 +64,9 @@ def claim_job(engine: Engine, worker: Worker, now: datetime) -> Engine:
|
|||
}
|
||||
)
|
||||
),
|
||||
).model_copy(
|
||||
update=MappingProxyType({"next_run_at": now + timedelta(minutes=engine.settings.interval_minutes)})
|
||||
)
|
||||
).model_copy(update=MappingProxyType({"next_run_at": now + timedelta(minutes=lens.settings.interval_minutes)}))
|
||||
return replace_job(
|
||||
engine,
|
||||
lens,
|
||||
job.model_copy(
|
||||
update=MappingProxyType(
|
||||
{
|
||||
|
|
@ -83,23 +81,23 @@ def claim_job(engine: Engine, worker: Worker, now: datetime) -> Engine:
|
|||
)
|
||||
|
||||
|
||||
def renew_budget(engine: Engine, now: datetime) -> Engine:
|
||||
def renew_budget(lens: Lens, now: datetime) -> Lens:
|
||||
month: Final = now.strftime("%Y-%m")
|
||||
if engine.budget_month == month:
|
||||
return engine
|
||||
return engine.model_copy(update=MappingProxyType({"budget_month": month, "spent": 0}))
|
||||
if lens.budget_month == month:
|
||||
return lens
|
||||
return lens.model_copy(update=MappingProxyType({"budget_month": month, "spent": 0}))
|
||||
|
||||
|
||||
def merge_finding(engine: Engine, draft: FindingDraft, revision: int, now: datetime) -> Finding:
|
||||
legacy_identity: Final = hashlib.sha256(f"{engine.id}:{draft.check_id}:{draft.title.lower()}".encode()).hexdigest()[
|
||||
def merge_finding(lens: Lens, draft: FindingDraft, revision: int, now: datetime) -> Finding:
|
||||
legacy_identity: Final = hashlib.sha256(f"{lens.id}:{draft.check_id}:{draft.title.lower()}".encode()).hexdigest()[
|
||||
:24
|
||||
]
|
||||
identity: Final = hashlib.sha256(
|
||||
f"{engine.id}:{draft.check_id}:{draft.kind}:{draft.title.lower()}".encode()
|
||||
f"{lens.id}:{draft.check_id}:{draft.kind}:{draft.title.lower()}".encode()
|
||||
).hexdigest()[:24]
|
||||
identities: Final = (draft.existing_finding_id, identity, legacy_identity)
|
||||
previous: Final = next(
|
||||
(f for f in engine.findings if f.id in identities and f.kind == draft.kind and f.check_id == draft.check_id),
|
||||
(f for f in lens.findings if f.id in identities and f.kind == draft.kind and f.check_id == draft.check_id),
|
||||
None,
|
||||
)
|
||||
occurrences: Final = tuple(sorted(frozenset(e.execution_id for e in draft.evidence if e.role == "support")))
|
||||
|
|
@ -137,8 +135,8 @@ def merge_finding(engine: Engine, draft: FindingDraft, revision: int, now: datet
|
|||
)
|
||||
|
||||
|
||||
def snapshot_finding(engine: Engine, draft: FindingDraft, revision: int, now: datetime) -> Finding:
|
||||
merged: Final = merge_finding(engine, draft, revision, now)
|
||||
def snapshot_finding(lens: Lens, draft: FindingDraft, revision: int, now: datetime) -> Finding:
|
||||
merged: Final = merge_finding(lens, draft, revision, now)
|
||||
return Finding.model_validate(
|
||||
MappingProxyType(
|
||||
{
|
||||
|
|
@ -12,10 +12,10 @@ import httpx
|
|||
from .analysis import analyze_sample
|
||||
from .models import Claim, Coverage, ExecutionContent, ModelRequest, ModelResult, Progress, Result, Sample
|
||||
|
||||
logger: Final = logging.getLogger("litellm.engine.worker")
|
||||
logger: Final = logging.getLogger("litellm.lens.worker")
|
||||
|
||||
|
||||
class EngineWorker:
|
||||
class LensWorker:
|
||||
def __init__(self, client: httpx.AsyncClient, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep) -> None:
|
||||
self.client: Final = client
|
||||
self.sleep: Final = sleep
|
||||
|
|
@ -38,14 +38,12 @@ class EngineWorker:
|
|||
return await self.model_request(path, body, attempt + 1)
|
||||
|
||||
async def run_once(self) -> bool:
|
||||
response: Final = await self.client.post(
|
||||
"/engine/worker/claim", params=MappingProxyType({"protocol_version": 2})
|
||||
)
|
||||
response: Final = await self.client.post("/lens/worker/claim", params=MappingProxyType({"protocol_version": 2}))
|
||||
response.raise_for_status()
|
||||
if response.json() is None:
|
||||
return False
|
||||
claim: Final = Claim.model_validate(response.json())
|
||||
prefix: Final = f"/engine/worker/{claim.engine_id}/{claim.job.id}"
|
||||
prefix: Final = f"/lens/worker/{claim.lens_id}/{claim.job.id}"
|
||||
|
||||
async def model(body: ModelRequest) -> ModelResult:
|
||||
return await self.model_request(prefix + "/model", body)
|
||||
|
|
@ -111,7 +109,7 @@ async def main() -> None:
|
|||
async with httpx.AsyncClient(
|
||||
base_url=url, headers=MappingProxyType({"Authorization": f"Bearer {token}"}), timeout=180
|
||||
) as client:
|
||||
worker: Final = EngineWorker(client)
|
||||
worker: Final = LensWorker(client)
|
||||
while True:
|
||||
try:
|
||||
await worker.run_once()
|
||||
|
|
@ -545,7 +545,6 @@ from litellm.proxy.discovery_endpoints import (
|
|||
agent_skills_discovery_router,
|
||||
ui_discovery_endpoints_router,
|
||||
)
|
||||
from litellm.proxy.engine.endpoints import router as engine_router
|
||||
from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router
|
||||
from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config
|
||||
from litellm.proxy.google_endpoints.endpoints import router as google_router
|
||||
|
|
@ -569,6 +568,7 @@ from litellm.proxy.hooks.prompt_injection_detection import (
|
|||
)
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_spend_event
|
||||
from litellm.proxy.image_endpoints.endpoints import router as image_router
|
||||
from litellm.proxy.lens.endpoints import router as lens_router
|
||||
from litellm.proxy.list_api.common import (
|
||||
ManagementProblem,
|
||||
problem_response,
|
||||
|
|
@ -20009,7 +20009,7 @@ app.include_router(auto_router_management_router)
|
|||
app.include_router(tag_management_router)
|
||||
app.include_router(workflow_management_router)
|
||||
app.include_router(memory_router)
|
||||
app.include_router(engine_router)
|
||||
app.include_router(lens_router)
|
||||
app.include_router(plugin_router)
|
||||
app.include_router(cost_tracking_settings_router)
|
||||
app.include_router(prompt_caching_requests_router)
|
||||
|
|
|
|||
|
|
@ -1895,22 +1895,22 @@ model LiteLLM_WorkflowMessage {
|
|||
@@index([run_id])
|
||||
}
|
||||
|
||||
model LiteLLM_Engine {
|
||||
model LiteLLM_Lens {
|
||||
id String @id
|
||||
version Int @default(0)
|
||||
data Json
|
||||
}
|
||||
|
||||
model LiteLLM_EngineRun {
|
||||
model LiteLLM_LensRun {
|
||||
id String @id
|
||||
engine_id String
|
||||
lens_id String
|
||||
created_at DateTime
|
||||
data Json
|
||||
|
||||
@@index([engine_id, created_at])
|
||||
@@index([lens_id, created_at])
|
||||
}
|
||||
|
||||
model LiteLLM_EngineWorker {
|
||||
model LiteLLM_LensWorker {
|
||||
id String @id
|
||||
token_hash String @unique
|
||||
data Json
|
||||
|
|
|
|||
|
|
@ -9,14 +9,23 @@ Pure functions, no I/O. Two steps:
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Mapping
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.constants import OTLP_MAX_ATTRIBUTE_VALUE_BYTES, OTLP_MAX_BODY_BYTES
|
||||
from litellm.rust_bridge.traces import DecodedSpan
|
||||
from litellm.rust_bridge.traces import decode_otlp as native_decode_otlp
|
||||
from litellm.tracing.types import SpanRow, SpanType
|
||||
from litellm.tracing.normalizers import select_normalizer
|
||||
from litellm.tracing.normalizers.base import to_int
|
||||
from litellm.tracing.types import SpanRow
|
||||
|
||||
_MESSAGE_LIST: Final = TypeAdapter(tuple[dict[str, JsonValue], ...])
|
||||
_MAX_JSON_ESCAPE_BYTES: Final = 6
|
||||
|
||||
# attributes whose content we lift into Input/Output and drop from SpanAttributes
|
||||
_HEAVY_ATTRIBUTES: Final = frozenset(
|
||||
|
|
@ -30,18 +39,6 @@ _HEAVY_ATTRIBUTES: Final = frozenset(
|
|||
"output.value",
|
||||
}
|
||||
)
|
||||
# LangChain / Deep Agents middleware wrappers: real spans, but noise in the UI
|
||||
_FRAMEWORK_SUFFIXES: Final = (
|
||||
".wrap_model_call",
|
||||
".wrap_tool_call",
|
||||
".before_agent",
|
||||
".after_agent",
|
||||
".before_model",
|
||||
".after_model",
|
||||
)
|
||||
_LLM_OPERATIONS: Final = frozenset({"chat", "text_completion", "generate_content"})
|
||||
_LC_ROLES: Final = MappingProxyType({"human": "user", "ai": "assistant", "system": "system", "tool": "tool"})
|
||||
_OPENINFERENCE_TYPES: Final[Mapping[str, SpanType]] = MappingProxyType({"AGENT": "agent", "LLM": "llm", "TOOL": "tool"})
|
||||
|
||||
|
||||
class InvalidOTLPPayloadError(ValueError):
|
||||
|
|
@ -60,6 +57,83 @@ def _truncate(value: str) -> str:
|
|||
return f"{kept}…[truncated {size - OTLP_MAX_ATTRIBUTE_VALUE_BYTES} bytes]"
|
||||
|
||||
|
||||
def _size(value: str) -> int:
|
||||
return len(value.encode("utf-8"))
|
||||
|
||||
|
||||
class _ElisionMarker(TypedDict):
|
||||
role: ReadOnly[str]
|
||||
content: ReadOnly[str]
|
||||
|
||||
|
||||
def _elided(count: int) -> str:
|
||||
marker: Final[_ElisionMarker] = {"role": "system", "content": f"…[{count} earlier messages truncated]"}
|
||||
return json.dumps(marker)
|
||||
|
||||
|
||||
def _with_content(message: Mapping[str, JsonValue], content: str) -> str:
|
||||
return json.dumps(MappingProxyType({**message, "content": content}), default=lambda proxy: proxy.copy())
|
||||
|
||||
|
||||
def _shrunk_message(message: Mapping[str, JsonValue], budget: int) -> str:
|
||||
"""One message cut to `budget` bytes, as valid JSON.
|
||||
|
||||
Shortens `content` first; if other fields (e.g. huge tool_calls) still don't fit, keeps only role + content.
|
||||
"""
|
||||
content: Final = message.get("content")
|
||||
text: Final = content if isinstance(content, str) else json.dumps(content)
|
||||
role_only: Final = MappingProxyType({"role": message.get("role", "user")})
|
||||
attempts: Final = (
|
||||
_cut_content(message, text, budget, 1),
|
||||
_cut_content(role_only, text, budget, 1),
|
||||
_cut_content(role_only, text, budget, _MAX_JSON_ESCAPE_BYTES),
|
||||
)
|
||||
return next((attempt for attempt in attempts if _size(attempt) <= budget), attempts[-1])
|
||||
|
||||
|
||||
def _cut_content(message: Mapping[str, JsonValue], text: str, budget: int, escape_factor: int) -> str:
|
||||
overhead: Final = _size(_with_content(message, ""))
|
||||
room: Final = max(0, budget - overhead - 48) // escape_factor
|
||||
kept: Final = text.encode("utf-8")[:room].decode("utf-8", "ignore")
|
||||
return _with_content(message, f"{kept}…[truncated {_size(text) - _size(kept)} bytes]")
|
||||
|
||||
|
||||
def _newest_that_fit(encoded: tuple[str, ...], budget: int) -> int:
|
||||
"""How many trailing messages fit in `budget` bytes (comma separators included), scanning newest first."""
|
||||
sizes: Final = tuple(_size(m) + 1 for m in reversed(encoded))
|
||||
totals: Final = tuple(accumulate(sizes))
|
||||
return next((count for count, total in enumerate(totals) if total > budget), len(totals))
|
||||
|
||||
|
||||
def _truncate_payload(value: str) -> str:
|
||||
"""Message arrays keep the first message, an elision marker and the newest messages that fit.
|
||||
|
||||
The result is always valid JSON: if even those don't fit, the first and last messages are shortened.
|
||||
Anything that isn't a message array is byte-truncated as before.
|
||||
"""
|
||||
if _size(value) <= OTLP_MAX_ATTRIBUTE_VALUE_BYTES or not value.startswith("["):
|
||||
return _truncate(value)
|
||||
try:
|
||||
messages: Final = _MESSAGE_LIST.validate_json(value)
|
||||
except ValidationError:
|
||||
return _truncate(value)
|
||||
if len(messages) < 2:
|
||||
return _truncate(value)
|
||||
encoded: Final = tuple(json.dumps(m) for m in messages)
|
||||
marker_budget: Final = _size(_elided(len(messages))) + 1
|
||||
budget: Final = OTLP_MAX_ATTRIBUTE_VALUE_BYTES - 2 - _size(encoded[0]) - 1 - marker_budget
|
||||
kept: Final = min(_newest_that_fit(encoded[1:], budget), len(messages) - 2)
|
||||
if kept > 0:
|
||||
tail: Final = encoded[len(encoded) - kept :]
|
||||
return "[" + ", ".join((encoded[0], _elided(len(messages) - 1 - kept), *tail)) + "]"
|
||||
half: Final = (OTLP_MAX_ATTRIBUTE_VALUE_BYTES - marker_budget - 4) // 2
|
||||
middle: Final = (_elided(len(messages) - 2),) if len(messages) > 2 else ()
|
||||
shrunk: Final = (
|
||||
"[" + ", ".join((_shrunk_message(messages[0], half), *middle, _shrunk_message(messages[-1], half))) + "]"
|
||||
)
|
||||
return shrunk if _size(shrunk) <= OTLP_MAX_ATTRIBUTE_VALUE_BYTES else "[" + _elided(len(messages)) + "]"
|
||||
|
||||
|
||||
def decode_otlp(
|
||||
body: bytes, content_type: str | None = None, content_encoding: str | None = None
|
||||
) -> tuple[SpanRow, ...]:
|
||||
|
|
@ -116,157 +190,17 @@ def _span_row(span: DecodedSpan) -> SpanRow:
|
|||
row["SpanAttributes"] = { # mutable-ok: the Rust JSON bridge requires a plain dict for span attributes
|
||||
k: _truncate(v) for k, v in attributes.items() if k not in _HEAVY_ATTRIBUTES
|
||||
}
|
||||
row["Input"], row["Output"] = _truncate(row["Input"]), _truncate(row["Output"])
|
||||
row["Input"], row["Output"] = _truncate_payload(row["Input"]), _truncate(row["Output"])
|
||||
return row
|
||||
|
||||
|
||||
def _loads(value: str) -> object:
|
||||
try:
|
||||
return json.loads(value)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _lc_message(message: Mapping[str, Any]) -> dict[str, Any]:
|
||||
"""LangChain serialized message (or plain {role, content}) -> {role, content, tool_calls?}."""
|
||||
kwargs = message.get("kwargs", message)
|
||||
role = _LC_ROLES.get(kwargs.get("type") or kwargs.get("role"), kwargs.get("role") or kwargs.get("type") or "")
|
||||
content = kwargs.get("content", "")
|
||||
out: dict[str, Any] = { # mutable-ok: the framework message is built for JSON serialization
|
||||
"role": role,
|
||||
"content": content if isinstance(content, str) else json.dumps(content),
|
||||
}
|
||||
if kwargs.get("tool_calls"):
|
||||
out["tool_calls"] = tuple(
|
||||
{"name": t.get("name"), "args": t.get("args")} # mutable-ok: JSON tool calls need object payloads
|
||||
for t in kwargs["tool_calls"]
|
||||
)
|
||||
if role == "tool" and kwargs.get("name"):
|
||||
out["name"] = kwargs["name"]
|
||||
return out
|
||||
|
||||
|
||||
def _langsmith_type(row: SpanRow, attributes: Mapping[str, str]) -> SpanType:
|
||||
kind = attributes.get("langsmith.span.kind", "chain")
|
||||
name = row["SpanName"]
|
||||
if not row["ParentSpanId"] or name == attributes.get("langsmith.metadata.lc_agent_name"):
|
||||
return "agent"
|
||||
if kind in ("llm", "tool"):
|
||||
return kind
|
||||
if name.endswith(_FRAMEWORK_SUFFIXES):
|
||||
return "framework"
|
||||
return "chain"
|
||||
|
||||
|
||||
def _langsmith_io(row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
prompt = _loads(attributes.get("gen_ai.prompt", ""))
|
||||
completion = _loads(attributes.get("gen_ai.completion", ""))
|
||||
prompt_payload = prompt if isinstance(prompt, dict) else MappingProxyType({})
|
||||
if row["ObservationType"] == "llm" and isinstance(completion, dict):
|
||||
messages = prompt_payload.get("messages") or ((),)
|
||||
batch = messages[0] if messages and isinstance(messages[0], list) else messages
|
||||
row["Input"] = (
|
||||
json.dumps(tuple(_lc_message(m) for m in batch if isinstance(m, dict)))
|
||||
if isinstance(batch, (list, tuple))
|
||||
else ""
|
||||
)
|
||||
generations: Final = completion.get("generations")
|
||||
first: Final = generations[0] if isinstance(generations, list) and generations else None
|
||||
item: Final = first[0] if isinstance(first, list) and first else None
|
||||
message: Final = item.get("message") if isinstance(item, dict) else None
|
||||
generation: Final = message.get("kwargs") if isinstance(message, dict) else None
|
||||
if isinstance(generation, dict):
|
||||
row["Output"] = json.dumps(_lc_message(generation))
|
||||
metadata: Final = generation.get("response_metadata")
|
||||
row["LiteLLMRequestId"] = metadata.get("id", "") if isinstance(metadata, dict) else ""
|
||||
else:
|
||||
row["Output"] = attributes.get("gen_ai.completion", "")
|
||||
return
|
||||
if row["ObservationType"] == "tool":
|
||||
output = completion.get("output", completion) if isinstance(completion, dict) else completion
|
||||
if isinstance(output, dict) and "update" in output: # LangGraph Command, e.g. Deep Agents `task`
|
||||
update: Final = output.get("update")
|
||||
update_messages = update.get("messages") or () if isinstance(update, dict) else ()
|
||||
output = update_messages[-1] if update_messages else output
|
||||
if isinstance(output, dict):
|
||||
output = output.get("content", output)
|
||||
row["Input"] = attributes.get("gen_ai.prompt", "")
|
||||
row["Output"] = output if isinstance(output, str) else json.dumps(output)
|
||||
return
|
||||
if row["ObservationType"] == "agent":
|
||||
input_messages = prompt.get("messages") if isinstance(prompt, dict) else None
|
||||
output_messages = completion.get("messages") if isinstance(completion, dict) else None
|
||||
# agents built with @traceable take arbitrary args, not a message list: keep the raw payload then
|
||||
row["Input"] = (
|
||||
json.dumps(tuple(_lc_message(m) for m in input_messages if isinstance(m, dict)))
|
||||
if input_messages
|
||||
else attributes.get("gen_ai.prompt", "")
|
||||
)
|
||||
row["Output"] = (
|
||||
json.dumps(_lc_message(output_messages[-1]))
|
||||
if output_messages and isinstance(output_messages[-1], dict)
|
||||
else attributes.get("gen_ai.completion", "")
|
||||
)
|
||||
return
|
||||
row["Input"] = attributes.get("gen_ai.prompt", "")
|
||||
row["Output"] = attributes.get("gen_ai.completion", "")
|
||||
|
||||
|
||||
def normalize_langsmith(row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
row["ObservationType"] = _langsmith_type(row, attributes)
|
||||
row["AgentName"] = attributes.get("langsmith.metadata.lc_agent_name", "")
|
||||
row["Model"] = attributes.get("gen_ai.request.model", "")
|
||||
_langsmith_io(row, attributes)
|
||||
|
||||
|
||||
def normalize_genai(row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
operation = attributes.get("gen_ai.operation.name", "")
|
||||
if operation == "invoke_agent" or not row["ParentSpanId"]:
|
||||
row["ObservationType"] = "agent"
|
||||
elif operation in _LLM_OPERATIONS:
|
||||
row["ObservationType"] = "llm"
|
||||
elif operation == "execute_tool":
|
||||
row["ObservationType"] = "tool"
|
||||
row["AgentName"] = attributes.get("gen_ai.agent.name", "")
|
||||
row["Model"] = attributes.get("gen_ai.request.model") or attributes.get("gen_ai.response.model", "")
|
||||
row["LiteLLMRequestId"] = attributes.get("gen_ai.response.id", "")
|
||||
row["Input"] = attributes.get("gen_ai.input.messages") or attributes.get("gen_ai.tool.call.arguments", "")
|
||||
row["Output"] = attributes.get("gen_ai.output.messages") or attributes.get("gen_ai.tool.call.result", "")
|
||||
|
||||
|
||||
def normalize_openinference(row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
kind = attributes.get("openinference.span.kind", "").upper()
|
||||
row["ObservationType"] = _OPENINFERENCE_TYPES.get(kind, "agent" if not row["ParentSpanId"] else "chain")
|
||||
row["AgentName"] = attributes.get("agent.name", "")
|
||||
row["Model"] = attributes.get("llm.model_name", "")
|
||||
row["Input"] = attributes.get("input.value", "")
|
||||
row["Output"] = attributes.get("output.value", "")
|
||||
row["InputTokens"] = _to_int(attributes.get("llm.token_count.prompt"))
|
||||
row["OutputTokens"] = _to_int(attributes.get("llm.token_count.completion"))
|
||||
|
||||
|
||||
def _set_tokens(row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
row["InputTokens"] = _to_int(attributes.get("gen_ai.usage.input_tokens"))
|
||||
row["OutputTokens"] = _to_int(attributes.get("gen_ai.usage.output_tokens"))
|
||||
|
||||
|
||||
def _to_int(value: str | None) -> int:
|
||||
try:
|
||||
return int(value) if value else 0
|
||||
except ValueError:
|
||||
return 0
|
||||
|
||||
|
||||
def select_normalizer(scope_name: str, attributes: Mapping[str, str]) -> Callable[[SpanRow, Mapping[str, str]], None]:
|
||||
if scope_name == "langsmith" or "langsmith.span.kind" in attributes:
|
||||
return normalize_langsmith
|
||||
if "openinference.span.kind" in attributes:
|
||||
return normalize_openinference
|
||||
return normalize_genai
|
||||
row["InputTokens"] = to_int(attributes.get("gen_ai.usage.input_tokens"))
|
||||
row["OutputTokens"] = to_int(attributes.get("gen_ai.usage.output_tokens"))
|
||||
|
||||
|
||||
def normalize(row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
select_normalizer(row["ScopeName"], attributes)(row, attributes)
|
||||
select_normalizer(row["ScopeName"], attributes).normalize(row, attributes)
|
||||
if not row["InputTokens"] and not row["OutputTokens"]:
|
||||
_set_tokens(row, attributes)
|
||||
|
||||
|
|
|
|||
32
litellm/tracing/normalizers/__init__.py
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
"""Per-convention span normalizers, tried in order: the first whose `matches()` is true wins."""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
from litellm.tracing.normalizers.base import SpanNormalizer
|
||||
from litellm.tracing.normalizers.genai import GenAISemconvNormalizer
|
||||
from litellm.tracing.normalizers.langsmith import LangSmithNormalizer
|
||||
from litellm.tracing.normalizers.openinference import OpenInferenceNormalizer
|
||||
|
||||
NORMALIZERS: Final[tuple[SpanNormalizer, ...]] = (
|
||||
LangSmithNormalizer(),
|
||||
OpenInferenceNormalizer(),
|
||||
GenAISemconvNormalizer(),
|
||||
)
|
||||
_FALLBACK: Final[SpanNormalizer] = GenAISemconvNormalizer()
|
||||
|
||||
|
||||
def select_normalizer(
|
||||
scope_name: str, attributes: Mapping[str, str], registry: Sequence[SpanNormalizer] = NORMALIZERS
|
||||
) -> SpanNormalizer:
|
||||
return next((n for n in registry if n.matches(scope_name, attributes)), _FALLBACK)
|
||||
|
||||
|
||||
__all__ = (
|
||||
"NORMALIZERS",
|
||||
"GenAISemconvNormalizer",
|
||||
"LangSmithNormalizer",
|
||||
"OpenInferenceNormalizer",
|
||||
"SpanNormalizer",
|
||||
"select_normalizer",
|
||||
)
|
||||
22
litellm/tracing/normalizers/base.py
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Protocol
|
||||
|
||||
from litellm.tracing.types import SpanRow
|
||||
|
||||
|
||||
class SpanNormalizer(Protocol):
|
||||
"""Maps one tracing convention's span attributes onto the LiteLLM `SpanRow` columns."""
|
||||
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
|
||||
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool: ...
|
||||
|
||||
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None: ...
|
||||
|
||||
|
||||
def to_int(value: str | None) -> int:
|
||||
try:
|
||||
return int(value) if value else 0
|
||||
except ValueError:
|
||||
return 0
|
||||
31
litellm/tracing/normalizers/genai.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm.tracing.types import SpanRow
|
||||
|
||||
_LLM_OPERATIONS: Final = frozenset({"chat", "text_completion", "generate_content"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GenAISemconvNormalizer:
|
||||
"""OTEL `gen_ai.*` semantic conventions. Matches every span, so it belongs last as the fallback."""
|
||||
|
||||
name: str = "genai"
|
||||
|
||||
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool:
|
||||
return True
|
||||
|
||||
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
operation: Final = attributes.get("gen_ai.operation.name", "")
|
||||
if operation == "invoke_agent" or not row["ParentSpanId"]:
|
||||
row["ObservationType"] = "agent"
|
||||
elif operation in _LLM_OPERATIONS:
|
||||
row["ObservationType"] = "llm"
|
||||
elif operation == "execute_tool":
|
||||
row["ObservationType"] = "tool"
|
||||
row["AgentName"] = attributes.get("gen_ai.agent.name", "")
|
||||
row["Model"] = attributes.get("gen_ai.request.model") or attributes.get("gen_ai.response.model", "")
|
||||
row["LiteLLMRequestId"] = attributes.get("gen_ai.response.id", "")
|
||||
row["Input"] = attributes.get("gen_ai.input.messages") or attributes.get("gen_ai.tool.call.arguments", "")
|
||||
row["Output"] = attributes.get("gen_ai.output.messages") or attributes.get("gen_ai.tool.call.result", "")
|
||||
115
litellm/tracing/normalizers/langsmith.py
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
"""LangSmith OTEL mode, which LangChain, LangGraph and Deep Agents export through."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.tracing.normalizers.messages import lc_message
|
||||
from litellm.tracing.types import SpanRow, SpanType
|
||||
|
||||
# LangChain / Deep Agents middleware wrappers: real spans, but noise in the UI
|
||||
_FRAMEWORK_SUFFIXES: Final = (
|
||||
".wrap_model_call",
|
||||
".wrap_tool_call",
|
||||
".before_agent",
|
||||
".after_agent",
|
||||
".before_model",
|
||||
".after_model",
|
||||
)
|
||||
|
||||
|
||||
def _loads(value: str) -> object:
|
||||
try:
|
||||
return json.loads(value)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _span_type(row: SpanRow, attributes: Mapping[str, str]) -> SpanType:
|
||||
kind: Final = attributes.get("langsmith.span.kind", "chain")
|
||||
name: Final = row["SpanName"]
|
||||
if not row["ParentSpanId"] or name == attributes.get("langsmith.metadata.lc_agent_name"):
|
||||
return "agent"
|
||||
if kind in ("llm", "tool"):
|
||||
return kind
|
||||
if name.endswith(_FRAMEWORK_SUFFIXES):
|
||||
return "framework"
|
||||
return "chain"
|
||||
|
||||
|
||||
def _tool_output(completion: object) -> object:
|
||||
raw: Final = completion.get("output", completion) if isinstance(completion, dict) else completion
|
||||
update: Final = raw.get("update") if isinstance(raw, dict) else None
|
||||
update_messages: Final = update.get("messages") or () if isinstance(update, dict) else ()
|
||||
is_command: Final = isinstance(raw, dict) and "update" in raw
|
||||
# LangGraph Command (e.g. the Deep Agents `task` tool): the result is the last update message
|
||||
output: Final = update_messages[-1] if is_command and update_messages else raw
|
||||
return output.get("content", output) if isinstance(output, dict) else output
|
||||
|
||||
|
||||
def _set_agent_io(row: SpanRow, attributes: Mapping[str, str], prompt: object, completion: object) -> None:
|
||||
input_messages: Final = prompt.get("messages") if isinstance(prompt, dict) else None
|
||||
output_messages: Final = completion.get("messages") if isinstance(completion, dict) else None
|
||||
# agents built with @traceable take arbitrary args, not a message list: keep the raw payload then
|
||||
row["Input"] = (
|
||||
json.dumps(tuple(lc_message(m) for m in input_messages if isinstance(m, dict)))
|
||||
if input_messages
|
||||
else attributes.get("gen_ai.prompt", "")
|
||||
)
|
||||
row["Output"] = (
|
||||
json.dumps(lc_message(output_messages[-1]))
|
||||
if output_messages and isinstance(output_messages[-1], dict)
|
||||
else attributes.get("gen_ai.completion", "")
|
||||
)
|
||||
|
||||
|
||||
def _set_io(row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
prompt: Final = _loads(attributes.get("gen_ai.prompt", ""))
|
||||
completion: Final = _loads(attributes.get("gen_ai.completion", ""))
|
||||
if row["ObservationType"] == "llm" and isinstance(completion, dict):
|
||||
prompt_payload: Final = prompt if isinstance(prompt, dict) else MappingProxyType({})
|
||||
messages: Final = prompt_payload.get("messages") or ((),)
|
||||
batch: Final = messages[0] if messages and isinstance(messages[0], list) else messages
|
||||
row["Input"] = (
|
||||
json.dumps(tuple(lc_message(m) for m in batch if isinstance(m, dict)))
|
||||
if isinstance(batch, (list, tuple))
|
||||
else ""
|
||||
)
|
||||
generations: Final = completion.get("generations")
|
||||
first: Final = generations[0] if isinstance(generations, list) and generations else None
|
||||
item: Final = first[0] if isinstance(first, list) and first else None
|
||||
message: Final = item.get("message") if isinstance(item, dict) else None
|
||||
generation: Final = message.get("kwargs") if isinstance(message, dict) else None
|
||||
if isinstance(generation, dict):
|
||||
row["Output"] = json.dumps(lc_message(generation))
|
||||
metadata: Final = generation.get("response_metadata")
|
||||
row["LiteLLMRequestId"] = metadata.get("id", "") if isinstance(metadata, dict) else ""
|
||||
return
|
||||
row["Output"] = attributes.get("gen_ai.completion", "")
|
||||
return
|
||||
if row["ObservationType"] == "tool":
|
||||
output: Final = _tool_output(completion)
|
||||
row["Input"] = attributes.get("gen_ai.prompt", "")
|
||||
row["Output"] = output if isinstance(output, str) else json.dumps(output)
|
||||
return
|
||||
if row["ObservationType"] == "agent":
|
||||
_set_agent_io(row, attributes, prompt, completion)
|
||||
return
|
||||
row["Input"] = attributes.get("gen_ai.prompt", "")
|
||||
row["Output"] = attributes.get("gen_ai.completion", "")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LangSmithNormalizer:
|
||||
name: str = "langsmith"
|
||||
|
||||
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool:
|
||||
return scope_name == "langsmith" or "langsmith.span.kind" in attributes
|
||||
|
||||
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
row["ObservationType"] = _span_type(row, attributes)
|
||||
row["AgentName"] = attributes.get("langsmith.metadata.lc_agent_name", "")
|
||||
row["Model"] = attributes.get("gen_ai.request.model", "")
|
||||
_set_io(row, attributes)
|
||||
59
litellm/tracing/normalizers/messages.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
ChatRole: TypeAlias = Literal["system", "user", "assistant", "tool"]
|
||||
|
||||
MESSAGE_ROLES: Final[Mapping[str, ChatRole]] = MappingProxyType(
|
||||
{"human": "user", "user": "user", "ai": "assistant", "assistant": "assistant", "system": "system", "tool": "tool"}
|
||||
)
|
||||
|
||||
|
||||
class _ContentBlock(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
type: str = ""
|
||||
text: str | None = None
|
||||
|
||||
|
||||
_CONTENT_BLOCKS: Final = TypeAdapter(tuple[_ContentBlock, ...])
|
||||
_NON_TEXT_BLOCKS: Final = frozenset(
|
||||
{"reasoning", "thinking", "redacted_thinking", "function_call", "tool_use", "tool_call"}
|
||||
)
|
||||
|
||||
|
||||
def content_text(content: object) -> str:
|
||||
"""Message content as display text: Responses-style block lists keep only their text blocks."""
|
||||
if content is None:
|
||||
return ""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
try:
|
||||
blocks: Final = _CONTENT_BLOCKS.validate_python(content)
|
||||
except ValidationError:
|
||||
return json.dumps(content)
|
||||
if not all(block.text is not None or block.type in _NON_TEXT_BLOCKS for block in blocks):
|
||||
return json.dumps(content)
|
||||
return "\n\n".join(block.text for block in blocks if block.text is not None)
|
||||
|
||||
|
||||
def lc_message(message: Mapping[str, Any]) -> dict[str, Any]:
|
||||
"""LangChain serialized message (or plain {role, content}) -> {role, content, tool_calls?}."""
|
||||
kwargs: Final = message.get("kwargs", message)
|
||||
role: Final = MESSAGE_ROLES.get(
|
||||
kwargs.get("type") or kwargs.get("role"), kwargs.get("role") or kwargs.get("type") or ""
|
||||
)
|
||||
out: Final[dict[str, Any]] = { # mutable-ok: the framework message is built for JSON serialization
|
||||
"role": role,
|
||||
"content": content_text(kwargs.get("content", "")),
|
||||
}
|
||||
if kwargs.get("tool_calls"):
|
||||
out["tool_calls"] = tuple(
|
||||
{"name": t.get("name"), "args": t.get("args")} # mutable-ok: JSON tool calls need object payloads
|
||||
for t in kwargs["tool_calls"]
|
||||
)
|
||||
if role == "tool" and kwargs.get("name"):
|
||||
out["name"] = kwargs["name"]
|
||||
return out
|
||||
27
litellm/tracing/normalizers/openinference.py
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.tracing.normalizers.base import to_int
|
||||
from litellm.tracing.types import SpanRow, SpanType
|
||||
|
||||
_OPENINFERENCE_TYPES: Final[Mapping[str, SpanType]] = MappingProxyType({"AGENT": "agent", "LLM": "llm", "TOOL": "tool"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OpenInferenceNormalizer:
|
||||
name: str = "openinference"
|
||||
|
||||
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool:
|
||||
return "openinference.span.kind" in attributes
|
||||
|
||||
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None:
|
||||
kind: Final = attributes.get("openinference.span.kind", "").upper()
|
||||
row["ObservationType"] = _OPENINFERENCE_TYPES.get(kind, "agent" if not row["ParentSpanId"] else "chain")
|
||||
row["AgentName"] = attributes.get("agent.name", "")
|
||||
row["Model"] = attributes.get("llm.model_name", "")
|
||||
row["Input"] = attributes.get("input.value", "")
|
||||
row["Output"] = attributes.get("output.value", "")
|
||||
row["InputTokens"] = to_int(attributes.get("llm.token_count.prompt"))
|
||||
row["OutputTokens"] = to_int(attributes.get("llm.token_count.completion"))
|
||||
|
|
@ -28,6 +28,7 @@ from litellm.tracing.types import (
|
|||
TraceScope,
|
||||
TraceSummary,
|
||||
)
|
||||
from litellm.tracing.ui_format import to_ui_content
|
||||
|
||||
NANOS_PER_MS: Final = 1_000_000
|
||||
SPEND_WINDOW_MS: Final = 30 * 60 * 1000
|
||||
|
|
@ -344,5 +345,7 @@ class ClickHouseTraceStore:
|
|||
span_id=rows[0]["span_id"],
|
||||
input=rows[0]["input"],
|
||||
output=rows[0]["output"],
|
||||
input_ui=to_ui_content(rows[0]["input"]),
|
||||
output_ui=to_ui_content(rows[0]["output"]),
|
||||
attributes=rows[0]["attributes"],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ from typing import Literal
|
|||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.tracing.ui_format import UIContent
|
||||
|
||||
SpanType = Literal["agent", "llm", "tool", "chain", "framework"]
|
||||
SpanStatus = Literal["ok", "error", "unset"]
|
||||
|
||||
|
|
@ -84,6 +86,8 @@ class SpanDetail(TypedDict):
|
|||
span_id: ReadOnly[str]
|
||||
input: ReadOnly[str]
|
||||
output: ReadOnly[str]
|
||||
input_ui: ReadOnly[UIContent]
|
||||
output_ui: ReadOnly[UIContent]
|
||||
attributes: ReadOnly[dict[str, str]]
|
||||
|
||||
|
||||
|
|
|
|||
158
litellm/tracing/ui_format.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
"""The LiteLLM UI content format: span input / output reduced to messages, key/value fields or plain text."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.tracing.normalizers.messages import MESSAGE_ROLES, ChatRole, content_text
|
||||
|
||||
|
||||
class UIToolCall(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
arguments: ReadOnly[str]
|
||||
|
||||
|
||||
class UIMessage(TypedDict):
|
||||
role: ReadOnly[ChatRole]
|
||||
content: ReadOnly[str]
|
||||
name: ReadOnly[NotRequired[str]]
|
||||
tool_calls: ReadOnly[NotRequired[tuple[UIToolCall, ...]]]
|
||||
|
||||
|
||||
class UIField(TypedDict):
|
||||
key: ReadOnly[str]
|
||||
value: ReadOnly[str]
|
||||
|
||||
|
||||
class UIMessages(TypedDict):
|
||||
kind: ReadOnly[Literal["messages"]]
|
||||
messages: ReadOnly[tuple[UIMessage, ...]]
|
||||
|
||||
|
||||
class UIFields(TypedDict):
|
||||
kind: ReadOnly[Literal["fields"]]
|
||||
fields: ReadOnly[tuple[UIField, ...]]
|
||||
|
||||
|
||||
class UIText(TypedDict):
|
||||
kind: ReadOnly[Literal["text"]]
|
||||
text: ReadOnly[str]
|
||||
|
||||
|
||||
UIContent: TypeAlias = UIMessages | UIFields | UIText
|
||||
|
||||
|
||||
class _ToolFunction(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
name: str = ""
|
||||
arguments: JsonValue = None
|
||||
|
||||
|
||||
class _RawToolCall(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
name: str = ""
|
||||
args: JsonValue = None
|
||||
arguments: JsonValue = None
|
||||
function: _ToolFunction | None = None
|
||||
|
||||
|
||||
class _RawMessage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
role: str | None = None
|
||||
type: str | None = None
|
||||
content: JsonValue = None
|
||||
name: str | None = None
|
||||
tool_calls: tuple[_RawToolCall, ...] | None = None
|
||||
kwargs: "_RawMessage | None" = None
|
||||
|
||||
|
||||
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_MESSAGE: Final = TypeAdapter(_RawMessage)
|
||||
_MESSAGES: Final = TypeAdapter(tuple[_RawMessage, ...])
|
||||
|
||||
|
||||
def _unwrapped(message: _RawMessage) -> _RawMessage:
|
||||
return message.kwargs if message.kwargs is not None else message
|
||||
|
||||
|
||||
def _is_message(message: _RawMessage) -> bool:
|
||||
has_role: Final = message.role is not None or message.type in MESSAGE_ROLES
|
||||
return has_role and ("content" in message.model_fields_set or bool(message.tool_calls))
|
||||
|
||||
|
||||
def _arguments_text(arguments: JsonValue) -> str:
|
||||
match arguments:
|
||||
case str():
|
||||
return arguments
|
||||
case None:
|
||||
return "{}"
|
||||
case _:
|
||||
return json.dumps(arguments)
|
||||
|
||||
|
||||
def _tool_call(call: _RawToolCall) -> UIToolCall:
|
||||
if call.function is not None:
|
||||
return UIToolCall(name=call.function.name or call.name, arguments=_arguments_text(call.function.arguments))
|
||||
return UIToolCall(name=call.name, arguments=_arguments_text(call.arguments if call.args is None else call.args))
|
||||
|
||||
|
||||
def _role(message: _RawMessage, has_tool_calls: bool) -> ChatRole:
|
||||
"""Known roles and LangChain types map directly; any other role is the assistant when it calls tools, else the user."""
|
||||
known: Final = MESSAGE_ROLES.get(message.role or message.type or "")
|
||||
if known is not None:
|
||||
return known
|
||||
return "assistant" if has_tool_calls else "user"
|
||||
|
||||
|
||||
def _ui_message(message: _RawMessage) -> UIMessage:
|
||||
calls: Final = tuple(_tool_call(call) for call in message.tool_calls or ())
|
||||
role: Final = _role(message, bool(calls))
|
||||
content: Final = content_text(message.content)
|
||||
match (message.name or None, calls):
|
||||
case (None, ()):
|
||||
return UIMessage(role=role, content=content)
|
||||
case (None, _):
|
||||
return UIMessage(role=role, content=content, tool_calls=calls)
|
||||
case (str() as name, ()):
|
||||
return UIMessage(role=role, content=content, name=name)
|
||||
case (str() as name, _):
|
||||
return UIMessage(role=role, content=content, name=name, tool_calls=calls)
|
||||
|
||||
|
||||
def _messages(parsed: Sequence[JsonValue] | Mapping[str, JsonValue]) -> tuple[_RawMessage, ...] | None:
|
||||
try:
|
||||
raw: Final = (
|
||||
(_MESSAGE.validate_python(parsed),) if isinstance(parsed, Mapping) else _MESSAGES.validate_python(parsed)
|
||||
)
|
||||
except ValidationError:
|
||||
return None
|
||||
unwrapped: Final = tuple(_unwrapped(message) for message in raw)
|
||||
return unwrapped if unwrapped and all(_is_message(message) for message in unwrapped) else None
|
||||
|
||||
|
||||
def _field_value(value: JsonValue) -> str:
|
||||
return value if isinstance(value, str) else json.dumps(value)
|
||||
|
||||
|
||||
def _parsed(raw: str) -> JsonValue:
|
||||
try:
|
||||
return _JSON.validate_json(raw)
|
||||
except ValidationError:
|
||||
return raw
|
||||
|
||||
|
||||
def to_ui_content(raw: str) -> UIContent:
|
||||
if not raw:
|
||||
return UIText(kind="text", text="")
|
||||
parsed: Final = _parsed(raw)
|
||||
if not isinstance(parsed, list | dict):
|
||||
return UIText(kind="text", text=parsed if isinstance(parsed, str) else raw)
|
||||
messages: Final = _messages(parsed)
|
||||
if messages is not None:
|
||||
return UIMessages(kind="messages", messages=tuple(_ui_message(message) for message in messages))
|
||||
if isinstance(parsed, dict):
|
||||
return UIFields(kind="fields", fields=tuple(UIField(key=k, value=_field_value(v)) for k, v in parsed.items()))
|
||||
return UIText(kind="text", text=raw)
|
||||
|
|
@ -1,5 +1,9 @@
|
|||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
S3PartitionGranularity = Literal["day", "hour"]
|
||||
|
||||
|
||||
class s3BatchLoggingElement(BaseModel):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -397,7 +397,7 @@ paths_to_mutate = [
|
|||
# a mutation score is only meaningful against the tests that claim to cover
|
||||
# the mutated code anyway.
|
||||
tests_dir = [
|
||||
"tests/test_litellm/proxy/management_endpoints/",
|
||||
"tests/unit/proxy/management_endpoints/",
|
||||
]
|
||||
also_copy = [
|
||||
"litellm/",
|
||||
|
|
@ -423,7 +423,7 @@ pytest_add_cli_args = [
|
|||
"-p", "no:pytest-retry",
|
||||
"-p", "no:rerunfailures",
|
||||
"-p", "no:xdist",
|
||||
"--ignore=tests/test_litellm/proxy/management_endpoints/test_saml_sso.py",
|
||||
"--ignore=tests/unit/proxy/management_endpoints/test_saml_sso.py",
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
|
|
|
|||
|
|
@ -1895,22 +1895,22 @@ model LiteLLM_WorkflowMessage {
|
|||
@@index([run_id])
|
||||
}
|
||||
|
||||
model LiteLLM_Engine {
|
||||
model LiteLLM_Lens {
|
||||
id String @id
|
||||
version Int @default(0)
|
||||
data Json
|
||||
}
|
||||
|
||||
model LiteLLM_EngineRun {
|
||||
model LiteLLM_LensRun {
|
||||
id String @id
|
||||
engine_id String
|
||||
lens_id String
|
||||
created_at DateTime
|
||||
data Json
|
||||
|
||||
@@index([engine_id, created_at])
|
||||
@@index([lens_id, created_at])
|
||||
}
|
||||
|
||||
model LiteLLM_EngineWorker {
|
||||
model LiteLLM_LensWorker {
|
||||
id String @id
|
||||
token_hash String @unique
|
||||
data Json
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ Run from repo root:
|
|||
uv run python scripts/test_tool_allowlist_script.py
|
||||
|
||||
Or run the unit tests:
|
||||
uv run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v
|
||||
uv run pytest tests/unit/proxy/test_tools_allowlist_enforcement.py -v
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -148,7 +148,7 @@ def main():
|
|||
asyncio.run(test_check_tools_allowlist())
|
||||
print("Done. For full unit tests run:")
|
||||
print(
|
||||
" uv run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v"
|
||||
" uv run pytest tests/unit/proxy/test_tools_allowlist_enforcement.py -v"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3,8 +3,8 @@ import os
|
|||
|
||||
ALLOWED_FILES = [
|
||||
# The standalone Lens process reuses one client for its entire lifetime, without importing the proxy SDK.
|
||||
"../../litellm/proxy/engine/worker.py",
|
||||
"./litellm/proxy/engine/worker.py",
|
||||
"../../litellm/proxy/lens/worker.py",
|
||||
"./litellm/proxy/lens/worker.py",
|
||||
# local files
|
||||
"../../litellm/__init__.py",
|
||||
"../../litellm/llms/custom_httpx/http_handler.py",
|
||||
|
|
|
|||
|
|
@ -185,6 +185,6 @@ never landed.
|
|||
Unified (managed) batch cost is owned by the hourly `CheckBatchCost` poller, and a
|
||||
terminal DB status short-circuits retrieve for those ids, so the terminal-state cell
|
||||
uses the encoded path; poller timing does not fit an e2e gate and belongs in a
|
||||
DI-stubbed proxy integration test under `tests/test_litellm/proxy/`. Gemini
|
||||
DI-stubbed proxy integration test under `tests/unit/proxy/`. Gemini
|
||||
(non-Vertex) file content raises `NotImplementedError` upstream and is not a
|
||||
coverage cell.
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Logging integration delivery (behavior features). Grounded in litellm/integrations/.
|
||||
- {id: logging.s3.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/s3_v2.py", rationale: "Primary audit trail; batch flush no-drop"}
|
||||
- {id: logging.s3.failure.writes_object, module: logging, tier: P0, event: failure, assertions: [writes_object], exercised_on: [chat_completions, messages], source: "integrations/s3_v2.py", rationale: "Failed calls persisted for compliance"}
|
||||
- {id: logging.s3.success.partition_layout, module: logging, tier: P1, event: success, assertions: [object_key_layout], exercised_on: [chat_completions], source: "integrations/s3_v2.py / LIT-8985", rationale: "s3_partition_granularity picks the date or date/hour folder every downstream query and lifecycle rule reads"}
|
||||
- {id: logging.gcs_bucket.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/gcs_bucket/gcs_bucket.py", rationale: "GCS parallel to S3"}
|
||||
- {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"}
|
||||
- {id: logging.datadog.stream.exports_metric, module: logging, tier: P0, event: stream, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses], source: "integrations/datadog/datadog.py", rationale: "Streaming aggregates usage after the last chunk; delivery and cost must survive that path"}
|
||||
|
|
|
|||
|
|
@ -106,6 +106,7 @@ UI_BASE_URL = os.environ.get("E2E_UI_BASE_URL", PROXY_BASE_URL).rstrip("/")
|
|||
|
||||
CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5")
|
||||
CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
|
||||
S3_PARTITION_GRANULARITY = os.environ.get("E2E_S3_PARTITION_GRANULARITY", "day")
|
||||
|
||||
LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp")
|
||||
LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "")
|
||||
|
|
|
|||
|
|
@ -1,12 +1,13 @@
|
|||
"""Live e2e: Together AI through the gateway on /chat/completions and /v1/messages.
|
||||
|
||||
The reasoning and tool-calling backend is the cheapest live ``together_ai/`` chat row
|
||||
in the proxy's own cost map that carries both capability flags; the structured-output
|
||||
and cache-pricing backends are likewise the cheapest rows carrying
|
||||
``supports_response_schema`` and a ``cache_read_input_token_cost``. Two backends are
|
||||
pinned because the registry has no flag for what they prove: ``enable_thinking`` and
|
||||
the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps
|
||||
to are Qwen hybrid-model contracts, and MiniMax-M3 is the serverless model whose
|
||||
in the proxy's own cost map that carries both capability flags; the cache-pricing
|
||||
backend is likewise the cheapest row carrying a ``cache_read_input_token_cost``. Two
|
||||
backends are pinned because the registry has no flag for what they prove: ``enable_thinking``
|
||||
and the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps
|
||||
to are Qwen hybrid-model contracts, and the structured-output case runs on that hybrid
|
||||
model with reasoning off, since a reasoning-only model can spend the whole token budget
|
||||
thinking and return no content. MiniMax-M3 is the serverless model whose
|
||||
template renders a replayed ``reasoning_content`` back into the prompt (Qwen and
|
||||
DeepSeek silently drop it). MiniMax-M3 honors that replayed field on nearly every call, not
|
||||
every call (one miss in dozens of otherwise identical calls), so the replay case asks
|
||||
|
|
@ -109,7 +110,6 @@ MESSAGES_WEATHER_TOOL = AnthropicCustomTool(
|
|||
class _Needs:
|
||||
function_calling: bool = False
|
||||
reasoning: bool = False
|
||||
response_schema: bool = False
|
||||
cache_read_pricing: bool = False
|
||||
|
||||
|
||||
|
|
@ -172,7 +172,6 @@ def _cheapest_together_chat_model(registry: Mapping[str, CostMapEntry], needs: _
|
|||
and (entry.output_cost_per_token or 0.0) > 0
|
||||
and (not needs.function_calling or bool(entry.supports_function_calling))
|
||||
and (not needs.reasoning or bool(entry.supports_reasoning))
|
||||
and (not needs.response_schema or bool(entry.supports_response_schema))
|
||||
and (not needs.cache_read_pricing or (entry.cache_read_input_token_cost or 0.0) > 0)
|
||||
)
|
||||
|
||||
|
|
@ -586,10 +585,9 @@ class TestTogetherChatCompletions:
|
|||
|
||||
@pytest.mark.covers("llm.chat_completions.together_ai.structured_output.nonstream.works")
|
||||
def test_response_format_json_schema_shapes_the_reply(
|
||||
self, client: PassthroughClient, resources: ResourceManager, registry: dict[str, CostMapEntry]
|
||||
self, client: PassthroughClient, resources: ResourceManager
|
||||
) -> None:
|
||||
backend = _cheapest_together_chat_model(registry, _Needs(response_schema=True))
|
||||
model, key = _register(client, resources, backend)
|
||||
model, key = _register(client, resources, HYBRID_REASONING_BACKEND)
|
||||
|
||||
message = _message(
|
||||
unwrap(
|
||||
|
|
@ -599,12 +597,13 @@ class TestTogetherChatCompletions:
|
|||
model=model,
|
||||
messages=[ChatMessage(role="user", content=PERSON_PROMPT)],
|
||||
max_tokens=1024,
|
||||
reasoning_effort="none",
|
||||
response_format=PERSON_RESPONSE_FORMAT,
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
assert message.content, f"{backend} returned no content: {message}"
|
||||
assert message.content, f"{HYBRID_REASONING_BACKEND} returned no content: {message}"
|
||||
person = _Person.model_validate_json(message.content)
|
||||
assert person.name, f"schema-shaped reply carries an empty name: {message.content!r}"
|
||||
|
||||
|
|
|
|||
|
|
@ -22,11 +22,12 @@ alias per test turns the poll into a cheap prefix listing.
|
|||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
|
||||
from e2e_config import CHEAP_ANTHROPIC_MODEL, S3_PARTITION_GRANULARITY, unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
from logging_client import (
|
||||
INVALID_UPSTREAM_API_KEY,
|
||||
|
|
@ -106,6 +107,46 @@ class TestS3LogDelivery:
|
|||
record.response_cost, outcome.response_cost, rel_tol=1e-9
|
||||
), f"payload response_cost {record.response_cost!r} must equal the header cost {outcome.response_cost}"
|
||||
|
||||
@pytest.mark.covers("logging.s3.success.partition_layout", exercised_on=["chat_completions"])
|
||||
def test_chat_completions_object_key_follows_the_partition_granularity(
|
||||
self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager
|
||||
) -> None:
|
||||
"""The one object a call writes must sit in the folder layout the proxy's
|
||||
s3_partition_granularity names: {alias}/{date}/ for day and
|
||||
{alias}/{date}/{HH}/ for hour, where HH is the hour the object's own
|
||||
time- file name records. E2E_S3_PARTITION_GRANULARITY tells the test
|
||||
which one the proxy under test runs."""
|
||||
_assert_s3_configured(client)
|
||||
|
||||
alias = f"s3-layout-{unique_marker()}"
|
||||
key = client.key_with_alias(alias, models=[CHEAP_ANTHROPIC_MODEL])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
outcome = first_ok(
|
||||
client,
|
||||
lambda: client.chat_raw(
|
||||
key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {unique_marker()}", max_tokens=16
|
||||
),
|
||||
)
|
||||
body_id = completion_response_id(outcome.body)
|
||||
assert body_id is not None, "the completion body must carry an id (it names the s3 object)"
|
||||
records = s3_logs.poll_records(prefix=f"{alias}/", predicate=lambda r: r.id == body_id)
|
||||
assert len(records) == 1, f"expected exactly ONE s3 object for response {body_id}, got {len(records)}"
|
||||
|
||||
file_id = body_id.replace("/", "_").replace(":", "_")
|
||||
hour_folder = r"(?P<folder_hour>\d{2})/" if S3_PARTITION_GRANULARITY == "hour" else ""
|
||||
layout = re.compile(
|
||||
rf"{re.escape(alias)}/\d{{4}}-\d{{2}}-\d{{2}}/{hour_folder}"
|
||||
rf"time-(?P<file_hour>\d{{2}})-\d{{2}}-\d{{2}}-\d{{6}}_{re.escape(file_id)}\.json"
|
||||
)
|
||||
keys = [object_key for object_key in s3_logs.list_keys(f"{alias}/") if file_id in object_key]
|
||||
assert len(keys) == 1, f"expected one object key for response {body_id}, got {keys}"
|
||||
match = layout.fullmatch(keys[0])
|
||||
assert match is not None, f"{keys[0]!r} is outside the {S3_PARTITION_GRANULARITY} layout {layout.pattern!r}"
|
||||
assert S3_PARTITION_GRANULARITY != "hour" or match.group("folder_hour") == match.group("file_hour"), (
|
||||
f"the hour folder must be the hour the object's file name records: {keys[0]!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("logging.s3.failure.writes_object", exercised_on=["chat_completions"])
|
||||
def test_chat_completions_failure_writes_one_object(
|
||||
self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ that would catch a regression.
|
|||
|
||||
Companion: live suite `test_spend_tracking_e2e.py` + route breadth
|
||||
`test_spend_routes.py` (this directory). Offline regression suite:
|
||||
`tests/test_litellm/proxy/spend_tracking/`. Reference PR: BerriAI/litellm#29956.
|
||||
`tests/unit/proxy/spend_tracking/`. Reference PR: BerriAI/litellm#29956.
|
||||
|
||||
Levels: `unit` mocked; `integration` real DB/cost-map; `live` real provider +
|
||||
proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`.
|
||||
|
|
|
|||
|
|
@ -17,9 +17,10 @@ sets rather than on whatever the model happens to do by default.
|
|||
The streaming cases pin the served-tier contract: OpenAI stamps the tier it actually
|
||||
used on every stream chunk, and that echo is what the caller sees and what the bill
|
||||
must be computed on. The request sets no service_tier, so the only place the tier
|
||||
can come from is the provider's response. The spend row must record the served tier
|
||||
and price input at that tier's rate, and every chunk the proxy relays must carry the
|
||||
same service_tier the provider sent.
|
||||
can come from is the provider's response. The spend row must record the tier the bill
|
||||
was priced on and price input at that tier's rate, and every chunk the proxy relays must
|
||||
carry the same service_tier the provider sent. A served `default` tier is base pricing,
|
||||
which the bill records as no tier
|
||||
"""
|
||||
|
||||
import json
|
||||
|
|
@ -61,7 +62,8 @@ PRIORITY_OUTPUT_RATE = 1.6e-04
|
|||
|
||||
REASONING_EFFORT = "high"
|
||||
|
||||
TIER_INPUT_RATES = {"default": INPUT_RATE, "priority": PRIORITY_INPUT_RATE}
|
||||
PRICING_BASIS_FOR_SERVED_TIER: dict[str, str | None] = {"default": None, "priority": "priority"}
|
||||
INPUT_RATE_FOR_PRICING_BASIS: dict[str | None, float] = {None: INPUT_RATE, "priority": PRIORITY_INPUT_RATE}
|
||||
|
||||
|
||||
class _StreamChunk(BaseModel):
|
||||
|
|
@ -213,17 +215,20 @@ class TestServiceTierPricing:
|
|||
)
|
||||
chunks = _stream_chunks(result.stream_events)
|
||||
served_tier = _served_tier(chunks)
|
||||
assert served_tier in TIER_INPUT_RATES, f"no custom rate registered for served tier {served_tier!r}"
|
||||
assert served_tier in PRICING_BASIS_FOR_SERVED_TIER, (
|
||||
f"no custom rate registered for served tier {served_tier!r}"
|
||||
)
|
||||
pricing_basis = PRICING_BASIS_FOR_SERVED_TIER[served_tier]
|
||||
stream_id = chunks[0].id
|
||||
assert stream_id, f"first stream chunk carried no id: {result.stream_events[0][:200]}"
|
||||
|
||||
row = poll_cost_row(client.proxy, stream_id)
|
||||
assert row is not None, f"no spend row with a cost breakdown landed for {stream_id}"
|
||||
assert row.breakdown.service_tier == served_tier, (
|
||||
f"the provider served tier {served_tier!r} on every chunk but the bill records "
|
||||
f"pricing basis {row.breakdown.service_tier!r}"
|
||||
assert row.breakdown.service_tier == pricing_basis, (
|
||||
f"the provider served tier {served_tier!r} on every chunk, so the bill should record pricing "
|
||||
f"basis {pricing_basis!r}, but it records {row.breakdown.service_tier!r}"
|
||||
)
|
||||
assert_fresh_tokens_billed_at(row, TIER_INPUT_RATES[served_tier])
|
||||
assert_fresh_tokens_billed_at(row, INPUT_RATE_FOR_PRICING_BASIS[pricing_basis])
|
||||
assert_total_is_sum_of_components(row)
|
||||
|
||||
@pytest.mark.covers("llm.chat_completions.openai.service_tier.stream.echoes_served_tier")
|
||||
|
|
@ -264,7 +269,14 @@ class TestServiceTierPricing:
|
|||
client.proxy,
|
||||
resources,
|
||||
"tier-responses-stream",
|
||||
LiteLLMParamsBody(model=STREAM_BACKEND, api_key=OPENAI_API_KEY),
|
||||
LiteLLMParamsBody(
|
||||
model=STREAM_BACKEND,
|
||||
api_key=OPENAI_API_KEY,
|
||||
input_cost_per_token=INPUT_RATE,
|
||||
output_cost_per_token=OUTPUT_RATE,
|
||||
input_cost_per_token_priority=PRIORITY_INPUT_RATE,
|
||||
output_cost_per_token_priority=PRIORITY_OUTPUT_RATE,
|
||||
),
|
||||
)
|
||||
|
||||
result = client.proxy.responses_stream(
|
||||
|
|
@ -282,14 +294,18 @@ class TestServiceTierPricing:
|
|||
)
|
||||
served_tier = completed.response.service_tier
|
||||
assert served_tier, f"response.completed carried no service_tier: {completed.response}"
|
||||
assert served_tier in TIER_INPUT_RATES, f"no custom rate registered for served tier {served_tier!r}"
|
||||
assert served_tier in PRICING_BASIS_FOR_SERVED_TIER, (
|
||||
f"no custom rate registered for served tier {served_tier!r}"
|
||||
)
|
||||
pricing_basis = PRICING_BASIS_FOR_SERVED_TIER[served_tier]
|
||||
|
||||
row = poll_cost_row_where(client.proxy, scoped_key, lambda r: r.spend is not None and r.spend > 0)
|
||||
assert row is not None, f"no spend row with a cost breakdown landed for the streamed responses call on {model}"
|
||||
assert row.breakdown.service_tier == served_tier, (
|
||||
f"response.completed served tier {served_tier!r} but the bill records "
|
||||
f"pricing basis {row.breakdown.service_tier!r}"
|
||||
assert row.breakdown.service_tier == pricing_basis, (
|
||||
f"response.completed served tier {served_tier!r}, so the bill should record pricing basis "
|
||||
f"{pricing_basis!r}, but it records {row.breakdown.service_tier!r}"
|
||||
)
|
||||
assert_fresh_tokens_billed_at(row, INPUT_RATE_FOR_PRICING_BASIS[pricing_basis])
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.service_tier_stream.messages_records_served_tier")
|
||||
def test_messages_stream_records_the_served_tier(
|
||||
|
|
@ -299,7 +315,14 @@ class TestServiceTierPricing:
|
|||
client.proxy,
|
||||
resources,
|
||||
"tier-messages-stream",
|
||||
LiteLLMParamsBody(model=STREAM_BACKEND, api_key=OPENAI_API_KEY),
|
||||
LiteLLMParamsBody(
|
||||
model=STREAM_BACKEND,
|
||||
api_key=OPENAI_API_KEY,
|
||||
input_cost_per_token=INPUT_RATE,
|
||||
output_cost_per_token=OUTPUT_RATE,
|
||||
input_cost_per_token_priority=PRIORITY_INPUT_RATE,
|
||||
output_cost_per_token_priority=PRIORITY_OUTPUT_RATE,
|
||||
),
|
||||
)
|
||||
|
||||
result = client.proxy.messages_stream(
|
||||
|
|
@ -322,8 +345,9 @@ class TestServiceTierPricing:
|
|||
|
||||
row = poll_cost_row_where(client.proxy, scoped_key, lambda r: r.spend is not None and r.spend > 0)
|
||||
assert row is not None, f"no spend row with a cost breakdown landed for the streamed messages call on {model}"
|
||||
served_tier = row.breakdown.service_tier
|
||||
assert served_tier in TIER_INPUT_RATES and served_tier is not None, (
|
||||
pricing_basis = row.breakdown.service_tier
|
||||
assert pricing_basis in INPUT_RATE_FOR_PRICING_BASIS, (
|
||||
"the anthropic wire format carries no service_tier, so the bill is the only record of "
|
||||
f"the tier OpenAI served; the row recorded pricing basis {served_tier!r}"
|
||||
f"the tier OpenAI served; the row recorded pricing basis {pricing_basis!r}"
|
||||
)
|
||||
assert_fresh_tokens_billed_at(row, INPUT_RATE_FOR_PRICING_BASIS[pricing_basis])
|
||||
|
|
|
|||
|
|
@ -111,7 +111,15 @@ test.describe("Logs page", () => {
|
|||
await dismissFeedbackPopup(page);
|
||||
const search = visibleTestId(page, "datatable-search");
|
||||
await expect(search).toBeVisible({ timeout: 20_000 });
|
||||
const searched = page.waitForResponse(
|
||||
(response) =>
|
||||
response.url().includes("/spend/logs/ui") &&
|
||||
new URL(response.url()).searchParams.get("search") === callId &&
|
||||
response.status() === 200,
|
||||
{ timeout: 20_000 },
|
||||
);
|
||||
await search.fill(callId);
|
||||
await searched;
|
||||
|
||||
const row = requestLogsRows(page).filter({ hasText: requestId });
|
||||
await expect(row, `no logs row for call id ${callId}`).toHaveCount(1, { timeout: 30_000 });
|
||||
|
|
|
|||
|
|
@ -22,11 +22,10 @@ test.describe("Tag management", () => {
|
|||
async () => {
|
||||
await navigateToPage(page, DashboardPage.TagManagement);
|
||||
await page.getByRole("button", { name: "+ Create New Tag" }).click();
|
||||
await expect(
|
||||
page.getByRole("dialog", { name: "Create New Tag" }),
|
||||
).toBeVisible();
|
||||
await page.getByLabel("Tag Name").fill(tagName);
|
||||
await page.getByLabel("Description").fill(description);
|
||||
const createDialog = page.getByRole("dialog", { name: "Create New Tag" });
|
||||
await expect(createDialog).toBeVisible();
|
||||
await createDialog.getByLabel("Tag Name").fill(tagName);
|
||||
await createDialog.getByLabel("Description").fill(description);
|
||||
await page.getByRole("button", { name: "Create Tag" }).click();
|
||||
|
||||
await expect
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from pydantic import JsonValue, TypeAdapter
|
|||
from tests.integration._support.database import read_rows
|
||||
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
GATEWAY_LIMITS: Final = httpx.Limits(keepalive_expiry=2)
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
|
|
@ -243,7 +244,7 @@ class Scenario:
|
|||
def gateway_from_environment() -> Iterator[Gateway]:
|
||||
url: Final = os.environ["INTEGRATION_PROXY_URL"]
|
||||
upstream: Final = os.environ["INTEGRATION_UPSTREAM_URL"]
|
||||
with httpx.Client(base_url=url, timeout=15, trust_env=False) as client:
|
||||
with httpx.Client(base_url=url, timeout=15, trust_env=False, limits=GATEWAY_LIMITS) as client:
|
||||
yield Gateway(client, os.environ["INTEGRATION_MASTER_KEY"], upstream)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from typing import Final
|
||||
|
|
@ -9,6 +10,7 @@ from urllib.parse import urlsplit, urlunsplit
|
|||
from pydantic import TypeAdapter
|
||||
|
||||
PORT: Final = TypeAdapter(int)
|
||||
OUTAGE_SECONDS: Final = 10.0
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
|
|
@ -27,6 +29,8 @@ class DatabaseRelay:
|
|||
self._armed: Final = threading.Event()
|
||||
self.tripped: Final = threading.Event()
|
||||
self.refused = 0
|
||||
self.reconnected: Final = threading.Event()
|
||||
self._tripped_at = 0.0
|
||||
self._writers: tuple[asyncio.StreamWriter, ...] = ()
|
||||
self._ready: Final = threading.Event()
|
||||
self._thread: Final = threading.Thread(target=self._run, daemon=True)
|
||||
|
|
@ -54,10 +58,12 @@ class DatabaseRelay:
|
|||
self._writers = ()
|
||||
|
||||
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
||||
if self.tripped.is_set() and self.refused < 5:
|
||||
if self.tripped.is_set() and time.monotonic() - self._tripped_at < OUTAGE_SECONDS:
|
||||
self.refused += 1
|
||||
client_writer.close()
|
||||
return
|
||||
if self.tripped.is_set():
|
||||
self.reconnected.set()
|
||||
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
||||
self._writers = (*self._writers, client_writer, server_writer)
|
||||
|
||||
|
|
@ -65,6 +71,7 @@ class DatabaseRelay:
|
|||
try:
|
||||
while chunk := await reader.read(65536):
|
||||
if inspect and self._armed.is_set() and not self.tripped.is_set() and self._trigger in chunk:
|
||||
self._tripped_at = time.monotonic()
|
||||
self.tripped.set()
|
||||
self._drop_all()
|
||||
return
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from typing import Final
|
|||
|
||||
import httpx
|
||||
import psutil
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.client import GATEWAY_LIMITS, Gateway
|
||||
|
||||
|
||||
def proxy_database_environment() -> Mapping[str, str]:
|
||||
|
|
@ -80,6 +80,88 @@ def owned_proxy(
|
|||
yield owned.gateway
|
||||
|
||||
|
||||
def _stop(process: subprocess.Popen[bytes]) -> None:
|
||||
root_stopped: Final = stop_root_process(process)
|
||||
residual: Final = group_members(process.pid)
|
||||
if residual:
|
||||
signal_group(process.pid, signal.SIGTERM)
|
||||
psutil.wait_procs(residual, timeout=5)
|
||||
remaining: Final = group_members(process.pid)
|
||||
if remaining:
|
||||
signal_group(process.pid, signal.SIGKILL)
|
||||
psutil.wait_procs(remaining, timeout=3)
|
||||
process.wait(timeout=3)
|
||||
survivors: Final = group_members(process.pid)
|
||||
assert not survivors, "Owned proxy child survived cleanup"
|
||||
assert root_stopped and not remaining, "Owned proxy required forced cleanup"
|
||||
|
||||
|
||||
_PORT_ATTEMPTS: Final = 3
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket() as reserve:
|
||||
reserve.bind(("127.0.0.1", 0))
|
||||
return reserve.getsockname()[1]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Launch:
|
||||
process: subprocess.Popen[bytes]
|
||||
port: int
|
||||
log: Path
|
||||
|
||||
|
||||
def _launch(command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path) -> _Launch:
|
||||
port: Final = _free_port()
|
||||
log_path: Final = output / f"owned-proxy-{uuid.uuid4().hex}.log"
|
||||
with log_path.open("w") as log:
|
||||
process: Final = subprocess.Popen(
|
||||
[*command, "--port", str(port)],
|
||||
cwd=root,
|
||||
env=environment,
|
||||
stdout=log,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
return _Launch(process, port, log_path)
|
||||
|
||||
|
||||
def _lost_port_race(launch: _Launch) -> bool:
|
||||
return launch.process.poll() is not None and "address already in use" in launch.log.read_text()
|
||||
|
||||
|
||||
def _wait_until_ready(launch: _Launch) -> None:
|
||||
with httpx.Client(base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False) as client:
|
||||
deadline: Final = time.monotonic() + 70
|
||||
while launch.process.poll() is None:
|
||||
try:
|
||||
if client.get("/health/readiness", timeout=2).status_code == 200:
|
||||
return
|
||||
except httpx.TransportError:
|
||||
pass
|
||||
assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded"
|
||||
time.sleep(0.1)
|
||||
|
||||
|
||||
def _launch_until_bound(
|
||||
command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path, attempts: int
|
||||
) -> _Launch:
|
||||
launch: Final = _launch(command, root, environment, output)
|
||||
try:
|
||||
_wait_until_ready(launch)
|
||||
assert launch.process.poll() is None or (attempts > 1 and _lost_port_race(launch)), (
|
||||
"Owned proxy exited before readiness"
|
||||
)
|
||||
except BaseException:
|
||||
_stop(launch.process)
|
||||
raise
|
||||
if launch.process.poll() is None:
|
||||
return launch
|
||||
_stop(launch.process)
|
||||
return _launch_until_bound(command, root, environment, output, attempts - 1)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def owned_proxy_process(
|
||||
gateway: Gateway,
|
||||
|
|
@ -90,9 +172,6 @@ def owned_proxy_process(
|
|||
remove_environment: tuple[str, ...] = (),
|
||||
workers: int = 1,
|
||||
) -> Iterator[OwnedProxy]:
|
||||
with socket.socket() as reserve:
|
||||
reserve.bind(("127.0.0.1", 0))
|
||||
port: Final = reserve.getsockname()[1]
|
||||
root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
|
||||
environment: Final = {
|
||||
**{
|
||||
|
|
@ -107,54 +186,25 @@ def owned_proxy_process(
|
|||
}
|
||||
output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
log_path: Final = output / f"owned-proxy-{uuid.uuid4().hex}.log"
|
||||
with log_path.open("w") as log:
|
||||
process: Final = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"integration._support.proxy",
|
||||
"--config",
|
||||
str(config or "tests/integration/proxy_config.yaml"),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(port),
|
||||
"--num_workers",
|
||||
str(workers),
|
||||
"--use_prisma_db_push",
|
||||
"--enforce_prisma_migration_check",
|
||||
],
|
||||
cwd=root,
|
||||
env=environment,
|
||||
stdout=log,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
try:
|
||||
with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client:
|
||||
deadline: Final = time.monotonic() + 70
|
||||
while True:
|
||||
assert process.poll() is None, "Owned proxy exited before readiness"
|
||||
try:
|
||||
if client.get("/health/readiness", timeout=2).status_code == 200:
|
||||
break
|
||||
except httpx.TransportError:
|
||||
pass
|
||||
assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded"
|
||||
time.sleep(0.1)
|
||||
yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, log_path)
|
||||
finally:
|
||||
root_stopped: Final = stop_root_process(process)
|
||||
residual: Final = group_members(process.pid)
|
||||
if residual:
|
||||
signal_group(process.pid, signal.SIGTERM)
|
||||
psutil.wait_procs(residual, timeout=5)
|
||||
remaining: Final = group_members(process.pid)
|
||||
if remaining:
|
||||
signal_group(process.pid, signal.SIGKILL)
|
||||
psutil.wait_procs(remaining, timeout=3)
|
||||
process.wait(timeout=3)
|
||||
survivors: Final = group_members(process.pid)
|
||||
assert not survivors, "Owned proxy child survived cleanup"
|
||||
assert root_stopped and not remaining, "Owned proxy required forced cleanup"
|
||||
command: Final = (
|
||||
sys.executable,
|
||||
"-m",
|
||||
"integration._support.proxy",
|
||||
"--config",
|
||||
str(config or "tests/integration/proxy_config.yaml"),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--num_workers",
|
||||
str(workers),
|
||||
"--use_prisma_db_push",
|
||||
"--enforce_prisma_migration_check",
|
||||
)
|
||||
launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS)
|
||||
process: Final = launch.process
|
||||
try:
|
||||
with httpx.Client(
|
||||
base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS
|
||||
) as client:
|
||||
yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log)
|
||||
finally:
|
||||
_stop(process)
|
||||
|
|
|
|||
|
|
@ -1,9 +1,6 @@
|
|||
"""Run the normal single-process CLI with the existing behavior-suite test entitlement."""
|
||||
|
||||
import signal
|
||||
import sys
|
||||
from types import FrameType
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm import run_server
|
||||
|
||||
|
|
@ -14,10 +11,7 @@ def _exit_on_reraised_term(signum: int, frame: FrameType | None) -> None:
|
|||
|
||||
def main() -> None:
|
||||
signal.signal(signal.SIGTERM, _exit_on_reraised_term)
|
||||
with patch( # test-quality-ok: route entitlement only; license validation is outside these HTTP/DB contracts
|
||||
"litellm.proxy.auth.litellm_license.LicenseCheck.is_premium", return_value=True
|
||||
):
|
||||
run_server()
|
||||
run_server()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -1,65 +0,0 @@
|
|||
import asyncio
|
||||
import os
|
||||
from collections.abc import AsyncIterator
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from prisma import Prisma
|
||||
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.engine.models import Check, Engine, EngineSettings, Scope, Worker
|
||||
from litellm.proxy.engine.repository import EngineRepository, WriterDatabase
|
||||
from litellm.proxy.engine.state import claim_job, queue_job
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def engine_db() -> AsyncIterator[Prisma]:
|
||||
async with Prisma(datasource={"url": os.environ["DATABASE_URL"]}) as db:
|
||||
yield db
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_workers_cannot_both_acquire_the_same_job(engine_db: Prisma) -> None:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
scope: Final = Scope(team_id=uuid4().hex)
|
||||
repo: Final = EngineRepository(WriterDatabase(PrismaWrapper(engine_db)))
|
||||
engine: Final = Engine(
|
||||
id=uuid4().hex,
|
||||
scope=scope,
|
||||
settings=EngineSettings(name="Lease test", model="test", checks=(Check(id="c", instruction="Find retries"),)),
|
||||
created_at=now,
|
||||
next_run_at=now,
|
||||
budget_month=now.strftime("%Y-%m"),
|
||||
)
|
||||
await repo.create(queue_job(engine, now, uuid4().hex))
|
||||
try:
|
||||
workers: Final = tuple(Worker(id=uuid4().hex, name="worker", scope=scope, last_seen=now) for _ in range(2))
|
||||
results: Final = await asyncio.gather(
|
||||
*(repo.update(engine.id, lambda e, w=w: claim_job(e, w, now)) for w in workers)
|
||||
)
|
||||
stored: Final = await repo.get(engine.id)
|
||||
assert stored is not None
|
||||
assert stored.jobs[0].attempts == 1
|
||||
assert stored.jobs[0].worker_id in tuple(w.id for w in workers)
|
||||
assert tuple(r.jobs[0].worker_id for r in results if r) == (stored.jobs[0].worker_id, stored.jobs[0].worker_id)
|
||||
finally:
|
||||
await engine_db.execute_raw('DELETE FROM "LiteLLM_Engine" WHERE id=$1', engine.id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_never_restores_revoked_access(engine_db: Prisma) -> None:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
repo: Final = EngineRepository(WriterDatabase(PrismaWrapper(engine_db)))
|
||||
worker: Final = Worker(id=uuid4().hex, name="worker", scope=Scope(team_id=uuid4().hex), last_seen=now)
|
||||
token_hash: Final = uuid4().hex
|
||||
await repo.save_worker(worker, token_hash)
|
||||
try:
|
||||
await repo.save_worker(worker.model_copy(update={"revoked": True}))
|
||||
await repo.heartbeat(worker.id, now.isoformat())
|
||||
stored: Final = await repo.worker(token_hash)
|
||||
assert stored is not None and stored.revoked is True
|
||||
finally:
|
||||
await engine_db.execute_raw('DELETE FROM "LiteLLM_EngineWorker" WHERE id=$1', worker.id)
|
||||
173
tests/integration/database/test_lens_repository.py
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
import asyncio
|
||||
import os
|
||||
from collections.abc import AsyncIterator
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
from uuid import uuid4
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from prisma import Prisma
|
||||
from psycopg import sql
|
||||
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.lens.models import Check, Lens, LensSettings, Scope, Worker
|
||||
from litellm.proxy.lens.repository import LensRepository, WriterDatabase
|
||||
from litellm.proxy.lens.state import claim_job, queue_job
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def lens_db() -> AsyncIterator[Prisma]:
|
||||
async with Prisma(datasource={"url": os.environ["DATABASE_URL"]}) as db:
|
||||
yield db
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_workers_cannot_both_acquire_the_same_job(lens_db: Prisma) -> None:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
scope: Final = Scope(team_id=uuid4().hex)
|
||||
repo: Final = LensRepository(WriterDatabase(PrismaWrapper(lens_db)))
|
||||
lens: Final = Lens(
|
||||
id=uuid4().hex,
|
||||
scope=scope,
|
||||
settings=LensSettings(name="Lease test", model="test", checks=(Check(id="c", instruction="Find retries"),)),
|
||||
created_at=now,
|
||||
next_run_at=now,
|
||||
budget_month=now.strftime("%Y-%m"),
|
||||
)
|
||||
await repo.create(queue_job(lens, now, uuid4().hex))
|
||||
try:
|
||||
workers: Final = tuple(Worker(id=uuid4().hex, name="worker", scope=scope, last_seen=now) for _ in range(2))
|
||||
results: Final = await asyncio.gather(
|
||||
*(repo.update(lens.id, lambda e, w=w: claim_job(e, w, now)) for w in workers)
|
||||
)
|
||||
stored: Final = await repo.get(lens.id)
|
||||
assert stored is not None
|
||||
assert stored.jobs[0].attempts == 1
|
||||
assert stored.jobs[0].worker_id in tuple(w.id for w in workers)
|
||||
assert tuple(r.jobs[0].worker_id for r in results if r) == (stored.jobs[0].worker_id, stored.jobs[0].worker_id)
|
||||
finally:
|
||||
await lens_db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', lens.id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_never_restores_revoked_access(lens_db: Prisma) -> None:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
repo: Final = LensRepository(WriterDatabase(PrismaWrapper(lens_db)))
|
||||
worker: Final = Worker(id=uuid4().hex, name="worker", scope=Scope(team_id=uuid4().hex), last_seen=now)
|
||||
token_hash: Final = uuid4().hex
|
||||
await repo.save_worker(worker, token_hash)
|
||||
try:
|
||||
await repo.save_worker(worker.model_copy(update={"revoked": True}))
|
||||
await repo.heartbeat(worker.id, now.isoformat())
|
||||
stored: Final = await repo.worker(token_hash)
|
||||
assert stored is not None and stored.revoked is True
|
||||
finally:
|
||||
await lens_db.execute_raw('DELETE FROM "LiteLLM_LensWorker" WHERE id=$1', worker.id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("populated", (False, True))
|
||||
@pytest.mark.parametrize("preceding_schema", (False, True))
|
||||
def test_lens_rename_preserves_saved_data_and_worker_credentials(populated: bool, preceding_schema: bool) -> None:
|
||||
migrations: Final = (
|
||||
Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
|
||||
)
|
||||
schema: Final = f"lens_migration_{uuid4().hex}"
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
|
||||
try:
|
||||
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
|
||||
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
|
||||
for name in ("20260930000000_agent_engine", "20261001000000_lens_run_history"):
|
||||
connection.execute(sql.SQL((migrations / name / "migration.sql").read_text()))
|
||||
if populated:
|
||||
connection.execute(
|
||||
"""INSERT INTO "LiteLLM_Engine" VALUES ('lens', 7, '{"findings":[{"id":"finding"}]}');
|
||||
INSERT INTO "LiteLLM_EngineWorker" VALUES ('worker', 'token-hash', '{"analysis_key_id":"key"}');
|
||||
INSERT INTO "LiteLLM_EngineRun" VALUES ('batch', 'lens', '2026-01-01', '{"cost":1.25}')"""
|
||||
)
|
||||
if preceding_schema:
|
||||
first_schema: Final = f"lens_first_{uuid4().hex}"
|
||||
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(first_schema)))
|
||||
connection.execute(
|
||||
sql.SQL("SET LOCAL search_path TO {}, {}").format(
|
||||
sql.Identifier(first_schema), sql.Identifier(schema)
|
||||
)
|
||||
)
|
||||
connection.execute(sql.SQL((migrations / "20261001100000_rename_lens" / "migration.sql").read_text()))
|
||||
connection.execute(sql.SQL((migrations / "20261001100000_rename_lens" / "migration.sql").read_text()))
|
||||
assert connection.execute('SELECT id, version, data FROM "LiteLLM_Lens"').fetchall() == (
|
||||
[("lens", 7, {"findings": [{"id": "finding"}]})] if populated else []
|
||||
)
|
||||
assert connection.execute('SELECT id, token_hash, data FROM "LiteLLM_LensWorker"').fetchall() == (
|
||||
[("worker", "token-hash", {"analysis_key_id": "key"})] if populated else []
|
||||
)
|
||||
assert connection.execute('SELECT id, lens_id, data FROM "LiteLLM_LensRun"').fetchall() == (
|
||||
[("batch", "lens", {"cost": 1.25})] if populated else []
|
||||
)
|
||||
finally:
|
||||
connection.rollback()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entrypoint", ("proxy", "extras-v1", "extras-v2"))
|
||||
@pytest.mark.parametrize("legacy_table", ("LiteLLM_Engine", "LiteLLM_EngineRun", "LiteLLM_EngineWorker"))
|
||||
def test_db_push_refuses_legacy_lens_data(monkeypatch: pytest.MonkeyPatch, entrypoint: str, legacy_table: str) -> None:
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
|
||||
from litellm.proxy.db.prisma_client import PrismaManager
|
||||
|
||||
database_url: Final = os.environ["DATABASE_URL"]
|
||||
schema: Final = f"lens_push_{uuid4().hex}"
|
||||
parsed: Final = urlsplit(database_url)
|
||||
scoped: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})))
|
||||
with psycopg.connect(database_url, autocommit=True) as connection:
|
||||
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
|
||||
try:
|
||||
connection.execute(
|
||||
sql.SQL("CREATE TABLE {} (id TEXT PRIMARY KEY, data JSONB)").format(
|
||||
sql.Identifier(schema, legacy_table)
|
||||
)
|
||||
)
|
||||
connection.execute(
|
||||
sql.SQL("INSERT INTO {} VALUES ('saved', '{{\"keep\":true}}')").format(
|
||||
sql.Identifier(schema, legacy_table)
|
||||
)
|
||||
)
|
||||
monkeypatch.setenv("DATABASE_URL", scoped)
|
||||
setup: Final = (
|
||||
PrismaManager.setup_database if entrypoint == "proxy" else ProxyExtrasDBManager.setup_database
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="Legacy Lens tables exist"):
|
||||
setup(use_migrate=False, use_v2_resolver=entrypoint == "extras-v2")
|
||||
assert connection.execute(
|
||||
sql.SQL("SELECT id, data FROM {}").format(sql.Identifier(schema, legacy_table))
|
||||
).fetchall() == [("saved", {"keep": True})]
|
||||
finally:
|
||||
connection.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
|
||||
|
||||
|
||||
def test_db_push_creates_fresh_lens_tables_and_preserves_them_on_restart(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.proxy.db.prisma_client import PrismaManager
|
||||
|
||||
database_url: Final = os.environ["DATABASE_URL"]
|
||||
schema: Final = f"lens_fresh_push_{uuid4().hex}"
|
||||
parsed: Final = urlsplit(database_url)
|
||||
scoped: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})))
|
||||
with psycopg.connect(database_url, autocommit=True) as connection:
|
||||
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
|
||||
try:
|
||||
monkeypatch.setenv("DATABASE_URL", scoped)
|
||||
assert PrismaManager.setup_database(use_migrate=False)
|
||||
connection.execute(
|
||||
sql.SQL("INSERT INTO {} (id, data) VALUES ('saved', '{{\"keep\":true}}')").format(
|
||||
sql.Identifier(schema, "LiteLLM_Lens")
|
||||
)
|
||||
)
|
||||
assert PrismaManager.setup_database(use_migrate=False)
|
||||
assert connection.execute(
|
||||
sql.SQL("SELECT id, data FROM {}").format(sql.Identifier(schema, "LiteLLM_Lens"))
|
||||
).fetchall() == [("saved", {"keep": True})]
|
||||
finally:
|
||||
connection.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
|
||||
|
|
@ -30,8 +30,10 @@ from integration._support.wire import Reply, Request, wire_server
|
|||
TENANT: Final = "00000000-0000-4000-8000-0000000a3650"
|
||||
REJECTED: Final = "Agent 365 guardrail rejected the tool call"
|
||||
GUARDRAIL_ROWS: Final = (
|
||||
"SELECT metadata->'guardrail_information' AS gi FROM \"LiteLLM_SpendLogs\" "
|
||||
'WHERE api_key = %s AND call_type = %s ORDER BY "startTime"'
|
||||
"SELECT COALESCE(jsonb_path_query_first(metadata, "
|
||||
"'$.guardrail_information[*] ? (@.guardrail_name == $name).guardrail_status', "
|
||||
"jsonb_build_object('name', %s::text)) #>> '{}', 'none') AS status "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s ORDER BY "startTime"'
|
||||
)
|
||||
FALLBACKS: Final = (None, "fail_open", "fail_closed")
|
||||
|
||||
|
|
@ -95,11 +97,11 @@ class Rig:
|
|||
|
||||
def guardrail_statuses(self, call_type: str, at_least: int) -> list[str]:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(GUARDRAIL_ROWS, (sha256(self.key.encode()).hexdigest(), call_type)),
|
||||
lambda: read_rows(GUARDRAIL_ROWS, (self.alias, sha256(self.key.encode()).hexdigest(), call_type)),
|
||||
lambda seen: len(seen) >= at_least,
|
||||
seconds=70,
|
||||
)
|
||||
return [row["gi"][0]["guardrail_status"] if row["gi"] else "none" for row in rows]
|
||||
return [str(row["status"]) for row in rows]
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
|
|
|||
1241
tests/integration/observability/test_s3_v2_partition_granularity.py
Normal file
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from typing import Final
|
||||
|
||||
import openai
|
||||
|
|
@ -69,8 +70,27 @@ def _sent_messages(request: Request) -> list[dict[str, JsonValue]]:
|
|||
return _MESSAGES.validate_python(_JSON_OBJECT.validate_json(request.body)["messages"])
|
||||
|
||||
|
||||
_DISCOVERY_PROBE: Final = ("GET", "/v1/models")
|
||||
|
||||
|
||||
def _is_discovery_probe(request: Request) -> bool:
|
||||
return (request.method, request.target) == _DISCOVERY_PROBE
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _vllm_server(respond: Callable[[Request], Reply]) -> Iterator[Wire]:
|
||||
with wire_server(
|
||||
lambda request: Reply(body=b'{"object":"list","data":[]}') if _is_discovery_probe(request) else respond(request)
|
||||
) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
def _provider_calls(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(request for request in wire.drain() if not _is_discovery_probe(request))
|
||||
|
||||
|
||||
def _only_request(wire: Wire) -> Request:
|
||||
received: Final = wire.drain()
|
||||
received: Final = _provider_calls(wire)
|
||||
assert [(request.method, request.target) for request in received] == [("POST", "/v1/chat/completions")]
|
||||
return received[0]
|
||||
|
||||
|
|
@ -139,7 +159,7 @@ def test_hosted_vllm_assistant_reasoning_content_reaches_the_wire(gateway: Gatew
|
|||
], body["messages"]
|
||||
return Reply(body=_completion(identity, "The totals differ by 42."))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
with _vllm_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
|
|
@ -167,13 +187,13 @@ def test_hosted_vllm_assistant_reasoning_content_reaches_the_wire(gateway: Gatew
|
|||
assert response.status_code == 200, response.text
|
||||
payload: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
assert payload["id"] == identity
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/chat/completions")]
|
||||
_only_request(wire)
|
||||
|
||||
|
||||
def test_openai_sdk_replayed_reasoning_reaches_hosted_vllm_and_is_billed_once(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
identity: Final = f"chatcmpl-sdk-{marker}"
|
||||
with wire_server(lambda _: Reply(body=_completion(identity, "They differ by 42."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(identity, "They differ by 42."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
completion: Final = _openai_client(gateway).chat.completions.create(
|
||||
|
|
@ -195,7 +215,7 @@ def test_openai_sdk_replayed_reasoning_reaches_hosted_vllm_and_is_billed_once(ga
|
|||
async def test_async_openai_sdk_stream_forwards_replayed_reasoning_to_hosted_vllm(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
identity: Final = f"chatcmpl-stream-{marker}"
|
||||
with wire_server(lambda _: _streamed_completion(identity, "They differ by 42.")) as wire:
|
||||
with _vllm_server(lambda _: _streamed_completion(identity, "They differ by 42.")) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
stream: Final = await _async_openai_client(gateway).chat.completions.create(
|
||||
|
|
@ -224,7 +244,7 @@ def test_each_replayed_turn_keeps_its_own_reasoning_in_order(gateway: Gateway) -
|
|||
{"role": "assistant", "content": "Step two.", "reasoning_content": f"second thought {marker}"},
|
||||
{"role": "user", "content": "Summarize."},
|
||||
]
|
||||
with wire_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
assert _post_chat(gateway, model, conversation)["id"] == f"chatcmpl-{marker}"
|
||||
|
|
@ -246,7 +266,7 @@ def test_only_string_reasoning_content_is_forwarded_to_hosted_vllm(
|
|||
gateway: Gateway, reasoning: JsonValue, forwarded: str | None
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
assert _post_chat(gateway, model, _replayed_conversation(reasoning, marker))["id"] == f"chatcmpl-{marker}"
|
||||
|
|
@ -264,7 +284,7 @@ def test_assistant_turn_without_reasoning_gets_no_reasoning_key(gateway: Gateway
|
|||
{"role": "assistant", "content": "Hi there."},
|
||||
{"role": "user", "content": "Again"},
|
||||
]
|
||||
with wire_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Hello again."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Hello again."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
_post_chat(gateway, model, conversation)
|
||||
|
|
@ -281,7 +301,7 @@ def test_same_reasoning_on_two_turns_is_forwarded_on_both(gateway: Gateway) -> N
|
|||
{"role": "assistant", "content": "Second.", "reasoning_content": reasoning},
|
||||
{"role": "user", "content": "Three"},
|
||||
]
|
||||
with wire_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Third."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Third."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
_post_chat(gateway, model, conversation)
|
||||
|
|
@ -290,7 +310,7 @@ def test_same_reasoning_on_two_turns_is_forwarded_on_both(gateway: Gateway) -> N
|
|||
|
||||
def test_thinking_blocks_are_stripped_while_reasoning_content_is_kept(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
_post_chat(
|
||||
|
|
@ -316,7 +336,7 @@ def test_thinking_blocks_are_stripped_while_reasoning_content_is_kept(gateway: G
|
|||
|
||||
def test_list_content_is_flattened_while_reasoning_content_is_kept(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
_post_chat(
|
||||
|
|
@ -341,7 +361,7 @@ def test_list_content_is_flattened_while_reasoning_content_is_kept(gateway: Gate
|
|||
|
||||
def test_unauthenticated_replay_is_rejected_before_hosted_vllm(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(f"chatcmpl-{marker}", "Done."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
|
|
@ -351,7 +371,7 @@ def test_unauthenticated_replay_is_rejected_before_hosted_vllm(gateway: Gateway)
|
|||
key=f"sk-not-a-key-{marker}",
|
||||
)
|
||||
assert response.status_code == 401, response.text
|
||||
assert wire.drain() == ()
|
||||
assert _provider_calls(wire) == ()
|
||||
|
||||
|
||||
def test_hosted_vllm_auth_error_reaches_the_caller_after_one_attempt_with_reasoning(gateway: Gateway) -> None:
|
||||
|
|
@ -361,7 +381,7 @@ def test_hosted_vllm_auth_error_reaches_the_caller_after_one_attempt_with_reason
|
|||
status=401,
|
||||
body=json.dumps({"error": {"message": error_message, "type": "authentication_error"}}).encode(),
|
||||
)
|
||||
with wire_server(lambda _: reply) as wire, gateway.scenario() as scenario:
|
||||
with _vllm_server(lambda _: reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
|
|
@ -381,7 +401,7 @@ def test_fallback_attempt_replays_reasoning_to_the_second_deployment(gateway: Ga
|
|||
return Reply(status=500, body=b'{"error": {"message": "primary deployment is down"}}')
|
||||
return Reply(body=_completion(f"chatcmpl-fallback-{marker}", "Recovered."))
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
with _vllm_server(respond) as wire, gateway.scenario() as scenario:
|
||||
primary: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
fallback: Final = scenario.model(
|
||||
model=f"hosted_vllm/{_FALLBACK_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY
|
||||
|
|
@ -399,7 +419,7 @@ def test_fallback_attempt_replays_reasoning_to_the_second_deployment(gateway: Ga
|
|||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["id"] == f"chatcmpl-fallback-{marker}"
|
||||
attempts: Final = wire.drain()
|
||||
attempts: Final = _provider_calls(wire)
|
||||
assert [_JSON_OBJECT.validate_json(attempt.body)["model"] for attempt in attempts] == [
|
||||
_BACKEND,
|
||||
_FALLBACK_BACKEND,
|
||||
|
|
@ -413,13 +433,13 @@ def test_fallback_attempt_replays_reasoning_to_the_second_deployment(gateway: Ga
|
|||
def test_identical_uncached_replays_are_each_forwarded_and_billed_once(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
identities: Final = iter((f"chatcmpl-first-{marker}", f"chatcmpl-second-{marker}"))
|
||||
with wire_server(lambda _: Reply(body=_completion(next(identities), "Done."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(next(identities), "Done."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
first: Final = _post_chat(gateway, model, _replayed_conversation(_REASONING, marker))
|
||||
second: Final = _post_chat(gateway, model, _replayed_conversation(_REASONING, marker))
|
||||
assert (first["id"], second["id"]) == (f"chatcmpl-first-{marker}", f"chatcmpl-second-{marker}")
|
||||
assert [_sent_messages(request) for request in wire.drain()] == [
|
||||
assert [_sent_messages(request) for request in _provider_calls(wire)] == [
|
||||
_replayed_conversation(_REASONING, marker),
|
||||
_replayed_conversation(_REASONING, marker),
|
||||
]
|
||||
|
|
@ -430,7 +450,7 @@ def test_identical_uncached_replays_are_each_forwarded_and_billed_once(gateway:
|
|||
def test_cached_replay_hits_only_for_the_same_reasoning(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
identities: Final = iter((f"chatcmpl-cached-{marker}", f"chatcmpl-other-{marker}"))
|
||||
with wire_server(lambda _: Reply(body=_completion(next(identities), "Done."))) as wire:
|
||||
with _vllm_server(lambda _: Reply(body=_completion(next(identities), "Done."))) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
|
||||
|
|
@ -446,7 +466,7 @@ def test_cached_replay_hits_only_for_the_same_reasoning(gateway: Gateway) -> Non
|
|||
assert ask(_REASONING)["id"] == f"chatcmpl-cached-{marker}"
|
||||
assert ask(_REASONING)["id"] == f"chatcmpl-cached-{marker}"
|
||||
assert ask(f"a different thought {marker}")["id"] == f"chatcmpl-other-{marker}"
|
||||
assert [_sent_messages(request)[1].get("reasoning_content") for request in wire.drain()] == [
|
||||
assert [_sent_messages(request)[1].get("reasoning_content") for request in _provider_calls(wire)] == [
|
||||
_REASONING,
|
||||
f"a different thought {marker}",
|
||||
]
|
||||
|
|
@ -514,14 +534,14 @@ def _responses_reply(identity: str, stream: bool) -> Reply:
|
|||
|
||||
|
||||
def _only_responses_body(wire: Wire) -> dict[str, JsonValue]:
|
||||
received: Final = wire.drain()
|
||||
received: Final = _provider_calls(wire)
|
||||
assert [(request.method, request.target) for request in received] == [("POST", "/v1/responses")]
|
||||
return _JSON_OBJECT.validate_json(received[0].body)
|
||||
|
||||
|
||||
def test_openai_sdk_responses_replay_reaches_hosted_vllm_with_its_reasoning_item(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(lambda _: _responses_reply(f"resp_upstream_{marker}", stream=False)) as wire:
|
||||
with _vllm_server(lambda _: _responses_reply(f"resp_upstream_{marker}", stream=False)) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
response: Final = _openai_client(gateway).responses.create(
|
||||
|
|
@ -542,7 +562,7 @@ async def test_async_openai_sdk_responses_stream_reaches_hosted_vllm_with_its_re
|
|||
gateway: Gateway,
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(lambda _: _responses_reply(f"resp_upstream_{marker}", stream=True)) as wire:
|
||||
with _vllm_server(lambda _: _responses_reply(f"resp_upstream_{marker}", stream=True)) as wire:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"hosted_vllm/{_BACKEND}", api_base=wire.url + "/v1", api_key=_API_KEY)
|
||||
stream: Final = await _async_openai_client(gateway).responses.create(
|
||||
|
|
|
|||
|
|
@ -12,10 +12,10 @@ from tests.integration._support.process import owned_proxy
|
|||
from tests.integration.pricing.test_off_peak_pricing import off_peak_window
|
||||
|
||||
|
||||
def delete_lens(engine_id: str) -> None:
|
||||
write_rows('DELETE FROM "LiteLLM_EngineRun" WHERE engine_id=%s', (engine_id,))
|
||||
write_rows('DELETE FROM "LiteLLM_Engine" WHERE id=%s', (engine_id,))
|
||||
assert read_rows('SELECT id FROM "LiteLLM_Engine" WHERE id=%s', (engine_id,)) == []
|
||||
def delete_lens(lens_id: str) -> None:
|
||||
write_rows('DELETE FROM "LiteLLM_LensRun" WHERE lens_id=%s', (lens_id,))
|
||||
write_rows('DELETE FROM "LiteLLM_Lens" WHERE id=%s', (lens_id,))
|
||||
assert read_rows('SELECT id FROM "LiteLLM_Lens" WHERE id=%s', (lens_id,)) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("off_peak", (False, True))
|
||||
|
|
@ -37,12 +37,12 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
|
|||
key: Final = scenario.key(models=[model], max_budget=1)
|
||||
key_id: Final = sha256(key.encode()).hexdigest()
|
||||
worker: Final = gateway.post(
|
||||
"/engine/workers/register", {"name": "Billing regression", "analysis_key_id": key_id}
|
||||
"/lens/workers/register", {"name": "Billing regression", "analysis_key_id": key_id}
|
||||
)
|
||||
worker_id: Final = string_value(object_value(worker["worker"])["id"])
|
||||
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_EngineWorker" WHERE id=%s', (worker_id,))
|
||||
engine: Final = gateway.post(
|
||||
"/engine",
|
||||
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_LensWorker" WHERE id=%s', (worker_id,))
|
||||
lens: Final = gateway.post(
|
||||
"/lens",
|
||||
{
|
||||
"name": "Billing regression",
|
||||
"model": model,
|
||||
|
|
@ -51,17 +51,17 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
|
|||
"source": "requests",
|
||||
},
|
||||
)
|
||||
engine_id: Final = string_value(engine["id"])
|
||||
scenario.cleanups.callback(delete_lens, engine_id)
|
||||
lens_id: Final = string_value(lens["id"])
|
||||
scenario.cleanups.callback(delete_lens, lens_id)
|
||||
worker_key: Final = string_value(worker["token"])
|
||||
unauthorized: Final = gateway.request(
|
||||
"POST", "/engine/workers/register", {"name": "Denied", "analysis_key_id": key_id}, key=key
|
||||
"POST", "/lens/workers/register", {"name": "Denied", "analysis_key_id": key_id}, key=key
|
||||
)
|
||||
assert unauthorized.status_code == 403, unauthorized.text
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
claims: Final = tuple(
|
||||
pool.map(
|
||||
lambda _: gateway.request("POST", "/engine/worker/claim?protocol_version=2", {}, key=worker_key),
|
||||
lambda _: gateway.request("POST", "/lens/worker/claim?protocol_version=2", {}, key=worker_key),
|
||||
range(8),
|
||||
)
|
||||
)
|
||||
|
|
@ -69,9 +69,9 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
|
|||
winners: Final = tuple(response.json() for response in claims if response.json() is not None)
|
||||
assert len(winners) == 1
|
||||
claim: Final = object_value(winners[0])
|
||||
assert claim["engine_id"] == engine_id
|
||||
assert claim["lens_id"] == lens_id
|
||||
job_id: Final = string_value(object_value(claim["job"])["id"])
|
||||
path: Final = f"/engine/worker/{engine_id}/{job_id}/model"
|
||||
path: Final = f"/lens/worker/{lens_id}/{job_id}/model"
|
||||
result: Final = gateway.post(path, {"prompt": "Inspect this run", "purpose": "extract"}, key=worker_key)
|
||||
expected: Final = (20 * 0.000001 + 20 * 0.000002) * (0.5 if off_peak else 1)
|
||||
assert result["cost"] == pytest.approx(expected)
|
||||
|
|
@ -81,7 +81,7 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
|
|||
seconds=70,
|
||||
)
|
||||
assert rows[0]["spend"] == pytest.approx(expected)
|
||||
assert gateway.get(f"/engine/{engine_id}")["spent"] == pytest.approx(expected)
|
||||
assert gateway.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected)
|
||||
raw_hash: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
|
|
@ -105,11 +105,11 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
|
|||
gateway.post("/key/block", {"key": key})
|
||||
blocked: Final = gateway.request("POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key)
|
||||
assert blocked.status_code == 400, blocked.text
|
||||
assert gateway.get(f"/engine/{engine_id}")["spent"] == pytest.approx(expected)
|
||||
assert gateway.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected)
|
||||
replacement: Final = scenario.key(models=[model], rpm_limit=1)
|
||||
replacement_id: Final = sha256(replacement.encode()).hexdigest()
|
||||
changed: Final = gateway.request(
|
||||
"PUT", f"/engine/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
|
||||
"PUT", f"/lens/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
|
||||
)
|
||||
assert changed.status_code == 200, changed.text
|
||||
billed_replacement: Final = gateway.post(
|
||||
|
|
@ -124,17 +124,17 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
|
|||
seconds=70,
|
||||
)
|
||||
assert second_rows[0]["spend"] == pytest.approx(expected)
|
||||
revoked: Final = gateway.request("DELETE", f"/engine/workers/{worker_id}")
|
||||
revoked: Final = gateway.request("DELETE", f"/lens/workers/{worker_id}")
|
||||
assert revoked.status_code == 200, revoked.text
|
||||
denied_worker: Final = gateway.request(
|
||||
"POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key
|
||||
)
|
||||
assert denied_worker.status_code == 401, denied_worker.text
|
||||
forbidden_change: Final = gateway.request(
|
||||
"PUT", f"/engine/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
|
||||
"PUT", f"/lens/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
|
||||
)
|
||||
assert forbidden_change.status_code == 409, forbidden_change.text
|
||||
gateway.post(f"/engine/{engine_id}/cancel", {})
|
||||
gateway.post(f"/lens/{lens_id}/cancel", {})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cancel_on_disconnect", (False, True))
|
||||
|
|
@ -174,19 +174,19 @@ def test_worker_spend_logs_do_not_expose_investigation_content(
|
|||
seconds=70,
|
||||
)
|
||||
assert marker in str(retained[0]), "Control must prove this proxy retains ordinary prompts"
|
||||
worker: Final = isolated.post("/engine/workers/register", {"analysis_key_id": key_id})
|
||||
worker: Final = isolated.post("/lens/workers/register", {"analysis_key_id": key_id})
|
||||
worker_id: Final = string_value(object_value(worker["worker"])["id"])
|
||||
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_EngineWorker" WHERE id=%s', (worker_id,))
|
||||
engine: Final = isolated.post(
|
||||
"/engine", {"name": "Log privacy", "model": model, "enabled": False, "context": "Find problems"}
|
||||
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_LensWorker" WHERE id=%s', (worker_id,))
|
||||
lens: Final = isolated.post(
|
||||
"/lens", {"name": "Log privacy", "model": model, "enabled": False, "context": "Find problems"}
|
||||
)
|
||||
engine_id: Final = string_value(engine["id"])
|
||||
scenario.cleanups.callback(delete_lens, engine_id)
|
||||
lens_id: Final = string_value(lens["id"])
|
||||
scenario.cleanups.callback(delete_lens, lens_id)
|
||||
worker_token: Final = string_value(worker["token"])
|
||||
claim: Final = isolated.post("/engine/worker/claim?protocol_version=2", {}, key=worker_token)
|
||||
claim: Final = isolated.post("/lens/worker/claim?protocol_version=2", {}, key=worker_token)
|
||||
job_id: Final = string_value(object_value(claim["job"])["id"])
|
||||
result: Final = isolated.post(
|
||||
f"/engine/worker/{engine_id}/{job_id}/model", {"prompt": marker, "purpose": "extract"}, key=worker_token
|
||||
f"/lens/worker/{lens_id}/{job_id}/model", {"prompt": marker, "purpose": "extract"}, key=worker_token
|
||||
)
|
||||
assert result["content"], "The worker must still receive model output"
|
||||
rows: Final = eventually(
|
||||
|
|
@ -200,4 +200,4 @@ def test_worker_spend_logs_do_not_expose_investigation_content(
|
|||
assert float(rows[0]["spend"]) == pytest.approx(result["cost"])
|
||||
assert marker not in str(rows[0])
|
||||
assert result["content"] not in str(rows[0]["response"])
|
||||
isolated.post(f"/engine/{engine_id}/cancel", {})
|
||||
isolated.post(f"/lens/{lens_id}/cancel", {})
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@
|
|||
"user": "",
|
||||
"team_id": "",
|
||||
"organization_id": "",
|
||||
"metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
|
||||
"metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
|
||||
"cache_key": "Cache OFF",
|
||||
"spend": 0.00022500000000000002,
|
||||
"total_tokens": 30,
|
||||
|
|
|
|||
|
|
@ -13,13 +13,13 @@ from typing import Final
|
|||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.proxy.engine.analysis import analyze_sample
|
||||
from litellm.proxy.engine.inference import _SYSTEM
|
||||
from litellm.proxy.engine.models import (
|
||||
from litellm.proxy.lens.analysis import analyze_sample
|
||||
from litellm.proxy.lens.inference import _SYSTEM
|
||||
from litellm.proxy.lens.models import (
|
||||
Check,
|
||||
Claim,
|
||||
Coverage,
|
||||
EngineSettings,
|
||||
LensSettings,
|
||||
Execution,
|
||||
ExecutionContent,
|
||||
Finding,
|
||||
|
|
@ -90,7 +90,7 @@ async def evaluate(
|
|||
feedback: tuple[Finding, ...] = (),
|
||||
) -> dict[str, object]:
|
||||
records: Final = MappingProxyType({case.name: fixtures(case) for case in cases})
|
||||
settings: Final = EngineSettings(
|
||||
settings: Final = LensSettings(
|
||||
name="Quality evaluation",
|
||||
model=model_name,
|
||||
checks=checks,
|
||||
|
|
@ -100,7 +100,7 @@ async def evaluate(
|
|||
)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
claim: Final = Claim(
|
||||
engine_id="evaluation",
|
||||
lens_id="evaluation",
|
||||
findings=feedback,
|
||||
job=Job(id="evaluation", created_at=now, start=now, end=now, settings=settings, revision=1),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -15,8 +15,8 @@ from litellm import Router
|
|||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.engine import endpoints
|
||||
from litellm.proxy.engine.models import Check, Coverage, EngineSettings, ModelRequest, Progress, Result, RunRequest
|
||||
from litellm.proxy.lens import endpoints
|
||||
from litellm.proxy.lens.models import Check, Coverage, LensSettings, ModelRequest, Progress, Result, RunRequest
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
||||
|
||||
|
|
@ -61,13 +61,13 @@ async def lens_database() -> AsyncIterator[PrismaClient]:
|
|||
@pytest.mark.asyncio
|
||||
async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database: PrismaClient) -> None:
|
||||
admin: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
settings: Final = EngineSettings(
|
||||
settings: Final = LensSettings(
|
||||
name="Lifecycle regression",
|
||||
model="lens-test-analysis",
|
||||
enabled=False,
|
||||
checks=(Check(id="retries", instruction="Find unrecovered retries"),),
|
||||
)
|
||||
engine: Final = await endpoints.create_engine(settings, admin)
|
||||
lens: Final = await endpoints.create_lens(settings, admin)
|
||||
key_id: Final = hashlib.sha256(uuid4().bytes).hexdigest()
|
||||
await lens_database.db.litellm_verificationtoken.create(data={"token": key_id, "models": ["lens-test-analysis"]})
|
||||
registration: Final = await endpoints.register_worker(
|
||||
|
|
@ -76,17 +76,17 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
credentials: Final = HTTPAuthorizationCredentials(scheme="Bearer", credentials=registration.token)
|
||||
worker: Final = await endpoints.worker_auth(credentials)
|
||||
try:
|
||||
assert engine.jobs[0].status == "queued"
|
||||
assert lens.jobs[0].status == "queued"
|
||||
stored_worker: Final = await endpoints.repository().worker(
|
||||
hashlib.sha256(registration.token.encode()).hexdigest()
|
||||
)
|
||||
assert stored_worker is not None and stored_worker.id == worker.id
|
||||
assert worker.id == registration.worker.id
|
||||
listing: Final = await endpoints.list_engines(admin)
|
||||
assert engine.id in tuple(e.id for e in listing.engines)
|
||||
listing: Final = await endpoints.list_lenses(admin)
|
||||
assert lens.id in tuple(e.id for e in listing.lenses)
|
||||
assert worker.id in tuple(w.id for w in listing.workers)
|
||||
claims: Final = await asyncio.gather(
|
||||
*(endpoints.claim_candidate(engine, worker, datetime.now(timezone.utc)) for _ in range(8))
|
||||
*(endpoints.claim_candidate(lens, worker, datetime.now(timezone.utc)) for _ in range(8))
|
||||
)
|
||||
winners: Final = tuple(claim for claim in claims if claim is not None)
|
||||
assert len(winners) == 1
|
||||
|
|
@ -94,16 +94,16 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
assert claimed.job.worker_id == worker.id
|
||||
assert (
|
||||
await endpoints.claim_candidate(
|
||||
await endpoints.get_engine(engine.id, worker.scope), worker, datetime.now(timezone.utc)
|
||||
await endpoints.get_lens(lens.id, worker.scope), worker, datetime.now(timezone.utc)
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert await endpoints.progress(
|
||||
engine.id, claimed.job.id, Progress(stage="Reviewing", coverage=Coverage(screened=2)), worker
|
||||
lens.id, claimed.job.id, Progress(stage="Reviewing", coverage=Coverage(screened=2)), worker
|
||||
)
|
||||
assert await endpoints.heartbeat(engine.id, claimed.job.id, worker)
|
||||
assert await endpoints.heartbeat(lens.id, claimed.job.id, worker)
|
||||
response: Final = await endpoints.model(
|
||||
engine.id,
|
||||
lens.id,
|
||||
claimed.job.id,
|
||||
ModelRequest(prompt="Return an empty observations list", purpose="extract"),
|
||||
worker,
|
||||
|
|
@ -111,7 +111,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
{
|
||||
"type": "http",
|
||||
"scheme": "http",
|
||||
"path": "/engine/worker/model",
|
||||
"path": "/lens/worker/model",
|
||||
"headers": [],
|
||||
"client": ("127.0.0.1", 1234),
|
||||
}
|
||||
|
|
@ -120,7 +120,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
assert '"observations"' in response.content
|
||||
with pytest.raises(HTTPException) as denied_ip:
|
||||
await endpoints.model(
|
||||
engine.id,
|
||||
lens.id,
|
||||
claimed.job.id,
|
||||
ModelRequest(prompt="Must not run", purpose="extract"),
|
||||
worker,
|
||||
|
|
@ -128,7 +128,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
{
|
||||
"type": "http",
|
||||
"scheme": "http",
|
||||
"path": "/engine/worker/model",
|
||||
"path": "/lens/worker/model",
|
||||
"headers": [(b"x-forwarded-for", b"127.0.0.1")],
|
||||
"client": ("192.0.2.1", 1234),
|
||||
}
|
||||
|
|
@ -136,7 +136,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
)
|
||||
assert denied_ip.value.status_code == 403
|
||||
forwarded: Final = await endpoints.model(
|
||||
engine.id,
|
||||
lens.id,
|
||||
claimed.job.id,
|
||||
ModelRequest(prompt="Return an empty observations list", purpose="extract"),
|
||||
worker,
|
||||
|
|
@ -144,7 +144,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
{
|
||||
"type": "http",
|
||||
"scheme": "http",
|
||||
"path": "/engine/worker/model",
|
||||
"path": "/lens/worker/model",
|
||||
"headers": [(b"x-forwarded-for", b"127.0.0.1")],
|
||||
"client": ("192.0.2.100", 1234),
|
||||
}
|
||||
|
|
@ -153,7 +153,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
assert '"observations"' in forwarded.content
|
||||
with pytest.raises(HTTPException) as spoofed_chain:
|
||||
await endpoints.model(
|
||||
engine.id,
|
||||
lens.id,
|
||||
claimed.job.id,
|
||||
ModelRequest(prompt="Must not run", purpose="extract"),
|
||||
worker,
|
||||
|
|
@ -161,14 +161,14 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
{
|
||||
"type": "http",
|
||||
"scheme": "http",
|
||||
"path": "/engine/worker/model",
|
||||
"path": "/lens/worker/model",
|
||||
"headers": [(b"x-forwarded-for", b"127.0.0.1, 192.0.2.1")],
|
||||
"client": ("192.0.2.100", 1234),
|
||||
}
|
||||
),
|
||||
)
|
||||
assert spoofed_chain.value.status_code == 403
|
||||
charged: Final = await endpoints.get_engine(engine.id, worker.scope)
|
||||
charged: Final = await endpoints.get_lens(lens.id, worker.scope)
|
||||
assert charged.spent == pytest.approx(response.cost + forwarded.cost)
|
||||
assert charged.jobs[0].cost == pytest.approx(response.cost + forwarded.cost)
|
||||
legacy: Final = worker.model_copy(update={"analysis_key_id": None})
|
||||
|
|
@ -178,37 +178,37 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
with pytest.raises(HTTPException) as needs_billing:
|
||||
await endpoints.claim(authenticated_legacy, protocol_version=2)
|
||||
assert needs_billing.value.status_code == 409
|
||||
assert await endpoints.heartbeat(engine.id, claimed.job.id, authenticated_legacy)
|
||||
assert await endpoints.heartbeat(lens.id, claimed.job.id, authenticated_legacy)
|
||||
finished: Final = await endpoints.result(
|
||||
engine.id, claimed.job.id, Result(coverage=Coverage(screened=2)), authenticated_legacy
|
||||
lens.id, claimed.job.id, Result(coverage=Coverage(screened=2)), authenticated_legacy
|
||||
)
|
||||
assert finished.jobs[0].status == "completed"
|
||||
assert finished.jobs[0].coverage.screened == 2
|
||||
assert finished.last_scan_at == claimed.job.end
|
||||
assert finished.next_run_at > finished.jobs[0].finished_at
|
||||
assert await endpoints.result(engine.id, claimed.job.id, Result(coverage=Coverage()), worker) == finished
|
||||
assert await endpoints.result(lens.id, claimed.job.id, Result(coverage=Coverage()), worker) == finished
|
||||
with pytest.raises(HTTPException) as stale:
|
||||
await endpoints.heartbeat(engine.id, claimed.job.id, worker)
|
||||
await endpoints.heartbeat(lens.id, claimed.job.id, worker)
|
||||
assert stale.value.status_code == 409
|
||||
edited: Final = await endpoints.update_engine(
|
||||
engine.id, settings.model_copy(update={"interval_minutes": 7}), admin
|
||||
edited: Final = await endpoints.update_lens(
|
||||
lens.id, settings.model_copy(update={"interval_minutes": 7}), admin
|
||||
)
|
||||
assert edited.revision == engine.revision + 1
|
||||
rerun: Final = await endpoints.run_engine(engine.id, RunRequest(lookback_hours=3), admin)
|
||||
assert edited.revision == lens.revision + 1
|
||||
rerun: Final = await endpoints.run_lens(lens.id, RunRequest(lookback_hours=3), admin)
|
||||
assert rerun.jobs[0].settings.interval_minutes == 7
|
||||
assert rerun.jobs[0].created_at - rerun.jobs[0].start == timedelta(hours=3)
|
||||
history: Final = await endpoints.list_runs(engine.id, admin, offset=0)
|
||||
history: Final = await endpoints.list_runs(lens.id, admin, offset=0)
|
||||
assert {job.id for job in history} == {claimed.job.id, rerun.jobs[0].id}
|
||||
archived: Final = await endpoints.read_run(engine.id, claimed.job.id, admin)
|
||||
archived: Final = await endpoints.read_run(lens.id, claimed.job.id, admin)
|
||||
assert archived == finished.jobs[0]
|
||||
assert archived.settings.interval_minutes == 15
|
||||
assert archived.findings == ()
|
||||
with pytest.raises(HTTPException) as foreign_history:
|
||||
await endpoints.read_run(engine.id, claimed.job.id, UserAPIKeyAuth(team_id="other"))
|
||||
await endpoints.read_run(lens.id, claimed.job.id, UserAPIKeyAuth(team_id="other"))
|
||||
assert foreign_history.value.status_code == 403
|
||||
cancelled: Final = await endpoints.cancel_engine(engine.id, admin)
|
||||
cancelled: Final = await endpoints.cancel_lens(lens.id, admin)
|
||||
assert cancelled.jobs[0].status == "cancelled"
|
||||
assert await endpoints.cancel_engine(engine.id, admin) == cancelled
|
||||
assert await endpoints.cancel_lens(lens.id, admin) == cancelled
|
||||
assert await endpoints.revoke_worker(worker.id, admin)
|
||||
assert await endpoints.repository().set_worker_billing(worker.id, key_id) is None
|
||||
with pytest.raises(HTTPException) as revoked_billing:
|
||||
|
|
@ -218,10 +218,10 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
|
|||
await endpoints.worker_auth(credentials)
|
||||
assert revoked.value.status_code == 401
|
||||
with pytest.raises(HTTPException) as foreign:
|
||||
await endpoints.get_engine(engine.id, endpoints.Scope(team_id="other"))
|
||||
await endpoints.get_lens(lens.id, endpoints.Scope(team_id="other"))
|
||||
assert foreign.value.status_code == 404
|
||||
finally:
|
||||
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_EngineRun" WHERE engine_id=$1', engine.id)
|
||||
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_Engine" WHERE id=$1', engine.id)
|
||||
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_EngineWorker" WHERE id=$1', worker.id)
|
||||
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_LensRun" WHERE lens_id=$1', lens.id)
|
||||
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', lens.id)
|
||||
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_LensWorker" WHERE id=$1', worker.id)
|
||||
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_VerificationToken" WHERE token=$1', key_id)
|
||||
|
|
|
|||
|
|
@ -6,9 +6,9 @@ from queue import SimpleQueue
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from engine.models import (
|
||||
from lens.models import (
|
||||
Claim,
|
||||
EngineSettings,
|
||||
LensSettings,
|
||||
Execution,
|
||||
ExecutionContent,
|
||||
Job,
|
||||
|
|
@ -17,7 +17,7 @@ from engine.models import (
|
|||
Sample,
|
||||
TracePart,
|
||||
)
|
||||
from engine.worker import EngineWorker
|
||||
from lens.worker import LensWorker
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
|
|
@ -25,7 +25,7 @@ async def main() -> None:
|
|||
claims: Final = iter(("full", "healthy"))
|
||||
saved: Final = SimpleQueue[Result]()
|
||||
pages: Final = SimpleQueue[str]()
|
||||
settings: Final = EngineSettings(name="Storage recovery", model="unused", context="Finish the task", concurrency=1)
|
||||
settings: Final = LensSettings(name="Storage recovery", model="unused", context="Finish the task", concurrency=1)
|
||||
execution: Final = Execution(
|
||||
id="run", source="traces", trace_id="trace", team_id="", name="Task", start_time="", span_count=10000
|
||||
)
|
||||
|
|
@ -34,7 +34,7 @@ async def main() -> None:
|
|||
path: Final = request.url.path
|
||||
if path.endswith("/claim"):
|
||||
claim: Final = Claim(
|
||||
engine_id="lens",
|
||||
lens_id="lens",
|
||||
job=Job(id=next(claims), created_at=now, start=now, end=now, settings=settings, revision=1),
|
||||
findings=(),
|
||||
)
|
||||
|
|
@ -71,7 +71,7 @@ async def main() -> None:
|
|||
return httpx.Response(200, json=True)
|
||||
|
||||
async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client:
|
||||
worker: Final = EngineWorker(client)
|
||||
worker: Final = LensWorker(client)
|
||||
assert await worker.run_once()
|
||||
failed: Final = saved.get_nowait()
|
||||
assert failed.error.startswith("Worker temporary storage failed.")
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
Behavior tests for the LiteLLM_AutoRouterSession conditional upsert and the benchmarks
|
||||
aggregate, against a real Postgres. The classification lives in SQL, so these tests are
|
||||
the ones that exercise it; the builder and flush contracts are unit-tested in
|
||||
tests/test_litellm/proxy/db/test_autorouter_session_rollup.py.
|
||||
tests/unit/proxy/db/test_autorouter_session_rollup.py.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
Behavior tests for the cache analytics queries against a real Postgres. The info-route
|
||||
exclusion and the Unknown grouping live in SQL, so these tests are the ones that exercise
|
||||
them; the endpoint wiring is unit-tested in
|
||||
tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py.
|
||||
tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
|
|
|||