mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
Merge remote-tracking branch 'origin/main' into HEAD
This commit is contained in:
commit
17822fd186
504 changed files with 31504 additions and 5645 deletions
|
|
@ -1,14 +1,14 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if [ "${GITHUB_ACTIONS:-}" = true ]; then
|
||||
echo "Integration contracts are owned by CircleCI" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
suite="${1:?integration suite required}"
|
||||
mode="${2:-standard}"
|
||||
side="${3:-}"
|
||||
requested_workers="${INTEGRATION_WORKERS:-}"
|
||||
license_env=()
|
||||
if [ -n "${LITELLM_LICENSE:-}" ]; then
|
||||
license_env=("LITELLM_LICENSE=$LITELLM_LICENSE")
|
||||
fi
|
||||
if [ "$mode" = replica ]; then
|
||||
results="test-results/integration-${suite}-replica"
|
||||
elif [ "$mode" = parity ]; then
|
||||
|
|
@ -84,7 +84,7 @@ if [ "$suite" = browser ]; then
|
|||
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
|
||||
test -f "$LITELLM_UI_PATH/index.html"
|
||||
fi
|
||||
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
|
||||
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256(((os.environ.get("CIRCLE_SHA1") or os.environ.get("GITHUB_SHA") or "local") + (os.environ.get("CIRCLE_WORKFLOW_ID") or os.environ.get("GITHUB_RUN_ID") or "local")).encode()).hexdigest()[:8],16))')"
|
||||
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
|
||||
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
|
||||
|
|
@ -107,7 +107,7 @@ sudo iptables -N integration_only
|
|||
guard_created=true
|
||||
sudo iptables -A integration_only -o lo -j ACCEPT
|
||||
sudo iptables -A integration_only -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
|
||||
for service in postgres-db redis-cache; do
|
||||
for service in "${INTEGRATION_POSTGRES_CONTAINER:-postgres-db}" "${INTEGRATION_REDIS_CONTAINER:-redis-cache}"; do
|
||||
address="$(docker inspect --format '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$service")"
|
||||
sudo iptables -A integration_only -d "$address" -j ACCEPT
|
||||
done
|
||||
|
|
@ -145,6 +145,9 @@ fi
|
|||
if [ "$suite" = mcp ]; then
|
||||
export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1
|
||||
fi
|
||||
if [ -n "$requested_workers" ]; then
|
||||
export INTEGRATION_WORKERS="$requested_workers"
|
||||
fi
|
||||
isolated_workers=false
|
||||
if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then
|
||||
isolated_workers=true
|
||||
|
|
@ -185,7 +188,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:-}" \
|
||||
"${license_env[@]}" \
|
||||
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True LITELLM_ENABLE_MCP_STDIO=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 \
|
||||
|
|
@ -257,7 +260,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
|
|||
"${shared_services[@]}" "${worker_coverage[@]}" \
|
||||
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
|
||||
LITELLM_MODE=PRODUCTION \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
"${license_env[@]}" \
|
||||
INTEGRATION_SEED="$INTEGRATION_SEED" \
|
||||
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
|
||||
INTEGRATION_CASE_SHARDS="${CIRCLE_NODE_TOTAL:-1}" INTEGRATION_CASE_SHARD_INDEX="${CIRCLE_NODE_INDEX:-0}" \
|
||||
|
|
|
|||
119
.github/scripts/assert_secret_absent.py
vendored
Normal file
119
.github/scripts/assert_secret_absent.py
vendored
Normal file
|
|
@ -0,0 +1,119 @@
|
|||
import base64
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import urllib.parse
|
||||
import zipfile
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TextIO
|
||||
|
||||
SECRET_VARIABLE: Final = "SCANNED_SECRET"
|
||||
MIN_ENCODED_LENGTH: Final = 12
|
||||
MAX_ARCHIVE_DEPTH: Final = 3
|
||||
MASK: Final = "***"
|
||||
|
||||
Forms = tuple[tuple[str, str], ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Hit:
|
||||
location: str
|
||||
form: str
|
||||
|
||||
|
||||
def _aligned_encodings(raw: bytes, name: str, encode: Callable[[bytes], bytes]) -> Iterator[tuple[str, str]]:
|
||||
for offset in range(3):
|
||||
chunk: Final = raw[offset:]
|
||||
yield f"{name} offset {offset}", encode(chunk[: len(chunk) - len(chunk) % 3]).decode()
|
||||
|
||||
|
||||
def secret_forms(secret: str) -> Forms:
|
||||
encoded: Final = (
|
||||
("url-encoded", urllib.parse.quote(secret, safe="")),
|
||||
*_aligned_encodings(secret.encode(), "base64", base64.b64encode),
|
||||
*_aligned_encodings(secret.encode(), "base64url", base64.urlsafe_b64encode),
|
||||
)
|
||||
by_value: Final = MappingProxyType(
|
||||
{value: name for name, value in encoded if value != secret and len(value) >= MIN_ENCODED_LENGTH}
|
||||
)
|
||||
return (("raw", secret), *((name, value) for value, name in by_value.items()))
|
||||
|
||||
|
||||
def mask(text: str, forms: Forms) -> str:
|
||||
return reduce(lambda masked, form: masked.replace(form[1], MASK), forms, text)
|
||||
|
||||
|
||||
def _name_hits(location: str, name: str, forms: Forms) -> Iterator[Hit]:
|
||||
for form, value in forms:
|
||||
if value in name:
|
||||
yield Hit(mask(location, forms), f"{form} in name")
|
||||
|
||||
|
||||
def _content_hits(location: str, data: bytes, forms: Forms, depth: int) -> Iterator[Hit]:
|
||||
for form, value in forms:
|
||||
if value.encode() in data:
|
||||
yield Hit(mask(location, forms), form)
|
||||
if depth >= MAX_ARCHIVE_DEPTH or not zipfile.is_zipfile(io.BytesIO(data)):
|
||||
return
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as archive:
|
||||
for member in archive.infolist():
|
||||
if member.is_dir():
|
||||
continue
|
||||
yield from _content_hits(f"{location}!{member.filename}", archive.read(member), forms, depth + 1)
|
||||
|
||||
|
||||
def _walk(root: Path) -> Iterator[tuple[Path, bool]]:
|
||||
for directory, subdirectories, files in os.walk(root):
|
||||
yield from ((Path(directory, name), True) for name in subdirectories if Path(directory, name).is_symlink())
|
||||
yield from ((Path(directory, name), False) for name in files)
|
||||
|
||||
|
||||
def _path_hits(root: Path, path: Path, linked_directory: bool, forms: Forms) -> Iterator[Hit]:
|
||||
location: Final = str(path.relative_to(root))
|
||||
if linked_directory:
|
||||
yield Hit(mask(location, forms), "directory symlink, not scanned")
|
||||
return
|
||||
yield from _name_hits(location, location, forms)
|
||||
yield from _content_hits(location, path.read_bytes(), forms, 0)
|
||||
|
||||
|
||||
def scan(root: Path, secret: str) -> tuple[int, tuple[Hit, ...]]:
|
||||
forms: Final = secret_forms(secret)
|
||||
entries: Final = tuple(_walk(root))
|
||||
hits: Final = tuple(chain.from_iterable(_path_hits(root, path, linked, forms) for path, linked in entries))
|
||||
return len(entries), hits
|
||||
|
||||
|
||||
def main(argv: tuple[str, ...], environ: Mapping[str, str], out: TextIO) -> int:
|
||||
if len(argv) != 1:
|
||||
print(f"usage: assert_secret_absent.py <directory>, with the secret in ${SECRET_VARIABLE}", file=out)
|
||||
return 2
|
||||
secret: Final = environ.get(SECRET_VARIABLE, "")
|
||||
root: Final = Path(argv[0])
|
||||
if not secret:
|
||||
print(f"{SECRET_VARIABLE} is empty, nothing to scan for", file=out)
|
||||
return 0
|
||||
if not root.is_dir():
|
||||
print(f"{root} is not a directory, nothing to scan", file=out)
|
||||
return 0
|
||||
try:
|
||||
scanned, hits = scan(root, secret)
|
||||
except Exception as error:
|
||||
print(f"scan failed with {type(error).__name__}, treating the results as unsafe", file=out)
|
||||
return 1
|
||||
for hit in hits:
|
||||
print(f"secret found: {hit.location} ({hit.form})", file=out)
|
||||
if hits:
|
||||
print(f"{len(hits)} copies of the secret under {root}, refusing to upload", file=out)
|
||||
return 1
|
||||
print(f"no copy of the secret in {scanned} files under {root}", file=out)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main(tuple(sys.argv[1:]), os.environ, sys.stdout))
|
||||
9
.github/scripts/smoke_test_native_wheel.py
vendored
9
.github/scripts/smoke_test_native_wheel.py
vendored
|
|
@ -7,6 +7,8 @@ import zipfile
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
COST_MAP: Final = Path(__file__).resolve().parents[2] / "model_prices_and_context_window.json"
|
||||
|
||||
CHILD_SCRIPT: Final = """
|
||||
from importlib.util import module_from_spec, spec_from_file_location
|
||||
from pathlib import Path
|
||||
|
|
@ -34,6 +36,9 @@ else:
|
|||
after = module.gil_stats()
|
||||
if not isinstance(after.get("releases"), int):
|
||||
raise AssertionError(f"native module unusable after panic: {after!r}")
|
||||
|
||||
if module.bundled_model_cost_map() != Path(sys.argv[2]).read_bytes():
|
||||
raise AssertionError("the wheel embeds a different model_prices_and_context_window.json than this checkout")
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -55,7 +60,9 @@ def main() -> int:
|
|||
|
||||
native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name
|
||||
native_path.write_bytes(archive.read(native_members[0]))
|
||||
result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False)
|
||||
result: Final = subprocess.run(
|
||||
(sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n")
|
||||
|
|
|
|||
21
.github/workflows/cost-map-guard.yml
vendored
21
.github/workflows/cost-map-guard.yml
vendored
|
|
@ -46,3 +46,24 @@ jobs:
|
|||
HEAD_REF: ${{ github.event.pull_request.head.ref }}
|
||||
run: |
|
||||
uv run --no-project python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"
|
||||
- name: Check whether the pull request changes the cost map
|
||||
id: changed
|
||||
env:
|
||||
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
if git diff --quiet "$MERGE_BASE" "$HEAD_SHA" -- model_prices_and_context_window.json; then
|
||||
echo "cost_map=false" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "cost_map=true" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
- if: steps.changed.outputs.cost_map == 'true'
|
||||
run: rustup toolchain install --no-self-update
|
||||
- name: Load the cost map through model-catalog
|
||||
if: steps.changed.outputs.cost_map == 'true'
|
||||
env:
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
git show "${HEAD_SHA}:model_prices_and_context_window.json" > "$RUNNER_TEMP/model_prices_and_context_window.json"
|
||||
cargo run --locked --quiet --manifest-path litellm-rust/Cargo.toml -p litellm-model-catalog --bin check-registry \
|
||||
-- "$RUNNER_TEMP/model_prices_and_context_window.json"
|
||||
|
|
|
|||
2
.github/workflows/image-scan.yml
vendored
2
.github/workflows/image-scan.yml
vendored
|
|
@ -13,8 +13,10 @@ on:
|
|||
- migrations/Dockerfile
|
||||
- migrations/run.py
|
||||
- gateway/Dockerfile
|
||||
- gateway/launch.py
|
||||
- gateway/main.py
|
||||
- gateway/routes/allowlist.py
|
||||
- litellm_cli/**
|
||||
- backend/Dockerfile
|
||||
- backend/main.py
|
||||
- litellm-rust/**
|
||||
|
|
|
|||
199
.github/workflows/test-integration.yml
vendored
Normal file
199
.github/workflows/test-integration.yml
vendored
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
name: integration
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [labeled]
|
||||
schedule:
|
||||
- cron: "5 8-20/3 * * 1-6"
|
||||
timezone: America/Los_Angeles
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
suite:
|
||||
description: Integration group to run
|
||||
type: choice
|
||||
options:
|
||||
- all
|
||||
- providers
|
||||
- extensions
|
||||
- accounting
|
||||
- management
|
||||
- management-replica
|
||||
- database
|
||||
- database-replica
|
||||
- mcp
|
||||
- security
|
||||
- cost
|
||||
- sdk
|
||||
- browser
|
||||
default: all
|
||||
workers:
|
||||
description: Workers per group, blank keeps each group's default
|
||||
type: string
|
||||
default: ""
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name }}-${{ inputs.suite }}-${{ inputs.workers }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
integration:
|
||||
name: integration ${{ matrix.suite }}
|
||||
if: github.event_name != 'pull_request' || github.event.label.name == 'run-ci'
|
||||
environment: integration-license
|
||||
runs-on: ${{ fromJSON('{"providers":"ubuntu-latest-64-cores","extensions":"ubuntu-latest-64-cores","accounting":"ubuntu-latest-32-cores","management":"ubuntu-latest-32-cores","management-replica":"ubuntu-latest-32-cores"}')[matrix.suite] || 'ubuntu-latest-16-cores' }}
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
suite: ${{ fromJSON((inputs.suite || 'all') == 'all' && '["providers","extensions","accounting","management","management-replica","database","database-replica","mcp","security","cost","sdk","browser"]' || format('["{0}"]', inputs.suite)) }}
|
||||
env:
|
||||
UV_PYTHON: "3.12"
|
||||
INTEGRATION_WORKERS: ${{ inputs.workers || fromJSON('{"providers":"32","extensions":"48","accounting":"16","management":"16","management-replica":"16","cost":"16","mcp":"8","security":"8","database":"4","database-replica":"4"}')[matrix.suite] || '' }}
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
command: -c max_connections=1000 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000
|
||||
env:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: circle_test
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--shm-size 1g
|
||||
--health-cmd "pg_isready -U postgres"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 20
|
||||
redis:
|
||||
image: redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6
|
||||
ports:
|
||||
- 6379:6379
|
||||
options: >-
|
||||
--health-cmd "redis-cli ping"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 20
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: ${{ env.UV_PYTHON }}
|
||||
|
||||
- name: Set up uv
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Install dependencies
|
||||
timeout-minutes: 20
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras
|
||||
|
||||
- name: Install redis-server 7.2.16
|
||||
timeout-minutes: 3
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/redis-bin"
|
||||
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
|
||||
docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server"
|
||||
docker rm "$cid"
|
||||
"$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16'
|
||||
echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH"
|
||||
|
||||
- name: Install MCP SDK1 peer
|
||||
if: matrix.suite == 'mcp'
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
uv venv --python 3.12 .venv-mcp-peer
|
||||
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
echo "MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Set up Node.js
|
||||
if: matrix.suite == 'browser'
|
||||
timeout-minutes: 5
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
cache: "npm"
|
||||
cache-dependency-path: |
|
||||
ui/litellm-dashboard/package-lock.json
|
||||
tests/e2e/ui/package-lock.json
|
||||
|
||||
- name: Install browser dependencies and build the dashboard
|
||||
if: matrix.suite == 'browser'
|
||||
timeout-minutes: 20
|
||||
run: |
|
||||
(cd ui/litellm-dashboard && npm ci)
|
||||
(cd tests/e2e/ui && npm ci)
|
||||
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive \
|
||||
node tests/e2e/ui/node_modules/@playwright/test/cli.js install-deps chromium
|
||||
node tests/e2e/ui/node_modules/@playwright/test/cli.js install chromium
|
||||
(cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build)
|
||||
|
||||
- name: Run integration contracts
|
||||
timeout-minutes: 45
|
||||
env:
|
||||
SUITE: ${{ matrix.suite }}
|
||||
INTEGRATION_POSTGRES_CONTAINER: ${{ job.services.postgres.id }}
|
||||
INTEGRATION_REDIS_CONTAINER: ${{ job.services.redis.id }}
|
||||
LITELLM_LICENSE: ${{ secrets.LITELLM_LICENSE }}
|
||||
run: |
|
||||
case "$SUITE" in
|
||||
*-replica) bash .circleci/scripts/run_integration.sh "${SUITE%-replica}" replica ;;
|
||||
*) bash .circleci/scripts/run_integration.sh "$SUITE" ;;
|
||||
esac
|
||||
|
||||
- name: Verify every case ran exactly once
|
||||
if: ${{ !cancelled() && matrix.suite != 'browser' }}
|
||||
timeout-minutes: 2
|
||||
run: |
|
||||
record="$(find test-results -mindepth 2 -maxdepth 2 -name execution.json -print -quit)"
|
||||
test -n "$record"
|
||||
mkdir -p "$RUNNER_TEMP/shards/0"
|
||||
cp "$record" "$RUNNER_TEMP/shards/0/execution.json"
|
||||
.venv/bin/python -m tests.integration._support.sharding "$RUNNER_TEMP/shards" --shards 1
|
||||
|
||||
- name: Report MCP coverage
|
||||
if: ${{ !cancelled() && matrix.suite == 'mcp' }}
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
coverage_dir=test-results/integration-mcp/coverage
|
||||
test -f "$coverage_dir/coverage.txt" && exit 0
|
||||
export COVERAGE_FILE="$PWD/$coverage_dir/data"
|
||||
.venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml
|
||||
.venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml > "$coverage_dir/coverage.txt"
|
||||
.venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml -d "$coverage_dir/html"
|
||||
tail -n 1 "$coverage_dir/coverage.txt"
|
||||
|
||||
- name: Check results for the license
|
||||
id: secret-scan
|
||||
if: ${{ !cancelled() }}
|
||||
timeout-minutes: 5
|
||||
env:
|
||||
SCANNED_SECRET: ${{ secrets.LITELLM_LICENSE }}
|
||||
run: python .github/scripts/assert_secret_absent.py test-results
|
||||
|
||||
- name: Upload results
|
||||
if: ${{ !cancelled() && steps.secret-scan.outcome == 'success' }}
|
||||
timeout-minutes: 10
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: integration-${{ matrix.suite }}-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
path: test-results/
|
||||
retention-days: 14
|
||||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -322,6 +322,7 @@ jobs:
|
|||
tests/unit/llms/openai_like
|
||||
tests/unit/llms/opencode
|
||||
tests/unit/llms/opensandbox
|
||||
tests/unit/llms/opensearch
|
||||
tests/unit/llms/openrouter
|
||||
tests/unit/llms/ovhcloud
|
||||
tests/unit/llms/parallel_ai
|
||||
|
|
|
|||
50
README.md
50
README.md
|
|
@ -141,12 +141,56 @@ response = client.chat.completions.create(
|
|||
)
|
||||
```
|
||||
|
||||
[**Docs: LLM Providers**](https://docs.litellm.ai/docs/providers)
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Auto Routing</b> - Choose a model for each prompt (Python SDK + AI Gateway)</summary>
|
||||
|
||||
### Python SDK
|
||||
|
||||
[**Docs: Autorouter SDK**](https://docs.litellm.ai/docs/autorouter_sdk)
|
||||
|
||||
Install with `uv add litellm` and set `OPENAI_API_KEY` for the models below.
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from litellm import Router
|
||||
|
||||
router = Router(model_list=[
|
||||
{"model_name": "fast", "litellm_params": {"model": "openai/gpt-6-luna"}},
|
||||
{"model_name": "capable", "litellm_params": {"model": "openai/gpt-6.1-sol"}},
|
||||
{"model_name": "auto", "litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"classifier_type": "heuristic",
|
||||
"tiers": {"SIMPLE": "fast", "MEDIUM": "fast",
|
||||
"COMPLEX": "capable", "REASONING": "capable"},
|
||||
"return_raw_model_name": True,
|
||||
},
|
||||
"complexity_router_default_model": "capable",
|
||||
}},
|
||||
])
|
||||
response = asyncio.run(router.acompletion(
|
||||
model="auto", messages=[{"role": "user", "content": "What is 2+2?"}],
|
||||
))
|
||||
print(response.model, response.choices[0].message.content)
|
||||
```
|
||||
|
||||
### AI Gateway (Proxy Server)
|
||||
|
||||
[**Create an Auto Router**](https://docs.litellm.ai/docs/auto_router/setup), then call its model name. Replace `auto` and the URL below with yours, and set `LITELLM_API_KEY` to your virtual key.
|
||||
|
||||
```shell
|
||||
curl http://localhost:4000/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LITELLM_API_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model":"auto","messages":[{"role":"user","content":"What is 2+2?"}]}'
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Agents</b> - Invoke A2A Agents (Python SDK + AI Gateway)</summary>
|
||||
<summary><b>A2A</b> - Invoke A2A Agents (Python SDK + AI Gateway)</summary>
|
||||
|
||||
[**Supported Providers**](https://docs.litellm.ai/docs/a2a#add-a2a-agents) - LangGraph, Vertex AI Agent Engine, Azure AI Foundry, Bedrock AgentCore, Pydantic AI
|
||||
|
||||
|
|
@ -442,7 +486,7 @@ You can use LiteLLM through either the Proxy Server or Python SDK. Both give you
|
|||
<tr>
|
||||
<th style={{width: '14%'}}></th>
|
||||
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/simple_proxy">LiteLLM AI Gateway</a></strong></th>
|
||||
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/">LiteLLM Python SDK</a></strong></th>
|
||||
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/python_sdk">LiteLLM Python SDK</a></strong></th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ ENV UV_PYTHON_DOWNLOADS=0
|
|||
COPY pyproject.toml README.md rust-toolchain.toml ./
|
||||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
COPY model_prices_and_context_window.json ./
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
|
|
|
|||
11
codecov.yaml
11
codecov.yaml
|
|
@ -15,9 +15,10 @@ ignore:
|
|||
flag_management:
|
||||
default_rules:
|
||||
carryforward: true
|
||||
# Dead flags no CI job uploads anymore: their carried-forward sessions were
|
||||
# measured against old revisions, and the stale line maps mark comment lines
|
||||
# of since-edited files as missed, sinking patch coverage on unrelated PRs.
|
||||
# Dead flags no CI job uploads anymore, and flags whose job is path-gated on
|
||||
# PRs: their carried-forward sessions were measured against old revisions, and
|
||||
# the stale line maps mark comment lines of since-edited files as missed,
|
||||
# sinking patch coverage on unrelated PRs.
|
||||
individual_flags:
|
||||
- name: proxy-mgmt-behavior
|
||||
carryforward: false
|
||||
|
|
@ -99,6 +100,10 @@ flag_management:
|
|||
carryforward: false
|
||||
- name: proxy-db-endpoints-and-responses
|
||||
carryforward: false
|
||||
- name: roi-postgres
|
||||
carryforward: false
|
||||
- name: redis-compat
|
||||
carryforward: false
|
||||
|
||||
component_management:
|
||||
individual_components:
|
||||
|
|
|
|||
|
|
@ -398,7 +398,7 @@ class CheckBatchCost:
|
|||
deployment. Once that deployment leaves the router, default fallbacks can
|
||||
silently send the retrieve to a provider that never saw the batch, so its
|
||||
404 must not retire a potentially recoverable row."""
|
||||
return self.llm_router.get_deployment(model_id=model_id) is not None
|
||||
return self.llm_router.registry().deployment(model_id) is not None
|
||||
|
||||
@staticmethod
|
||||
def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool:
|
||||
|
|
@ -564,7 +564,7 @@ class CheckBatchCost:
|
|||
def _get_deployment_id_for_bare_model(
|
||||
self, bare_model_name: str, llm_provider: str
|
||||
) -> Optional[str]:
|
||||
model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name)
|
||||
model_group = self.llm_router.registry().resolve(bare_model_name, kind="from_id")
|
||||
deployment_id = (
|
||||
self._get_deployment_id_for_provider(model_group, llm_provider) if model_group else None
|
||||
)
|
||||
|
|
@ -580,7 +580,7 @@ class CheckBatchCost:
|
|||
) -> Optional[str]:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
for deployment in self.llm_router.get_model_list(model_name=None) or []:
|
||||
for deployment in self.llm_router.registry().deployments() or []:
|
||||
litellm_params = deployment.get("litellm_params") or {}
|
||||
actual_model = litellm_params.get("model")
|
||||
if not isinstance(actual_model, str):
|
||||
|
|
@ -623,8 +623,9 @@ class CheckBatchCost:
|
|||
"""
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
for deployment_id in self.llm_router.get_model_ids(model_name=model_group):
|
||||
deployment_info = self.llm_router.get_deployment(model_id=deployment_id)
|
||||
registry = self.llm_router.registry()
|
||||
for deployment_id in registry.names(model_group, kind="ids"):
|
||||
deployment_info = registry.deployment(deployment_id)
|
||||
if deployment_info is None:
|
||||
continue
|
||||
try:
|
||||
|
|
@ -743,7 +744,7 @@ class CheckBatchCost:
|
|||
except (IndexError, AttributeError):
|
||||
pass
|
||||
|
||||
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
_file_content = await afile_content(
|
||||
file_id=cast(str, raw_output_file_id),
|
||||
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
|
||||
|
|
@ -766,7 +767,7 @@ class CheckBatchCost:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
deployment_info = self.llm_router.get_deployment(model_id=model_id)
|
||||
deployment_info = self.llm_router.registry().deployment(model_id)
|
||||
if deployment_info is None:
|
||||
verbose_proxy_logger.info(
|
||||
f"Skipping job {job.unified_object_id} because it is not a valid deployment info"
|
||||
|
|
@ -996,7 +997,7 @@ class CheckBatchCost:
|
|||
)
|
||||
|
||||
try:
|
||||
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
|
||||
deployment_info: Final = self.llm_router.registry().deployment(model_id)
|
||||
if deployment_info is None:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
|
||||
|
|
@ -1009,7 +1010,7 @@ class CheckBatchCost:
|
|||
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
|
||||
)
|
||||
cost_provider: Final = cast(BatchFileProvider, llm_provider)
|
||||
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
file_access: Final = {
|
||||
**credentials,
|
||||
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
|
||||
|
|
|
|||
|
|
@ -273,7 +273,7 @@ class CheckResponsesCost:
|
|||
|
||||
def _resolve_deployment(self, response_id: str) -> Deployment | None:
|
||||
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
|
||||
return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
|
||||
return self.llm_router.registry().deployment(model_id) if model_id is not None else None
|
||||
|
||||
async def _get_response(
|
||||
self,
|
||||
|
|
@ -415,7 +415,7 @@ class CheckResponsesCost:
|
|||
if target.deployment is None or not isinstance(requested, str):
|
||||
return MappingProxyType({})
|
||||
routed: Final = _ROUTED_DEPLOYMENTS.validate_python(
|
||||
self.llm_router.get_model_list(model_name=requested, team_id=team_id) or ()
|
||||
self.llm_router.registry().deployments(requested, team_id=team_id) or ()
|
||||
)
|
||||
serving_id: Final = target.deployment.model_info.id
|
||||
return MappingProxyType(
|
||||
|
|
|
|||
|
|
@ -245,7 +245,7 @@ def _provider_file_retrieve_credentials(
|
|||
if llm_router is None or model_id is None:
|
||||
return None
|
||||
try:
|
||||
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id)
|
||||
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
except Exception as error:
|
||||
verbose_logger.warning(
|
||||
"Failed to retrieve credentials for provider file "
|
||||
|
|
@ -2650,7 +2650,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
|
||||
}
|
||||
for model_id, model_file_id in specific_model_file_id_mapping.items():
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data = {
|
||||
**filtered_data,
|
||||
**(
|
||||
|
|
@ -2694,7 +2694,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
# against the deployment's configured bucket, which they only
|
||||
# trust from this immutable server-side snapshot, never from
|
||||
# request params.
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
if credentials is not None:
|
||||
data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -290,7 +290,7 @@ def _check_team_project_limits(
|
|||
|
||||
|
||||
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
|
||||
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
|
||||
return frozenset(llm_router.registry().access_groups()) if llm_router is not None else frozenset()
|
||||
|
||||
|
||||
def _project_models_expanding_at_request_time(
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ ENV UV_PYTHON_DOWNLOADS=0
|
|||
COPY pyproject.toml README.md rust-toolchain.toml ./
|
||||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
COPY model_prices_and_context_window.json ./
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
|
|
@ -144,6 +145,7 @@ COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
|
|||
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/litellm_cli /app/litellm_cli
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
|
||||
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
|
|
|
|||
|
|
@ -8,6 +8,11 @@ handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
|
|||
wins under token auth too, so exporting it here is enough for every worker to
|
||||
pick the pooled URL up unchanged.
|
||||
|
||||
LiteLLM imports wait for ``main``: every uvicorn worker imports this module again
|
||||
before it can answer its first 5 s health ping, so a slow host would kill each one.
|
||||
Workers then load the gateway app through ``worker_app`` without garbage collection
|
||||
pauses, which hold the GIL long enough to miss a ping.
|
||||
|
||||
Run with:
|
||||
python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
|
||||
"""
|
||||
|
|
@ -15,19 +20,19 @@ Run with:
|
|||
import os
|
||||
import sys
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from uvicorn.main import main as uvicorn_main
|
||||
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import (
|
||||
PgBouncerError,
|
||||
PgBouncerSettings,
|
||||
export_pooled_database_url,
|
||||
start_in_container_pgbouncer,
|
||||
)
|
||||
from litellm_cli.main import without_gc_pauses
|
||||
|
||||
GATEWAY_APP: Final = "gateway.main:app"
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import FastAPI
|
||||
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings
|
||||
|
||||
GATEWAY_APP: Final = "gateway.launch:worker_app"
|
||||
KEEPALIVE_FLAG: Final = "--timeout-keep-alive"
|
||||
|
||||
|
||||
|
|
@ -35,21 +40,23 @@ def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str,
|
|||
"""Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly."""
|
||||
keepalive: Final = environ.get("KEEPALIVE_TIMEOUT")
|
||||
if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv):
|
||||
return (GATEWAY_APP, *argv)
|
||||
return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive)
|
||||
return (GATEWAY_APP, "--factory", *argv)
|
||||
return (GATEWAY_APP, "--factory", *argv, KEEPALIVE_FLAG, keepalive)
|
||||
|
||||
|
||||
def pool_database_url(
|
||||
settings: DatabaseURLSettings,
|
||||
pgbouncer: PgBouncerSettings,
|
||||
settings: "DatabaseURLSettings",
|
||||
pgbouncer: "PgBouncerSettings",
|
||||
environ: Mapping[str, str],
|
||||
) -> str | PgBouncerError | None:
|
||||
) -> "str | PgBouncerError | None":
|
||||
"""Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.
|
||||
|
||||
The upstream URL is whatever ``apply_to_env`` assembled from the discrete
|
||||
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token
|
||||
auth the pooler mints and renews the upstream token itself.
|
||||
"""
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, start_in_container_pgbouncer
|
||||
|
||||
if not pgbouncer.enabled:
|
||||
return None
|
||||
upstream_url: Final = environ.get("DATABASE_URL")
|
||||
|
|
@ -58,11 +65,20 @@ def pool_database_url(
|
|||
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth())
|
||||
|
||||
|
||||
def worker_app() -> "FastAPI":
|
||||
with without_gc_pauses():
|
||||
from gateway.main import app
|
||||
return app
|
||||
|
||||
|
||||
def _serve(argv: Sequence[str]) -> None:
|
||||
uvicorn_main(tuple(argv), prog_name="uvicorn")
|
||||
|
||||
|
||||
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url
|
||||
|
||||
settings: Final = DatabaseURLSettings.from_env()
|
||||
settings.apply_to_env()
|
||||
pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)
|
||||
|
|
|
|||
|
|
@ -259,3 +259,9 @@ itself.
|
|||
`lensWorker.mode` selects `disabled`, `bundled` or `external`. Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
|
||||
|
||||
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for complete values, verification, GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
|
||||
|
||||
## Shutdown accounting
|
||||
|
||||
`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged
|
||||
|
||||
Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml`
|
||||
|
|
|
|||
|
|
@ -14,11 +14,9 @@ either, so the migration keeps today's behavior for it.
|
|||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable, Iterator
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from itertools import accumulate
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
|
@ -43,7 +41,7 @@ MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIME
|
|||
DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0
|
||||
REFUSAL: Final = "Existing rows would make a pending schema change fail:"
|
||||
|
||||
Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"]
|
||||
Kind = Literal["too_long"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -54,9 +52,6 @@ class Check:
|
|||
statement: str
|
||||
source: str
|
||||
index: str | None = None
|
||||
referenced_table: str | None = None
|
||||
referenced_columns: tuple[str, ...] = ()
|
||||
expression: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -100,33 +95,6 @@ class ColumnType:
|
|||
|
||||
ColumnTypeLookup = Callable[[str, str], ColumnType | None]
|
||||
|
||||
_IDENT: Final = r'"?(\w+)"?'
|
||||
_QUALIFIED: Final = r'(?:"?\w+"?\s*\.\s*)?"?(\w+)"?'
|
||||
_CREATE_INDEX: Final = re.compile(
|
||||
rf"^\s*CREATE\s+(UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?{_IDENT}\s+ON\s+(?:ONLY\s+)?"
|
||||
rf"{_QUALIFIED}\s*(?:USING\s+(\w+)\s*)?\((.*?)\)",
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
_ALTER_TABLE: Final = re.compile(
|
||||
rf"^\s*ALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?{_QUALIFIED}\s+(.*)$", re.IGNORECASE | re.DOTALL
|
||||
)
|
||||
_SET_NOT_NULL: Final = re.compile(rf"^ALTER\s+(?:COLUMN\s+)?{_IDENT}\s+SET\s+NOT\s+NULL$", re.IGNORECASE)
|
||||
_ADD_COLUMN: Final = re.compile(
|
||||
rf"^ADD\s+(?:COLUMN\s+)?(?:IF\s+NOT\s+EXISTS\s+)?(?!(?:CONSTRAINT|FOREIGN|UNIQUE|CHECK|PRIMARY|EXCLUDE)\b){_IDENT}\s+(.*)$",
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
_CONSTRAINT: Final = rf"^ADD\s+(?:CONSTRAINT\s+{_IDENT}\s+)?"
|
||||
_ADD_FOREIGN_KEY: Final = re.compile(
|
||||
rf"{_CONSTRAINT}FOREIGN\s+KEY\s*\(([^)]*)\)\s*REFERENCES\s+{_QUALIFIED}\s*\(([^)]*)\)(.*)$",
|
||||
re.IGNORECASE | re.DOTALL,
|
||||
)
|
||||
_ADD_UNIQUE: Final = re.compile(rf"{_CONSTRAINT}UNIQUE\s*\(([^)]*)\)", re.IGNORECASE | re.DOTALL)
|
||||
_ADD_CHECK: Final = re.compile(rf"{_CONSTRAINT}CHECK\s*\((.*)\)\s*(NOT\s+VALID)?\s*$", re.IGNORECASE | re.DOTALL)
|
||||
_PLAIN_COLUMN: Final = re.compile(rf"^{_IDENT}(?:\s+(?:ASC|DESC|NULLS\s+(?:FIRST|LAST)))*$", re.IGNORECASE)
|
||||
_COMMENT: Final = re.compile(r"--[^\n]*")
|
||||
_NOT_NULL: Final = re.compile(r"\bNOT\s+NULL\b", re.IGNORECASE)
|
||||
_HAS_DEFAULT: Final = re.compile(r"\bDEFAULT\b|\bGENERATED\b", re.IGNORECASE)
|
||||
|
||||
|
||||
def migration_data_check_timeout_seconds() -> float:
|
||||
"""Seconds one count query may run before the preflight gives up on verifying it."""
|
||||
|
|
@ -145,117 +113,6 @@ def migration_data_check_timeout_seconds() -> float:
|
|||
return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
|
||||
|
||||
|
||||
def statements(migration_sql: str) -> tuple[str, ...]:
|
||||
"""The statements of one migration file, comments stripped, empty ones dropped."""
|
||||
parts: Final = _COMMENT.sub("", migration_sql).split(";")
|
||||
return tuple(part.strip() for part in parts if part.strip())
|
||||
|
||||
|
||||
def _split_top_level(text: str) -> tuple[str, ...]:
|
||||
"""Split on the commas outside parentheses."""
|
||||
depth: Final = accumulate((1 if char == "(" else -1 if char == ")" else 0) for char in text)
|
||||
cuts: Final = tuple(
|
||||
position for position, (char, level) in enumerate(zip(text, depth, strict=True)) if char == "," and level == 0
|
||||
)
|
||||
starts: Final = (0, *(cut + 1 for cut in cuts))
|
||||
stops: Final = (*cuts, len(text))
|
||||
pieces: Final = tuple(text[start:stop].strip() for start, stop in zip(starts, stops, strict=True))
|
||||
return tuple(piece for piece in pieces if piece)
|
||||
|
||||
|
||||
def _plain_columns(column_list: str) -> tuple[str, ...] | None:
|
||||
"""The column names of an index or constraint column list, None when one entry is an
|
||||
expression the derivation cannot turn into a count."""
|
||||
matches: Final = tuple(_PLAIN_COLUMN.match(entry) for entry in _split_top_level(column_list))
|
||||
if any(match is None for match in matches):
|
||||
return None
|
||||
return tuple(match.group(1) for match in matches if match is not None)
|
||||
|
||||
|
||||
def _index_checks(match: "re.Match[str]", statement_text: str, source: str) -> Iterator[Check]:
|
||||
unique, name, table, method, column_list = match.groups()
|
||||
columns: Final = _plain_columns(column_list)
|
||||
if columns is None:
|
||||
return
|
||||
label: Final = f'CREATE {"UNIQUE " if unique else ""}INDEX "{name}"'
|
||||
if method is None or method.lower() == "btree":
|
||||
yield Check("too_long", table, columns, label, source, index=name)
|
||||
if unique:
|
||||
yield Check("duplicates", table, columns, label, source, index=name)
|
||||
|
||||
|
||||
def _alter_table_checks(table: str, action: str, source: str) -> Iterator[Check]:
|
||||
if (set_not_null := _SET_NOT_NULL.match(action)) is not None:
|
||||
yield Check("nulls", table, (set_not_null.group(1),), f'ALTER TABLE "{table}" {action}', source)
|
||||
return
|
||||
if (add_column := _ADD_COLUMN.match(action)) is not None:
|
||||
column, body = add_column.groups()
|
||||
if _NOT_NULL.search(body) is not None and _HAS_DEFAULT.search(body) is None:
|
||||
yield Check(
|
||||
"any_row", table, (column,), f'ALTER TABLE "{table}" ADD COLUMN "{column}" ... NOT NULL', source
|
||||
)
|
||||
return
|
||||
if (foreign_key := _ADD_FOREIGN_KEY.match(action)) is not None:
|
||||
constraint, columns, referenced_table, referenced_columns, tail = foreign_key.groups()
|
||||
own: Final = _plain_columns(columns)
|
||||
theirs: Final = _plain_columns(referenced_columns)
|
||||
if own is None or theirs is None or len(own) != len(theirs) or "NOT VALID" in tail.upper():
|
||||
return
|
||||
yield Check(
|
||||
"orphans",
|
||||
table,
|
||||
own,
|
||||
_constraint_label(table, constraint, "FOREIGN KEY"),
|
||||
source,
|
||||
referenced_table=referenced_table,
|
||||
referenced_columns=theirs,
|
||||
)
|
||||
return
|
||||
if (unique := _ADD_UNIQUE.match(action)) is not None:
|
||||
constraint, columns = unique.groups()
|
||||
unique_columns: Final = _plain_columns(columns)
|
||||
if unique_columns is not None:
|
||||
yield Check("duplicates", table, unique_columns, _constraint_label(table, constraint, "UNIQUE"), source)
|
||||
return
|
||||
if (check := _ADD_CHECK.match(action)) is not None:
|
||||
constraint, expression, not_valid = check.groups()
|
||||
if not_valid is None:
|
||||
yield Check(
|
||||
"check_violations",
|
||||
table,
|
||||
(),
|
||||
_constraint_label(table, constraint, "CHECK"),
|
||||
source,
|
||||
expression=expression,
|
||||
)
|
||||
|
||||
|
||||
def _constraint_label(table: str, constraint: str | None, kind: str) -> str:
|
||||
named: Final = f'CONSTRAINT "{constraint}" ' if constraint else ""
|
||||
return f'ALTER TABLE "{table}" ADD {named}{kind}'
|
||||
|
||||
|
||||
def checks_for_statement(statement_text: str, source: str) -> tuple[Check, ...]:
|
||||
"""The count queries one SQL statement needs before it can be applied, empty for a
|
||||
statement that cannot fail on existing rows or that the derivation cannot read."""
|
||||
if (index := _CREATE_INDEX.match(statement_text)) is not None:
|
||||
return tuple(_index_checks(index, statement_text, source))
|
||||
if (alter := _ALTER_TABLE.match(statement_text)) is not None:
|
||||
table, actions = alter.groups()
|
||||
return tuple(
|
||||
check for action in _split_top_level(actions) for check in _alter_table_checks(table, action, source)
|
||||
) # comprehension-ok: flattening the per-action checks of one statement
|
||||
return ()
|
||||
|
||||
|
||||
def checks_for_migration(name: str, migration_sql: str) -> tuple[Check, ...]:
|
||||
return tuple(
|
||||
check
|
||||
for statement_text in statements(migration_sql)
|
||||
for check in checks_for_statement(statement_text, f"migration {name}")
|
||||
) # comprehension-ok: flattening the per-statement checks of one migration
|
||||
|
||||
|
||||
def checks_for_index(table: str, name: str, columns: tuple[str, ...]) -> tuple[Check, ...]:
|
||||
"""The entry-size check a request-log index needs before `CREATE INDEX CONCURRENTLY`."""
|
||||
return (Check("too_long", table, columns, f'CREATE INDEX "{name}"', "request-log index", index=name),)
|
||||
|
|
@ -424,42 +281,7 @@ def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...],
|
|||
|
||||
|
||||
def _count_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
|
||||
from psycopg import sql
|
||||
|
||||
table: Final = sql.Identifier(schema, check.table)
|
||||
columns: Final = sql.SQL(", ").join(sql.Identifier(column) for column in check.columns)
|
||||
match check.kind:
|
||||
case "too_long":
|
||||
return _entry_size_query(check, schema, types, limit_bytes)
|
||||
case "duplicates":
|
||||
not_null: Final = sql.SQL(" AND ").join(
|
||||
sql.SQL("{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
|
||||
)
|
||||
return sql.SQL(
|
||||
"SELECT count(*) FROM (SELECT 1 FROM {t} WHERE {nn} GROUP BY {cols} HAVING count(*) > 1) AS d"
|
||||
).format(t=table, nn=not_null, cols=columns)
|
||||
case "nulls":
|
||||
return sql.SQL("SELECT count(*) FROM {t} WHERE {col} IS NULL").format(
|
||||
t=table, col=sql.Identifier(check.columns[0])
|
||||
)
|
||||
case "any_row":
|
||||
return sql.SQL("SELECT count(*) FROM {t}").format(t=table)
|
||||
case "orphans":
|
||||
pairs: Final = sql.SQL(" AND ").join(
|
||||
sql.SQL("parent.{} = child.{}").format(sql.Identifier(theirs), sql.Identifier(own))
|
||||
for own, theirs in zip(check.columns, check.referenced_columns, strict=True)
|
||||
)
|
||||
present: Final = sql.SQL(" AND ").join(
|
||||
sql.SQL("child.{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
|
||||
)
|
||||
return sql.SQL(
|
||||
"SELECT count(*) FROM {t} AS child WHERE {present} AND NOT EXISTS (SELECT 1 FROM {ref} AS parent WHERE {pairs})"
|
||||
).format(t=table, present=present, ref=sql.Identifier(schema, check.referenced_table or ""), pairs=pairs)
|
||||
case "check_violations":
|
||||
return sql.SQL("SELECT count(*) FROM {t} WHERE NOT ({expr})").format(
|
||||
t=table,
|
||||
expr=sql.SQL(check.expression or "true"), # pyright: ignore[reportArgumentType] # copied from our own migration file
|
||||
)
|
||||
return _entry_size_query(check, schema, types, limit_bytes)
|
||||
|
||||
|
||||
def _applies(
|
||||
|
|
@ -474,15 +296,7 @@ def _applies(
|
|||
return False
|
||||
if check.index is not None and _index_is_valid(connection, schema, check.index):
|
||||
return False
|
||||
if check.kind == "any_row":
|
||||
return types is None
|
||||
if types is None:
|
||||
return False
|
||||
if check.kind == "too_long":
|
||||
return any(column.can_overflow_btree for column in types)
|
||||
if check.kind == "orphans":
|
||||
return check.referenced_table is not None and _table_exists(connection, schema, check.referenced_table)
|
||||
return True
|
||||
return types is not None and any(column.can_overflow_btree for column in types)
|
||||
|
||||
|
||||
def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str:
|
||||
|
|
@ -515,21 +329,7 @@ def _run_check(
|
|||
rows: Final = _int(values[0], 0) if values else 0
|
||||
if rows == 0:
|
||||
return None
|
||||
match check.kind:
|
||||
case "too_long":
|
||||
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
|
||||
case "duplicates":
|
||||
return BlockedChange(check, rows, "duplicated value groups")
|
||||
case "nulls":
|
||||
return BlockedChange(check, rows, "NULL")
|
||||
case "any_row":
|
||||
return BlockedChange(
|
||||
check, rows, "rows, and a NOT NULL column with no DEFAULT cannot be added to a table with rows"
|
||||
)
|
||||
case "orphans":
|
||||
return BlockedChange(check, rows, f'with no matching row in "{check.referenced_table}"')
|
||||
case "check_violations":
|
||||
return BlockedChange(check, rows, f"failing CHECK ({check.expression})")
|
||||
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
|
||||
|
||||
|
||||
def find_blockers(
|
||||
|
|
@ -571,19 +371,7 @@ def find_blockers(
|
|||
|
||||
def _remedy(check: Check) -> str:
|
||||
joined: Final = ", ".join(check.columns)
|
||||
match check.kind:
|
||||
case "too_long":
|
||||
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
|
||||
case "duplicates":
|
||||
return f'Remove or merge the duplicate "{check.table}" ({joined}) values.'
|
||||
case "nulls":
|
||||
return f'Fill the NULL "{check.table}"."{joined}" values, or give the column a default first.'
|
||||
case "any_row":
|
||||
return "Give the new column a DEFAULT in the migration, or add it nullable and fill it afterwards."
|
||||
case "orphans":
|
||||
return f'Delete or repoint the "{check.table}" ({joined}) values with no "{check.referenced_table}" row.'
|
||||
case "check_violations":
|
||||
return f'Fix the "{check.table}" rows that fail the expression before adding the constraint.'
|
||||
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
|
||||
|
||||
|
||||
def _describe(finding: Finding) -> str:
|
||||
|
|
@ -607,30 +395,4 @@ def _describe(finding: Finding) -> str:
|
|||
def report(findings: tuple[Finding, ...]) -> str:
|
||||
"""The operator-facing refusal: one block per finding, then what to do next."""
|
||||
lines: Final = "\n".join(_describe(finding) for finding in findings)
|
||||
return (
|
||||
f"{REFUSAL}\n{lines}\n"
|
||||
"No schema change was made. Fix the rows and start again, or run `litellm --check_migrations_against_existing_data` with the new image to check again."
|
||||
)
|
||||
|
||||
|
||||
def pending_change_checks(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
migrations_dir: Path | None,
|
||||
indexes: tuple[tuple[str, str, tuple[str, ...]], ...],
|
||||
) -> tuple[Check, ...]:
|
||||
"""Every check the pending migrations and the request-log indexes need. `migrations_dir`
|
||||
is None when the schema is pushed rather than migrated, so only the indexes are checked."""
|
||||
migration_checks: Final = (
|
||||
tuple(
|
||||
check
|
||||
for name, migration_sql in pending_migrations(connection, schema, migrations_dir)
|
||||
for check in checks_for_migration(name, migration_sql)
|
||||
) # comprehension-ok: flattening the per-migration checks
|
||||
if migrations_dir is not None
|
||||
else ()
|
||||
)
|
||||
index_checks: Final = tuple(
|
||||
check for table, name, columns in indexes for check in checks_for_index(table, name, columns)
|
||||
) # comprehension-ok: flattening the per-index checks
|
||||
return migration_checks + index_checks
|
||||
return f"{REFUSAL}\n{lines}\nNo index was built. Fix the rows and the next index build picks it up."
|
||||
|
|
|
|||
184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py
Normal file
184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
"""`litellm --check_migrations_against_existing_data`: run the pending migrations for real, each
|
||||
inside its own transaction that is always rolled back, so Postgres itself says whether this
|
||||
database's rows and schema take them. Each migration holds the locks the real upgrade will,
|
||||
and releases them before the next one runs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm_proxy_extras.migration_data_checks import pending_migrations
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
||||
_CONCURRENTLY: Final = re.compile(
|
||||
r"\b((?:CREATE\s+(?:UNIQUE\s+)?INDEX|DROP\s+INDEX|REINDEX\s+\w+)\s+)CONCURRENTLY\s+", re.IGNORECASE
|
||||
)
|
||||
UNDEFINED_SQLSTATES: Final = frozenset({"42703", "42P01", "42704", "42883"})
|
||||
|
||||
Migration = tuple[str, str]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Applied:
|
||||
name: str
|
||||
seconds: float
|
||||
replayed: tuple[str, ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rejected:
|
||||
name: str
|
||||
error: str
|
||||
seconds: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Baseline:
|
||||
"""A database with tables but no `_prisma_migrations`: the boot baselines it instead of replaying
|
||||
history, so the dry run asks the same question the boot does. `refusal` is None when the boot
|
||||
will baseline, otherwise its reason (a custom migration folder, or the schema diff)."""
|
||||
|
||||
refusal: str | None
|
||||
|
||||
@property
|
||||
def matches(self) -> bool:
|
||||
return self.refusal is None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DryRun:
|
||||
applied: tuple[Applied, ...]
|
||||
rejected: Rejected | None
|
||||
baseline: Baseline | None = None
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.rejected is None and (self.baseline is None or self.baseline.matches)
|
||||
|
||||
|
||||
def transactional(migration_sql: str) -> str:
|
||||
"""The migration as it can run inside a transaction: `CONCURRENTLY` index builds become plain ones."""
|
||||
return _CONCURRENTLY.sub(r"\1", migration_sql)
|
||||
|
||||
|
||||
def _attempt(
|
||||
connection: psycopg.Connection[tuple[object, ...]],
|
||||
schema: str,
|
||||
chain: tuple[Migration, ...],
|
||||
clock: Callable[[], float],
|
||||
) -> tuple[float, psycopg.Error | None]:
|
||||
import psycopg
|
||||
from psycopg import sql
|
||||
|
||||
started: Final = clock()
|
||||
try:
|
||||
with connection.transaction(force_rollback=True):
|
||||
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
|
||||
for _, migration_sql in chain:
|
||||
connection.execute(transactional(migration_sql))
|
||||
except psycopg.Error as exc:
|
||||
return clock() - started, exc
|
||||
return clock() - started, None
|
||||
|
||||
|
||||
def _run(
|
||||
connection: psycopg.Connection[tuple[object, ...]],
|
||||
schema: str,
|
||||
pending: tuple[Migration, ...],
|
||||
position: int,
|
||||
done: DryRun,
|
||||
clock: Callable[[], float],
|
||||
) -> DryRun:
|
||||
"""One migration in its own rolled-back transaction. A migration that fails because something
|
||||
an earlier pending migration creates does not exist yet is rerun with those replayed in front."""
|
||||
if done.rejected is not None:
|
||||
return done
|
||||
name: Final = pending[position][0]
|
||||
seconds, error = _attempt(connection, schema, (pending[position],), clock)
|
||||
if error is None:
|
||||
return DryRun(done.applied + (Applied(name, seconds),), None, done.baseline)
|
||||
if error.sqlstate not in UNDEFINED_SQLSTATES or position == 0:
|
||||
return DryRun(done.applied, Rejected(name, str(error).strip(), seconds), done.baseline)
|
||||
earlier: Final = pending[:position]
|
||||
chained_seconds, chained_error = _attempt(connection, schema, earlier + (pending[position],), clock)
|
||||
if chained_error is None:
|
||||
replayed: Final = tuple(earlier_name for earlier_name, _ in earlier)
|
||||
return DryRun(done.applied + (Applied(name, chained_seconds, replayed),), None, done.baseline)
|
||||
return DryRun(done.applied, Rejected(name, str(chained_error).strip(), chained_seconds), done.baseline)
|
||||
|
||||
|
||||
def _has_ledger(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
|
||||
from psycopg import sql
|
||||
|
||||
row: Final = connection.execute(
|
||||
"SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(connection),)
|
||||
).fetchone()
|
||||
return row is not None and row[0] is not None
|
||||
|
||||
|
||||
def _has_litellm_tables(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
|
||||
row: Final = connection.execute(
|
||||
"SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name LIKE 'LiteLLM\\_%%' LIMIT 1",
|
||||
(schema,),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
|
||||
def dry_run(
|
||||
connection: psycopg.Connection[tuple[object, ...]],
|
||||
schema: str,
|
||||
migrations_dir: Path,
|
||||
baseline_check: Callable[[], str | None] | None = None,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> DryRun:
|
||||
"""Apply every pending migration in order, each in its own transaction that is rolled back
|
||||
before the next one starts. The first migration Postgres rejects ends the run with its error.
|
||||
`baseline_check` (None when the boot would baseline, otherwise why it refuses) stands in for
|
||||
the replay on a database that has tables but no migration history, which the boot baselines."""
|
||||
if baseline_check is not None and not _has_ledger(connection, schema) and _has_litellm_tables(connection, schema):
|
||||
return DryRun((), None, Baseline(baseline_check()))
|
||||
pending: Final = pending_migrations(connection, schema, migrations_dir)
|
||||
return reduce(
|
||||
lambda done, position: _run(connection, schema, pending, position, done, clock),
|
||||
range(len(pending)),
|
||||
DryRun((), None),
|
||||
)
|
||||
|
||||
|
||||
def _timing(applied: Applied) -> str:
|
||||
replay: Final = f" (depends on {', '.join(applied.replayed)}, replayed in front of it)" if applied.replayed else ""
|
||||
return f"\n {applied.name}: {applied.seconds:.1f} s{replay}"
|
||||
|
||||
|
||||
def report(result: DryRun) -> str:
|
||||
"""What the operator reads: the per-migration time the real upgrade will hold its locks,
|
||||
or the migration Postgres rejected with Postgres's own words. Nothing was applied either way."""
|
||||
if result.baseline is not None:
|
||||
if result.baseline.matches:
|
||||
return (
|
||||
"No migration history and the schema matches this build, "
|
||||
"the proxy baselines it at boot, nothing to dry-run."
|
||||
)
|
||||
return f"No migration history and the proxy cannot baseline it at boot: {result.baseline.refusal}"
|
||||
timings: Final = "".join(_timing(applied) for applied in result.applied)
|
||||
if result.rejected is None:
|
||||
count: Final = len(result.applied)
|
||||
if count == 0:
|
||||
return "No pending migrations, this database is on this version's schema."
|
||||
return (
|
||||
f"{count} pending migration(s) applied one at a time in a dry run, each rolled back before the next "
|
||||
f"so its locks are released, nothing changed:{timings}"
|
||||
)
|
||||
return (
|
||||
f"Migration {result.rejected.name} fails on this database after {result.rejected.seconds:.1f} s, "
|
||||
f"nothing was applied.\n{result.rejected.error}"
|
||||
+ (f"\nApplied before it in the dry run:{timings}" if timings else "")
|
||||
)
|
||||
|
|
@ -86,6 +86,18 @@ def migration_files(directory: Path) -> tuple[tuple[str, str], ...]:
|
|||
)
|
||||
|
||||
|
||||
def has_packaged_history(migrations_dir: Path) -> bool:
|
||||
"""Whether `migrations_dir` carries exactly the migration files and schema.prisma this package ships,
|
||||
the only history the boot baselines an existing database onto."""
|
||||
packaged_dir: Final = Path(__file__).parent
|
||||
migrations: Final = migration_files(migrations_dir)
|
||||
return (
|
||||
bool(migrations)
|
||||
and migrations == migration_files(packaged_dir)
|
||||
and (migrations_dir / "schema.prisma").read_bytes() == (packaged_dir / "schema.prisma").read_bytes()
|
||||
)
|
||||
|
||||
|
||||
def baseline_current_schema(
|
||||
coordinator: MigrationCoordinator,
|
||||
schema: str,
|
||||
|
|
@ -95,14 +107,9 @@ def baseline_current_schema(
|
|||
) -> None:
|
||||
from psycopg import sql
|
||||
|
||||
packaged_dir: Final = Path(__file__).parent
|
||||
migrations: Final = migration_files(migrations_dir)
|
||||
if (
|
||||
not migrations
|
||||
or migrations != migration_files(packaged_dir)
|
||||
or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes()
|
||||
):
|
||||
if not has_packaged_history(migrations_dir):
|
||||
raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
|
||||
migrations: Final = migration_files(migrations_dir)
|
||||
|
||||
coordinator.acquire_prisma_lock()
|
||||
existing: Final = coordinator.connection.execute(
|
||||
|
|
|
|||
|
|
@ -11,18 +11,13 @@ from collections.abc import Callable
|
|||
from contextlib import AbstractContextManager, nullcontext
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras._redaction import _redact_credentials
|
||||
from litellm_proxy_extras.migration_data_checks import (
|
||||
Finding,
|
||||
find_blockers,
|
||||
migration_data_check_timeout_seconds,
|
||||
pending_change_checks,
|
||||
report,
|
||||
)
|
||||
from litellm_proxy_extras.migration_dry_run import dry_run
|
||||
from litellm_proxy_extras.migration_dry_run import report as dry_run_report
|
||||
from litellm_proxy_extras.migration_lock import held_migration_lock
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
|
|
@ -36,7 +31,6 @@ from litellm_proxy_extras.replica_identity import (
|
|||
apply_replica_identity_full,
|
||||
)
|
||||
from litellm_proxy_extras.request_log_indexes import (
|
||||
REQUEST_LOG_INDEXES,
|
||||
ensure_request_log_indexes,
|
||||
filter_request_log_index_diff,
|
||||
)
|
||||
|
|
@ -52,7 +46,7 @@ class LensCheckConnect(Protocol):
|
|||
) -> "psycopg.Connection[tuple[object, ...]]": ...
|
||||
|
||||
|
||||
def str_to_bool(value: Optional[str]) -> bool:
|
||||
def str_to_bool(value: str | None) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
return value.lower() in ("true", "1", "t", "y", "yes")
|
||||
|
|
@ -169,7 +163,7 @@ def _without_sql_comments(statement: str) -> str:
|
|||
).strip()
|
||||
|
||||
|
||||
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
|
||||
def _without_spend_logs_pk_clauses(statement: str) -> str | None:
|
||||
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
|
||||
if not prefix_match:
|
||||
return statement
|
||||
|
|
@ -217,7 +211,7 @@ def _max_migration_timestamp(names) -> int:
|
|||
return max(_migration_timestamp(n) for n in names)
|
||||
|
||||
|
||||
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
|
||||
def _redacted_command(command: object) -> str | tuple[str, ...] | list[str]:
|
||||
if isinstance(command, tuple):
|
||||
return tuple(_redact_credentials(str(argument)) for argument in command)
|
||||
if isinstance(command, list):
|
||||
|
|
@ -404,7 +398,7 @@ class ProxyExtrasDBManager:
|
|||
return tuple(row) if row is not None else ()
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]:
|
||||
def _failed_migration_logs(migration_name: str, started_at: str) -> str | None:
|
||||
row: Final = ProxyExtrasDBManager._read_migration_ledger(
|
||||
"SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
|
||||
"AND finished_at IS NULL AND rolled_back_at IS NULL",
|
||||
|
|
@ -721,10 +715,10 @@ class ProxyExtrasDBManager:
|
|||
return row is not None
|
||||
|
||||
@staticmethod
|
||||
def _prisma_schema_param(url: str) -> Optional[str]:
|
||||
def _prisma_schema_param(url: str) -> str | None:
|
||||
"""The `schema` query param Prisma uses to pick its target schema,
|
||||
or None when the URL does not set one."""
|
||||
from urllib.parse import urlparse, parse_qsl
|
||||
from urllib.parse import parse_qsl, urlparse
|
||||
|
||||
return next(
|
||||
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
|
||||
|
|
@ -1424,7 +1418,6 @@ class ProxyExtrasDBManager:
|
|||
bool: True if setup was successful, False otherwise
|
||||
"""
|
||||
ensure_prisma_toolchain(prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env())
|
||||
ProxyExtrasDBManager.raise_if_rows_block_pending_changes(use_migrate=use_migrate)
|
||||
migrated = ProxyExtrasDBManager._run_migrations(use_migrate=use_migrate, use_v2_resolver=use_v2_resolver)
|
||||
if not migrated:
|
||||
return False
|
||||
|
|
@ -1433,46 +1426,69 @@ class ProxyExtrasDBManager:
|
|||
return True
|
||||
|
||||
@staticmethod
|
||||
def migration_data_check_findings(use_migrate: bool = True) -> tuple[Finding, ...]:
|
||||
"""What in the existing rows would make a pending migration or request-log index
|
||||
fail, read with count queries only: no write, no table lock. Empty when nothing
|
||||
blocks, when no database is configured, or when it cannot be reached, so the
|
||||
migration step that follows reports the connection problem the way it does today."""
|
||||
def _baseline_refusal(prisma_dir: Path) -> str | None:
|
||||
"""Why `baseline_current_schema` would refuse this database at boot, None when it would baseline:
|
||||
the migrations folder must be the packaged one, then `prisma migrate diff` must find the live
|
||||
database and this build's schema.prisma identical."""
|
||||
from litellm_proxy_extras.migration_lock import migration_environment
|
||||
from litellm_proxy_extras.migration_recovery import has_packaged_history
|
||||
|
||||
if not has_packaged_history(prisma_dir):
|
||||
return (
|
||||
f"{prisma_dir} is a custom migration history (LITELLM_MIGRATION_DIR), it differs from the "
|
||||
"migrations this build ships, and the proxy only baselines onto the packaged history."
|
||||
)
|
||||
try:
|
||||
prisma_toolchain.run_prisma(
|
||||
(
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
"diff",
|
||||
"--from-schema-datasource",
|
||||
str(prisma_dir / "schema.prisma"),
|
||||
"--to-schema-datamodel",
|
||||
str(prisma_dir / "schema.prisma"),
|
||||
"--exit-code",
|
||||
),
|
||||
timeout=prisma_toolchain.prisma_command_timeout(),
|
||||
env=migration_environment(_get_prisma_env()),
|
||||
)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
diff: Final = str(exc.stdout or exc.stderr or exc).strip()
|
||||
return f"the schema differs from this build. prisma migrate diff:\n{diff}"
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
return f"prisma migrate diff did not finish within {exc.timeout} s"
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]:
|
||||
"""`--check_migrations_against_existing_data`: the pending migrations applied for real on
|
||||
DIRECT_URL or DATABASE_URL inside one transaction that is rolled back, so Postgres reports
|
||||
what the upgrade would hit. (True, the per-migration timings) when they all apply,
|
||||
(False, Postgres's own error) when one does not."""
|
||||
import psycopg
|
||||
|
||||
if not use_migrate:
|
||||
return (True, "The schema is pushed rather than migrated, there is nothing to dry-run.")
|
||||
database_url: Final = os.environ.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
return ()
|
||||
return (False, "DATABASE_URL is not set.")
|
||||
direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
|
||||
os.environ.get("DIRECT_URL") or database_url
|
||||
)
|
||||
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
|
||||
migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" if use_migrate else None
|
||||
indexes: Final = tuple((index.table, index.name, index.columns) for index in REQUEST_LOG_INDEXES)
|
||||
prisma_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir())
|
||||
try:
|
||||
with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection:
|
||||
checks: Final = pending_change_checks(connection, schema, migrations_dir, indexes)
|
||||
return find_blockers(connection, schema, checks, migration_data_check_timeout_seconds())
|
||||
result: Final = dry_run(
|
||||
connection,
|
||||
schema,
|
||||
prisma_dir / "migrations",
|
||||
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(prisma_dir),
|
||||
)
|
||||
except psycopg.OperationalError as exc:
|
||||
logger.warning(
|
||||
"Could not reach the database for the pre-migration row check: %s", _redact_credentials(str(exc))
|
||||
)
|
||||
return ()
|
||||
|
||||
@staticmethod
|
||||
def migration_data_check_report(use_migrate: bool = True) -> Optional[str]:
|
||||
"""The operator-facing refusal for `migration_data_check_findings`, None when nothing blocks."""
|
||||
findings: Final = ProxyExtrasDBManager.migration_data_check_findings(use_migrate=use_migrate)
|
||||
return report(findings) if findings else None
|
||||
|
||||
@staticmethod
|
||||
def raise_if_rows_block_pending_changes(use_migrate: bool) -> None:
|
||||
"""Refuse to run migrations that the existing rows would make fail, before any is
|
||||
applied. The RuntimeError carries the report and reaches the operator through the
|
||||
proxy's "Database migration cannot proceed" exit."""
|
||||
refusal: Final = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate)
|
||||
if refusal is not None:
|
||||
raise RuntimeError(refusal)
|
||||
return (False, f"Could not reach the database: {_redact_credentials(str(exc))}")
|
||||
return (result.ok, dry_run_report(result))
|
||||
|
||||
@staticmethod
|
||||
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
|
||||
|
|
@ -1724,5 +1740,4 @@ class ProxyExtrasDBManager:
|
|||
time.sleep(random.randrange(5, 15))
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
pass
|
||||
return False
|
||||
|
|
|
|||
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -4345,6 +4345,7 @@ dependencies = [
|
|||
"litellm-inference-transcription",
|
||||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-model-catalog",
|
||||
"litellm-router-types",
|
||||
"litellm-secrets",
|
||||
"litellm-secrets-aws",
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ pub(super) const RULES: &[Rule] = &[
|
|||
|mapping| {
|
||||
mapping.error_str.contains("invalid_request_error")
|
||||
&& !mapping.error_str.contains("Incorrect API key provided")
|
||||
&& !matches!(mapping.status, Some(401 | 404))
|
||||
},
|
||||
PublicError::BadRequest,
|
||||
),
|
||||
|
|
@ -132,6 +133,23 @@ mod tests {
|
|||
assert_eq!(classified(Some(400), text), Some(expected));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::unauthorized(Some(401), None)]
|
||||
#[case::not_found(Some(404), None)]
|
||||
#[case::forbidden(Some(403), Some(PublicError::BadRequest))]
|
||||
#[case::rate_limited(Some(429), Some(PublicError::BadRequest))]
|
||||
#[case::server_error(Some(500), Some(PublicError::BadRequest))]
|
||||
#[case::no_status(None, Some(PublicError::BadRequest))]
|
||||
fn invalid_request_errors_leave_401_and_404_to_the_status_mapping(
|
||||
#[case] status: Option<u16>,
|
||||
#[case] expected: Option<PublicError>,
|
||||
) {
|
||||
assert_eq!(
|
||||
classified(status, "invalid_request_error bad field"),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::rate_limit_before_context_window(
|
||||
"rate limit and This model's maximum context length is 10",
|
||||
|
|
|
|||
|
|
@ -25,3 +25,5 @@ A new `N`, image size or resolution in an existing key shape parses with no Rust
|
|||
## Tests
|
||||
|
||||
`BUNDLED_REGISTRY` embeds the root `model_prices_and_context_window.json`, so the tests over it fail a JSON change with a bad value or a key no type reads
|
||||
|
||||
`Catalog::problems` is the one definition of a mergeable cost map. The bundled test asserts it is empty, and the cost map guard workflow runs it on a pull request's JSON through `cargo run --bin check-registry <path>`
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ schema = ["dep:schemars", "litellm-cost/schema"]
|
|||
|
||||
[dependencies]
|
||||
litellm-cost = { workspace = true, features = ["serde"] }
|
||||
litellm-error.workspace = true
|
||||
|
||||
indexmap = { version = "2.14.0", features = ["serde"] }
|
||||
jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] }
|
||||
|
|
@ -24,5 +25,4 @@ time.workspace = true
|
|||
|
||||
[dev-dependencies]
|
||||
jsonschema = { version = "0.55.1", default-features = false }
|
||||
litellm-error.workspace = true
|
||||
rstest.workspace = true
|
||||
|
|
|
|||
26
litellm-rust/crates/model-catalog/src/bin/check-registry.rs
Normal file
26
litellm-rust/crates/model-catalog/src/bin/check-registry.rs
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
use std::{env, fs, path::PathBuf, process::ExitCode};
|
||||
|
||||
use litellm_model_catalog::Catalog;
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let Some(path) = env::args_os().nth(1).map(PathBuf::from) else {
|
||||
eprintln!("usage: check-registry <model_prices_and_context_window.json>");
|
||||
return ExitCode::from(2);
|
||||
};
|
||||
let body = match fs::read(&path) {
|
||||
Ok(body) => body,
|
||||
Err(error) => {
|
||||
eprintln!("{}: {error}", path.display());
|
||||
return ExitCode::FAILURE;
|
||||
}
|
||||
};
|
||||
let problems = Catalog::problems(&body);
|
||||
for problem in &problems {
|
||||
eprintln!("{}", litellm_error::Report::new(problem));
|
||||
}
|
||||
if problems.is_empty() {
|
||||
ExitCode::SUCCESS
|
||||
} else {
|
||||
ExitCode::FAILURE
|
||||
}
|
||||
}
|
||||
|
|
@ -5,7 +5,7 @@ use indexmap::IndexMap;
|
|||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::error::{Error, Loaded, Rejects, Repr};
|
||||
use crate::error::{Error, Loaded, Rejects, Repr, pointer_child, pointer_for_key};
|
||||
use crate::fallback::{FallbackGeneralizations, FallbackRule};
|
||||
use crate::model::Model;
|
||||
use crate::store::{Rows, Store};
|
||||
|
|
@ -211,6 +211,36 @@ impl Catalog {
|
|||
})
|
||||
}
|
||||
|
||||
pub fn problems(body: &[u8]) -> Vec<Error> {
|
||||
let Loaded {
|
||||
value: catalog,
|
||||
errors,
|
||||
} = match Self::parse(body) {
|
||||
Ok(loaded) => loaded,
|
||||
Err(fatal) => return vec![fatal],
|
||||
};
|
||||
let unread = catalog.rows().entries.iter().flat_map(|(model, entry)| {
|
||||
EntryWire::read(entry.fields())
|
||||
.unknown
|
||||
.into_iter()
|
||||
.map(move |(key, _)| {
|
||||
Error::new(
|
||||
pointer_child(&pointer_for_key(model), &key),
|
||||
Repr::UnreadField,
|
||||
)
|
||||
})
|
||||
});
|
||||
let no_fallback_rules = catalog
|
||||
.fallback_rules()
|
||||
.is_none_or(<[_]>::is_empty)
|
||||
.then(|| Error::from_key(FALLBACK_GENERALIZATIONS, Repr::NoFallbackRules));
|
||||
errors
|
||||
.into_iter()
|
||||
.chain(unread)
|
||||
.chain(no_fallback_rules)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn row_errors(model: &str, fields: &Fields) -> Vec<Error> {
|
||||
let (_, _, errors) = read_row(fields);
|
||||
errors
|
||||
|
|
@ -404,3 +434,85 @@ fn unknown_timezone(wire: &EntryWire) -> Option<Error> {
|
|||
)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
const RULES: &str = r#""fallback_generalizations":{"rules":[{"name":"family","pattern":"^new-","model_info":{"mode":"chat"}}]}"#;
|
||||
|
||||
#[rstest]
|
||||
#[case::clean(format!(r#"{{{RULES},"m":{{"litellm_provider":"test"}}}}"#), vec![])]
|
||||
#[case::not_json("{".to_owned(), vec![("", "json")])]
|
||||
#[case::row_error(format!(r#"{{{RULES},"m":{{}}}}"#), vec![("/m/litellm_provider", "missing_provider")])]
|
||||
#[case::unread_key(
|
||||
format!(r#"{{{RULES},"m":{{"litellm_provider":"test","typo_cost":1}}}}"#),
|
||||
vec![("/m/typo_cost", "unread_field")]
|
||||
)]
|
||||
#[case::empty_model_key(
|
||||
format!(r#"{{{RULES},"":{{"litellm_provider":"test","future_field":true}}}}"#),
|
||||
vec![("//future_field", "unread_field")]
|
||||
)]
|
||||
#[case::unread_key_on_invalid_row(
|
||||
format!(r#"{{{RULES},"m":{{"typo_cost":1}}}}"#),
|
||||
vec![("/m/litellm_provider", "missing_provider"), ("/m/typo_cost", "unread_field")]
|
||||
)]
|
||||
#[case::escaped_unread_key_on_invalid_row(
|
||||
format!(r#"{{{RULES},"open/router~v1":{{"typo/cost~":1}}}}"#),
|
||||
vec![("/open~1router~0v1/litellm_provider", "missing_provider"), ("/open~1router~0v1/typo~1cost~0", "unread_field")]
|
||||
)]
|
||||
#[case::no_rules(
|
||||
r#"{"m":{"litellm_provider":"test"}}"#.to_owned(),
|
||||
vec![("/fallback_generalizations", "no_fallback_rules")]
|
||||
)]
|
||||
#[case::empty_rules(
|
||||
r#"{"fallback_generalizations":{"rules":[]},"m":{"litellm_provider":"test"}}"#.to_owned(),
|
||||
vec![("/fallback_generalizations", "no_fallback_rules")]
|
||||
)]
|
||||
fn problems_list_what_a_cost_map_change_must_fix(
|
||||
#[case] body: String,
|
||||
#[case] expected: Vec<(&str, &str)>,
|
||||
) {
|
||||
let problems: Vec<_> = Catalog::problems(body.as_bytes())
|
||||
.iter()
|
||||
.map(|problem| (problem.path().to_owned(), <&str>::from(problem.kind())))
|
||||
.collect();
|
||||
let expected: Vec<_> = expected
|
||||
.into_iter()
|
||||
.map(|(path, kind)| (path.to_owned(), kind))
|
||||
.collect();
|
||||
assert_eq!(problems, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn runtime_parsing_preserves_custom_fields_and_recovers_valid_rows() {
|
||||
let body = br#"{
|
||||
"good":{"litellm_provider":"test","custom_metadata":true},
|
||||
"bad":{"litellm_provider":"test","input_cost_per_token":"wrong"}
|
||||
}"#;
|
||||
let loaded = Catalog::parse(body).unwrap();
|
||||
assert_eq!(loaded.errors.len(), 1);
|
||||
assert_eq!(loaded.errors[0].path(), "/bad/input_cost_per_token");
|
||||
let (catalog, rejected) = loaded.value.without_rejected(&loaded.errors);
|
||||
assert_eq!(rejected.len(), 1);
|
||||
assert!(catalog.lookup("bad").is_none());
|
||||
let good = catalog.lookup("good").unwrap();
|
||||
assert!(good.entry.model().is_some());
|
||||
assert_eq!(
|
||||
good.entry.field("custom_metadata"),
|
||||
Some(&Value::Bool(true))
|
||||
);
|
||||
assert!(
|
||||
Catalog::problems(body)
|
||||
.iter()
|
||||
.any(|problem| problem.kind() == crate::ErrorKind::UnreadField)
|
||||
);
|
||||
assert!(
|
||||
Catalog::problems(body)
|
||||
.iter()
|
||||
.any(|problem| problem.kind() == crate::ErrorKind::NoFallbackRules)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -109,6 +109,12 @@ pub(crate) enum Repr {
|
|||
#[source]
|
||||
cause: Option<DateCause>,
|
||||
},
|
||||
#[error("not a field litellm reads")]
|
||||
#[strum_discriminants(strum(serialize = "unread_field"))]
|
||||
UnreadField,
|
||||
#[error("has no rules")]
|
||||
#[strum_discriminants(strum(serialize = "no_fallback_rules"))]
|
||||
NoFallbackRules,
|
||||
}
|
||||
|
||||
/// What a problem takes out of the served map
|
||||
|
|
@ -148,7 +154,7 @@ impl Error {
|
|||
| Repr::NoOffPeakHours
|
||||
| Repr::EmptyList
|
||||
| Repr::InvalidDate { .. } => Rejects::Row,
|
||||
Repr::MissingProvider => Rejects::Nothing,
|
||||
Repr::MissingProvider | Repr::UnreadField | Repr::NoFallbackRules => Rejects::Nothing,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -334,6 +340,8 @@ mod tests {
|
|||
#[case::no_off_peak_hours(Repr::NoOffPeakHours, Rejects::Row)]
|
||||
#[case::empty_list(Repr::EmptyList, Rejects::Row)]
|
||||
#[case::invalid_date(Repr::InvalidDate { cause: None }, Rejects::Row)]
|
||||
#[case::unread_field(Repr::UnreadField, Rejects::Nothing)]
|
||||
#[case::no_fallback_rules(Repr::NoFallbackRules, Rejects::Nothing)]
|
||||
fn policy_classifies_each_error(#[case] error: Repr, #[case] expected: Rejects) {
|
||||
assert_eq!(Error::from(error).rejects(), expected);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -366,18 +366,14 @@ fn written_is_the_document_in_order_including_sections_and_bad_entries() {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_bundled_registry_parses_with_no_problems_and_every_key_understood() {
|
||||
let loaded = Catalog::parse(BUNDLED_REGISTRY).unwrap();
|
||||
let catalog = loaded.value;
|
||||
assert!(loaded.errors.is_empty(), "{:#?}", loaded.errors);
|
||||
let unrecognized: Vec<_> = catalog.unrecognized().collect();
|
||||
assert!(unrecognized.is_empty(), "{unrecognized:#?}");
|
||||
assert!(
|
||||
catalog
|
||||
.fallback_rules()
|
||||
.is_some_and(|rules| !rules.is_empty())
|
||||
);
|
||||
fn the_bundled_registry_has_no_problems() {
|
||||
let problems: Vec<_> = Catalog::problems(BUNDLED_REGISTRY)
|
||||
.iter()
|
||||
.map(report)
|
||||
.collect();
|
||||
assert!(problems.is_empty(), "{problems:#?}");
|
||||
const REFERENCE_MODEL_COUNT: usize = 4303;
|
||||
let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value;
|
||||
assert!(catalog.model_count() as f64 >= REFERENCE_MODEL_COUNT as f64 * 0.9);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -106,6 +106,18 @@ GIL handling to `litellm-host-python`.
|
|||
|
||||
## Data Handling
|
||||
|
||||
- Reuse large immutable bundled data and parsed catalogs across calls, including in
|
||||
replacement APIs. Share Rust storage and cache repeated Python projections instead
|
||||
of reparsing, serializing or copying the full payload on every access
|
||||
- Cache owned Python objects with `PyOnceLock<Py<T>>`, then bind and clone references
|
||||
while attached. Do not use blocking `OnceLock::get_or_init` while attached
|
||||
([PyO3 initialization](https://pyo3.rs/v0.29.2/faq.html),
|
||||
[reference ownership](https://pyo3.rs/v0.29.2/types.html))
|
||||
- Never reuse Python objects across interpreters or interpreter reinitialization.
|
||||
Once-only caches are for immutable data. Reloadable/custom data needs versioned
|
||||
entries or invalidation, and caller-mutable Python containers must not be shared
|
||||
- Test unchanged contents and allocation reuse, not just equal values. Keep this
|
||||
coverage when replacing a bundled-data accessor with a catalog API
|
||||
- OCR payloads can contain personal data and large base64 images. Do not log
|
||||
payloads or provider responses.
|
||||
- Avoid copying large payloads more than needed. The current JSON round-trip is
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ litellm-secrets-types.workspace = true
|
|||
litellm-llms-types.workspace = true
|
||||
litellm-router-types.workspace = true
|
||||
litellm-host-python.workspace = true
|
||||
litellm-model-catalog.workspace = true
|
||||
litellm-token-counter = { path = "../token-counter", default-features = false }
|
||||
pyo3.workspace = true
|
||||
prost.workspace = true
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ mod execution;
|
|||
mod http;
|
||||
mod logger;
|
||||
mod marshal;
|
||||
mod model_registry;
|
||||
mod preflight;
|
||||
mod python_settings;
|
||||
mod routes;
|
||||
|
|
@ -30,6 +31,8 @@ mod _native {
|
|||
#[pymodule_export]
|
||||
use crate::logger::NativeDiagnosticProcessor;
|
||||
#[pymodule_export]
|
||||
use crate::model_registry::bundled_model_cost_map;
|
||||
#[pymodule_export]
|
||||
use crate::routes::audio_transcription::{atranscription, transcription};
|
||||
#[pymodule_export]
|
||||
use crate::routes::chat_completions::{acompletion, completion};
|
||||
|
|
@ -110,6 +113,7 @@ mod tests {
|
|||
"NativeClickHouseSpendConfig",
|
||||
"NativeClickHouseSpendStorage",
|
||||
"trace_encode_error",
|
||||
"bundled_model_cost_map",
|
||||
"TokenCounter",
|
||||
"Tokenizer",
|
||||
"gil_stats",
|
||||
|
|
|
|||
30
litellm-rust/crates/python-bridge/src/model_registry.rs
Normal file
30
litellm-rust/crates/python-bridge/src/model_registry.rs
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
use pyo3::prelude::*;
|
||||
use pyo3::sync::PyOnceLock;
|
||||
use pyo3::types::PyBytes;
|
||||
|
||||
#[pyfunction]
|
||||
pub fn bundled_model_cost_map(py: Python<'_>) -> Bound<'_, PyBytes> {
|
||||
static BYTES: PyOnceLock<Py<PyBytes>> = PyOnceLock::new();
|
||||
BYTES
|
||||
.get_or_init(py, || {
|
||||
PyBytes::new(py, litellm_model_catalog::BUNDLED_REGISTRY).unbind()
|
||||
})
|
||||
.bind(py)
|
||||
.clone()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[rstest::rstest]
|
||||
fn repeated_calls_reuse_the_bundled_registry_bytes() {
|
||||
Python::initialize();
|
||||
let first = Python::attach(|py| bundled_model_cost_map(py).unbind());
|
||||
Python::attach(|py| {
|
||||
let second = bundled_model_cost_map(py);
|
||||
assert!(first.bind(py).is(&second));
|
||||
assert_eq!(second.as_bytes(), litellm_model_catalog::BUNDLED_REGISTRY);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -520,7 +520,7 @@ aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings
|
|||
disable_aiohttp_transport: bool = False # Set this to true to use httpx instead
|
||||
disable_aiohttp_trust_env: bool = False # When False, aiohttp will respect HTTP(S)_PROXY env vars
|
||||
force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6.
|
||||
http2: bool = False
|
||||
http2: bool | None = None
|
||||
network_mock: bool = False # When True, use mock transport — no real network calls
|
||||
|
||||
####### STOP SEQUENCE LIMIT #######
|
||||
|
|
|
|||
|
|
@ -57,7 +57,7 @@ def resolve_embedding_max_input_tokens(
|
|||
return configured_max_input_tokens
|
||||
if router is None:
|
||||
return None
|
||||
deployment_max_input_tokens, _ = router.get_configured_token_limits(embedding_model)
|
||||
deployment_max_input_tokens, _ = router.registry().group(embedding_model, field="token_limits")
|
||||
return deployment_max_input_tokens
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -830,6 +830,13 @@ class LLMCachingHandler:
|
|||
if is_internal_call.get():
|
||||
return
|
||||
|
||||
if logging_obj.defer_async_logging and not logging_obj.stream:
|
||||
from litellm.utils import dispatch_success_logging
|
||||
|
||||
logging_obj.model_call_details["cache_hit"] = cache_hit # rebind-ok: capture cache status before deferring
|
||||
dispatch_success_logging(logging_obj, cached_result, start_time, end_time, False, False)
|
||||
return
|
||||
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
|
|
|
|||
|
|
@ -29,6 +29,9 @@ it does not belong to. Queued clients are therefore bucketed by what it takes to
|
|||
close them, and each bucket is ordered by deadline, so a reap walks the entries
|
||||
that are due rather than the whole queue.
|
||||
|
||||
A bucket whose loop has been collected is dropped, and its clients are left to the
|
||||
collector, as they are when the queue is full.
|
||||
|
||||
The queue holds its clients weakly, so waiting out a grace window never keeps
|
||||
alive anything the collector would have reclaimed first.
|
||||
"""
|
||||
|
|
@ -36,6 +39,7 @@ alive anything the collector would have reclaimed first.
|
|||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import itertools
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
|
|
@ -44,6 +48,8 @@ from collections.abc import Awaitable, Callable, Iterator
|
|||
from dataclasses import dataclass, replace
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import (
|
||||
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
|
||||
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
|
||||
|
|
@ -53,6 +59,9 @@ _CLOSABLE_ANYWHERE: Final = "closable-anywhere"
|
|||
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
|
||||
|
||||
_BucketKey = str | int
|
||||
_LOOP_TOKENS: Final[weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, int]] = weakref.WeakKeyDictionary()
|
||||
_LOOP_TOKEN_COUNTER: Final = itertools.count()
|
||||
_LOOP_TOKEN_LOCK: Final = threading.Lock()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -63,12 +72,12 @@ class _PendingClose:
|
|||
collector would otherwise have reclaimed first.
|
||||
|
||||
``needs_loop`` is set for a client whose close is a coroutine; those can only
|
||||
be closed from the event loop they were evicted on, recorded in ``loop_id``.
|
||||
be closed from the event loop they were evicted on, recorded in ``loop_token``.
|
||||
A client that closes synchronously carries neither constraint.
|
||||
"""
|
||||
|
||||
client_ref: "weakref.ref[object]"
|
||||
loop_id: int | None
|
||||
loop_token: int | None
|
||||
needs_loop: bool
|
||||
close_after: float
|
||||
|
||||
|
|
@ -77,16 +86,24 @@ def _bucket_key(pending: _PendingClose) -> _BucketKey:
|
|||
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
|
||||
if not pending.needs_loop:
|
||||
return _CLOSABLE_ANYWHERE
|
||||
if pending.loop_id is None:
|
||||
if pending.loop_token is None:
|
||||
return _CLOSABLE_ON_ANY_LOOP
|
||||
return pending.loop_id
|
||||
return pending.loop_token
|
||||
|
||||
|
||||
def _running_loop_id() -> int | None:
|
||||
def running_loop_token() -> int | None:
|
||||
try:
|
||||
return id(asyncio.get_running_loop())
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
try:
|
||||
token: Final[int | None] = _LOOP_TOKENS.get(loop)
|
||||
if token is not None:
|
||||
return token
|
||||
with _LOOP_TOKEN_LOCK:
|
||||
return _LOOP_TOKENS.setdefault(loop, next(_LOOP_TOKEN_COUNTER))
|
||||
except TypeError:
|
||||
return id(loop)
|
||||
|
||||
|
||||
def _close_function(client: object) -> Callable[[], object] | None:
|
||||
|
|
@ -94,13 +111,24 @@ def _close_function(client: object) -> Callable[[], object] | None:
|
|||
return close_fn
|
||||
|
||||
|
||||
def _transport_of(client: object) -> object:
|
||||
"""The httpx transport behind an SDK wrapper, a litellm handler, or a bare client."""
|
||||
def _mounted_transports(holder: object) -> tuple[object, ...]:
|
||||
mounts: Final[object] = getattr(holder, "_mounts", None)
|
||||
if not isinstance(mounts, dict):
|
||||
return ()
|
||||
return tuple(
|
||||
mounted
|
||||
for mounted in mounts.values() # pyright: ignore[reportUnknownVariableType] # untyped httpx mount map
|
||||
if isinstance(mounted, (httpx.AsyncBaseTransport, httpx.BaseTransport))
|
||||
)
|
||||
|
||||
|
||||
def _transports_of(client: object) -> tuple[object, ...]:
|
||||
"""The httpx transports behind an SDK wrapper, a litellm handler, or a bare client, proxy mounts included."""
|
||||
for holder in (getattr(client, "_client", None), getattr(client, "client", None), client):
|
||||
transport: object = getattr(holder, "_transport", None)
|
||||
if transport is not None:
|
||||
return transport
|
||||
return None
|
||||
return (transport, *_mounted_transports(holder))
|
||||
return ()
|
||||
|
||||
|
||||
def _connection_is_idle(connection: object) -> bool:
|
||||
|
|
@ -123,6 +151,17 @@ def _pool_has_busy_connection(transport: object) -> bool | None:
|
|||
)
|
||||
|
||||
|
||||
def _transport_in_flight(transport: object) -> bool:
|
||||
spread_in_flight: Final[object] = getattr(transport, "has_request_in_flight", None)
|
||||
if callable(spread_in_flight):
|
||||
return bool(spread_in_flight())
|
||||
pooled_busy: Final = _pool_has_busy_connection(transport)
|
||||
if pooled_busy is not None:
|
||||
return pooled_busy
|
||||
session: Final[object] = getattr(transport, "client", None)
|
||||
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
|
||||
|
||||
|
||||
def _has_connection_in_flight(client: object) -> bool:
|
||||
"""Whether the client is servicing a request right now.
|
||||
|
||||
|
|
@ -138,12 +177,7 @@ def _has_connection_in_flight(client: object) -> bool:
|
|||
try:
|
||||
if getattr(getattr(client, "connection_pool", None), "_in_use_connections", None):
|
||||
return True
|
||||
transport: Final = _transport_of(client)
|
||||
pooled_busy: Final = _pool_has_busy_connection(transport)
|
||||
if pooled_busy is not None:
|
||||
return pooled_busy
|
||||
session: Final[object] = getattr(transport, "client", None)
|
||||
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
|
||||
return any(_transport_in_flight(transport) for transport in _transports_of(client))
|
||||
except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle
|
||||
return False
|
||||
|
||||
|
|
@ -167,6 +201,8 @@ class EvictedClientCloser:
|
|||
self._clock = clock
|
||||
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
|
||||
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
|
||||
self._dead_loop_tokens: deque[int] = deque() # mutable-ok: loop finalizers append without taking the queue lock
|
||||
self._watched_loop_tokens: set[int] = set() # mutable-ok: tracks tokens until their dead buckets are discarded
|
||||
self._pending_count = 0
|
||||
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
|
||||
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
|
||||
|
|
@ -197,13 +233,19 @@ class EvictedClientCloser:
|
|||
close_fn: Final = _close_function(client)
|
||||
if close_fn is None:
|
||||
return
|
||||
loop_token: Final = running_loop_token()
|
||||
needs_loop: Final = inspect.iscoroutinefunction(close_fn)
|
||||
if self._pending_count >= self._max_pending:
|
||||
return
|
||||
self._discard_dead_loops()
|
||||
if self._pending_count >= self._max_pending:
|
||||
return
|
||||
if needs_loop and loop_token is not None:
|
||||
self._watch_loop(loop_token)
|
||||
self._enqueue(
|
||||
_PendingClose(
|
||||
client_ref=weakref.ref(client),
|
||||
loop_id=_running_loop_id(),
|
||||
needs_loop=inspect.iscoroutinefunction(close_fn),
|
||||
loop_token=loop_token,
|
||||
needs_loop=needs_loop,
|
||||
close_after=self._clock() + self._grace_seconds,
|
||||
)
|
||||
)
|
||||
|
|
@ -217,7 +259,7 @@ class EvictedClientCloser:
|
|||
if not self._pending_count:
|
||||
return
|
||||
now: Final = self._clock()
|
||||
for pending in self._take_due(_running_loop_id(), now):
|
||||
for pending in self._take_due(running_loop_token(), now):
|
||||
client = pending.client_ref()
|
||||
if client is None:
|
||||
continue
|
||||
|
|
@ -230,6 +272,29 @@ class EvictedClientCloser:
|
|||
def pending_count(self) -> int:
|
||||
return self._pending_count
|
||||
|
||||
def _watch_loop(self, loop_token: int) -> None:
|
||||
if loop_token in self._watched_loop_tokens:
|
||||
return
|
||||
try:
|
||||
weakref.finalize(asyncio.get_running_loop(), self._dead_loop_tokens.append, loop_token)
|
||||
except TypeError:
|
||||
return
|
||||
self._watched_loop_tokens.add(loop_token)
|
||||
|
||||
def _discard_dead_loops(self) -> None:
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
|
||||
def _discard_dead_loops_locked(self) -> None:
|
||||
while self._dead_loop_tokens:
|
||||
self._discard_dead_loop_token_locked(self._dead_loop_tokens.popleft())
|
||||
|
||||
def _discard_dead_loop_token_locked(self, loop_token: int) -> None:
|
||||
self._watched_loop_tokens.discard(loop_token)
|
||||
bucket: Final = self._buckets.pop(loop_token, None)
|
||||
if bucket is not None:
|
||||
self._pending_count -= len(bucket)
|
||||
|
||||
def _enqueue(self, pending: _PendingClose) -> None:
|
||||
"""Append to the entry's bucket, dropping any dead entries it queues behind.
|
||||
|
||||
|
|
@ -238,6 +303,7 @@ class EvictedClientCloser:
|
|||
the front rather than having to be searched for.
|
||||
"""
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque())
|
||||
while bucket and bucket[0].client_ref() is None:
|
||||
bucket.popleft()
|
||||
|
|
@ -245,9 +311,12 @@ class EvictedClientCloser:
|
|||
bucket.append(pending)
|
||||
self._pending_count += 1
|
||||
|
||||
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
|
||||
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
|
||||
def _take_due(self, loop_token: int | None, now: float) -> tuple[_PendingClose, ...]:
|
||||
buckets = (
|
||||
(_CLOSABLE_ANYWHERE,) if loop_token is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_token)
|
||||
)
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
|
||||
|
||||
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:
|
||||
|
|
|
|||
|
|
@ -2,10 +2,9 @@
|
|||
Add the event loop to the cache key, to prevent event loop closed errors.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Final
|
||||
|
||||
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
|
||||
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer, running_loop_token
|
||||
from .in_memory_cache import InMemoryCache
|
||||
|
||||
|
||||
|
|
@ -43,24 +42,22 @@ class LLMClientCache(InMemoryCache):
|
|||
self.evicted_client_closer.schedule(evicted)
|
||||
self.evicted_client_closer.reap()
|
||||
|
||||
def update_cache_key_with_event_loop(self, key):
|
||||
def update_cache_key_with_event_loop(self, key: str) -> str:
|
||||
"""
|
||||
Add the event loop to the cache key, to prevent event loop closed errors.
|
||||
If none, use the key as is.
|
||||
"""
|
||||
try:
|
||||
event_loop: Final = asyncio.get_running_loop()
|
||||
stringified_event_loop: Final = str(id(event_loop))
|
||||
return f"{key}-{stringified_event_loop}"
|
||||
except RuntimeError: # handle no current running event loop
|
||||
return key
|
||||
token: Final = running_loop_token()
|
||||
return f"{key}-{token}" if token is not None else key
|
||||
|
||||
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
|
||||
def _scoped_key(self, key: str, loop_bound: bool) -> str:
|
||||
return self.update_cache_key_with_event_loop(key) if loop_bound else key
|
||||
|
||||
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, loop_bound: bool = True, **kwargs):
|
||||
"""``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
|
||||
if litellm_owned_client:
|
||||
self.evicted_client_closer.mark_owned(value)
|
||||
key = self.update_cache_key_with_event_loop(key)
|
||||
return super().set_cache(key, value, **kwargs)
|
||||
return super().set_cache(self._scoped_key(key, loop_bound), value, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
|
||||
if litellm_owned_client:
|
||||
|
|
@ -68,13 +65,13 @@ class LLMClientCache(InMemoryCache):
|
|||
key = self.update_cache_key_with_event_loop(key)
|
||||
return await super().async_set_cache(key, value, **kwargs)
|
||||
|
||||
def get_cache(self, key, **kwargs):
|
||||
key = self.update_cache_key_with_event_loop(key)
|
||||
def get_cache(self, key: str, loop_bound: bool = True, **kwargs):
|
||||
scoped_key: Final = self._scoped_key(key, loop_bound)
|
||||
self.evicted_client_closer.reap()
|
||||
|
||||
return super().get_cache(key, **kwargs)
|
||||
return super().get_cache(scoped_key, **kwargs)
|
||||
|
||||
async def async_get_cache(self, key, **kwargs):
|
||||
async def async_get_cache(self, key: str, **kwargs):
|
||||
key = self.update_cache_key_with_event_loop(key)
|
||||
|
||||
return await super().async_get_cache(key, **kwargs)
|
||||
|
|
|
|||
|
|
@ -139,6 +139,7 @@ class AuthenticationError(openai.AuthenticationError):
|
|||
litellm_debug_info: str | None = None,
|
||||
max_retries: int | None = None,
|
||||
num_retries: int | None = None,
|
||||
body: Mapping[str, object] | None = None,
|
||||
):
|
||||
self.status_code = 401
|
||||
self.message = f"litellm.AuthenticationError: {message}"
|
||||
|
|
@ -152,7 +153,7 @@ class AuthenticationError(openai.AuthenticationError):
|
|||
request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object
|
||||
)
|
||||
super().__init__(
|
||||
self.message, response=self.response, body=None
|
||||
self.message, response=self.response, body=body
|
||||
) # Call the base class constructor with the parameters it needs
|
||||
|
||||
def __str__(self):
|
||||
|
|
@ -189,6 +190,7 @@ class NotFoundError(openai.NotFoundError):
|
|||
litellm_debug_info: str | None = None,
|
||||
max_retries: int | None = None,
|
||||
num_retries: int | None = None,
|
||||
body: Mapping[str, object] | None = None,
|
||||
):
|
||||
self.status_code = 404
|
||||
self.message = f"litellm.NotFoundError: {message}"
|
||||
|
|
@ -202,7 +204,7 @@ class NotFoundError(openai.NotFoundError):
|
|||
request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object
|
||||
)
|
||||
super().__init__(
|
||||
self.message, response=self.response, body=None
|
||||
self.message, response=self.response, body=body
|
||||
) # Call the base class constructor with the parameters it needs
|
||||
|
||||
def __str__(self):
|
||||
|
|
@ -321,6 +323,17 @@ class UnsupportedEndpointError(BadRequestError):
|
|||
pass
|
||||
|
||||
|
||||
class EndpointNotSupportedByModelError(BadRequestError):
|
||||
def __init__(self, message: str, model: str, llm_provider: str = "") -> None:
|
||||
super().__init__( # pyright: ignore[reportUnknownMemberType] # inherited public exception constructor is untyped
|
||||
message=message, model=model, llm_provider=llm_provider
|
||||
)
|
||||
self.message = message
|
||||
self.type = "invalid_request_error"
|
||||
self.param = "model"
|
||||
self.code = "400"
|
||||
|
||||
|
||||
class UnprocessableEntityError(openai.UnprocessableEntityError):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -68,6 +68,11 @@ from litellm.utils import (
|
|||
)
|
||||
|
||||
_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
|
||||
_PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE: Final = (
|
||||
LlmProviders.DASHSCOPE,
|
||||
LlmProviders.QWENCLOUD,
|
||||
LlmProviders.QWEN_AI_PLATFORM,
|
||||
)
|
||||
|
||||
# Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls
|
||||
_ImageEditRequestUtils_cache: type["ImageEditRequestUtils"] | None = None
|
||||
|
|
@ -271,6 +276,7 @@ def image_generation(
|
|||
"""
|
||||
try:
|
||||
args: Final = locals()
|
||||
requested_api_base: Final = api_base
|
||||
aimg_generation: Final = kwargs.get("aimg_generation", False)
|
||||
litellm_call_id: Final = kwargs.get("litellm_call_id", None)
|
||||
logger_fn: Final = kwargs.get("logger_fn", None)
|
||||
|
|
@ -468,7 +474,11 @@ def image_generation(
|
|||
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
|
||||
|
||||
# Resolve api_base from litellm.api_base if not explicitly provided
|
||||
_api_base: Final = api_base or litellm.api_base
|
||||
_api_base: Final = (
|
||||
requested_api_base
|
||||
if custom_llm_provider in _PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE
|
||||
else api_base or litellm.api_base
|
||||
)
|
||||
litellm_params_dict["api_base"] = _api_base
|
||||
|
||||
return llm_http_handler.image_generation_handler(
|
||||
|
|
|
|||
|
|
@ -349,7 +349,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
except Exception:
|
||||
return 0
|
||||
|
||||
async def send_daily_reports(self, router) -> bool:
|
||||
async def send_daily_reports(self, router: Router) -> bool:
|
||||
"""
|
||||
Send a daily report on:
|
||||
- Top 5 deployments with most failed requests
|
||||
|
|
@ -365,7 +365,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
False -> if not sent
|
||||
"""
|
||||
|
||||
ids: Final = router.get_model_ids()
|
||||
ids: Final = router.registry().names(kind="ids")
|
||||
|
||||
# get keys
|
||||
failed_request_keys: Final = [f"{id}:{SlackAlertingCacheKeys.failed_requests_key.value}" for id in ids]
|
||||
|
|
@ -430,7 +430,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
message += "\tNone\n"
|
||||
for i in range(len(top_5_failed)):
|
||||
key = failed_request_keys[top_5_failed[i]].split(":")[0]
|
||||
_deployment = router.get_model_info(key)
|
||||
_deployment = router.registry().deployment_dict(key)
|
||||
if isinstance(_deployment, dict):
|
||||
deployment_name = _deployment["litellm_params"].get("model", "")
|
||||
else:
|
||||
|
|
@ -452,7 +452,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
message += "\tNone\n"
|
||||
for i in range(len(top_5_slowest)):
|
||||
key = latency_keys[top_5_slowest[i]].split(":")[0]
|
||||
_deployment = router.get_model_info(key)
|
||||
_deployment = router.registry().deployment_dict(key)
|
||||
if _deployment is not None:
|
||||
deployment_name = _deployment["litellm_params"].get("model", "")
|
||||
else:
|
||||
|
|
@ -802,7 +802,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
if self.llm_router is None:
|
||||
return
|
||||
|
||||
deployment: Final = self.llm_router.get_deployment(model_id=deployment_id)
|
||||
deployment: Final = self.llm_router.registry().deployment(deployment_id)
|
||||
|
||||
if deployment is None:
|
||||
return
|
||||
|
|
@ -953,7 +953,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
return
|
||||
|
||||
### EXTRACT MODEL DETAILS ###
|
||||
deployment: Final = self.llm_router.get_deployment(model_id=deployment_id)
|
||||
deployment: Final = self.llm_router.registry().deployment(deployment_id)
|
||||
if deployment is None:
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -102,18 +102,14 @@ class AgentOps(OpenTelemetry):
|
|||
"Connection": "keep-alive",
|
||||
}
|
||||
|
||||
client: Final = get_httpx_client()
|
||||
try:
|
||||
response: Final = client.post(
|
||||
url=auth_endpoint,
|
||||
headers=headers,
|
||||
json={"api_key": api_key},
|
||||
timeout=10,
|
||||
)
|
||||
response: Final = get_httpx_client().post(
|
||||
url=auth_endpoint,
|
||||
headers=headers,
|
||||
json={"api_key": api_key},
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch auth token: {response.text}")
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch auth token: {response.text}")
|
||||
|
||||
return _DECODED_JSON.validate_python(response.json())
|
||||
finally:
|
||||
client.close()
|
||||
return _DECODED_JSON.validate_python(response.json())
|
||||
|
|
|
|||
|
|
@ -35,6 +35,8 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
Span = _Span
|
||||
else:
|
||||
Span = Any
|
||||
|
|
@ -126,6 +128,26 @@ def _starts_its_own_trace(span: object, kwargs: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _a_destination_records_the_call(kwargs: Mapping[str, object]) -> bool:
|
||||
"""Whether a key's or team's ``langfuse_otel`` destination gets this call from the OTel v2 fan-out, which records
|
||||
every call but one whose logging payload shows it never reached a provider."""
|
||||
from litellm.integrations.otel.model.metadata import LLMCallEvent
|
||||
|
||||
if not _a_destination_carries_the_trace():
|
||||
return False
|
||||
call: Final = LLMCallEvent.from_dict(kwargs)
|
||||
return call.payload is None or call.reached_a_provider
|
||||
|
||||
|
||||
def _a_destination_carries_the_trace() -> bool:
|
||||
"""Whether this request resolved a key's or team's ``langfuse_otel`` destination, whose fan-out sends it the whole
|
||||
trace."""
|
||||
from litellm.integrations.otel.model.config import ExporterOwner
|
||||
from litellm.integrations.otel.plumbing.context import destination_backends
|
||||
|
||||
return ExporterOwner.LANGFUSE_OTEL in destination_backends()
|
||||
|
||||
|
||||
class LangfuseOtelLogger(OpenTelemetry):
|
||||
def __init__(self, config=None, *args, **kwargs):
|
||||
# Prevent LangfuseOtelLogger from modifying global environment variables by constructing config manually
|
||||
|
|
@ -541,6 +563,27 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
"""
|
||||
return None
|
||||
|
||||
def _emit_once(
|
||||
self,
|
||||
kwargs: dict[str, object], # mutable-ok: the v1 dedupe marker is written into the request's kwargs
|
||||
*scope: object,
|
||||
) -> bool:
|
||||
"""Whether to emit this span: once per request, and never for a call a Langfuse destination records."""
|
||||
return not _a_destination_records_the_call(kwargs) and super()._emit_once(kwargs, *scope) # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object], # mutable-ok: the v1 hook contract
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
"""Record the proxy failure of a request no Langfuse destination carries."""
|
||||
if not _a_destination_carries_the_trace():
|
||||
await super().async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
|
||||
request_data, original_exception, user_api_key_dict, traceback_str
|
||||
)
|
||||
|
||||
async def async_service_success_hook(self, *args, **kwargs):
|
||||
"""
|
||||
Langfuse should not receive service success logs.
|
||||
|
|
|
|||
|
|
@ -13,9 +13,9 @@ from opentelemetry.trace.status import Status, StatusCode
|
|||
|
||||
from litellm.integrations.otel.mappers import resolve_mappers
|
||||
from litellm.integrations.otel.mappers.base import AttributeMapper, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES
|
||||
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES, LangfuseMapper
|
||||
from litellm.integrations.otel.mappers.openinference import fit_indexed_messages
|
||||
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
|
|
@ -34,6 +34,7 @@ from litellm.integrations.otel.model.spans import (
|
|||
mcp_tool_call_span_name,
|
||||
service_span_name,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.context import destination_backends
|
||||
from litellm.integrations.otel.plumbing.events import GenAIEventRecorder
|
||||
from litellm.integrations.otel.plumbing.providers import to_otel_span_kind
|
||||
|
||||
|
|
@ -168,6 +169,9 @@ class SpanEmitter:
|
|||
self._mappers: list[AttributeMapper] = (
|
||||
list(mappers) if mappers is not None else resolve_mappers(config.mapper_names)
|
||||
)
|
||||
self._langfuse: Final = (
|
||||
None if any(isinstance(mapper, LangfuseMapper) for mapper in self._mappers) else LangfuseMapper()
|
||||
)
|
||||
# Bounded LRU (ordered by insertion / most-recent touch). Storing keys
|
||||
# only — the value is unused — so it behaves like a capped set.
|
||||
self._emitted: OrderedDict[tuple[str, SpanRole], None] = OrderedDict()
|
||||
|
|
@ -295,17 +299,27 @@ class SpanEmitter:
|
|||
)
|
||||
else None
|
||||
)
|
||||
mapped: Final = MappingProxyType(
|
||||
{key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
|
||||
configured: Final = {key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
|
||||
for_langfuse: Final = (
|
||||
dict(
|
||||
sorted(
|
||||
self._langfuse.map(data).items(),
|
||||
key=lambda item: item[0] not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
|
||||
)
|
||||
)
|
||||
if self._langfuse is not None and ExporterOwner.LANGFUSE_OTEL in destination_backends()
|
||||
else _NO_ATTRIBUTES
|
||||
)
|
||||
mapped: Final = MappingProxyType({**configured, **for_langfuse})
|
||||
stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES
|
||||
reserved: Final = len(stamped_later.keys() - mapped.keys())
|
||||
overwrites: Final = carried_keys(span) & mapped.keys()
|
||||
# The Langfuse request metadata sheds first, so a span at its attribute limit loses nothing else to it.
|
||||
# The Langfuse request metadata, and the Langfuse vocabulary only a destination asked for, shed first,
|
||||
# so a span at its attribute limit loses nothing else to them.
|
||||
fitted: Final = fit_indexed_messages(
|
||||
MappingProxyType({key: value for key, value in mapped.items() if key not in overwrites}),
|
||||
attribute_budget(span, reserved),
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES.union(for_langfuse),
|
||||
)
|
||||
for key, value in mapped.items():
|
||||
if key in overwrites or key in fitted:
|
||||
|
|
|
|||
|
|
@ -198,6 +198,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
tracer_provider: TracerProvider | None = None,
|
||||
logger_provider: LoggerProvider | None = None,
|
||||
meter_provider: "MeterProvider | None" = None,
|
||||
console_default: bool = False,
|
||||
**kwargs: Unpack[_CustomLoggerOptions],
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
|
|
@ -206,7 +207,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
self._tracer_provider: TracerProvider = (
|
||||
tracer_provider
|
||||
if tracer_provider is not None
|
||||
else build_tracer_provider(self.config, tenant_overrides=True)
|
||||
else build_tracer_provider(self.config, tenant_overrides=True, console_default=console_default)
|
||||
)
|
||||
self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME)
|
||||
self._metrics_recorder = self._init_metrics(meter_provider)
|
||||
|
|
@ -217,7 +218,9 @@ class OpenTelemetryV2(CustomLogger):
|
|||
mappers=resolve_mappers(self.config.mapper_names),
|
||||
event_recorder=self._init_events(logger_provider),
|
||||
)
|
||||
self._tenant_tracers = TenantTracerCache(self.config, callback_name, LITELLM_TRACER_NAME)
|
||||
self._tenant_tracers = TenantTracerCache(
|
||||
self.config, callback_name, LITELLM_TRACER_NAME, console_default=console_default
|
||||
)
|
||||
self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict()
|
||||
self._init_otel_logger_on_litellm_proxy()
|
||||
|
||||
|
|
@ -541,9 +544,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
# gate rejection carries ``is_no_upstream_call`` and gets no span. A cache
|
||||
# hit never hands off to a provider, so ``cache_hit`` substitutes for the
|
||||
# upstream stamp, since the payload proves the call completed.
|
||||
if carrier is None and (
|
||||
call.is_no_upstream_call or call.payload is None or (not call.upstream_started and not call.cache_hit)
|
||||
):
|
||||
if carrier is None and not call.reached_a_provider:
|
||||
return None
|
||||
try:
|
||||
return self._finish_carrier(carrier, call, start_time, end_time)
|
||||
|
|
@ -1091,6 +1092,7 @@ def build_otel_v2_logger(
|
|||
logger_provider: LoggerProvider | None = None,
|
||||
meter_provider: "MeterProvider | None" = None,
|
||||
settings: Mapping[str, object] = MappingProxyType({}),
|
||||
console_default: bool = False,
|
||||
) -> OpenTelemetryV2:
|
||||
return _logger_class(config)(
|
||||
config=config,
|
||||
|
|
@ -1098,6 +1100,7 @@ def build_otel_v2_logger(
|
|||
tracer_provider=tracer_provider,
|
||||
logger_provider=logger_provider,
|
||||
meter_provider=meter_provider,
|
||||
console_default=console_default,
|
||||
**settings,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -103,6 +103,18 @@ class LangfuseMapper:
|
|||
),
|
||||
}
|
||||
|
||||
#: Every attribute this mapper can write on a model call.
|
||||
LLM_CALL_ATTRIBUTES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
*_LLM_CALL_ATTRS,
|
||||
*_BLOB_ATTRS,
|
||||
LANGFUSE_TRACE_NAME,
|
||||
LANGFUSE_TRACE_USER_ID,
|
||||
LANGFUSE_TRACE_SESSION_ID,
|
||||
LANGFUSE_TRACE_TAGS,
|
||||
}
|
||||
)
|
||||
|
||||
def map(self, data: SpanData) -> AttributeMap:
|
||||
match data:
|
||||
case LLMCallSpanData():
|
||||
|
|
|
|||
|
|
@ -121,6 +121,17 @@ class ExporterSpec(LiteLLMBaseModel):
|
|||
)
|
||||
|
||||
|
||||
def is_unconfigured_placeholder(spec: ExporterSpec) -> bool:
|
||||
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
|
||||
|
||||
No field set is what says the operator asked for nothing: an exporter they did
|
||||
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
|
||||
and so does the gated spec ``credential_gated_exporters`` appends, which would
|
||||
otherwise eat itself when one preset layers onto another.
|
||||
"""
|
||||
return not spec.model_fields_set
|
||||
|
||||
|
||||
class _EnvWithoutBareExcludedServices(PydanticBaseSettingsSource):
|
||||
def __init__(self, settings_cls: type[BaseSettings], env_settings: PydanticBaseSettingsSource) -> None:
|
||||
super().__init__(settings_cls)
|
||||
|
|
|
|||
|
|
@ -272,6 +272,11 @@ class LLMCallEvent:
|
|||
trace: TraceControls
|
||||
session_id: str | None
|
||||
|
||||
@property
|
||||
def reached_a_provider(self) -> bool:
|
||||
"""Whether the call was handed to a provider or answered from the response cache."""
|
||||
return not self.is_no_upstream_call and self.payload is not None and (self.upstream_started or self.cache_hit)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent:
|
||||
raw_payload: Final = kwargs.get("standard_logging_object")
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ from litellm.integrations.otel.mappers.langfuse import (
|
|||
LANGFUSE_OBSERVATION_OUTPUT,
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
|
||||
LANGFUSE_TRACE_NAME,
|
||||
LangfuseMapper,
|
||||
)
|
||||
from litellm.integrations.otel.model.config import (
|
||||
CaptureMessageContent,
|
||||
|
|
@ -55,6 +56,7 @@ from litellm.integrations.otel.model.config import (
|
|||
OpenTelemetryV2Config,
|
||||
capture_message_content_from_env,
|
||||
captures_span_content,
|
||||
is_unconfigured_placeholder,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
DB,
|
||||
|
|
@ -490,19 +492,25 @@ def _is_litellm_span(span: ReadableSpan) -> bool:
|
|||
return scope is not None and scope.name == LITELLM_TRACER_NAME
|
||||
|
||||
|
||||
def without_langfuse_request_metadata(span: ReadableSpan) -> ReadableSpan:
|
||||
"""``span`` without the request metadata the Langfuse mapper writes for a Langfuse destination.
|
||||
def without_destination_langfuse_attributes(span: ReadableSpan, langfuse_mapped: bool) -> ReadableSpan:
|
||||
"""``span`` without the Langfuse attributes only a Langfuse destination asked for.
|
||||
|
||||
Only LiteLLM's own spans lose it: any other span keeps every key its own code set.
|
||||
That is the request metadata the Langfuse mapper writes and, on a model call written by
|
||||
a logger whose own mappers leave the Langfuse vocabulary out (``langfuse_mapped`` off),
|
||||
every Langfuse attribute, since that logger wrote them for a key's or team's Langfuse
|
||||
destination alone. Only LiteLLM's own spans lose them: any other span keeps every key
|
||||
its own code set.
|
||||
"""
|
||||
attributes: Final = span.attributes or _NO_ATTRIBUTES
|
||||
if not (_is_litellm_span(span) and any(key in LANGFUSE_REQUEST_METADATA_ATTRIBUTES for key in attributes)):
|
||||
dropped: Final = (
|
||||
LANGFUSE_REQUEST_METADATA_ATTRIBUTES
|
||||
if langfuse_mapped or not is_llm_call_span(span)
|
||||
else LangfuseMapper.LLM_CALL_ATTRIBUTES
|
||||
)
|
||||
if not (_is_litellm_span(span) and any(key in dropped for key in attributes)):
|
||||
return span
|
||||
return _CountedSpanView(
|
||||
span,
|
||||
MappingProxyType(
|
||||
{key: value for key, value in attributes.items() if key not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES}
|
||||
),
|
||||
span, MappingProxyType({key: value for key, value in attributes.items() if key not in dropped})
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -658,6 +666,7 @@ def _for_destination(
|
|||
span: ReadableSpan,
|
||||
destination: "OtelDestination",
|
||||
default_capture: str,
|
||||
langfuse_mapped: bool,
|
||||
) -> ReadableSpan:
|
||||
"""The view of ``span`` a tenant destination receives.
|
||||
|
||||
|
|
@ -672,13 +681,16 @@ def _for_destination(
|
|||
and the query string of the request URL, which can hold the same key. A
|
||||
destination whose effective setting disables span content loses the prompt and
|
||||
response content too, and a destination for any backend but Langfuse loses the
|
||||
request metadata the Langfuse mapper writes on LiteLLM's own spans. The span itself
|
||||
stays, so the tenant still gets the whole trace tree.
|
||||
Langfuse attributes only a Langfuse destination asked for (see
|
||||
:func:`without_destination_langfuse_attributes`). The span itself stays, so the
|
||||
tenant still gets the whole trace tree.
|
||||
"""
|
||||
extra: Final = destination.resource_attributes
|
||||
defaults: Final = destination.resource_defaults
|
||||
source: Final = (
|
||||
span if destination.callback_name == ExporterOwner.LANGFUSE_OTEL else without_langfuse_request_metadata(span)
|
||||
span
|
||||
if destination.callback_name == ExporterOwner.LANGFUSE_OTEL
|
||||
else without_destination_langfuse_attributes(span, langfuse_mapped)
|
||||
)
|
||||
attributes: Final = source.attributes or _NO_ATTRIBUTES
|
||||
database: Final = _is_database_span(attributes)
|
||||
|
|
@ -728,8 +740,10 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
default_capture_by_backend: Mapping[str, str] = MappingProxyType({}),
|
||||
pending_drains: int = _MAX_PENDING_DRAINS,
|
||||
drain_pool: _DrainPool | None = None,
|
||||
langfuse_mapped: bool = True,
|
||||
) -> None:
|
||||
self._operator_sinks: Final = operator_sinks
|
||||
self._langfuse_mapped: Final = langfuse_mapped
|
||||
self._excluded_db_systems: Final = excluded_db_systems
|
||||
self._default_capture: Final = default_capture
|
||||
self._default_capture_by_backend: Final = default_capture_by_backend
|
||||
|
|
@ -822,7 +836,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
try:
|
||||
processor.on_end(
|
||||
_scoped(
|
||||
_for_destination(span, destination, self._default_capture_for(destination)), destination.span_scope
|
||||
_for_destination(span, destination, self._default_capture_for(destination), self._langfuse_mapped),
|
||||
destination.span_scope,
|
||||
)
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span
|
||||
|
|
@ -1198,17 +1213,19 @@ class _OverriddenBackendFilter(SpanProcessor):
|
|||
return self._inner.force_flush(timeout_millis)
|
||||
|
||||
|
||||
class _LangfuseRequestMetadataFilter(SpanProcessor):
|
||||
"""Strip the Langfuse request metadata from LiteLLM's own spans before ``inner`` exports them."""
|
||||
class _DestinationLangfuseAttributeFilter(SpanProcessor):
|
||||
"""Strip the Langfuse attributes only a Langfuse destination asked for from LiteLLM's own spans before ``inner``
|
||||
exports them."""
|
||||
|
||||
def __init__(self, inner: SpanProcessor) -> None:
|
||||
def __init__(self, inner: SpanProcessor, langfuse_mapped: bool) -> None:
|
||||
self._inner: Final = inner
|
||||
self._langfuse_mapped: Final = langfuse_mapped
|
||||
|
||||
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
|
||||
self._inner.on_start(span, parent_context)
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
self._inner.on_end(without_langfuse_request_metadata(span))
|
||||
self._inner.on_end(without_destination_langfuse_attributes(span, self._langfuse_mapped))
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._inner.shutdown()
|
||||
|
|
@ -1491,13 +1508,18 @@ def _spec_processor(
|
|||
spec: ExporterSpec,
|
||||
use_simple_processor: bool | None,
|
||||
strip_message_content: bool,
|
||||
strip_langfuse_request_metadata: bool,
|
||||
strip_destination_langfuse_attributes: bool,
|
||||
langfuse_mapped: bool,
|
||||
) -> SpanProcessor:
|
||||
exporting: Final = _processor_for(
|
||||
_exporter_from_spec(spec),
|
||||
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
|
||||
)
|
||||
scoped: Final = _LangfuseRequestMetadataFilter(exporting) if strip_langfuse_request_metadata else exporting
|
||||
scoped: Final = (
|
||||
_DestinationLangfuseAttributeFilter(exporting, langfuse_mapped)
|
||||
if strip_destination_langfuse_attributes
|
||||
else exporting
|
||||
)
|
||||
return _MessageContentFilter(scoped) if strip_message_content else scoped
|
||||
|
||||
|
||||
|
|
@ -1530,6 +1552,7 @@ def build_tracer_provider(
|
|||
use_simple_processor: bool | None = None,
|
||||
tenant_overrides: bool = False,
|
||||
strip_message_content: bool = False,
|
||||
console_default: bool = False,
|
||||
) -> TracerProvider:
|
||||
"""Build the shared :class:`TracerProvider`.
|
||||
|
||||
|
|
@ -1548,11 +1571,16 @@ def build_tracer_provider(
|
|||
|
||||
``config.langfuse_span_scope`` narrows the exporter owned by ``langfuse_otel``
|
||||
alone; a collector or any other backend in the same config keeps the full tree.
|
||||
In a config that writes the Langfuse vocabulary, every exporter but ``langfuse_otel``'s
|
||||
goes without the request metadata the Langfuse mapper writes on LiteLLM's own spans.
|
||||
Every exporter but ``langfuse_otel``'s goes without the Langfuse attributes only a
|
||||
Langfuse destination asked for (see :func:`without_destination_langfuse_attributes`),
|
||||
which any logger writes for a request with a Langfuse destination.
|
||||
``strip_message_content`` removes prompt and response content before every
|
||||
exporter, for a routed provider whose spans can carry content another
|
||||
destination asked for.
|
||||
|
||||
The console spec ``config._normalize`` folds in when nothing was configured
|
||||
prints only under ``console_default``, which the ``otel`` callback passes: a
|
||||
vendor preset or the proxy's background global logger was never asked to print.
|
||||
"""
|
||||
provider: Final = TracerProvider(resource=build_resource(config))
|
||||
provider.sampler = _RequestExclusionSampler(provider.sampler)
|
||||
|
|
@ -1564,19 +1592,20 @@ def build_tracer_provider(
|
|||
provider.add_span_processor(_processor_for(exporter, use_simple_processor))
|
||||
return provider
|
||||
|
||||
langfuse_mapped: Final = _maps_langfuse(config)
|
||||
# ``config._normalize`` guarantees at least one spec (it folds the top-level
|
||||
# ``exporter``/``endpoint``/``headers`` fields in when ``exporters`` is empty).
|
||||
writes_langfuse: Final = "langfuse" in config.mapper_names or any(
|
||||
spec.owner is ExporterOwner.LANGFUSE_OTEL for spec in config.exporters
|
||||
)
|
||||
for spec in config.exporters:
|
||||
if spec.requires_headers and not spec.headers:
|
||||
continue
|
||||
if is_unconfigured_placeholder(spec) and not console_default:
|
||||
continue
|
||||
processor = _spec_processor(
|
||||
spec,
|
||||
use_simple_processor,
|
||||
strip_message_content,
|
||||
strip_langfuse_request_metadata=writes_langfuse and spec.owner is not ExporterOwner.LANGFUSE_OTEL,
|
||||
strip_destination_langfuse_attributes=spec.owner is not ExporterOwner.LANGFUSE_OTEL,
|
||||
langfuse_mapped=langfuse_mapped,
|
||||
)
|
||||
owner = spec.owner.value if tenant_overrides and spec.owner is not None else None
|
||||
scope = _operator_scope(config, spec)
|
||||
|
|
@ -1613,8 +1642,9 @@ def attach_tenant_fan_out(
|
|||
so exactly one fan-out lands. ``configs`` name the operator's own exporters, one
|
||||
config per v2 logger since each keeps its own provider and still writes its
|
||||
account, so an additive destination pointing at any of them is delivered once
|
||||
rather than twice. ``excluded_db_systems`` only filters what the fan-out
|
||||
delivers, never the operator's own exporters.
|
||||
rather than twice. The first is the config of the logger that owns ``provider``,
|
||||
whose model calls the fan-out carries. ``excluded_db_systems`` only filters what
|
||||
the fan-out delivers, never the operator's own exporters.
|
||||
"""
|
||||
with _FAN_OUT_ATTACH_LOCK:
|
||||
if any(isinstance(processor, TenantFanOutSpanProcessor) for processor in _attached_processors(provider)):
|
||||
|
|
@ -1625,10 +1655,16 @@ def attach_tenant_fan_out(
|
|||
excluded_db_systems=excluded_db_systems,
|
||||
default_capture=_operator_capture_mode(configs),
|
||||
default_capture_by_backend=MappingProxyType(dict(_owned_capture_modes(configs))),
|
||||
langfuse_mapped=any(_maps_langfuse(config) for config in configs[:1]),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _maps_langfuse(config: OpenTelemetryV2Config) -> bool:
|
||||
"""Whether a logger built from ``config`` writes the Langfuse vocabulary on every model call."""
|
||||
return "langfuse" in config.mapper_names
|
||||
|
||||
|
||||
def _operator_capture_mode(configs: Sequence[OpenTelemetryV2Config]) -> str:
|
||||
"""The mode of the operator's generic ``otel`` callback, so ``callback_settings.otel`` counts as the global default"""
|
||||
return next(
|
||||
|
|
|
|||
|
|
@ -163,10 +163,12 @@ class TenantTracerCache:
|
|||
config: OpenTelemetryV2Config,
|
||||
callback_name: str | None,
|
||||
tracer_name: str,
|
||||
console_default: bool = False,
|
||||
) -> None:
|
||||
self._config = config
|
||||
self._callback_name = callback_name
|
||||
self._tracer_name = tracer_name
|
||||
self._console_default = console_default
|
||||
# Guards the three mutable structures below: ``pre_call`` can run on
|
||||
# thread-pool workers concurrently with the event loop, so cache
|
||||
# updates, span counts, and retirement must be atomic.
|
||||
|
|
@ -284,6 +286,7 @@ class TenantTracerCache:
|
|||
built: Final = build_tracer_provider(
|
||||
self._routed_config(credential_headers, project_headers, endpoint, service_name),
|
||||
strip_message_content=not self._config.capture_span_content,
|
||||
console_default=self._console_default,
|
||||
)
|
||||
self._providers[cache_key] = built
|
||||
return built
|
||||
|
|
|
|||
|
|
@ -111,7 +111,7 @@ _REQUIRED_HEADERS_BY_CALLBACK: Final[Mapping[str, frozenset[str]]] = MappingProx
|
|||
_NO_ATTRS: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _arize_project() -> Mapping[str, str]:
|
||||
def _arize_project(_params: StandardCallbackDynamicParams) -> Mapping[str, str]:
|
||||
from litellm.integrations.arize.arize import ArizeLogger
|
||||
|
||||
project: Final = ArizeLogger.get_arize_config().project_name
|
||||
|
|
@ -124,10 +124,27 @@ def _arize_default_project() -> Mapping[str, str]:
|
|||
return MappingProxyType({"model_id": ARIZE_DEFAULT_PROJECT})
|
||||
|
||||
|
||||
#: Resource attributes the operator chose for a backend, applied to every export the
|
||||
#: team receives over the span's own resource: ``ARIZE_PROJECT_NAME`` for Arize.
|
||||
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str]]]] = MappingProxyType(
|
||||
{"arize": _arize_project}
|
||||
def _stands_in_for_the_legacy_langfuse_logger() -> bool:
|
||||
"""Whether a Langfuse destination replaces the older Langfuse logger's copy, which a proxy without Langfuse keys
|
||||
of its own sends for a key's or team's ``langfuse_otel``."""
|
||||
return not (os.environ.get("LANGFUSE_PUBLIC_KEY") and os.environ.get("LANGFUSE_SECRET_KEY"))
|
||||
|
||||
|
||||
def _langfuse_environment(params: StandardCallbackDynamicParams) -> Mapping[str, str]:
|
||||
"""The key's or team's ``langfuse_environment``, else ``LANGFUSE_TRACING_ENVIRONMENT``, where the older Langfuse
|
||||
logger wrote it."""
|
||||
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
|
||||
|
||||
if not _stands_in_for_the_legacy_langfuse_logger():
|
||||
return _NO_ATTRS
|
||||
environment: Final = params.get("langfuse_environment") or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
|
||||
return MappingProxyType({LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value: environment} if environment else {})
|
||||
|
||||
|
||||
#: Resource attributes a backend gets on every export the team receives, over the
|
||||
#: span's own resource: ``ARIZE_PROJECT_NAME`` for Arize, the Langfuse environment.
|
||||
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], Mapping[str, str]]]] = (
|
||||
MappingProxyType({"arize": _arize_project, "langfuse_otel": _langfuse_environment})
|
||||
)
|
||||
|
||||
#: Resource attributes a backend needs on every export to accept it, filled only where
|
||||
|
|
@ -137,12 +154,14 @@ _RESOURCE_DEFAULTS_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str
|
|||
)
|
||||
|
||||
|
||||
def _resource_attributes(callback_name: str, service_name: str | None) -> Mapping[str, str]:
|
||||
def _resource_attributes(
|
||||
callback_name: str, service_name: str | None, params: StandardCallbackDynamicParams
|
||||
) -> Mapping[str, str]:
|
||||
backend: Final = _RESOURCE_ATTRIBUTES_BY_CALLBACK.get(callback_name)
|
||||
return MappingProxyType(
|
||||
{
|
||||
**({"service.name": service_name} if service_name else {}),
|
||||
**(backend() if backend is not None else {}),
|
||||
**(backend(params) if backend is not None else {}),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -214,7 +233,7 @@ def destination_for(
|
|||
return OtelDestination(
|
||||
endpoint=endpoint,
|
||||
headers=MappingProxyType(dict(headers)),
|
||||
resource_attributes=_resource_attributes(callback_name, service_name),
|
||||
resource_attributes=_resource_attributes(callback_name, service_name, params),
|
||||
resource_defaults=_resource_defaults(callback_name),
|
||||
callback_name=callback_name,
|
||||
protocol=protocol,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
from collections.abc import Iterable
|
||||
from typing import Final
|
||||
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, is_unconfigured_placeholder
|
||||
|
||||
|
||||
def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]:
|
||||
|
|
@ -35,14 +35,3 @@ def credential_gated_exporters(
|
|||
*(spec for spec in exporters if not is_unconfigured_placeholder(spec)),
|
||||
ExporterSpec(owner=owner, requires_headers=True),
|
||||
)
|
||||
|
||||
|
||||
def is_unconfigured_placeholder(spec: "ExporterSpec") -> bool:
|
||||
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
|
||||
|
||||
No field set is what says the operator asked for nothing: an exporter they did
|
||||
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
|
||||
and so does the gated spec this module appends, which would otherwise eat itself
|
||||
when one preset layers onto another.
|
||||
"""
|
||||
return not spec.model_fields_set
|
||||
|
|
|
|||
|
|
@ -468,15 +468,16 @@ def _bounded_requested_model_label(requested_model: object, router_originated: b
|
|||
llm_router: Final = _get_proxy_llm_router()
|
||||
if llm_router is None:
|
||||
return requested_model if router_originated else UNRECOGNIZED_REQUESTED_MODEL_LABEL
|
||||
if llm_router.is_recognized_model(requested_model):
|
||||
registry: Final = llm_router.registry()
|
||||
if registry.resolve(requested_model, kind="recognized"):
|
||||
return requested_model
|
||||
if requested_model in llm_router.team_public_model_names:
|
||||
if requested_model in registry.names(kind="team_public"):
|
||||
return requested_model
|
||||
if llm_router.pattern_router.route(requested_model) is not None:
|
||||
if registry.patterns().route(requested_model) is not None:
|
||||
return requested_model
|
||||
if any(
|
||||
team_pattern_router.route(requested_model) is not None
|
||||
for team_pattern_router in llm_router.team_pattern_routers.values()
|
||||
for team_pattern_router in registry.patterns(team=True).values()
|
||||
):
|
||||
return requested_model
|
||||
return UNRECOGNIZED_REQUESTED_MODEL_LABEL
|
||||
|
|
@ -3410,7 +3411,7 @@ class PrometheusLogger(CustomLogger):
|
|||
OpenAI / Anthropic / Azure return remaining tokens/requests in response
|
||||
headers, but Bedrock and Vertex AI do not. This fallback computes
|
||||
``configured_limit - current_usage`` via
|
||||
``Router.get_remaining_model_group_usage`` so the same gauges are
|
||||
``Router.health().usage(model_group, kind="remaining")`` so the same gauges are
|
||||
emitted for every provider when tpm/rpm is configured on the
|
||||
deployment.
|
||||
"""
|
||||
|
|
@ -3437,7 +3438,7 @@ class PrometheusLogger(CustomLogger):
|
|||
return
|
||||
|
||||
try:
|
||||
remaining_usage: Final = await llm_router.get_remaining_model_group_usage(model_group)
|
||||
remaining_usage: Final = await llm_router.health().usage(model_group, kind="remaining")
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"Prometheus: get_remaining_model_group_usage failed for model_group=%s: %s",
|
||||
|
|
|
|||
|
|
@ -40,7 +40,6 @@ from litellm.litellm_core_utils.llm_judge import (
|
|||
)
|
||||
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
|
||||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.management_endpoints.auto_router_endpoints import ShadowEvalDirection
|
||||
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
|
||||
|
|
@ -790,9 +789,7 @@ class ActiveShadowEvalJob(LiteLLMBaseModel):
|
|||
def _canonical_group(router: "Router | None", model_group: str) -> str:
|
||||
"""A model group in the one spelling both a job's scope and a request's model compare
|
||||
under: an alias resolves to its target so the two never fail to match on spelling."""
|
||||
return (
|
||||
resolve_model_group_alias(router.model_group_alias, model_group) if router is not None else None
|
||||
) or model_group
|
||||
return (router.registry().resolve(model_group, kind="alias") if router is not None else None) or model_group
|
||||
|
||||
|
||||
def _scope_admits(router: "Router | None", job: "ActiveShadowEvalJob", model_group: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ if TYPE_CHECKING:
|
|||
from litellm.llms.base_llm.anthropic_messages.transformation import (
|
||||
BaseAnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
|
@ -132,7 +133,7 @@ class _SearchToolLitellmParams(TypedDict, total=False):
|
|||
|
||||
|
||||
class _SearchToolConfig(TypedDict, total=False):
|
||||
search_tool_name: str
|
||||
search_tool_name: ReadOnly[str]
|
||||
litellm_params: ReadOnly[_SearchToolLitellmParams | None]
|
||||
|
||||
|
||||
|
|
@ -1768,10 +1769,13 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
return None
|
||||
|
||||
def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None":
|
||||
if llm_router is None or not hasattr(llm_router, "search_tools"):
|
||||
def _select_search_tool_from_router(self, llm_router: "Router | None") -> "_SearchToolConfig | None":
|
||||
if llm_router is None:
|
||||
return None
|
||||
search_tools: Final = tuple(getattr(llm_router, "search_tools", None) or ())
|
||||
settings: Final = llm_router.settings()
|
||||
if not settings.is_set("search_tools"):
|
||||
return None
|
||||
search_tools: Final = tuple(settings.search_tools or ())
|
||||
return self._select_search_tool_from_list(search_tools=search_tools, source="router")
|
||||
|
||||
def _select_search_tool_from_list(
|
||||
|
|
|
|||
|
|
@ -413,6 +413,25 @@ def _map_openai_exception(
|
|||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str:
|
||||
invalid_request_status: Final = getattr(original_exception, "status_code", None)
|
||||
if invalid_request_status == 401:
|
||||
raise AuthenticationError(
|
||||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
if invalid_request_status == 404:
|
||||
raise NotFoundError(
|
||||
message=f"{exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
raise BadRequestError(
|
||||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
|
|
|
|||
|
|
@ -31,7 +31,13 @@ EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = Mapping
|
|||
)
|
||||
|
||||
|
||||
def native_health_check_mode(model: str, custom_llm_provider: str | None) -> Literal["anthropic_messages"] | None:
|
||||
def native_health_check_mode(
|
||||
model: str, custom_llm_provider: str | None
|
||||
) -> Literal["anthropic_messages", "realtime"] | None:
|
||||
if custom_llm_provider == "soniox":
|
||||
from litellm.llms.soniox.common_utils import soniox_health_check_mode
|
||||
|
||||
return soniox_health_check_mode(model)
|
||||
if custom_llm_provider != "bedrock_mantle":
|
||||
return None
|
||||
from litellm.llms.bedrock_mantle.common_utils import mantle_health_check_mode
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# What is this?
|
||||
## Common Utility file for Logging handler
|
||||
# Logging function -> log the exact model details + what's being sent | Non-Blocking
|
||||
import asyncio
|
||||
import copy
|
||||
import datetime
|
||||
import functools
|
||||
|
|
@ -96,6 +97,7 @@ from litellm.litellm_core_utils.internal_call_metadata import (
|
|||
MODEL_ACCESS_GROUP_METADATA_KEY,
|
||||
is_unbilled_non_inference_call,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost, PreparedCost
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
|
||||
bedrock_guardrail_trace,
|
||||
cost_breakdown_with_guardrail,
|
||||
|
|
@ -813,6 +815,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
# Initialize cost breakdown field
|
||||
self.cost_breakdown: CostBreakdown | None = None
|
||||
self.finalized_cost: FinalizedCost | None = None
|
||||
self.prepared_cost: PreparedCost | None = None
|
||||
self.cost_allowance_discount: float | None = None
|
||||
self.cost_finalization_lock = asyncio.Lock()
|
||||
self.deferred_success_result: tuple[object, datetime.datetime, datetime.datetime] | None = None
|
||||
self.billed_token_rates: BilledTokenRates | None = None
|
||||
self._provider_reported_cost: float | None = None
|
||||
|
||||
|
|
@ -874,6 +881,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"current_attempt_api_call_start_time": None,
|
||||
"litellm_params": litellm_params,
|
||||
}
|
||||
new_call.finalized_cost = None
|
||||
new_call.prepared_cost = None
|
||||
new_call.cost_allowance_discount = None
|
||||
new_call.cost_finalization_lock = asyncio.Lock()
|
||||
new_call.deferred_success_result = None
|
||||
new_call.streaming_chunks = []
|
||||
new_call.sync_streaming_chunks = []
|
||||
return new_call
|
||||
|
|
@ -2462,6 +2474,55 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
**kwargs,
|
||||
)
|
||||
|
||||
def prepare_response_cost(self, response: object, now: datetime.datetime) -> PreparedCost | None:
|
||||
if self.prepared_cost is not None:
|
||||
return self.prepared_cost
|
||||
recorded_cache_hit: Final[object] = self.model_call_details.get("cache_hit")
|
||||
deferred: Final = self.deferred_success_result
|
||||
prepared: Final = self.prepare_success_logging(
|
||||
result=deferred[0] if deferred is not None else response,
|
||||
start_time=deferred[1] if deferred is not None else self.start_time,
|
||||
end_time=deferred[2] if deferred is not None else now.replace(tzinfo=None),
|
||||
cache_hit=recorded_cache_hit if isinstance(recorded_cache_hit, bool) else None,
|
||||
standard_logging_object=None,
|
||||
build_logging_payload=True,
|
||||
publish_logging_payload=False,
|
||||
)
|
||||
payload: Final = prepared[3]
|
||||
raw_cost: Final[object] = self.model_call_details.get("response_cost")
|
||||
if payload is None or not isinstance(raw_cost, (int, float)):
|
||||
return None
|
||||
recorded_usage: Final[object] = self.model_call_details.get("combined_usage_object")
|
||||
response_usage: Final[object] = getattr(prepared[2], "usage", None)
|
||||
normalized_usage: Final = (
|
||||
recorded_usage
|
||||
if isinstance(recorded_usage, Usage)
|
||||
else response_usage
|
||||
if isinstance(response_usage, Usage)
|
||||
else None
|
||||
)
|
||||
usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(
|
||||
_extract_response_obj_and_hidden_params(prepared[2], None)[0],
|
||||
combined_usage_object=normalized_usage,
|
||||
).model_copy(deep=True)
|
||||
self.prepared_cost = PreparedCost(
|
||||
raw_cost,
|
||||
{
|
||||
**payload,
|
||||
"response": None,
|
||||
"messages": None,
|
||||
"model_map_information": copy.deepcopy(payload["model_map_information"]),
|
||||
"cost_breakdown": copy.deepcopy(payload["cost_breakdown"]),
|
||||
},
|
||||
usage,
|
||||
)
|
||||
self.record_partial_usage_for_failure(usage, raw_cost)
|
||||
return self.prepared_cost
|
||||
|
||||
def record_finalized_cost(self, finalized: FinalizedCost) -> None:
|
||||
self.finalized_cost = finalized
|
||||
self.model_call_details["response_cost"] = finalized.amount
|
||||
|
||||
def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None:
|
||||
"""Stash what an interrupted stream already consumed so the failure log bills it instead of zero."""
|
||||
self.model_call_details["combined_usage_object"] = usage
|
||||
|
|
@ -2812,12 +2873,16 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
logging_result, "_hidden_params", {}
|
||||
)
|
||||
|
||||
if self.model_call_details.get("cache_hit") is True:
|
||||
if self.finalized_cost is not None:
|
||||
self.model_call_details["response_cost"] = self.finalized_cost.amount
|
||||
elif self.model_call_details.get("cache_hit") is True:
|
||||
self.model_call_details["response_cost"] = 0.0
|
||||
elif "response_cost" in hidden_params:
|
||||
self.model_call_details["response_cost"] = hidden_params["response_cost"]
|
||||
self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"])
|
||||
elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0:
|
||||
elif (existing_cost := self.model_call_details.get("response_cost")) is not None and (
|
||||
existing_cost != 0 or self.call_type == CallTypes.pass_through.value # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
|
||||
):
|
||||
# Preserve response_cost if already calculated (e.g., by pass-through
|
||||
# handlers like Gemini/Vertex which call completion_cost directly).
|
||||
# Do not preserve 0 from failure_handler on intermediate router retries.
|
||||
|
|
@ -2936,7 +3001,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
standard_logging_object: StandardLoggingPayload | None = None,
|
||||
build_logging_payload: bool = True,
|
||||
) -> tuple[datetime.datetime, datetime.datetime, object]:
|
||||
prepared: Final = self._prepare_success_logging(
|
||||
prepared: Final = self.prepare_success_logging(
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
|
|
@ -2947,7 +3012,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
return prepared[0], prepared[1], prepared[2]
|
||||
|
||||
def _prepare_success_logging(
|
||||
def prepare_success_logging(
|
||||
self,
|
||||
result: object,
|
||||
start_time: datetime.datetime | None,
|
||||
|
|
@ -2968,7 +3033,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
self.model_call_details["log_event_type"] = "successful_api_call"
|
||||
self.model_call_details["end_time"] = end_time
|
||||
self.model_call_details["cache_hit"] = cache_hit
|
||||
self.model_call_details["cache_hit"] = (
|
||||
self.prepared_cost.payload["cache_hit"]
|
||||
if cache_hit is None and self.prepared_cost is not None
|
||||
else cache_hit
|
||||
)
|
||||
|
||||
if self.call_type == CallTypes.anthropic_messages.value:
|
||||
result = self._anthropic_messages_logged_response(result=result)
|
||||
|
|
@ -3007,6 +3076,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
if (
|
||||
litellm.max_budget
|
||||
and publish_logging_payload
|
||||
and self.stream is False
|
||||
and result is not None
|
||||
and isinstance(result, dict)
|
||||
|
|
@ -3037,10 +3107,17 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
publish_logging_payload: bool,
|
||||
) -> StandardLoggingPayload | None:
|
||||
if standard_logging_object is not None:
|
||||
effective_payload: Final = (
|
||||
self.finalized_cost.logging_payload(standard_logging_object)
|
||||
if self.finalized_cost is not None
|
||||
else standard_logging_object
|
||||
)
|
||||
if self.finalized_cost is not None:
|
||||
self.model_call_details["response_cost"] = self.finalized_cost.amount
|
||||
if not publish_logging_payload:
|
||||
return standard_logging_object.copy()
|
||||
self.model_call_details["standard_logging_object"] = standard_logging_object
|
||||
return standard_logging_object
|
||||
return effective_payload.copy()
|
||||
self.model_call_details["standard_logging_object"] = effective_payload
|
||||
return effective_payload
|
||||
if result is None or self.stream is True:
|
||||
# Streaming reaches here before its cost is known, so the cost
|
||||
# is seeded to None, but only when nothing has already
|
||||
|
|
@ -3281,7 +3358,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit)
|
||||
if not self.should_run_logging(event_type="sync_success"): # prevent double logging
|
||||
return
|
||||
start_time, end_time, result, request_logging_payload = self._prepare_success_logging(
|
||||
start_time, end_time, result, request_logging_payload = self.prepare_success_logging(
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
result=result,
|
||||
|
|
@ -4204,7 +4281,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
# chunks already delivered; the router stashes that recovered usage as
|
||||
# ``combined_usage_object`` and pre-computes its cost, so preserve it
|
||||
# here instead of zeroing the spend on an otherwise-failed request.
|
||||
if self.model_call_details.get("combined_usage_object") is None:
|
||||
if self.finalized_cost is not None:
|
||||
self.model_call_details["response_cost"] = self.finalized_cost.amount
|
||||
elif self.model_call_details.get("combined_usage_object") is None:
|
||||
self.model_call_details["response_cost"] = 0
|
||||
|
||||
if hasattr(exception, "headers") and isinstance(exception.headers, dict):
|
||||
|
|
@ -5544,7 +5623,7 @@ def build_custom_logger_compatible_class(
|
|||
return callback
|
||||
otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
|
||||
otel_logger_v2: Final = build_otel_v2_logger(
|
||||
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings
|
||||
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings, console_default=True
|
||||
)
|
||||
_in_memory_loggers.append(otel_logger_v2)
|
||||
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
|
||||
|
|
@ -7367,6 +7446,12 @@ def _cache_hit_cost_breakdown(served: CostBreakdown | None) -> CostBreakdown:
|
|||
)
|
||||
|
||||
|
||||
def _finalized_cost_of(logging_obj: object) -> FinalizedCost | None:
|
||||
"""Duck-typed logging objects (and test doubles) carry no finalized cost."""
|
||||
finalized: Final = getattr(logging_obj, "finalized_cost", None)
|
||||
return finalized if isinstance(finalized, FinalizedCost) else None
|
||||
|
||||
|
||||
def get_standard_logging_object_payload(
|
||||
kwargs: dict | None,
|
||||
init_response_obj: object,
|
||||
|
|
@ -7470,7 +7555,10 @@ def get_standard_logging_object_payload(
|
|||
raw_response_cost: Final = kwargs.get("response_cost")
|
||||
llm_response_cost: Final[float] = raw_response_cost or 0.0
|
||||
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))
|
||||
response_cost: Final[float] = llm_response_cost + guardrail_cost
|
||||
finalized_cost: Final = _finalized_cost_of(logging_obj)
|
||||
response_cost: Final[float] = (
|
||||
finalized_cost.amount if finalized_cost is not None else llm_response_cost + guardrail_cost
|
||||
)
|
||||
|
||||
# clean up litellm hidden params
|
||||
clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params)
|
||||
|
|
@ -7537,13 +7625,17 @@ def get_standard_logging_object_payload(
|
|||
):
|
||||
model_name = response_model_name
|
||||
|
||||
request_cost_breakdown: Final = cost_breakdown_with_guardrail(
|
||||
(
|
||||
_cache_hit_cost_breakdown(uncached_price.cost_breakdown)
|
||||
if uncached_price is not None
|
||||
else logging_obj.cost_breakdown
|
||||
),
|
||||
guardrail_cost,
|
||||
request_cost_breakdown: Final = (
|
||||
finalized_cost.logging_breakdown()
|
||||
if finalized_cost is not None
|
||||
else cost_breakdown_with_guardrail(
|
||||
(
|
||||
_cache_hit_cost_breakdown(uncached_price.cost_breakdown)
|
||||
if uncached_price is not None
|
||||
else logging_obj.cost_breakdown
|
||||
),
|
||||
guardrail_cost,
|
||||
)
|
||||
)
|
||||
captured_baseline: Final = logging_obj.baseline_observation
|
||||
autorouter_savings: Final = (
|
||||
|
|
@ -7600,6 +7692,7 @@ def get_standard_logging_object_payload(
|
|||
model=model_name,
|
||||
metadata=clean_metadata,
|
||||
cache_key=clean_hidden_params["cache_key"],
|
||||
cost_finalized=finalized_cost is not None,
|
||||
response_cost=response_cost,
|
||||
cost_breakdown=request_cost_breakdown,
|
||||
autorouter_savings=autorouter_savings,
|
||||
|
|
|
|||
61
litellm/litellm_core_utils/llm_cost_calc/finalization.py
Normal file
61
litellm/litellm_core_utils/llm_cost_calc/finalization.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.types.utils import CostBreakdown, StandardLoggingPayload, Usage
|
||||
|
||||
BreakdownValue: TypeAlias = float | str | None | Mapping[str, float]
|
||||
_BREAKDOWN_FIELDS: Final = TypeAdapter(Mapping[str, BreakdownValue])
|
||||
_COST_BREAKDOWN: Final = TypeAdapter(CostBreakdown)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PreparedCost:
|
||||
amount: float
|
||||
payload: StandardLoggingPayload
|
||||
usage: Usage
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FinalizedCost:
|
||||
amount: float
|
||||
breakdown: Mapping[str, BreakdownValue] | None
|
||||
allowance_discount: float
|
||||
|
||||
def logging_payload(self, payload: StandardLoggingPayload) -> StandardLoggingPayload:
|
||||
return {
|
||||
**payload,
|
||||
"response_cost": self.amount,
|
||||
"cost_breakdown": self.logging_breakdown(),
|
||||
"cost_finalized": True,
|
||||
}
|
||||
|
||||
def logging_breakdown(self) -> CostBreakdown | None:
|
||||
return _COST_BREAKDOWN.validate_python(self.breakdown) if self.breakdown is not None else None
|
||||
|
||||
|
||||
def finalize_cost(amount: float, breakdown: CostBreakdown | None, allowance_discount: float = 0.0) -> FinalizedCost:
|
||||
discount: Final = min(max(amount, 0.0), max(allowance_discount, 0.0))
|
||||
total: Final = amount - discount
|
||||
existing: Final = _BREAKDOWN_FIELDS.validate_python(breakdown) if breakdown is not None else None
|
||||
adjusted: Final = (
|
||||
{
|
||||
**(existing or {}),
|
||||
"total_cost": total,
|
||||
"original_cost": (breakdown or {}).get("original_cost", amount),
|
||||
"discount_amount": (breakdown or {}).get("discount_amount", 0.0) + discount,
|
||||
}
|
||||
if discount > 0
|
||||
else existing
|
||||
)
|
||||
frozen: Final = (
|
||||
MappingProxyType(
|
||||
{key: MappingProxyType(value) if isinstance(value, Mapping) else value for key, value in adjusted.items()}
|
||||
)
|
||||
if adjusted is not None
|
||||
else None
|
||||
)
|
||||
return FinalizedCost(amount=total, breakdown=frozen, allowance_discount=discount)
|
||||
|
|
@ -108,7 +108,7 @@ def judge_target(router: Router | None, model: str, team_id: str | None = None)
|
|||
only for its own team and a team's own deployment resolves for nobody else, so asking
|
||||
without it answers for a caller who does not exist.
|
||||
"""
|
||||
served: Final = router.resolved_litellm_models(model, team_id=team_id) if router is not None else ()
|
||||
served: Final = router.registry().names(model, kind="litellm_models", team_id=team_id) if router is not None else ()
|
||||
if served:
|
||||
return JudgeTarget("router", frozenset(_provider_qualified(m) or m for m in served))
|
||||
qualified: Final = _provider_qualified(model)
|
||||
|
|
@ -148,7 +148,7 @@ def attached_auth(request_data: object) -> UserAPIKeyAuth | None:
|
|||
|
||||
|
||||
def _named_route(router: Router, deployment_id: str) -> bool:
|
||||
deployment: Final = router.get_deployment(model_id=deployment_id)
|
||||
deployment: Final = router.registry().deployment(deployment_id)
|
||||
if deployment is None:
|
||||
return False
|
||||
public_names: Final = (deployment.model_name, deployment.model_info.team_public_model_name or "")
|
||||
|
|
@ -167,7 +167,7 @@ def caller_team_for(router: Router | None, model: str, auth: UserAPIKeyAuth | No
|
|||
team_id: Final = auth.team_id if auth is not None else None
|
||||
if team_id is None or router is None or judge_target(router, model).via != "nothing":
|
||||
return None
|
||||
team_route: Final = router.get_candidate_model_ids_for_route(model, team_id)
|
||||
team_route: Final = router.registry().names(model, kind="route_candidates", team_id=team_id)
|
||||
return team_id if team_route and all(_named_route(router, deployment_id) for deployment_id in team_route) else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import json
|
|||
import re
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Generator, Iterator, Mapping, Sequence
|
||||
from typing import Final, Literal, cast
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
|
||||
|
|
@ -19,7 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
|
||||
from litellm.types.llms.databricks import DatabricksTool
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionReasoningItem,
|
||||
ChatCompletionAnnotation,
|
||||
ChatCompletionThinkingBlock,
|
||||
ImageURLListItem,
|
||||
OpenAIModerationResponse,
|
||||
|
|
@ -59,6 +59,31 @@ _CHOICES_FIELDS: Final[frozenset] = frozenset(Choices.model_fields.keys())
|
|||
_MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields.keys()) | {"usage"}
|
||||
|
||||
|
||||
_ANNOTATION_ENTRIES: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def _is_provider_specific_annotation(annotation: object) -> bool:
|
||||
match annotation:
|
||||
case {"type": "url_citation"}:
|
||||
return False
|
||||
case {"type": _}:
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def _annotation_entries(annotations: object) -> tuple[object, ...]:
|
||||
return _ANNOTATION_ENTRIES.validate_python(annotations) if isinstance(annotations, list) else ()
|
||||
|
||||
|
||||
def partition_annotations(annotations: object) -> tuple[object, tuple[object, ...]]:
|
||||
entries: Final = _annotation_entries(annotations)
|
||||
others: Final = tuple(annotation for annotation in entries if _is_provider_specific_annotation(annotation))
|
||||
if not others:
|
||||
return annotations, ()
|
||||
return [annotation for annotation in entries if not _is_provider_specific_annotation(annotation)] or None, others
|
||||
|
||||
|
||||
def _normalize_images_for_message(
|
||||
images: list[dict] | None,
|
||||
) -> list[ImageURLListItem] | None:
|
||||
|
|
@ -145,31 +170,89 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]:
|
|||
return _REPLAY_CONTENT_SLICE_RE.findall(content)
|
||||
|
||||
|
||||
def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
|
||||
# Rebuild the delta as content-only so every accumulate-able field (role,
|
||||
# tool_calls, reasoning_content, thinking_blocks, audio, images,
|
||||
# annotations, ...) is dropped on later slices instead of an enumerated
|
||||
# subset; repeating any of them makes downstream handlers that accumulate
|
||||
# streamed deltas collect it once per slice, and a field added to Delta
|
||||
# later can't silently re-introduce the duplication.
|
||||
choice.delta = Delta(content=choice.delta.content)
|
||||
choice.logprobs = None
|
||||
if hasattr(choice, "enhancements"):
|
||||
del choice.enhancements
|
||||
_REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"})
|
||||
_REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"}
|
||||
_REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"})
|
||||
_REPLAY_TEXT_FIELDS: Final = frozenset({"content", "annotations"})
|
||||
|
||||
|
||||
def _first_replay_delta(
|
||||
message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None
|
||||
) -> Delta:
|
||||
return Delta(
|
||||
content=message.get("content"),
|
||||
role=message["role"],
|
||||
function_call=message.get("function_call"),
|
||||
tool_calls=message.get("tool_calls"),
|
||||
reasoning_items=None if reasoning_items is None else list(reasoning_items),
|
||||
provider_specific_fields=_cached_provider_specific_fields(
|
||||
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
|
||||
),
|
||||
def _fields_in(fields: Mapping[str, object], keys: frozenset[str]) -> dict[str, object]:
|
||||
return {key: value for key, value in fields.items() if key in keys}
|
||||
|
||||
|
||||
def _replay_deltas(delta: Delta) -> tuple[Delta, ...]:
|
||||
fields: Final[dict[str, object]] = delta.model_dump(exclude_none=True)
|
||||
lead: Final = _fields_in(fields, _REPLAY_LEAD_FIELDS) if _REPLAY_REASONING_FIELDS.intersection(fields) else {}
|
||||
pieces: Final = _split_assembled_content_for_replay(delta.content) or ([delta.content] if delta.content else [])
|
||||
if not pieces and not lead:
|
||||
return (delta,)
|
||||
if not pieces:
|
||||
held_back: Final = lead.keys() | {"content"}
|
||||
rest: Final = {key: value for key, value in fields.items() if key not in held_back}
|
||||
return (Delta.model_validate(lead),) + ((Delta.model_validate(rest),) if rest else ())
|
||||
tools: Final = _fields_in(fields, _REPLAY_TOOL_FIELDS)
|
||||
annotations: Final = _fields_in(fields, frozenset({"annotations"}))
|
||||
moved: Final = lead.keys() | _REPLAY_TOOL_FIELDS | _REPLAY_TEXT_FIELDS
|
||||
first_slice: Final = {key: value for key, value in fields.items() if key not in moved}
|
||||
last_index: Final = len(pieces) - 1
|
||||
return (
|
||||
((Delta.model_validate(lead),) if lead else ())
|
||||
+ tuple(
|
||||
Delta.model_validate(
|
||||
{
|
||||
**(first_slice if index == 0 else {}),
|
||||
"content": piece,
|
||||
**(annotations if index == last_index else {}),
|
||||
}
|
||||
)
|
||||
for index, piece in enumerate(pieces)
|
||||
)
|
||||
+ ((Delta.model_validate(tools),) if tools else ())
|
||||
)
|
||||
|
||||
|
||||
def _replay_frame(
|
||||
response: ModelResponseStream, delta: Delta, finish_reason: str | None, keeps_choice_metadata: bool
|
||||
) -> ModelResponseStream:
|
||||
frame: Final = response.model_copy(deep=True)
|
||||
choice: Final = frame.choices[0]
|
||||
choice.delta = delta
|
||||
choice.finish_reason = finish_reason
|
||||
if not keeps_choice_metadata:
|
||||
choice.logprobs = None
|
||||
if hasattr(choice, "enhancements"):
|
||||
del choice.enhancements
|
||||
return frame
|
||||
|
||||
|
||||
def _replay_frames(response: ModelResponseStream) -> Iterator[ModelResponseStream]:
|
||||
if len(response.choices) != 1:
|
||||
yield response
|
||||
return
|
||||
usage: Final = getattr(response, "usage", None)
|
||||
if usage is not None:
|
||||
delattr(response, "usage")
|
||||
choice: Final = response.choices[0]
|
||||
deltas: Final = _replay_deltas(choice.delta)
|
||||
first_text: Final = next((index for index, delta in enumerate(deltas) if delta.content), 0)
|
||||
for index, delta in enumerate(deltas):
|
||||
yield _replay_frame(response, delta, finish_reason=None, keeps_choice_metadata=index == first_text)
|
||||
finish_frame: Final = _replay_frame(
|
||||
response, Delta(), finish_reason=choice.finish_reason, keeps_choice_metadata=False
|
||||
)
|
||||
if usage is not None:
|
||||
setattr(finish_frame, "usage", usage)
|
||||
yield finish_frame
|
||||
|
||||
|
||||
def _first_replay_delta(message: Mapping[str, object]) -> Delta:
|
||||
return Delta.model_validate(
|
||||
{
|
||||
**message,
|
||||
"provider_specific_fields": _cached_provider_specific_fields(
|
||||
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -234,7 +317,7 @@ async def convert_to_streaming_response_async(
|
|||
t["index"] = index
|
||||
pydantic_tool_calls.append(ChatCompletionDeltaToolCall(**t))
|
||||
choice["message"]["tool_calls"] = pydantic_tool_calls
|
||||
delta = _first_replay_delta(choice["message"], choice["message"].get("reasoning_items"))
|
||||
delta = _first_replay_delta(choice["message"])
|
||||
finish_reason = choice.get("finish_reason", None)
|
||||
|
||||
if finish_reason is None:
|
||||
|
|
@ -266,41 +349,14 @@ async def convert_to_streaming_response_async(
|
|||
if "model" in response_object:
|
||||
model_response_object.model = response_object["model"]
|
||||
|
||||
# Replay cached content with per-word cadence so stream=true cache hits
|
||||
# don't arrive as a single SSE frame. Multi-choice (n>1) responses and
|
||||
# unsplittable content (None/empty/whitespace-free) keep the original
|
||||
# single-yield behavior.
|
||||
slices: list[str] = []
|
||||
if len(model_response_object.choices) == 1:
|
||||
slices = _split_assembled_content_for_replay(model_response_object.choices[0].delta.content)
|
||||
if len(slices) <= 1:
|
||||
yield model_response_object
|
||||
await asyncio.sleep(0)
|
||||
return
|
||||
|
||||
# Detach usage from the base object so we can re-attach it only to the
|
||||
# final slice chunk. A non-None usage always lives in __pydantic_extra__
|
||||
# here (set via setattr above), so delattr cannot fail.
|
||||
original_usage: Final = getattr(model_response_object, "usage", None)
|
||||
if original_usage is not None:
|
||||
delattr(model_response_object, "usage")
|
||||
original_finish_reason: Final = model_response_object.choices[0].finish_reason
|
||||
last_idx: Final = len(slices) - 1
|
||||
for i, piece in enumerate(slices):
|
||||
slice_chunk = model_response_object.model_copy(deep=True)
|
||||
slice_chunk.choices[0].delta.content = piece
|
||||
if i > 0:
|
||||
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
|
||||
slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
|
||||
if i == last_idx and original_usage is not None:
|
||||
setattr(slice_chunk, "usage", original_usage)
|
||||
yield slice_chunk
|
||||
for frame in _replay_frames(model_response_object):
|
||||
yield frame
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
def convert_to_streaming_response(
|
||||
response_object: dict | None = None,
|
||||
):
|
||||
) -> Generator[ModelResponseStream, None, None]:
|
||||
# used for yielding Cache hits when stream == True
|
||||
if response_object is None:
|
||||
raise Exception("Error in response object format")
|
||||
|
|
@ -356,30 +412,7 @@ def convert_to_streaming_response(
|
|||
if "model" in response_object:
|
||||
model_response_object.model = response_object["model"]
|
||||
|
||||
# Replay cached content with per-word cadence on sync cache-hit paths
|
||||
# (S3Cache, sync completion()). See convert_to_streaming_response_async
|
||||
# for the full rationale — this mirrors its tail.
|
||||
slices: list[str] = []
|
||||
if len(model_response_object.choices) == 1:
|
||||
slices = _split_assembled_content_for_replay(model_response_object.choices[0].delta.content)
|
||||
if len(slices) <= 1:
|
||||
yield model_response_object
|
||||
return
|
||||
|
||||
original_usage: Final = getattr(model_response_object, "usage", None)
|
||||
if original_usage is not None:
|
||||
delattr(model_response_object, "usage")
|
||||
original_finish_reason: Final = model_response_object.choices[0].finish_reason
|
||||
last_idx: Final = len(slices) - 1
|
||||
for i, piece in enumerate(slices):
|
||||
slice_chunk = model_response_object.model_copy(deep=True)
|
||||
slice_chunk.choices[0].delta.content = piece
|
||||
if i > 0:
|
||||
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
|
||||
slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
|
||||
if i == last_idx and original_usage is not None:
|
||||
setattr(slice_chunk, "usage", original_usage)
|
||||
yield slice_chunk
|
||||
yield from _replay_frames(model_response_object)
|
||||
|
||||
|
||||
from collections import defaultdict
|
||||
|
|
@ -787,6 +820,9 @@ def convert_to_model_response_object(
|
|||
provider_specific_fields = dict(choice["message"].get("provider_specific_fields", None) or {})
|
||||
for f in choice["message"].keys() - _MESSAGE_FIELDS:
|
||||
provider_specific_fields[f] = choice["message"][f]
|
||||
annotations, other_annotations = partition_annotations(choice["message"].get("annotations", None))
|
||||
if other_annotations:
|
||||
provider_specific_fields["annotations"] = list(other_annotations)
|
||||
|
||||
# Handle reasoning models that display `reasoning_content` within `content`
|
||||
reasoning_content, content = extract_reasoning_content(choice["message"])
|
||||
|
|
@ -809,7 +845,9 @@ def convert_to_model_response_object(
|
|||
reasoning_content=reasoning_content,
|
||||
thinking_blocks=thinking_blocks,
|
||||
reasoning_items=choice["message"].get("reasoning_items", None),
|
||||
annotations=choice["message"].get("annotations", None),
|
||||
annotations=cast( # cast-ok: provider annotation dicts pass through whole, keeping keys such as url_citation.content
|
||||
list[ChatCompletionAnnotation] | None, annotations
|
||||
),
|
||||
images=_normalize_images_for_message(choice["message"].get("images", None)),
|
||||
)
|
||||
finish_reason = choice.get("finish_reason", None)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import atexit
|
|||
import contextvars
|
||||
import inspect
|
||||
import logging
|
||||
from collections.abc import Coroutine, Iterator
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Iterator
|
||||
from typing import Final
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -66,14 +66,19 @@ class LoggingWorker:
|
|||
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
|
||||
concurrency: int = LOGGING_WORKER_CONCURRENCY,
|
||||
timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS,
|
||||
*,
|
||||
retry_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
):
|
||||
self.timeout = timeout
|
||||
self.max_queue_size = max_queue_size
|
||||
self.concurrency = concurrency
|
||||
self.timeout_summary_window = timeout_summary_window
|
||||
self._queue: asyncio.Queue[LoggingTask] | None = None
|
||||
self._worker_task: asyncio.Task | None = None
|
||||
self._running_tasks: set[asyncio.Task] = set()
|
||||
self._worker_task: asyncio.Task[None] | None = None
|
||||
self._running_tasks: set[asyncio.Task[None]] = set()
|
||||
self._overflow_tasks: set[asyncio.Task[None]] = set() # mutable-ok: done callbacks discard refs
|
||||
self._pending_enqueues: dict[int, LoggingTask] = {} # mutable-ok: retries and shutdown share callback ownership
|
||||
self._retry_sleep = retry_sleep
|
||||
self._dequeued_tasks: dict[int, LoggingTask] = {} # mutable-ok: refs so flush can rescue never-started tasks
|
||||
self._sem: asyncio.Semaphore | None = None
|
||||
self._bound_loop: asyncio.AbstractEventLoop | None = None
|
||||
|
|
@ -168,6 +173,7 @@ class LoggingWorker:
|
|||
self._sem = None
|
||||
self._worker_task = None
|
||||
self._running_tasks.clear()
|
||||
self._overflow_tasks.clear()
|
||||
self._timeout_summary_task = None
|
||||
self._timeout_burst_count = 0
|
||||
self._queue = new_queue
|
||||
|
|
@ -185,6 +191,10 @@ class LoggingWorker:
|
|||
self._sem = asyncio.Semaphore(self.concurrency)
|
||||
if self._worker_task is None or self._worker_task.done():
|
||||
self._worker_task = asyncio.create_task(self._worker_loop())
|
||||
if not self._overflow_tasks:
|
||||
self._aggressive_clear_in_progress = False
|
||||
for pending in tuple(self._pending_enqueues.values()):
|
||||
self._schedule_delayed_enqueue_retry(pending)
|
||||
|
||||
async def _process_log_task(
|
||||
self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]"
|
||||
|
|
@ -278,12 +288,24 @@ class LoggingWorker:
|
|||
# Capture the current context when enqueueing
|
||||
task: Final = LoggingTask(coroutine=coroutine, context=contextvars.copy_context())
|
||||
|
||||
self._enqueue_task(task)
|
||||
|
||||
def _enqueue_task(self, task: LoggingTask) -> None:
|
||||
if self._queue is None:
|
||||
return
|
||||
try:
|
||||
self._queue.put_nowait(task)
|
||||
except asyncio.QueueFull:
|
||||
# Queue is full - handle it appropriately
|
||||
verbose_logger.exception("LoggingWorker queue is full")
|
||||
self._pending_enqueues[id(task)] = task
|
||||
self._handle_queue_full(task)
|
||||
else:
|
||||
self._pending_enqueues.pop(id(task), None)
|
||||
|
||||
def _start_overflow_task(self, coroutine: Coroutine[object, object, None]) -> None:
|
||||
task: Final = asyncio.create_task(coroutine)
|
||||
self._overflow_tasks.add(task)
|
||||
task.add_done_callback(self._overflow_tasks.discard)
|
||||
|
||||
def _should_start_aggressive_clear(self) -> bool:
|
||||
"""
|
||||
|
|
@ -326,7 +348,7 @@ class LoggingWorker:
|
|||
if self._should_start_aggressive_clear():
|
||||
self._mark_aggressive_clear_started()
|
||||
# Schedule clearing as async task so enqueue returns immediately (non-blocking)
|
||||
asyncio.create_task(self._aggressively_clear_queue_async(task))
|
||||
self._start_overflow_task(self._aggressively_clear_queue_async(task))
|
||||
else:
|
||||
# Cooldown active or clear in progress, schedule a delayed retry
|
||||
self._schedule_delayed_enqueue_retry(task)
|
||||
|
|
@ -363,7 +385,7 @@ class LoggingWorker:
|
|||
delay: Final = self._calculate_retry_delay()
|
||||
|
||||
# Schedule the retry as a background task
|
||||
asyncio.create_task(self._retry_enqueue_task(task, delay))
|
||||
self._start_overflow_task(self._retry_enqueue_task(task, delay))
|
||||
except RuntimeError:
|
||||
# No event loop, drop the task as we can't schedule a retry
|
||||
pass
|
||||
|
|
@ -373,17 +395,8 @@ class LoggingWorker:
|
|||
Retry enqueueing the task after delay, preserving original context.
|
||||
This is called as a background task from _schedule_delayed_enqueue_retry.
|
||||
"""
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# Try to enqueue the task directly, preserving its original context
|
||||
if self._queue is None:
|
||||
return
|
||||
|
||||
try:
|
||||
self._queue.put_nowait(task)
|
||||
except asyncio.QueueFull:
|
||||
# Still full - handle it appropriately (clear or retry again)
|
||||
self._handle_queue_full(task)
|
||||
await self._retry_sleep(delay)
|
||||
self._enqueue_task(task)
|
||||
|
||||
def _extract_tasks_from_queue(self) -> list[LoggingTask]:
|
||||
"""
|
||||
|
|
@ -423,15 +436,10 @@ class LoggingWorker:
|
|||
return
|
||||
|
||||
extracted_tasks: Final = self._extract_tasks_from_queue()
|
||||
|
||||
# Add new task to extracted tasks to process directly
|
||||
if new_task is not None:
|
||||
self._track_dequeued(new_task)
|
||||
extracted_tasks.append(new_task)
|
||||
|
||||
# Process extracted tasks directly
|
||||
self._enqueue_task(new_task)
|
||||
if extracted_tasks:
|
||||
await self._process_extracted_tasks(extracted_tasks)
|
||||
await asyncio.gather(*(self._process_single_task(task) for task in extracted_tasks))
|
||||
except Exception as e:
|
||||
verbose_logger.exception("LoggingWorker error during aggressive clear: %s", e)
|
||||
finally:
|
||||
|
|
@ -439,7 +447,7 @@ class LoggingWorker:
|
|||
self._aggressive_clear_in_progress = False
|
||||
|
||||
async def _process_single_task(self, task: LoggingTask) -> None:
|
||||
"""Process a single task and mark it done."""
|
||||
"""Process a task extracted from the queue and mark it done."""
|
||||
if self._queue is None:
|
||||
return
|
||||
|
||||
|
|
@ -455,17 +463,6 @@ class LoggingWorker:
|
|||
self._untrack_dequeued(task)
|
||||
self._queue.task_done()
|
||||
|
||||
async def _process_extracted_tasks(self, tasks: list[LoggingTask]) -> None:
|
||||
"""
|
||||
Process tasks that were extracted from the queue to make room.
|
||||
Processes them concurrently without semaphore limits for maximum speed.
|
||||
"""
|
||||
if not tasks or self._queue is None:
|
||||
return
|
||||
|
||||
# Process all tasks concurrently for maximum speed
|
||||
await asyncio.gather(*[self._process_single_task(task) for task in tasks])
|
||||
|
||||
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None:
|
||||
"""
|
||||
Ensure the logging worker is initialized and enqueue the coroutine.
|
||||
|
|
@ -473,32 +470,63 @@ class LoggingWorker:
|
|||
self.start()
|
||||
self.enqueue(async_coroutine)
|
||||
|
||||
async def stop(self) -> None:
|
||||
@property
|
||||
def pending_count(self) -> int:
|
||||
return (
|
||||
(self._queue.qsize() if self._queue is not None else 0)
|
||||
+ len(self._dequeued_tasks)
|
||||
+ len(self._pending_enqueues)
|
||||
)
|
||||
|
||||
async def stop(self, *, discard_pending: bool = False, timeout: float | None = None) -> None:
|
||||
"""Stop the logging worker and clean up resources."""
|
||||
if discard_pending and self._queue is not None:
|
||||
discarded: Final = self._drain_pending(self._queue)
|
||||
for task in discarded:
|
||||
task["coroutine"].close()
|
||||
self._queue.task_done()
|
||||
if discarded:
|
||||
verbose_logger.warning("LoggingWorker shutdown: %d queued callbacks discarded", len(discarded))
|
||||
if self._timeout_summary_task is not None:
|
||||
self._timeout_summary_task.cancel()
|
||||
self._timeout_summary_task = None
|
||||
self._emit_timeout_summary()
|
||||
|
||||
if self._worker_task is None and not self._running_tasks:
|
||||
# No worker launched and no in-flight tasks to drain.
|
||||
return
|
||||
if discard_pending:
|
||||
rejected: Final = tuple(self._pending_enqueues.values())
|
||||
self._pending_enqueues.clear()
|
||||
for task in rejected:
|
||||
task["coroutine"].close()
|
||||
if rejected:
|
||||
verbose_logger.warning("LoggingWorker shutdown: %d overflow callbacks discarded", len(rejected))
|
||||
|
||||
tasks_to_cancel: Final[list[asyncio.Task]] = list(self._running_tasks)
|
||||
if self._worker_task:
|
||||
# Include the main worker loop so it stops fetching work.
|
||||
tasks_to_cancel.append(self._worker_task)
|
||||
tasks_to_cancel: Final = (
|
||||
*self._running_tasks,
|
||||
*self._overflow_tasks,
|
||||
*((self._worker_task,) if self._worker_task is not None else ()),
|
||||
)
|
||||
if not tasks_to_cancel:
|
||||
return
|
||||
|
||||
for task in tasks_to_cancel:
|
||||
# Propagate cancellation to every pending task.
|
||||
task.cancel()
|
||||
|
||||
# Wait for cancellation to settle; ignore errors raised during shutdown.
|
||||
await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
|
||||
done, pending = await asyncio.wait(tasks_to_cancel, timeout=timeout)
|
||||
await asyncio.gather(*done, return_exceptions=True)
|
||||
if pending:
|
||||
verbose_logger.warning("LoggingWorker shutdown: %d tasks still alive after cancellation", len(pending))
|
||||
|
||||
self._worker_task = None
|
||||
if discard_pending and self._queue is not None:
|
||||
for task in self._unstarted_dequeued_tasks():
|
||||
task["coroutine"].close()
|
||||
self._untrack_dequeued(task)
|
||||
self._queue.task_done()
|
||||
if self._worker_task is not None and self._worker_task.done():
|
||||
self._worker_task = None
|
||||
# Drop references to completed tasks so we can restart cleanly.
|
||||
self._running_tasks.clear()
|
||||
self._running_tasks.difference_update(done)
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Flush the logging queue.
|
||||
|
|
@ -516,6 +544,9 @@ class LoggingWorker:
|
|||
return
|
||||
self.start()
|
||||
await self._queue.join()
|
||||
while self._overflow_tasks:
|
||||
await asyncio.wait(tuple(self._overflow_tasks))
|
||||
await self._queue.join()
|
||||
|
||||
async def clear_queue(self):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -424,6 +424,109 @@ def anthropic_image_source_to_openai_url(image_source: Mapping[str, object]) ->
|
|||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InlineDocument:
|
||||
filename: str
|
||||
file_data: str
|
||||
|
||||
|
||||
def _document_media_type(source: object) -> str:
|
||||
match source:
|
||||
case {"media_type": str() as media_type} if media_type:
|
||||
return media_type
|
||||
case _:
|
||||
return "application/pdf"
|
||||
|
||||
|
||||
def anthropic_inline_document(document: Mapping[str, object]) -> InlineDocument | None:
|
||||
source: Final = document.get("source")
|
||||
match source:
|
||||
case {"type": "base64", "data": str() as data} if data:
|
||||
title: Final = document.get("title")
|
||||
return InlineDocument(
|
||||
filename=title if isinstance(title, str) and title else "document.pdf",
|
||||
file_data=f"data:{_document_media_type(source)};base64,{data}",
|
||||
)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def is_pdf_file_part(part: object) -> bool:
|
||||
match part:
|
||||
case {"type": "file", "file": {"file_data": str() as file_data}} if file_data.startswith("data:"):
|
||||
return file_data.startswith("data:application/pdf")
|
||||
case {"type": "file", "file": {"format": str() as file_format}}:
|
||||
return file_format == "application/pdf"
|
||||
case {"type": "file", "file": {"filename": str() as filename}}:
|
||||
return filename.lower().endswith(".pdf")
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
_CONTENT_PARTS: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def content_parts(message: Mapping[str, object]) -> Sequence[object]:
|
||||
content: Final = message.get("content")
|
||||
return _CONTENT_PARTS.validate_python(content) if isinstance(content, list) else ()
|
||||
|
||||
|
||||
def _has_pdf_file_part(message: Mapping[str, object]) -> bool:
|
||||
return any(is_pdf_file_part(part) for part in content_parts(message))
|
||||
|
||||
|
||||
def messages_have_pdf_file_part(messages: Sequence[Mapping[str, object]]) -> bool:
|
||||
return any(_has_pdf_file_part(message) for message in messages)
|
||||
|
||||
|
||||
def _message_without_pdf_file_parts(message: AllMessageValues) -> AllMessageValues:
|
||||
if not _has_pdf_file_part(message):
|
||||
return message
|
||||
kept: Final = [part for part in content_parts(message) if not is_pdf_file_part(part)]
|
||||
return cast(AllMessageValues, {**message, "content": kept}) # cast-ok: same message with fewer content parts
|
||||
|
||||
|
||||
def _is_pdf_only_user_turn(message: AllMessageValues) -> bool:
|
||||
parts: Final = content_parts(message)
|
||||
return message["role"] == "user" and bool(parts) and all(is_pdf_file_part(part) for part in parts)
|
||||
|
||||
|
||||
def messages_without_pdf_file_parts(messages: Sequence[AllMessageValues]) -> tuple[AllMessageValues, ...]:
|
||||
return tuple(
|
||||
_message_without_pdf_file_parts(message) for message in messages if not _is_pdf_only_user_turn(message)
|
||||
)
|
||||
|
||||
|
||||
_CONTENT_PART: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _file_part_as_image_url(part: object) -> object:
|
||||
match part:
|
||||
case {"type": "file", "file": {"file_data": str() as file_data}} if file_data:
|
||||
fields: Final = _CONTENT_PART.validate_python(part)
|
||||
return {key: value for key, value in fields.items() if key != "file"} | {
|
||||
"type": "image_url",
|
||||
"image_url": {"url": file_data},
|
||||
}
|
||||
case _:
|
||||
return part
|
||||
|
||||
|
||||
def message_with_file_parts_as_image_urls(message: Mapping[str, object]) -> Mapping[str, object]:
|
||||
parts: Final = content_parts(message)
|
||||
if not parts:
|
||||
return message
|
||||
return {**message, "content": [_file_part_as_image_url(part) for part in parts]}
|
||||
|
||||
|
||||
def _typed_message_with_file_parts_as_image_urls(message: AllMessageValues) -> AllMessageValues:
|
||||
return cast(AllMessageValues, message_with_file_parts_as_image_urls(message)) # cast-ok: same message, image parts
|
||||
|
||||
|
||||
def messages_with_file_parts_as_image_urls(messages: Sequence[AllMessageValues]) -> tuple[AllMessageValues, ...]:
|
||||
return tuple(_typed_message_with_file_parts_as_image_urls(message) for message in messages)
|
||||
|
||||
|
||||
def _image_part_url(part: Mapping[str, object]) -> str | None:
|
||||
"""The image URL carried by one content part, whichever of the three dialects wrote it."""
|
||||
part_type: Final = part.get("type")
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.constants import (
|
|||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
|
||||
from litellm.types.llms.openai import (
|
||||
OpenAIRealtimeEvents,
|
||||
|
|
@ -1281,24 +1282,29 @@ class RealTimeStreaming:
|
|||
verbose_logger.exception("Connection closed in backend to client send messages - %s", e)
|
||||
close: Final = backend_close_from(e)
|
||||
self._flush_unbilled_transcription_usage()
|
||||
if self._backend_refused_session(close):
|
||||
await self.log_backend_refusal(e)
|
||||
else:
|
||||
await self.log_messages()
|
||||
await self._log_session_end(e if self._backend_refused_session(close) else None)
|
||||
return close
|
||||
except asyncio.CancelledError:
|
||||
self._flush_unbilled_transcription_usage()
|
||||
await self.log_messages()
|
||||
await self._log_session_end(None)
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error in backend to client send messages: %s", e)
|
||||
self._flush_unbilled_transcription_usage()
|
||||
await self.log_messages()
|
||||
await self._log_session_end(None)
|
||||
return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket")
|
||||
|
||||
def _backend_refused_session(self, close: BackendClose) -> bool:
|
||||
return close.code != 1000 and not self._backend_sent_frames
|
||||
|
||||
async def _log_session_end(self, transport_refusal: Exception | None) -> None:
|
||||
provider_refusal: Final = None if self.provider_config is None else self.provider_config.session_refusal()
|
||||
refusal: Final = provider_refusal or transport_refusal
|
||||
if refusal is None:
|
||||
await self.log_messages()
|
||||
return
|
||||
await self.log_backend_refusal(refusal)
|
||||
|
||||
async def log_backend_refusal(self, error: Exception) -> None:
|
||||
if not self.logging_obj:
|
||||
return
|
||||
|
|
@ -1691,7 +1697,13 @@ class RealTimeStreaming:
|
|||
# ``_maybe_send_guardrail_turn_detection_update`` would retry.
|
||||
if msg_type == "session.update" and self.provider_config is None:
|
||||
await self._send_session_gate_once()
|
||||
sent = await self._send_to_backend(message)
|
||||
try:
|
||||
sent = await self._send_to_backend(message)
|
||||
except RealtimeTranscriptionProtocolError as e:
|
||||
await self.websocket.send_text(realtime_error_event(str(e), error_type="invalid_request_error"))
|
||||
if self._uses_deferred_backend_setup():
|
||||
raise
|
||||
continue
|
||||
if guardrail_turn_detection_injected and sent:
|
||||
self._guardrail_turn_detection_update_sent = True
|
||||
|
||||
|
|
|
|||
|
|
@ -4,12 +4,13 @@ post-call guardrail rewrote rather than the provider response the proxy assemble
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from itertools import chain
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream, TextCompletionResponse
|
||||
|
||||
SERVED_OUTPUT_TEXTS_KEY: Final = "served_output_texts"
|
||||
|
||||
|
|
@ -44,6 +45,19 @@ class _ChatChoices(LiteLLMBaseModel):
|
|||
choices: list[object]
|
||||
|
||||
|
||||
class _TextCompletionChoice(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
index: int | None = None
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class _TextCompletionChunk(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
choices: tuple[_TextCompletionChoice, ...]
|
||||
|
||||
|
||||
def _as_json_object(response: object) -> dict[str, object] | None:
|
||||
candidate: Final = response.model_dump() if isinstance(response, BaseModel) else response
|
||||
try:
|
||||
|
|
@ -108,8 +122,13 @@ def served_output_texts(response: object) -> ServedTexts | None:
|
|||
|
||||
|
||||
def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
|
||||
if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks):
|
||||
return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)))
|
||||
chat_chunks: Final = tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream))
|
||||
if chunks and len(chat_chunks) == len(chunks):
|
||||
return _indexed_texts_by_choice(tuple(chain.from_iterable(_chat_chunk_texts(chunk) for chunk in chat_chunks)))
|
||||
if chunks and all(isinstance(chunk, TextCompletionResponse) for chunk in chunks):
|
||||
return _indexed_texts_by_choice(
|
||||
tuple(chain.from_iterable(_text_completion_chunk_texts(chunk) for chunk in chunks))
|
||||
)
|
||||
from litellm.llms.anthropic.messages.stream_assembly import (
|
||||
assemble_anthropic_sse_stream,
|
||||
is_anthropic_sse_stream,
|
||||
|
|
@ -121,19 +140,26 @@ def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
|
|||
return _chat_texts(assembled) if assembled is not None else None
|
||||
|
||||
|
||||
def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) -> str | None:
|
||||
contents: Final = tuple(
|
||||
content
|
||||
for chunk in chunks
|
||||
def _chat_chunk_texts(chunk: ModelResponseStream) -> tuple[tuple[int, str | None], ...]:
|
||||
return tuple(
|
||||
(choice.index or 0, content if isinstance(content := choice.delta.content, str) else None)
|
||||
for choice in chunk.choices
|
||||
if (choice.index or 0) == index and isinstance(content := choice.delta.content, str)
|
||||
)
|
||||
return "".join(contents) if contents else None
|
||||
|
||||
|
||||
def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None:
|
||||
choice_count: Final = max(((choice.index or 0) + 1 for chunk in chunks for choice in chunk.choices), default=0)
|
||||
texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count))
|
||||
def _text_completion_chunk_texts(chunk: object) -> tuple[tuple[int, str | None], ...]:
|
||||
parsed: Final = _TextCompletionChunk.model_validate(_as_json_object(chunk))
|
||||
return tuple((choice.index or 0, choice.text) for choice in parsed.choices)
|
||||
|
||||
|
||||
def _joined_choice_text(indexed_texts: Sequence[tuple[int, str | None]], index: int) -> str | None:
|
||||
texts: Final = tuple(text for text_index, text in indexed_texts if text_index == index and text is not None)
|
||||
return "".join(texts) if texts else None
|
||||
|
||||
|
||||
def _indexed_texts_by_choice(indexed_texts: Sequence[tuple[int, str | None]]) -> ServedTexts | None:
|
||||
choice_count: Final = max((index + 1 for index, _ in indexed_texts), default=0)
|
||||
texts: Final = tuple(_joined_choice_text(indexed_texts, index) for index in range(choice_count))
|
||||
return texts if any(text is not None for text in texts) else None
|
||||
|
||||
|
||||
|
|
@ -170,6 +196,8 @@ def _choice_with_text(choice: object, text: str | None) -> object:
|
|||
if choice_obj is None or text is None:
|
||||
return choice
|
||||
message: Final = _as_json_object(choice_obj.get("message"))
|
||||
if message is None or message.get("content") == text:
|
||||
return choice
|
||||
return {**choice_obj, "message": {**message, "content": text}}
|
||||
if message is not None:
|
||||
return choice if message.get("content") == text else {**choice_obj, "message": {**message, "content": text}}
|
||||
if isinstance(choice_obj.get("text"), str):
|
||||
return {**choice_obj, "text": text}
|
||||
return choice
|
||||
|
|
|
|||
27
litellm/litellm_core_utils/stream_usage_chunks.py
Normal file
27
litellm/litellm_core_utils/stream_usage_chunks.py
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
from collections.abc import Iterable
|
||||
from typing import Final
|
||||
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices
|
||||
|
||||
|
||||
def is_empty_streaming_choice(choice: StreamingChoices) -> bool:
|
||||
if choice.finish_reason is not None:
|
||||
return False
|
||||
if getattr(choice, "logprobs", None) is not None:
|
||||
return False
|
||||
delta: Final = choice.delta if hasattr(choice, "delta") else None
|
||||
if delta is None:
|
||||
return True
|
||||
return _all_values_none(delta.model_dump().values())
|
||||
|
||||
|
||||
def _all_values_none(values: Iterable[object]) -> bool:
|
||||
return all(value is None for value in values)
|
||||
|
||||
|
||||
def is_injected_stream_usage_artifact(chunk: object) -> bool:
|
||||
if not isinstance(chunk, ModelResponseStream):
|
||||
return False
|
||||
if chunk.provider_specific_fields is not None:
|
||||
return False
|
||||
return all(is_empty_streaming_choice(choice) for choice in chunk.choices or [])
|
||||
|
|
@ -19,6 +19,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
)
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import partition_annotations
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict
|
||||
)
|
||||
|
|
@ -1700,10 +1701,17 @@ def stream_chunk_builder(
|
|||
# Merge annotations from ALL chunks — providers may spread
|
||||
# them across multiple streaming chunks or send them only in
|
||||
# the final chunk.
|
||||
all_annotations: Final[list] = []
|
||||
all_annotations: Final[list[object]] = []
|
||||
for ac in annotation_chunks:
|
||||
all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
|
||||
response["choices"][0]["message"]["annotations"] = all_annotations
|
||||
message_annotations, provider_annotations = partition_annotations(all_annotations)
|
||||
if message_annotations is not None:
|
||||
response["choices"][0]["message"]["annotations"] = message_annotations
|
||||
if provider_annotations:
|
||||
response.choices[0].message.provider_specific_fields = {
|
||||
**(response.choices[0].message.provider_specific_fields or {}),
|
||||
"annotations": list(provider_annotations),
|
||||
}
|
||||
|
||||
audio_chunks: Final = [
|
||||
chunk
|
||||
|
|
@ -1766,7 +1774,10 @@ def stream_chunk_builder(
|
|||
|
||||
if combined_provider_fields:
|
||||
_choice = response.choices[0]
|
||||
_choice.message.provider_specific_fields = combined_provider_fields
|
||||
_choice.message.provider_specific_fields = {
|
||||
**(_choice.message.provider_specific_fields or {}),
|
||||
**combined_provider_fields,
|
||||
}
|
||||
|
||||
completion_output = get_content_from_model_response(response)
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import copy
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Container, Mapping, MutableMapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
|
||||
|
|
@ -1774,6 +1775,48 @@ def anthropic_request_errors_on_thinking_binding_mismatch(request_body: object)
|
|||
return binding is not None and binding.get("prefix_mismatch_behavior") == "error"
|
||||
|
||||
|
||||
_CLAUDE_STRICT_TOOLS_REJECTION: Final = re.compile(
|
||||
r"Too many strict tools \(\d+\)"
|
||||
r"|Schemas contains too many (?:optional parameters|parameters with union types)"
|
||||
r"|[Cc]ompiled grammar (?:is too large|size)"
|
||||
r"|tools\.\d+\.custom(?:\.strict)?: "
|
||||
r"|Grammar compilation timed out"
|
||||
)
|
||||
_CLAUDE_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
|
||||
|
||||
def is_claude_strict_tools_rejection(status_code: int, error_text: str) -> bool:
|
||||
return status_code == 400 and _CLAUDE_STRICT_TOOLS_REJECTION.search(error_text) is not None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ToolsWithoutStrict:
|
||||
tools: tuple[Mapping[str, object], ...]
|
||||
dropped_count: int
|
||||
|
||||
|
||||
def tools_without_strict(tools: object) -> ToolsWithoutStrict | None:
|
||||
try:
|
||||
validated: Final = _CLAUDE_TOOLS_ADAPTER.validate_python(tools)
|
||||
except ValidationError:
|
||||
return None
|
||||
dropped_count: Final = sum(tool.get("strict") is True for tool in validated)
|
||||
if dropped_count == 0:
|
||||
return None
|
||||
return ToolsWithoutStrict(
|
||||
tools=tuple({key: value for key, value in tool.items() if key != "strict"} for tool in validated),
|
||||
dropped_count=dropped_count,
|
||||
)
|
||||
|
||||
|
||||
def warn_claude_strict_tools_dropped(dropped_count: int) -> None:
|
||||
litellm.verbose_logger.warning(
|
||||
"Claude rejected the request's strict schemas, so litellm removed `strict` from every tool (%d had "
|
||||
"strict true) and sent the request again because drop_params is set.",
|
||||
dropped_count,
|
||||
)
|
||||
|
||||
|
||||
def strip_thinking_blocks_from_anthropic_messages(messages: Sequence[object]) -> list[object]:
|
||||
"""
|
||||
Return a new message list with thinking / redacted_thinking content blocks removed
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# What is this?
|
||||
## Translates OpenAI call to Anthropic `/v1/messages` format
|
||||
import copy
|
||||
import io
|
||||
import itertools
|
||||
import json
|
||||
import traceback
|
||||
|
|
@ -181,6 +182,44 @@ class _ToolCallPart:
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ThinkingPiece:
|
||||
text: str
|
||||
signature: str
|
||||
index: int | None
|
||||
|
||||
@staticmethod
|
||||
def of(chunk: "ModelResponseStream") -> "_ThinkingPiece | None":
|
||||
choices: Final = _optional_attr_sequence(chunk, "choices")
|
||||
if len(choices) != 1 or _optional_attr(choices[0], "finish_reason") is not None:
|
||||
return None
|
||||
delta: Final = _optional_attr(choices[0], "delta")
|
||||
if _optional_attr(delta, "content") or _optional_attr(delta, "tool_calls"):
|
||||
return None
|
||||
blocks: Final = _optional_attr_sequence(delta, "thinking_blocks")
|
||||
if not blocks:
|
||||
reasoning: Final = _optional_str_attr(delta, "reasoning_content")
|
||||
return _ThinkingPiece(text=reasoning, signature="", index=None) if reasoning else None
|
||||
block: Final = as_str_mapping(blocks[0]) if len(blocks) == 1 else None
|
||||
if block is None or block.get("type") != "thinking":
|
||||
return None
|
||||
text: Final = block.get("thinking")
|
||||
signature: Final = block.get("signature")
|
||||
index: Final = block.get("index")
|
||||
return _ThinkingPiece(
|
||||
text=text if isinstance(text, str) else "",
|
||||
signature=signature if isinstance(signature, str) else "",
|
||||
index=index if isinstance(index, int) else None,
|
||||
)
|
||||
|
||||
def block(self, *, text: str, signature: str) -> ChatCompletionThinkingBlock:
|
||||
if self.index is None:
|
||||
return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature or None)
|
||||
return ChatCompletionThinkingBlock(
|
||||
type="thinking", thinking=text, signature=signature or None, index=self.index
|
||||
)
|
||||
|
||||
|
||||
class _CombinedChunkSplitter:
|
||||
"""
|
||||
Splits a streaming chunk that carries BOTH response content and a
|
||||
|
|
@ -209,6 +248,10 @@ class _CombinedChunkSplitter:
|
|||
self._buffer: deque[ModelResponseStream] = deque()
|
||||
self._held_calls: Mapping[int, tuple[ModelResponseStream, _ToolCallPart]] = MappingProxyType({})
|
||||
self._live_call_index: int | None = None
|
||||
self._held_signature: ModelResponseStream | None = None
|
||||
self._unsigned_thinking = io.StringIO()
|
||||
self._unsigned_thinking_length = 0
|
||||
self._thinking_index: int | None = None
|
||||
self._closed = False
|
||||
|
||||
async def aclose(self) -> None:
|
||||
|
|
@ -458,7 +501,45 @@ class _CombinedChunkSplitter:
|
|||
for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits
|
||||
)
|
||||
block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks)))
|
||||
return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, block_chunks)))
|
||||
signed_chunks: Final = tuple(itertools.chain.from_iterable(map(self._defer_signature, block_chunks)))
|
||||
return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, signed_chunks)))
|
||||
|
||||
def _end_thinking_run(self) -> "tuple[ModelResponseStream, ...]":
|
||||
held: Final = self._held_signature
|
||||
self._held_signature = None
|
||||
self._unsigned_thinking = io.StringIO()
|
||||
self._unsigned_thinking_length = 0
|
||||
self._thinking_index = None
|
||||
return (held,) if held is not None else ()
|
||||
|
||||
def _defer_signature(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]":
|
||||
"""Stream a signed thinking chunk's new text as thinking and send its signature once, when the block ends.
|
||||
|
||||
Anthropic-shaped producers stream unsigned text and then sign it, either with no text or
|
||||
restating the unsigned text, which passes through untouched. Mistral signs every chunk while
|
||||
each one still carries new text, which would otherwise collapse into ``signature_delta``
|
||||
events and lose the text.
|
||||
"""
|
||||
piece: Final = _ThinkingPiece.of(chunk)
|
||||
released: Final = self._end_thinking_run() if piece is None or piece.index != self._thinking_index else ()
|
||||
if piece is None:
|
||||
return (*released, chunk)
|
||||
self._thinking_index = piece.index
|
||||
if not piece.signature:
|
||||
self._unsigned_thinking_length += self._unsigned_thinking.write(piece.text)
|
||||
return (*released, chunk)
|
||||
self._held_signature = None
|
||||
if not piece.text or self._restates_unsigned_thinking(piece.text):
|
||||
return (*released, chunk)
|
||||
signature_chunk: Final = self._with_thinking_block(
|
||||
chunk, piece.block(text="", signature=piece.signature), keep_role=False
|
||||
)
|
||||
self._clear_usage(signature_chunk)
|
||||
self._held_signature = signature_chunk
|
||||
return (*released, self._with_thinking_block(chunk, piece.block(text=piece.text, signature=""), keep_role=True))
|
||||
|
||||
def _restates_unsigned_thinking(self, text: str) -> bool:
|
||||
return len(text) == self._unsigned_thinking_length and text == self._unsigned_thinking.getvalue()
|
||||
|
||||
def __iter__(self) -> "Iterator[ModelResponseStream]":
|
||||
return self
|
||||
|
|
@ -493,7 +574,7 @@ class _CombinedChunkSplitter:
|
|||
try:
|
||||
chunk = next(self._sync_iter)
|
||||
except StopIteration:
|
||||
self._buffer.extend(self._release_held_calls())
|
||||
self._buffer.extend((*self._end_thinking_run(), *self._release_held_calls()))
|
||||
if not self._buffer:
|
||||
raise
|
||||
break
|
||||
|
|
@ -512,7 +593,7 @@ class _CombinedChunkSplitter:
|
|||
try:
|
||||
chunk = await self._async_iter.__anext__()
|
||||
except StopAsyncIteration:
|
||||
self._buffer.extend(self._release_held_calls())
|
||||
self._buffer.extend((*self._end_thinking_run(), *self._release_held_calls()))
|
||||
if not self._buffer:
|
||||
raise
|
||||
break
|
||||
|
|
|
|||
|
|
@ -123,6 +123,7 @@ from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingCho
|
|||
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
anthropic_image_source_to_openai_url,
|
||||
anthropic_inline_document,
|
||||
parse_tool_call_arguments,
|
||||
reasoning_content_from_thinking_blocks,
|
||||
with_prompt_cache_breakpoint,
|
||||
|
|
@ -185,6 +186,8 @@ from litellm.types.llms.openai import (
|
|||
AllMessageValues,
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionFileObject,
|
||||
ChatCompletionFileObjectFile,
|
||||
ChatCompletionImageObject,
|
||||
ChatCompletionImageUrlObject,
|
||||
ChatCompletionRedactedThinkingBlock,
|
||||
|
|
@ -498,7 +501,9 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
for m in ordered_messages:
|
||||
user_message: ChatCompletionUserMessage | None = None
|
||||
tool_message_list: list[ChatCompletionToolMessage] = []
|
||||
new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = []
|
||||
new_user_content_list: list[
|
||||
ChatCompletionTextObject | ChatCompletionImageObject | ChatCompletionFileObject
|
||||
] = []
|
||||
if m["role"] == "system":
|
||||
system_message = self._translate_midturn_system_message_to_openai(m, model)
|
||||
if system_message is not None:
|
||||
|
|
@ -531,13 +536,8 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
self._add_prompt_cache_breakpoint_if_present(content, image_obj)
|
||||
)
|
||||
elif content.get("type") == "document":
|
||||
# Convert Anthropic document format (PDF, etc.) to OpenAI format
|
||||
source = content.get("source", {})
|
||||
openai_image_url = self._translate_anthropic_image_to_openai(cast(dict, source))
|
||||
|
||||
if openai_image_url:
|
||||
image_url_obj = ChatCompletionImageUrlObject(url=openai_image_url)
|
||||
doc_obj = ChatCompletionImageObject(type="image_url", image_url=image_url_obj)
|
||||
doc_obj = self._translate_anthropic_document_to_openai(content)
|
||||
if doc_obj is not None:
|
||||
self._add_cache_control_if_applicable(content, doc_obj, model)
|
||||
new_user_content_list.append(doc_obj)
|
||||
elif content.get("type") == "tool_result":
|
||||
|
|
@ -1357,7 +1357,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
text: Final = tool_result_text_part(block)
|
||||
return None if text is None else ChatCompletionTextObject(type="text", text=text)
|
||||
case "image" | "document":
|
||||
return self._tool_result_image_part(block.get("source"))
|
||||
return self._source_image_url_part(block.get("source"))
|
||||
case "tool_reference":
|
||||
return ChatCompletionToolReferenceObject(
|
||||
type="tool_reference", tool_name=str(block.get("tool_name") or "")
|
||||
|
|
@ -1365,7 +1365,18 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
case _:
|
||||
return None
|
||||
|
||||
def _tool_result_image_part(self, image_source: object) -> ChatCompletionImageObject | None:
|
||||
def _translate_anthropic_document_to_openai(
|
||||
self, document: Mapping[str, object]
|
||||
) -> ChatCompletionFileObject | ChatCompletionImageObject | None:
|
||||
inline_document: Final = anthropic_inline_document(document)
|
||||
if inline_document is None:
|
||||
return self._source_image_url_part(document.get("source"))
|
||||
return ChatCompletionFileObject(
|
||||
type="file",
|
||||
file=ChatCompletionFileObjectFile(filename=inline_document.filename, file_data=inline_document.file_data),
|
||||
)
|
||||
|
||||
def _source_image_url_part(self, image_source: object) -> ChatCompletionImageObject | None:
|
||||
if not isinstance(image_source, dict):
|
||||
return None
|
||||
openai_image_url = self._translate_anthropic_image_to_openai(image_source)
|
||||
|
|
|
|||
|
|
@ -396,7 +396,8 @@ def _resolve_advisor_router(advisor_model: str) -> "Router | None":
|
|||
return None
|
||||
if llm_router is None:
|
||||
return None
|
||||
if llm_router.is_recognized_model(advisor_model) or llm_router.pattern_router.route(advisor_model):
|
||||
registry: Final = llm_router.registry()
|
||||
if registry.resolve(advisor_model, kind="recognized") or registry.patterns().route(advisor_model):
|
||||
return llm_router
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from typing import Any, Final, cast
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
TOOL_RESULT_IMAGE_BOUNDARY,
|
||||
TOOL_RESULT_IMAGE_PLACEHOLDER,
|
||||
anthropic_inline_document,
|
||||
encrypted_reasoning_signature,
|
||||
responses_reasoning_items_from_thinking_blocks,
|
||||
with_prompt_cache_breakpoint,
|
||||
|
|
@ -106,27 +107,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
block: Mapping[str, object],
|
||||
) -> dict[str, str] | None: # mutable-ok: API message payload
|
||||
"""Convert an Anthropic document block to a Responses input_file part."""
|
||||
inline_document: Final = anthropic_inline_document(block)
|
||||
if inline_document is not None:
|
||||
return {
|
||||
"type": "input_file",
|
||||
"filename": inline_document.filename,
|
||||
"file_data": inline_document.file_data,
|
||||
}
|
||||
raw_source: Final = block.get("source")
|
||||
if not isinstance(raw_source, Mapping):
|
||||
return None
|
||||
source: Final = cast(Mapping[str, object], raw_source) # cast-ok: untrusted client payload
|
||||
source_type: Final = source.get("type")
|
||||
if source_type == "base64":
|
||||
data: Final = source.get("data")
|
||||
if not isinstance(data, str) or not data:
|
||||
return None
|
||||
raw_media_type: Final = source.get("media_type")
|
||||
media_type: Final = (
|
||||
raw_media_type if isinstance(raw_media_type, str) and raw_media_type else "application/pdf"
|
||||
)
|
||||
raw_title: Final = block.get("title")
|
||||
filename: Final = raw_title if isinstance(raw_title, str) and raw_title else "document.pdf"
|
||||
return {
|
||||
"type": "input_file",
|
||||
"filename": filename,
|
||||
"file_data": f"data:{media_type};base64,{data}",
|
||||
}
|
||||
if source_type == "url":
|
||||
if source.get("type") == "url":
|
||||
url: Final = source.get("url")
|
||||
if not isinstance(url, str) or not url:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import copy
|
||||
import enum
|
||||
import re
|
||||
from collections.abc import AsyncIterator, Collection, Iterator
|
||||
from collections.abc import AsyncIterator, Collection, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -10,10 +10,11 @@ from httpx import Response
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
audio_or_image_in_message_content,
|
||||
content_parts,
|
||||
convert_content_list_to_str,
|
||||
drop_tool_reference_parts_from_tool_messages,
|
||||
filter_value_from_dict,
|
||||
is_pdf_file_part,
|
||||
)
|
||||
from litellm.llms.azure.common_utils import BaseAzureLLM
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
|
|
@ -46,6 +47,20 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = (
|
|||
"cache_control",
|
||||
)
|
||||
|
||||
LIST_ONLY_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_audio"})
|
||||
|
||||
|
||||
def _is_list_only_content_part(part: object) -> bool:
|
||||
match part:
|
||||
case {"type": str() as part_type} if part_type in LIST_ONLY_CONTENT_PART_TYPES:
|
||||
return True
|
||||
case _:
|
||||
return is_pdf_file_part(part)
|
||||
|
||||
|
||||
def _has_list_only_content_part(message: Mapping[str, object]) -> bool:
|
||||
return any(_is_list_only_content_part(part) for part in content_parts(message))
|
||||
|
||||
|
||||
class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator):
|
||||
def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract
|
||||
|
|
@ -246,7 +261,7 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
these with "Extra inputs are not permitted", which breaks multi-turn
|
||||
Anthropic-format clients that echo thinking blocks back as history.
|
||||
2. Transforms list content to a string.
|
||||
3. If message contains an image or audio, send as is (user-intended)
|
||||
3. If message contains an image, audio, or PDF file part, send as is (user-intended)
|
||||
|
||||
Operates on a deep copy so the caller's messages keep their thinking blocks
|
||||
and provider metadata, which a fallback to another provider still needs.
|
||||
|
|
@ -257,8 +272,7 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
for field in NON_OPENAI_SPEC_MESSAGE_FIELDS:
|
||||
filter_value_from_dict(message_dict, field)
|
||||
|
||||
# Do nothing if the message contains an image or audio
|
||||
if audio_or_image_in_message_content(message):
|
||||
if _has_list_only_content_part(message):
|
||||
continue
|
||||
|
||||
texts = convert_content_list_to_str(message=message)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
|
@ -201,11 +202,23 @@ class BaseAnthropicMessagesConfig(ABC):
|
|||
anthropic_request_errors_on_thinking_binding_mismatch,
|
||||
is_anthropic_invalid_thinking_block_error,
|
||||
is_anthropic_thinking_binding_mismatch_error,
|
||||
is_claude_strict_tools_rejection,
|
||||
tools_without_strict,
|
||||
validated_object_mapping,
|
||||
)
|
||||
|
||||
if e.response.status_code != 400:
|
||||
return False
|
||||
error_text: Final[str] = e.response.text
|
||||
body: Final = validated_object_mapping(request_body)
|
||||
drops_rejected_strict_tools: Final = (
|
||||
drops_unsupported_params(litellm_params.get("drop_params"))
|
||||
and is_claude_strict_tools_rejection(e.response.status_code, error_text)
|
||||
and body is not None
|
||||
and tools_without_strict(body.get("tools")) is not None
|
||||
)
|
||||
if drops_rejected_strict_tools:
|
||||
return True
|
||||
caller_asked_for_this_error: Final = anthropic_request_errors_on_thinking_binding_mismatch(
|
||||
request_body
|
||||
) and is_anthropic_thinking_binding_mismatch_error(error_text)
|
||||
|
|
@ -219,11 +232,22 @@ class BaseAnthropicMessagesConfig(ABC):
|
|||
"""
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
is_anthropic_invalid_thinking_block_error,
|
||||
is_claude_strict_tools_rejection,
|
||||
strip_thinking_blocks_from_anthropic_messages_request_dict,
|
||||
tools_without_strict,
|
||||
warn_claude_strict_tools_dropped,
|
||||
)
|
||||
|
||||
if e.response.status_code == 400 and is_anthropic_invalid_thinking_block_error(e.response.text):
|
||||
strip_thinking_blocks_from_anthropic_messages_request_dict(request_data)
|
||||
without_strict: Final = (
|
||||
tools_without_strict(request_data.get("tools"))
|
||||
if is_claude_strict_tools_rejection(e.response.status_code, e.response.text)
|
||||
else None
|
||||
)
|
||||
if without_strict is not None:
|
||||
warn_claude_strict_tools_dropped(without_strict.dropped_count)
|
||||
request_data["tools"] = list(without_strict.tools) # rebind-ok: the retry loop re-signs this same dict
|
||||
return request_data
|
||||
|
||||
def native_output_fallback_litellm_params(
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ from types import TracebackType
|
|||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
|
||||
import httpx
|
||||
from typing_extensions import Self
|
||||
|
||||
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
|
||||
from litellm.types.realtime import (
|
||||
|
|
@ -24,7 +23,7 @@ else:
|
|||
|
||||
|
||||
class RealtimeBackend(Protocol):
|
||||
async def __aenter__(self) -> Self: ...
|
||||
async def __aenter__(self) -> object: ...
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
|
|
@ -97,9 +96,15 @@ class BaseRealtimeConfig(ABC):
|
|||
def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None:
|
||||
return None
|
||||
|
||||
def session_refusal(self) -> Exception | None:
|
||||
return None
|
||||
|
||||
async def open_backend(self, url: str, headers: Mapping[str, str]) -> RealtimeBackend | None:
|
||||
return None
|
||||
|
||||
def wrap_backend(self, backend: RealtimeBackend) -> RealtimeBackend:
|
||||
return backend
|
||||
|
||||
def transform_session_created_event(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -93,9 +93,12 @@ class RouterVectorStoreEmbeddingExecutor:
|
|||
|
||||
def _router_serves(self, model: str) -> bool:
|
||||
team_id: Final = self.metadata.get("user_api_key_team_id")
|
||||
resolved: Final = self.router.resolved_litellm_models(model, team_id if isinstance(team_id, str) else None)
|
||||
registry: Final = self.router.registry()
|
||||
resolved: Final = registry.names(
|
||||
model, kind="litellm_models", team_id=team_id if isinstance(team_id, str) else None
|
||||
)
|
||||
deployment_models: Final = (
|
||||
deployment.get("litellm_params", {}).get("model") for deployment in self.router.get_model_list() or ()
|
||||
deployment.get("litellm_params", {}).get("model") for deployment in registry.deployments() or ()
|
||||
)
|
||||
return bool(resolved) or model in deployment_models
|
||||
|
||||
|
|
@ -124,6 +127,9 @@ class BaseVectorStoreConfig:
|
|||
def validate_create_vector_store(self) -> None:
|
||||
return None
|
||||
|
||||
def get_async_httpx_client_params(self) -> Mapping[str, object]:
|
||||
return MappingProxyType({})
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]:
|
||||
return []
|
||||
|
||||
|
|
@ -226,6 +232,9 @@ class BaseVectorStoreConfig:
|
|||
``get_complete_url`` adds a query string, like Azure, override this."""
|
||||
return f"{api_base}/{encoded_vector_store_id}"
|
||||
|
||||
def get_vector_store_list_url(self, api_base: str) -> str:
|
||||
return api_base
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
|
||||
from ..chat.transformation import BaseLLMException
|
||||
|
||||
|
|
|
|||
|
|
@ -2,11 +2,11 @@ from collections.abc import Callable, Coroutine
|
|||
from functools import partial
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from litellm.llms.anthropic.common_utils import is_claude_strict_tools_rejection
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
NativeOutputFallback,
|
||||
is_bedrock_native_output_rejection,
|
||||
is_bedrock_strict_tools_rejection,
|
||||
warn_bedrock_native_output_rejected,
|
||||
warn_bedrock_strict_tools_dropped,
|
||||
)
|
||||
|
|
@ -34,7 +34,7 @@ def _native_output_retry(
|
|||
def _strict_tools_retry(
|
||||
error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str
|
||||
) -> NativeOutputFallback | None:
|
||||
if fallback is None or not is_bedrock_strict_tools_rejection(error.status_code, error.message):
|
||||
if fallback is None or not is_claude_strict_tools_rejection(error.status_code, error.message):
|
||||
return None
|
||||
retry: Final = fallback()
|
||||
if retry is not None:
|
||||
|
|
|
|||
|
|
@ -366,21 +366,10 @@ class NativeOutputFallback:
|
|||
litellm_params: Mapping[str, object]
|
||||
|
||||
|
||||
_BEDROCK_STRICT_TOOLS_REJECTION: Final = re.compile(
|
||||
r"Too many strict tools \(\d+\)"
|
||||
r"|Schemas contains too many (?:optional parameters|parameters with union types)"
|
||||
r"|[Cc]ompiled grammar (?:is too large|size)"
|
||||
r"|tools\.\d+\.custom(?:\.strict)?: "
|
||||
r"|Grammar compilation timed out"
|
||||
)
|
||||
_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
_TOOL_SPEC_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def is_bedrock_strict_tools_rejection(status_code: int, error_text: str) -> bool:
|
||||
return status_code == 400 and _BEDROCK_STRICT_TOOLS_REJECTION.search(error_text) is not None
|
||||
|
||||
|
||||
def warn_bedrock_strict_tools_dropped(model: str) -> None:
|
||||
verbose_logger.warning(
|
||||
"Bedrock rejected the strict tool schemas for model=%s, so litellm sent the request again without `strict` "
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ import httpx
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items
|
||||
|
|
@ -196,7 +197,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
|
||||
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
|
||||
)
|
||||
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
|
||||
hoisted: Final = hoist_additional_tools(
|
||||
input, params.get("tools"), drop_params=drops_unsupported_params(litellm_params.drop_params)
|
||||
)
|
||||
normalized_input, rewritten_types = normalize_codex_input_items(hoisted.input)
|
||||
if rewritten_types:
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
128
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
128
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
import weakref
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
HTTP2_STREAMS_PER_CONNECTION: Final = 100
|
||||
|
||||
|
||||
class _Lane:
|
||||
__slots__ = ("in_flight", "transport")
|
||||
|
||||
def __init__(self, transport: httpx.AsyncBaseTransport) -> None:
|
||||
self.transport: Final = transport
|
||||
self.in_flight = 0
|
||||
|
||||
def acquire(self) -> None:
|
||||
self.in_flight += 1
|
||||
|
||||
def release(self) -> None:
|
||||
self.in_flight -= 1
|
||||
|
||||
|
||||
class _Lanes:
|
||||
def __init__(
|
||||
self,
|
||||
new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport],
|
||||
limits: httpx.Limits,
|
||||
streams_per_connection: int,
|
||||
) -> None:
|
||||
self._new_transport: Final = new_transport
|
||||
self._limits: Final = limits
|
||||
self._streams_per_connection: Final = streams_per_connection
|
||||
self.lanes: tuple[_Lane, ...] = ()
|
||||
|
||||
def acquire(self) -> _Lane:
|
||||
free: Final = next((lane for lane in self.lanes if lane.in_flight < self._streams_per_connection), None)
|
||||
lane: Final = free if free is not None else self._open_or_share()
|
||||
lane.acquire()
|
||||
return lane
|
||||
|
||||
def _open_or_share(self) -> _Lane:
|
||||
cap: Final = self._limits.max_connections
|
||||
if cap is not None and len(self.lanes) * self._streams_per_connection >= cap:
|
||||
return min(self.lanes, key=lambda lane: lane.in_flight)
|
||||
return self._open()
|
||||
|
||||
def _open(self) -> _Lane:
|
||||
lane: Final = _Lane(self._new_transport(self._next_lane_limits()))
|
||||
self.lanes = (*self.lanes, lane)
|
||||
return lane
|
||||
|
||||
def _next_lane_limits(self) -> httpx.Limits:
|
||||
cap: Final = self._limits.max_connections
|
||||
if cap is None:
|
||||
return self._limits
|
||||
return httpx.Limits(
|
||||
max_connections=min(self._streams_per_connection, cap - len(self.lanes) * self._streams_per_connection),
|
||||
max_keepalive_connections=self._limits.max_keepalive_connections,
|
||||
keepalive_expiry=self._limits.keepalive_expiry,
|
||||
)
|
||||
|
||||
|
||||
class _AsyncReleasingStream(httpx.AsyncByteStream):
|
||||
def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None:
|
||||
self._stream: Final = stream
|
||||
self._release: Final = weakref.finalize(self, release)
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
try:
|
||||
async for chunk in self._stream:
|
||||
yield chunk
|
||||
except BaseException:
|
||||
self._release()
|
||||
raise
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._release()
|
||||
await self._stream.aclose()
|
||||
|
||||
|
||||
def _allowed_over_http2(name: bytes, value: bytes) -> bool:
|
||||
return name.lower() != b"te" or value.strip().lower() == b"trailers"
|
||||
|
||||
|
||||
def _without_te_codings(request: httpx.Request) -> httpx.Request:
|
||||
if request.url.scheme != "https":
|
||||
return request
|
||||
kept: Final = tuple((name, value) for name, value in request.headers.raw if _allowed_over_http2(name, value))
|
||||
if len(kept) == len(request.headers.raw):
|
||||
return request
|
||||
return httpx.Request(
|
||||
request.method, request.url, headers=kept, stream=request.stream, extensions=request.extensions
|
||||
)
|
||||
|
||||
|
||||
def _with_stream(response: httpx.Response, stream: httpx.AsyncByteStream) -> httpx.Response:
|
||||
return httpx.Response(response.status_code, headers=response.headers, stream=stream, extensions=response.extensions)
|
||||
|
||||
|
||||
class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport):
|
||||
def __init__(
|
||||
self,
|
||||
new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport],
|
||||
limits: httpx.Limits,
|
||||
streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION,
|
||||
) -> None:
|
||||
self._lanes: Final = _Lanes(new_transport, limits, streams_per_connection)
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
lane: Final = self._lanes.acquire()
|
||||
try:
|
||||
response: Final = await lane.transport.handle_async_request(_without_te_codings(request))
|
||||
except BaseException:
|
||||
lane.release()
|
||||
raise
|
||||
stream: Final = response.stream
|
||||
if not isinstance(stream, httpx.AsyncByteStream):
|
||||
lane.release()
|
||||
return response
|
||||
return _with_stream(response, _AsyncReleasingStream(stream, lane.release))
|
||||
|
||||
def has_request_in_flight(self) -> bool:
|
||||
return any(lane.in_flight > 0 for lane in self._lanes.lanes)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
for lane in self._lanes.lanes:
|
||||
await lane.transport.aclose()
|
||||
|
|
@ -18,6 +18,7 @@ import certifi
|
|||
import httpx
|
||||
from aiohttp import ClientSession, DummyCookieJar, TCPConnector
|
||||
from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport
|
||||
from httpx._config import DEFAULT_LIMITS
|
||||
from httpx._types import CertTypes, RequestFiles
|
||||
from httpx._utils import get_environment_proxies
|
||||
from pydantic import TypeAdapter
|
||||
|
|
@ -43,6 +44,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
|||
from litellm.litellm_core_utils.request_timeout_resolver import (
|
||||
get_configured_request_timeout,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread
|
||||
from litellm.types.llms.custom_http import *
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -73,7 +75,7 @@ _RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes
|
|||
|
||||
_IPV4_LOCAL_ADDRESS: Final = "0.0.0.0"
|
||||
|
||||
_HttpxTransportT = TypeVar("_HttpxTransportT", HTTPTransport, AsyncHTTPTransport)
|
||||
_HttpxTransportT: Final = TypeVar("_HttpxTransportT", httpx.BaseTransport, httpx.AsyncBaseTransport)
|
||||
|
||||
HTTPEventHooks: TypeAlias = Mapping[str, list[Callable[..., object]]]
|
||||
|
||||
|
|
@ -84,6 +86,36 @@ def http2_enabled() -> bool:
|
|||
return litellm.http2 is True or str_to_bool(os.getenv("LITELLM_HTTP2", "False")) is True
|
||||
|
||||
|
||||
_HTTP2_OFF_VALUES: Final = frozenset({"false", "0", "no", "off"})
|
||||
|
||||
|
||||
def _is_http2_off(value: bool | int | str | None) -> bool:
|
||||
return str(value).strip().lower() in _HTTP2_OFF_VALUES
|
||||
|
||||
|
||||
def http2_disabled() -> bool:
|
||||
return _is_http2_off(litellm.http2) or _is_http2_off(os.getenv("LITELLM_HTTP2"))
|
||||
|
||||
|
||||
def _resolve_http2(provider_default: bool, ssl_verify: VerifyTypes | None) -> bool:
|
||||
if http2_enabled():
|
||||
return True
|
||||
return provider_default and not http2_disabled() and not isinstance(ssl_verify, ssl.SSLContext)
|
||||
|
||||
|
||||
def _aiohttp_env_proxies_disabled() -> bool:
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
return litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) is True
|
||||
|
||||
|
||||
_AIOHTTP_EQUIVALENT_LIMITS: Final = httpx.Limits(
|
||||
max_connections=AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT > 0 else None,
|
||||
max_keepalive_connections=None,
|
||||
keepalive_expiry=AIOHTTP_KEEPALIVE_TIMEOUT,
|
||||
)
|
||||
|
||||
|
||||
def _environment_proxy_mounts(
|
||||
build_proxy_transport: Callable[[str], _HttpxTransportT],
|
||||
) -> Mapping[str, _HttpxTransportT | None]:
|
||||
|
|
@ -145,7 +177,7 @@ def build_aiohttp_keepalive_socket_factory() -> Callable[[_AddrInfo], socket.soc
|
|||
_build_aiohttp_keepalive_socket_factory = build_aiohttp_keepalive_socket_factory
|
||||
|
||||
|
||||
def get_default_headers() -> dict:
|
||||
def get_default_headers() -> dict[str, str]:
|
||||
"""
|
||||
Get default headers for HTTP requests.
|
||||
|
||||
|
|
@ -280,9 +312,7 @@ def _prepare_request_data_and_content(
|
|||
|
||||
|
||||
# Cache for SSL contexts to avoid creating duplicate contexts with the same configuration
|
||||
# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve)
|
||||
# Value: ssl.SSLContext
|
||||
_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None], ssl.SSLContext]] = {}
|
||||
_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None, bool], ssl.SSLContext]] = {}
|
||||
|
||||
|
||||
def _create_ssl_context(
|
||||
|
|
@ -375,6 +405,7 @@ def get_ssl_verify(
|
|||
|
||||
def get_ssl_configuration(
|
||||
ssl_verify: VerifyTypes | None = None,
|
||||
http2: bool = False,
|
||||
) -> bool | str | ssl.SSLContext:
|
||||
"""
|
||||
Unified SSL configuration function that handles ssl_context and ssl_verify logic.
|
||||
|
|
@ -397,6 +428,8 @@ def get_ssl_configuration(
|
|||
- False: Disable SSL verification
|
||||
- True: Enable SSL verification
|
||||
- str: Path to CA bundle file
|
||||
http2: Whether the context is for an HTTP/2 client. httpx sets ALPN on the context it is given,
|
||||
so HTTP/2 clients get their own cached context instead of sharing the HTTP/1.1 one
|
||||
|
||||
Returns:
|
||||
Union[bool, str, ssl.SSLContext]: Appropriate SSL configuration
|
||||
|
|
@ -423,7 +456,7 @@ def get_ssl_configuration(
|
|||
|
||||
if ssl_verify is not False:
|
||||
# Create cache key from configuration parameters
|
||||
cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve)
|
||||
cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve, http2)
|
||||
|
||||
# Check if we have a cached SSL context for this configuration
|
||||
if cache_key not in _ssl_context_cache:
|
||||
|
|
@ -439,6 +472,11 @@ def get_ssl_configuration(
|
|||
return ssl_verify
|
||||
|
||||
|
||||
def _httpx_default_verify_unless_configured(verify: VerifyTypes) -> VerifyTypes:
|
||||
litellm_defaults: Final = (get_ssl_configuration(True), get_ssl_configuration(True, http2=True))
|
||||
return True if any(verify is default for default in litellm_defaults) else verify
|
||||
|
||||
|
||||
_shared_realtime_ssl_context: bool | str | ssl.SSLContext | None = None
|
||||
|
||||
|
||||
|
|
@ -662,6 +700,7 @@ class AsyncHTTPHandler:
|
|||
shared_session: Optional["ClientSession"] = None,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
follow_redirects: bool = True,
|
||||
http2: bool = False,
|
||||
):
|
||||
self.timeout = timeout
|
||||
self.event_hooks = event_hooks
|
||||
|
|
@ -669,6 +708,7 @@ class AsyncHTTPHandler:
|
|||
self.shared_session = shared_session
|
||||
self.transport = transport
|
||||
self.follow_redirects = follow_redirects
|
||||
self.http2 = http2
|
||||
self._owns_client = True
|
||||
self._client = self.create_client(
|
||||
timeout=timeout,
|
||||
|
|
@ -711,8 +751,11 @@ class AsyncHTTPHandler:
|
|||
follow_redirects=self.follow_redirects,
|
||||
trust_env=False,
|
||||
)
|
||||
use_http2: Final = _resolve_http2(self.http2, ssl_verify)
|
||||
replaces_aiohttp: Final = use_http2 and AsyncHTTPHandler.should_use_aiohttp_transport()
|
||||
limits: Final = _AIOHTTP_EQUIVALENT_LIMITS if replaces_aiohttp else DEFAULT_LIMITS
|
||||
# Get unified SSL configuration
|
||||
ssl_config: Final = get_ssl_configuration(ssl_verify)
|
||||
ssl_config: Final = get_ssl_configuration(ssl_verify, http2=use_http2)
|
||||
|
||||
# An SSL certificate used by the requested host to authenticate the client.
|
||||
# /path/to/client.pem
|
||||
|
|
@ -726,14 +769,23 @@ class AsyncHTTPHandler:
|
|||
ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None,
|
||||
ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
|
||||
shared_session=shared_session,
|
||||
http2=use_http2,
|
||||
cert=cert,
|
||||
limits=limits,
|
||||
)
|
||||
|
||||
# Get default headers (User-Agent, overridable via LITELLM_USER_AGENT)
|
||||
default_headers: Final = get_default_headers()
|
||||
trust_env: Final = not (replaces_aiohttp and _aiohttp_env_proxies_disabled())
|
||||
|
||||
return httpx.AsyncClient(
|
||||
transport=transport,
|
||||
mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert),
|
||||
mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(
|
||||
transport, verify=ssl_config, cert=cert, http2=use_http2, limits=limits
|
||||
)
|
||||
if trust_env
|
||||
else None,
|
||||
limits=limits,
|
||||
event_hooks=event_hooks,
|
||||
timeout=timeout,
|
||||
verify=ssl_config,
|
||||
|
|
@ -741,7 +793,8 @@ class AsyncHTTPHandler:
|
|||
headers=default_headers,
|
||||
cookies=blocked_cookie_jar(),
|
||||
follow_redirects=self.follow_redirects,
|
||||
http2=http2_enabled(),
|
||||
http2=use_http2,
|
||||
trust_env=trust_env,
|
||||
)
|
||||
|
||||
async def close(self):
|
||||
|
|
@ -1152,6 +1205,17 @@ class AsyncHTTPHandler:
|
|||
# Mirrors LiteLLMAiohttpTransport._background_close_tasks.
|
||||
_finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
|
||||
|
||||
@classmethod
|
||||
async def _close_client_at_finalization(cls, client: httpx.AsyncClient) -> None:
|
||||
try:
|
||||
await client.aclose()
|
||||
except (OSError, RuntimeError) as exc:
|
||||
verbose_logger.debug("Error closing client at finalization: %s", exc)
|
||||
finally:
|
||||
task: Final = asyncio.current_task()
|
||||
if task is not None:
|
||||
cls._finalizer_close_tasks.discard(task)
|
||||
|
||||
@classmethod
|
||||
def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
|
||||
cls._finalizer_close_tasks.discard(task)
|
||||
|
|
@ -1219,8 +1283,8 @@ class AsyncHTTPHandler:
|
|||
# here is the cross-loop path the transport refuses.
|
||||
self._dispose_wrapped_aiohttp_session()
|
||||
return
|
||||
task: Final = loop.create_task(self._client.aclose())
|
||||
cls: Final = type(self)
|
||||
task: Final = loop.create_task(cls._close_client_at_finalization(self._client))
|
||||
cls._finalizer_close_tasks.add(task)
|
||||
task.add_done_callback(cls._on_finalizer_close_done)
|
||||
except Exception:
|
||||
|
|
@ -1231,7 +1295,10 @@ class AsyncHTTPHandler:
|
|||
ssl_context: ssl.SSLContext | None = None,
|
||||
ssl_verify: bool | None = None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None:
|
||||
http2: bool = False,
|
||||
cert: CertTypes | None = None,
|
||||
limits: httpx.Limits = DEFAULT_LIMITS,
|
||||
) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None:
|
||||
"""
|
||||
- Creates a transport for httpx.AsyncClient
|
||||
- if litellm.force_ipv4 is True, it will return AsyncHTTPTransport with local_address="0.0.0.0"
|
||||
|
|
@ -1249,7 +1316,7 @@ class AsyncHTTPHandler:
|
|||
#########################################################
|
||||
# AIOHTTP TRANSPORT is off by default
|
||||
#########################################################
|
||||
if AsyncHTTPHandler.should_use_aiohttp_transport():
|
||||
if not http2 and AsyncHTTPHandler.should_use_aiohttp_transport():
|
||||
return AsyncHTTPHandler.create_aiohttp_transport(
|
||||
ssl_context=ssl_context,
|
||||
ssl_verify=ssl_verify,
|
||||
|
|
@ -1259,7 +1326,12 @@ class AsyncHTTPHandler:
|
|||
#########################################################
|
||||
# HTTPX TRANSPORT is used when aiohttp is not installed
|
||||
#########################################################
|
||||
return AsyncHTTPHandler._create_httpx_transport()
|
||||
return AsyncHTTPHandler._create_httpx_transport(
|
||||
http2=http2,
|
||||
verify=ssl_context if ssl_context is not None else ssl_verify is not False,
|
||||
cert=cert,
|
||||
limits=limits,
|
||||
)
|
||||
|
||||
_create_async_transport = create_async_transport
|
||||
|
||||
|
|
@ -1406,28 +1478,57 @@ class AsyncHTTPHandler:
|
|||
_create_aiohttp_transport = create_aiohttp_transport
|
||||
|
||||
@staticmethod
|
||||
def _create_httpx_transport() -> AsyncHTTPTransport | None:
|
||||
def _create_httpx_transport(
|
||||
http2: bool = False,
|
||||
verify: ssl.SSLContext | bool = True,
|
||||
cert: CertTypes | None = None,
|
||||
limits: httpx.Limits = DEFAULT_LIMITS,
|
||||
) -> AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None:
|
||||
"""
|
||||
Creates an AsyncHTTPTransport
|
||||
|
||||
- If HTTP/2 is on, it spreads requests over as many HTTP/2 connections as the stream limit needs
|
||||
- If force_ipv4 is True, it will create an AsyncHTTPTransport with local_address set to "0.0.0.0"
|
||||
- [Default] If force_ipv4 is False, it will return None
|
||||
- [Default] Otherwise it will return None
|
||||
"""
|
||||
local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None
|
||||
transport_verify: Final = (
|
||||
_httpx_default_verify_unless_configured(verify)
|
||||
if litellm.force_ipv4 and not AsyncHTTPHandler.should_use_aiohttp_transport()
|
||||
else verify
|
||||
)
|
||||
if http2 or http2_enabled():
|
||||
return AsyncHTTP2ConnectionSpread(
|
||||
lambda lane_limits: AsyncHTTPTransport(
|
||||
local_address=local_address, verify=transport_verify, cert=cert, http2=True, limits=lane_limits
|
||||
),
|
||||
limits,
|
||||
)
|
||||
if litellm.force_ipv4:
|
||||
return AsyncHTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2_enabled())
|
||||
else:
|
||||
return None
|
||||
return AsyncHTTPTransport(local_address=local_address, verify=transport_verify, cert=cert, limits=limits)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def create_httpx_proxy_mounts(
|
||||
transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None,
|
||||
transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None,
|
||||
verify: VerifyTypes,
|
||||
cert: CertTypes | None,
|
||||
) -> Mapping[str, AsyncHTTPTransport | None] | None:
|
||||
if not isinstance(transport, AsyncHTTPTransport):
|
||||
http2: bool = False,
|
||||
limits: httpx.Limits = DEFAULT_LIMITS,
|
||||
) -> Mapping[str, httpx.AsyncBaseTransport | None] | None:
|
||||
if not isinstance(transport, (AsyncHTTPTransport, AsyncHTTP2ConnectionSpread)):
|
||||
return None
|
||||
if http2 or http2_enabled():
|
||||
return _environment_proxy_mounts(
|
||||
lambda proxy_url: AsyncHTTP2ConnectionSpread(
|
||||
lambda lane_limits: AsyncHTTPTransport(
|
||||
proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=lane_limits
|
||||
),
|
||||
limits,
|
||||
)
|
||||
)
|
||||
return _environment_proxy_mounts(
|
||||
lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2_enabled())
|
||||
lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, limits=limits)
|
||||
)
|
||||
|
||||
_create_httpx_proxy_mounts = create_httpx_proxy_mounts
|
||||
|
|
@ -1453,8 +1554,9 @@ class HTTPHandler:
|
|||
self._client = self.create_client() if client is None else client
|
||||
|
||||
def create_client(self) -> httpx.Client:
|
||||
use_http2: Final = http2_enabled()
|
||||
# Get unified SSL configuration
|
||||
ssl_config: Final = get_ssl_configuration(self.ssl_verify)
|
||||
ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=use_http2)
|
||||
|
||||
# An SSL certificate used by the requested host to authenticate the client.
|
||||
# /path/to/client.pem
|
||||
|
|
@ -1463,17 +1565,19 @@ class HTTPHandler:
|
|||
# Get default headers (User-Agent, overridable via LITELLM_USER_AGENT)
|
||||
default_headers: Final = get_default_headers() if not self.disable_default_headers else None
|
||||
|
||||
transport: Final = self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2)
|
||||
|
||||
# Create a client with a connection pool
|
||||
return httpx.Client(
|
||||
transport=self._create_sync_transport(),
|
||||
mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert),
|
||||
transport=transport,
|
||||
mounts=self._create_sync_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=use_http2),
|
||||
timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT,
|
||||
verify=ssl_config,
|
||||
cert=cert,
|
||||
headers=default_headers,
|
||||
cookies=blocked_cookie_jar(),
|
||||
follow_redirects=self.follow_redirects,
|
||||
http2=http2_enabled(),
|
||||
http2=use_http2,
|
||||
)
|
||||
|
||||
@property
|
||||
|
|
@ -1769,27 +1873,34 @@ class HTTPHandler:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def _create_sync_transport(self) -> HTTPTransport | None:
|
||||
@staticmethod
|
||||
def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None:
|
||||
"""
|
||||
Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True.
|
||||
Otherwise, return None.
|
||||
Otherwise, return litellm.sync_transport or None.
|
||||
|
||||
Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them
|
||||
"""
|
||||
if litellm.force_ipv4:
|
||||
return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2_enabled())
|
||||
else:
|
||||
return getattr(litellm, "sync_transport", None)
|
||||
return HTTPTransport(
|
||||
local_address=_IPV4_LOCAL_ADDRESS,
|
||||
verify=_httpx_default_verify_unless_configured(verify),
|
||||
cert=cert,
|
||||
http2=http2,
|
||||
)
|
||||
return getattr(litellm, "sync_transport", None)
|
||||
|
||||
@staticmethod
|
||||
def _create_sync_proxy_mounts(
|
||||
transport: httpx.BaseTransport | None,
|
||||
verify: VerifyTypes,
|
||||
cert: CertTypes | None,
|
||||
) -> Mapping[str, HTTPTransport | None] | None:
|
||||
if not litellm.force_ipv4:
|
||||
http2: bool,
|
||||
) -> Mapping[str, httpx.BaseTransport | None] | None:
|
||||
if transport is None or transport is getattr(litellm, "sync_transport", None):
|
||||
return None
|
||||
return _environment_proxy_mounts(
|
||||
lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2_enabled())
|
||||
lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1873,7 +1984,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
|
|||
cache = LLMClientCache()
|
||||
setattr(litellm, "in_memory_llm_clients_cache", cache)
|
||||
|
||||
_cached_client: Final = cache.get_cache(_cache_key_name)
|
||||
_cached_client: Final = cache.get_cache(_cache_key_name, loop_bound=False)
|
||||
if _cached_client:
|
||||
return _cached_client
|
||||
|
||||
|
|
@ -1889,6 +2000,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
|
|||
value=_new_client,
|
||||
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
litellm_owned_client=True,
|
||||
loop_bound=False,
|
||||
)
|
||||
return _new_client
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from typing import (
|
|||
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from h2.events import ConnectionTerminated
|
||||
from httpx import USE_CLIENT_DEFAULT
|
||||
from httpx._types import FileContent
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
|
@ -371,6 +372,34 @@ class _SignedBodyKwargs(TypedDict, total=False):
|
|||
json: ReadOnly[dict[str, object]]
|
||||
|
||||
|
||||
_GOAWAY_RETRY_MESSAGE: Final = "Vector store search connection was closed by an HTTP/2 GOAWAY, retrying once"
|
||||
|
||||
|
||||
def _closed_by_http2_goaway(error: httpx.RemoteProtocolError) -> bool:
|
||||
cause: Final = error.__cause__
|
||||
return cause is not None and any(isinstance(arg, ConnectionTerminated) for arg in cause.args)
|
||||
|
||||
|
||||
async def _asend_with_goaway_retry(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
|
||||
try:
|
||||
return await send()
|
||||
except httpx.RemoteProtocolError as error:
|
||||
if not _closed_by_http2_goaway(error):
|
||||
raise
|
||||
verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
|
||||
return await send()
|
||||
|
||||
|
||||
def _send_with_goaway_retry(send: Callable[[], httpx.Response]) -> httpx.Response:
|
||||
try:
|
||||
return send()
|
||||
except httpx.RemoteProtocolError as error:
|
||||
if not _closed_by_http2_goaway(error):
|
||||
raise
|
||||
verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
|
||||
return send()
|
||||
|
||||
|
||||
def _is_sized_json_body(response: httpx.Response) -> bool:
|
||||
headers: Final = httpx.Headers(response.headers)
|
||||
media_types: Final = (value.split(";")[0].strip().lower() for value in headers.get_list("content-type"))
|
||||
|
|
@ -2438,8 +2467,7 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"Anthropic /v1/messages: invalid thinking signature; "
|
||||
"stripping thinking blocks and retrying (attempt %s/%s).",
|
||||
"Anthropic /v1/messages: retrying after a recoverable HTTP error (attempt %s/%s).",
|
||||
len(recovered_errors) + 2,
|
||||
max_attempts,
|
||||
)
|
||||
|
|
@ -6923,7 +6951,7 @@ class BaseLLMHTTPHandler:
|
|||
try:
|
||||
ssl_context: Final = realtime_ssl_for_url(url)
|
||||
provider_backend: Final = await provider_config.open_backend(url, headers)
|
||||
backend_ws: Final = (
|
||||
backend_ws: Final = provider_config.wrap_backend(
|
||||
provider_backend
|
||||
if provider_backend is not None
|
||||
else await self._open_realtime_backend_ws(websockets, url, headers, ssl_context)
|
||||
|
|
@ -6961,9 +6989,7 @@ class BaseLLMHTTPHandler:
|
|||
await backend_ws.send(_session_config)
|
||||
realtime_streaming.session_configuration_request = _session_config
|
||||
|
||||
# For providers that defer setup until client session.update, optionally
|
||||
# send synthetic session.created to unblock clients waiting on connect.
|
||||
if not provider_config.requires_session_configuration():
|
||||
if not _session_config:
|
||||
synthetic_session: Final = provider_config.transform_session_created_event(
|
||||
model=model,
|
||||
logging_session_id=logging_obj.litellm_trace_id,
|
||||
|
|
@ -10629,7 +10655,10 @@ class BaseLLMHTTPHandler:
|
|||
if client is None or not isinstance(client, AsyncHTTPHandler):
|
||||
async_httpx_client = get_async_httpx_client(
|
||||
llm_provider=litellm.LlmProviders(custom_llm_provider),
|
||||
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
|
||||
params={
|
||||
"ssl_verify": litellm_params.get("ssl_verify", None),
|
||||
**vector_store_provider_config.get_async_httpx_client_params(),
|
||||
},
|
||||
)
|
||||
else:
|
||||
async_httpx_client = client
|
||||
|
|
@ -10701,11 +10730,8 @@ class BaseLLMHTTPHandler:
|
|||
request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
|
||||
|
||||
try:
|
||||
response: Final = await async_httpx_client.post(
|
||||
url=url,
|
||||
headers=headers,
|
||||
data=request_data,
|
||||
timeout=timeout,
|
||||
response: Final = await _asend_with_goaway_retry(
|
||||
partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise vector_store_provider_config.get_error_class(
|
||||
|
|
@ -10838,11 +10864,8 @@ class BaseLLMHTTPHandler:
|
|||
request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
|
||||
|
||||
try:
|
||||
response: Final = sync_httpx_client.post(
|
||||
url=url,
|
||||
headers=headers,
|
||||
data=request_data,
|
||||
timeout=timeout,
|
||||
response: Final = _send_with_goaway_retry(
|
||||
partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise vector_store_provider_config.get_error_class(
|
||||
|
|
@ -11152,7 +11175,7 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
|
||||
url: Final = api_base
|
||||
url: Final = vector_store_provider_config.get_vector_store_list_url(api_base)
|
||||
|
||||
params: Final[dict[str, object]] = {}
|
||||
if after is not None:
|
||||
|
|
@ -11231,7 +11254,7 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
|
||||
url: Final = api_base
|
||||
url: Final = vector_store_provider_config.get_vector_store_list_url(api_base)
|
||||
|
||||
params: Final[dict[str, object]] = {}
|
||||
if after is not None:
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from litellm.secret_managers.main import get_secret_str
|
|||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from ..common_utils import resolve_dashscope_openai_api_base
|
||||
|
||||
|
||||
class DashScopeChatConfig(OpenAIGPTConfig):
|
||||
|
|
@ -55,11 +56,9 @@ class DashScopeChatConfig(OpenAIGPTConfig):
|
|||
def _get_openai_compatible_provider_info(
|
||||
self, api_base: str | None, api_key: str | None
|
||||
) -> tuple[str | None, str | None]:
|
||||
api_base = (
|
||||
api_base or get_secret_str("DASHSCOPE_API_BASE") or "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
)
|
||||
resolved_api_base: Final = resolve_dashscope_openai_api_base(api_base)
|
||||
dynamic_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
return api_base, dynamic_api_key
|
||||
return resolved_api_base, dynamic_api_key
|
||||
|
||||
def get_openai_compatible_provider_info(
|
||||
self,
|
||||
|
|
@ -69,7 +68,7 @@ class DashScopeChatConfig(OpenAIGPTConfig):
|
|||
return self._get_openai_compatible_provider_info(api_base, api_key)
|
||||
|
||||
def _resolve_chat_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
return resolve_dashscope_openai_api_base(api_base)
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
Common utilities for the DashScope LLM provider.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -19,10 +20,123 @@ if TYPE_CHECKING:
|
|||
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
|
||||
|
||||
DASHSCOPE_CHAT_COMPATIBLE_PATH: Final = "/compatible-mode/v1"
|
||||
DASHSCOPE_RERANK_PATH: Final = "/compatible-api/v1/reranks"
|
||||
DASHSCOPE_ANTHROPIC_PATH: Final = "/apps/anthropic"
|
||||
DASHSCOPE_NATIVE_V1_PATH: Final = "/api/v1"
|
||||
DASHSCOPE_COMPATIBLE_API_PATH: Final = "/compatible-api/v1"
|
||||
DASHSCOPE_RERANK_PATH: Final = f"{DASHSCOPE_COMPATIBLE_API_PATH}/reranks"
|
||||
DASHSCOPE_IMAGE_PATH: Final = "/api/v1/services/aigc/multimodal-generation/generation"
|
||||
DASHSCOPE_DEFAULT_ORIGIN: Final = "https://dashscope.aliyuncs.com"
|
||||
|
||||
_STANDARD_DASHSCOPE_HOSTS: Final = frozenset(
|
||||
{
|
||||
"dashscope.aliyuncs.com",
|
||||
"dashscope-intl.aliyuncs.com",
|
||||
"dashscope-us.aliyuncs.com",
|
||||
}
|
||||
)
|
||||
_WORKSPACE_HOST_SUFFIX: Final = ".maas.aliyuncs.com"
|
||||
_EXCLUDED_WORKSPACE_PREFIXES: Final = ("trial.", "token-plan.")
|
||||
_PROTOCOL_PATH_PREFIXES: Final = (
|
||||
DASHSCOPE_CHAT_COMPATIBLE_PATH,
|
||||
DASHSCOPE_ANTHROPIC_PATH,
|
||||
DASHSCOPE_NATIVE_V1_PATH,
|
||||
DASHSCOPE_COMPATIBLE_API_PATH,
|
||||
)
|
||||
|
||||
|
||||
def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None:
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DashScopeEndpoints:
|
||||
origin: str
|
||||
openai: str
|
||||
anthropic: str
|
||||
native: str
|
||||
|
||||
|
||||
def _is_dashscope_host(host: str) -> bool:
|
||||
if host in _STANDARD_DASHSCOPE_HOSTS:
|
||||
return True
|
||||
return host.endswith(_WORKSPACE_HOST_SUFFIX) and not host.startswith(_EXCLUDED_WORKSPACE_PREFIXES)
|
||||
|
||||
|
||||
def _dashscope_origin(api_base: str | None) -> str | None:
|
||||
if not api_base:
|
||||
return None
|
||||
parsed: Final = urlparse(api_base)
|
||||
if parsed.scheme not in ("http", "https") or not parsed.netloc:
|
||||
return None
|
||||
if parsed.query or parsed.fragment:
|
||||
return None
|
||||
host: Final = (parsed.hostname or "").lower()
|
||||
if not _is_dashscope_host(host):
|
||||
return None
|
||||
path: Final = parsed.path.rstrip("/")
|
||||
if path and not any(path == prefix or path.startswith(f"{prefix}/") for prefix in _PROTOCOL_PATH_PREFIXES):
|
||||
return None
|
||||
return f"{parsed.scheme}://{parsed.netloc}"
|
||||
|
||||
|
||||
def resolve_dashscope_endpoints(api_base: str | None) -> DashScopeEndpoints | None:
|
||||
"""
|
||||
Derive DashScope's protocol-specific bases from any standard DashScope base.
|
||||
|
||||
Accepts the domain root, an OpenAI-compatible base, an Anthropic Messages
|
||||
base, or a native `/api/v1` base for `dashscope.aliyuncs.com`,
|
||||
`dashscope-intl.aliyuncs.com`, `dashscope-us.aliyuncs.com`, and
|
||||
`{workspace}.{region}.maas.aliyuncs.com`. Returns None for any other host
|
||||
or path so callers keep the caller-supplied base untouched.
|
||||
"""
|
||||
origin: Final = _dashscope_origin(api_base)
|
||||
if origin is None:
|
||||
return None
|
||||
return DashScopeEndpoints(
|
||||
origin=origin,
|
||||
openai=f"{origin}{DASHSCOPE_CHAT_COMPATIBLE_PATH}",
|
||||
anthropic=f"{origin}{DASHSCOPE_ANTHROPIC_PATH}",
|
||||
native=f"{origin}{DASHSCOPE_NATIVE_V1_PATH}",
|
||||
)
|
||||
|
||||
|
||||
def is_dashscope_native_endpoint_url(api_base: str) -> bool:
|
||||
return urlparse(api_base).path.rstrip("/").startswith(f"{DASHSCOPE_NATIVE_V1_PATH}/")
|
||||
|
||||
|
||||
def resolve_dashscope_openai_api_base(
|
||||
api_base: str | None,
|
||||
env_var: str = "DASHSCOPE_API_BASE",
|
||||
default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_CHAT_COMPATIBLE_PATH}",
|
||||
) -> str:
|
||||
candidate: Final = api_base or get_secret_str(env_var)
|
||||
if candidate and is_dashscope_native_endpoint_url(candidate):
|
||||
return candidate
|
||||
endpoints: Final = resolve_dashscope_endpoints(candidate)
|
||||
if endpoints is not None:
|
||||
return endpoints.openai
|
||||
return candidate or default
|
||||
|
||||
|
||||
def resolve_dashscope_anthropic_api_base(
|
||||
api_base: str | None,
|
||||
env_var: str = "DASHSCOPE_API_BASE",
|
||||
default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_ANTHROPIC_PATH}",
|
||||
) -> str:
|
||||
candidate: Final = api_base or get_secret_str(env_var)
|
||||
endpoints: Final = resolve_dashscope_endpoints(candidate)
|
||||
if endpoints is not None:
|
||||
return endpoints.anthropic
|
||||
return candidate or default
|
||||
|
||||
|
||||
def dashscope_base_selects_anthropic_messages(api_base: str | None) -> bool:
|
||||
candidate: Final = api_base or get_secret_str("DASHSCOPE_API_BASE")
|
||||
if not candidate:
|
||||
return False
|
||||
path: Final = urlparse(candidate).path.rstrip("/")
|
||||
if path.endswith(DASHSCOPE_ANTHROPIC_PATH) or f"{DASHSCOPE_ANTHROPIC_PATH}/" in path:
|
||||
return True
|
||||
return not path and resolve_dashscope_endpoints(candidate) is not None
|
||||
|
||||
|
||||
def _chat_shaped_aliyun_origin(api_base: str | None) -> str | None:
|
||||
if api_base is None:
|
||||
return None
|
||||
parsed: Final = urlparse(api_base)
|
||||
|
|
@ -30,14 +144,24 @@ def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None:
|
|||
on_aliyun_host: Final = host == "aliyuncs.com" or host.endswith(".aliyuncs.com")
|
||||
if not on_aliyun_host or parsed.path.rstrip("/") != DASHSCOPE_CHAT_COMPATIBLE_PATH:
|
||||
return None
|
||||
return f"{parsed.scheme}://{parsed.netloc}{DASHSCOPE_RERANK_PATH}"
|
||||
return f"{parsed.scheme}://{parsed.netloc}"
|
||||
|
||||
|
||||
def _is_compatible_api_base(api_base: str) -> bool:
|
||||
path: Final = urlparse(api_base).path.rstrip("/")
|
||||
return path == DASHSCOPE_COMPATIBLE_API_PATH or path.startswith(f"{DASHSCOPE_COMPATIBLE_API_PATH}/")
|
||||
|
||||
|
||||
def resolve_dashscope_family_rerank_api_base(api_base: str | None, env_var: str, default_rerank_base: str) -> str:
|
||||
remapped: Final = _rerank_base_for_chat_shaped_base(api_base)
|
||||
if api_base is not None and remapped is None:
|
||||
if api_base is None:
|
||||
return get_secret_str(env_var) or default_rerank_base
|
||||
if _is_compatible_api_base(api_base):
|
||||
return api_base
|
||||
return get_secret_str(env_var) or remapped or default_rerank_base
|
||||
chat_origin: Final = _chat_shaped_aliyun_origin(api_base)
|
||||
if chat_origin is not None:
|
||||
return get_secret_str(env_var) or f"{chat_origin}{DASHSCOPE_RERANK_PATH}"
|
||||
endpoints: Final = resolve_dashscope_endpoints(api_base)
|
||||
return f"{endpoints.origin}{DASHSCOPE_RERANK_PATH}" if endpoints is not None else api_base
|
||||
|
||||
|
||||
def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig":
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ from litellm.secret_managers.main import get_secret_str
|
|||
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
|
||||
from litellm.types.utils import EmbeddingResponse, Usage
|
||||
|
||||
from ..common_utils import DashScopeError
|
||||
from ..common_utils import DashScopeError, resolve_dashscope_openai_api_base
|
||||
|
||||
DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
|
||||
|
|
@ -91,7 +91,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
|||
return resolved_api_key
|
||||
|
||||
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE
|
||||
return resolve_dashscope_openai_api_base(api_base, default=DEFAULT_API_BASE)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -39,6 +39,13 @@ from litellm.types.llms.openai import (
|
|||
)
|
||||
from litellm.types.utils import ImageObject, ImageResponse
|
||||
|
||||
from ..common_utils import (
|
||||
DASHSCOPE_CHAT_COMPATIBLE_PATH,
|
||||
DASHSCOPE_IMAGE_PATH,
|
||||
is_dashscope_native_endpoint_url,
|
||||
resolve_dashscope_endpoints,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
|
@ -53,8 +60,6 @@ _DECODED_JSON: Final = TypeAdapter(object)
|
|||
|
||||
DEFAULT_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
|
||||
|
||||
CHAT_COMPATIBLE_MODE_PATH: Final = "/compatible-mode/v1"
|
||||
|
||||
# Maps OpenAI size strings (WxH) to DashScope size strings (W*H)
|
||||
OPENAI_TO_DASHSCOPE_SIZE: Final[dict] = {
|
||||
"256x256": "256*256",
|
||||
|
|
@ -102,8 +107,14 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
raise ValueError("DASHSCOPE_API_KEY is not set")
|
||||
return resolved_api_key
|
||||
|
||||
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
|
||||
return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE
|
||||
def _image_api_base_env(self) -> str | None:
|
||||
return get_secret_str("DASHSCOPE_API_BASE_IMAGE")
|
||||
|
||||
def _default_image_api_base(self) -> str:
|
||||
return DEFAULT_API_BASE
|
||||
|
||||
def _configured_chat_api_base(self) -> str | None:
|
||||
return get_secret_str("DASHSCOPE_API_BASE")
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
@ -114,10 +125,18 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
litellm_params: dict,
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
image_api_base: Final = (
|
||||
api_base if api_base and not api_base.rstrip("/").endswith(CHAT_COMPATIBLE_MODE_PATH) else None
|
||||
)
|
||||
return self._resolve_image_api_base(image_api_base)
|
||||
image_env_api_base: Final = self._image_api_base_env()
|
||||
configured_api_base: Final = api_base or self._configured_chat_api_base()
|
||||
if not configured_api_base or configured_api_base.rstrip("/").endswith(DASHSCOPE_CHAT_COMPATIBLE_PATH):
|
||||
return image_env_api_base or self._default_image_api_base()
|
||||
if is_dashscope_native_endpoint_url(configured_api_base):
|
||||
return configured_api_base
|
||||
endpoints: Final = resolve_dashscope_endpoints(configured_api_base)
|
||||
if endpoints is None:
|
||||
return configured_api_base
|
||||
if not api_base and image_env_api_base:
|
||||
return image_env_api_base
|
||||
return f"{endpoints.origin}{DASHSCOPE_IMAGE_PATH}"
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
|
|
|
|||
7
litellm/llms/dashscope/messages/__init__.py
Normal file
7
litellm/llms/dashscope/messages/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""
|
||||
DashScope Anthropic-compatible Messages Module
|
||||
"""
|
||||
|
||||
from .transformation import DashScopeAnthropicMessagesConfig
|
||||
|
||||
__all__ = ["DashScopeAnthropicMessagesConfig"]
|
||||
96
litellm/llms/dashscope/messages/transformation.py
Normal file
96
litellm/llms/dashscope/messages/transformation.py
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
"""
|
||||
DashScope Anthropic-compatible Messages transformation config.
|
||||
"""
|
||||
|
||||
from typing import ClassVar, Final
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.litellm_core_utils.url_utils import ensure_api_route
|
||||
from litellm.llms.anthropic.messages.native.transformation import (
|
||||
AnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
|
||||
|
||||
from ..common_utils import (
|
||||
dashscope_base_selects_anthropic_messages,
|
||||
resolve_dashscope_anthropic_api_base,
|
||||
)
|
||||
|
||||
|
||||
class DashScopeAnthropicMessagesConfig(AnthropicMessagesConfig):
|
||||
"""
|
||||
DashScope exposes an Anthropic-compatible Messages API at
|
||||
``{origin}/apps/anthropic`` for every DashScope domain, including
|
||||
``dashscope.aliyuncs.com`` and ``{workspace}.{region}.maas.aliyuncs.com``.
|
||||
"""
|
||||
|
||||
_workload_identity_eligible: ClassVar[bool] = False
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> str | None:
|
||||
return "dashscope"
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
return True
|
||||
|
||||
def handles_web_search_natively(self) -> bool:
|
||||
return False
|
||||
|
||||
def serves_api_base(self, model: str, api_base: str | None) -> bool:
|
||||
return dashscope_base_selects_anthropic_messages(api_base)
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(api_key: str | None = None) -> str | None:
|
||||
return (
|
||||
normalize_nonempty_secret_str(api_key)
|
||||
or normalize_nonempty_secret_str(get_secret_str("DASHSCOPE_API_KEY"))
|
||||
or normalize_nonempty_secret_str(litellm.api_key)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_api_base(api_base: str | None = None) -> str:
|
||||
return resolve_dashscope_anthropic_api_base(api_base)
|
||||
|
||||
def validate_anthropic_messages_environment(
|
||||
self,
|
||||
headers: dict[str, str], # mutable-ok: mirrors the legacy base override signature
|
||||
model: str,
|
||||
messages: list[object], # mutable-ok: mirrors the legacy base override signature
|
||||
optional_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature
|
||||
litellm_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors the base signature
|
||||
resolved_api_key: Final = self.get_api_key(api_key=api_key)
|
||||
if resolved_api_key is None:
|
||||
present_headers: Final = {name.lower() for name in headers}
|
||||
if "x-api-key" not in present_headers and "authorization" not in present_headers:
|
||||
raise AuthenticationError(
|
||||
message=(
|
||||
"Missing DashScope API key. Set DASHSCOPE_API_KEY or pass api_key explicitly. "
|
||||
"Anthropic credentials are never sent to a DashScope endpoint."
|
||||
),
|
||||
llm_provider="dashscope",
|
||||
model=model,
|
||||
)
|
||||
return super().validate_anthropic_messages_environment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportReturnType] # base signature uses an unparameterized dict
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=resolved_api_key,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature
|
||||
litellm_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
return ensure_api_route(self.get_api_base(api_base=api_base), "v1/messages")
|
||||
|
|
@ -3,7 +3,7 @@ from typing import Final
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
from .chat.transformation import DashScopeChatConfig
|
||||
from .common_utils import resolve_dashscope_family_rerank_api_base
|
||||
from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base
|
||||
from .embed.transformation import DashScopeEmbeddingConfig
|
||||
from .image_generation.transformation import DashScopeImageGenerationConfig
|
||||
from .rerank.transformation import DashScopeRerankConfig
|
||||
|
|
@ -38,7 +38,7 @@ class QwenAIPlatformChatConfig(DashScopeChatConfig):
|
|||
return self._resolve_chat_api_base(api_base), _resolve_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_chat_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
|
||||
return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE)
|
||||
|
||||
|
||||
class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig):
|
||||
|
|
@ -46,7 +46,7 @@ class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig):
|
|||
return _require_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
|
||||
return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE)
|
||||
|
||||
|
||||
class QwenAIPlatformRerankConfig(DashScopeRerankConfig):
|
||||
|
|
@ -63,8 +63,14 @@ class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig):
|
|||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
|
||||
return image_api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE") or QWEN_AI_PLATFORM_IMAGE_API_BASE
|
||||
def _image_api_base_env(self) -> str | None:
|
||||
return get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE")
|
||||
|
||||
def _default_image_api_base(self) -> str:
|
||||
return QWEN_AI_PLATFORM_IMAGE_API_BASE
|
||||
|
||||
def _configured_chat_api_base(self) -> str | None:
|
||||
return get_secret_str("QWEN_AI_PLATFORM_API_BASE")
|
||||
|
||||
|
||||
class QwenAIPlatformVideoConfig(DashScopeVideoConfig):
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from typing import Final
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
from .chat.transformation import DashScopeChatConfig
|
||||
from .common_utils import resolve_dashscope_family_rerank_api_base
|
||||
from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base
|
||||
from .embed.transformation import DashScopeEmbeddingConfig
|
||||
from .image_generation.transformation import DashScopeImageGenerationConfig
|
||||
from .rerank.transformation import DashScopeRerankConfig
|
||||
|
|
@ -38,7 +38,7 @@ class QwenCloudChatConfig(DashScopeChatConfig):
|
|||
return self._resolve_chat_api_base(api_base), _resolve_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_chat_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
|
||||
return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE)
|
||||
|
||||
|
||||
class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig):
|
||||
|
|
@ -46,7 +46,7 @@ class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig):
|
|||
return _require_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
|
||||
return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE)
|
||||
|
||||
|
||||
class QwenCloudRerankConfig(DashScopeRerankConfig):
|
||||
|
|
@ -63,8 +63,14 @@ class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig):
|
|||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
|
||||
return image_api_base or get_secret_str("QWENCLOUD_API_BASE_IMAGE") or QWENCLOUD_IMAGE_API_BASE
|
||||
def _image_api_base_env(self) -> str | None:
|
||||
return get_secret_str("QWENCLOUD_API_BASE_IMAGE")
|
||||
|
||||
def _default_image_api_base(self) -> str:
|
||||
return QWENCLOUD_IMAGE_API_BASE
|
||||
|
||||
def _configured_chat_api_base(self) -> str | None:
|
||||
return get_secret_str("QWENCLOUD_API_BASE")
|
||||
|
||||
|
||||
class QwenCloudVideoConfig(DashScopeVideoConfig):
|
||||
|
|
|
|||
|
|
@ -38,6 +38,8 @@ from litellm.types.videos.utils import (
|
|||
)
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
from ..common_utils import resolve_dashscope_endpoints
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
|
@ -206,6 +208,9 @@ def _parse_task_response(raw_response: httpx.Response) -> _TaskResponse:
|
|||
|
||||
|
||||
def _normalized_api_base(api_base: str) -> str:
|
||||
endpoints: Final = resolve_dashscope_endpoints(api_base)
|
||||
if endpoints is not None:
|
||||
return endpoints.origin
|
||||
trimmed: Final = api_base.rstrip("/")
|
||||
if trimmed.endswith(DASHSCOPE_COMPATIBLE_MODE_PATH):
|
||||
return trimmed[: -len(DASHSCOPE_COMPATIBLE_MODE_PATH)]
|
||||
|
|
|
|||
|
|
@ -38,6 +38,15 @@ else:
|
|||
|
||||
_EMPTY_FUNCTION: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_JSON_VALUE: Final = TypeAdapter(object)
|
||||
_CONTENT_PART: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _file_part_data(part: Mapping[str, object]) -> str | None:
|
||||
match part:
|
||||
case {"type": "file", "file": {"file_data": str() as file_data}} if file_data:
|
||||
return file_data
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def is_valid_json(value: str) -> bool:
|
||||
|
|
@ -295,7 +304,7 @@ class GigaChatConfig(BaseConfig):
|
|||
|
||||
def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[object]]:
|
||||
"""
|
||||
Extract text and image attachments from a multimodal message content list.
|
||||
Extract text, image and file attachments from a multimodal message content list.
|
||||
|
||||
Args:
|
||||
content: List of content parts (OpenAI multimodal format)
|
||||
|
|
@ -320,6 +329,10 @@ class GigaChatConfig(BaseConfig):
|
|||
file_id = self._upload_image(upload_url)
|
||||
if file_id:
|
||||
attachments.append(file_id)
|
||||
elif (file_data := _file_part_data(_CONTENT_PART.validate_python(part))) and (
|
||||
uploaded_file_id := self._upload_image(file_data)
|
||||
):
|
||||
attachments.append(uploaded_file_id)
|
||||
text: Final = "\n".join(texts) if texts else ""
|
||||
return text, attachments
|
||||
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS
|
||||
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
|
|
@ -681,7 +682,9 @@ class OCIChatConfig(BaseConfig):
|
|||
servingMode=serving_mode,
|
||||
chatRequest=OCIChatRequestPayload(
|
||||
apiFormat=vendor.value,
|
||||
messages=adapt_messages_to_generic_oci_standard(messages),
|
||||
messages=adapt_messages_to_generic_oci_standard(
|
||||
list(messages_with_file_parts_as_image_urls(messages))
|
||||
),
|
||||
**self._get_optional_params(vendor, optional_params, model),
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -275,6 +275,10 @@ def build_output_token_limit_response(
|
|||
return e.response.headers, (_aiter_once(chunk) if is_async else _iter_once(chunk))
|
||||
|
||||
|
||||
def _is_loop_bound(client_initialization_params: Mapping[str, object]) -> bool:
|
||||
return client_initialization_params.get("is_async") is not False
|
||||
|
||||
|
||||
class BaseOpenAILLM:
|
||||
"""
|
||||
Base class for OpenAI LLMs for getting their httpx clients and SSL verification settings
|
||||
|
|
@ -289,7 +293,9 @@ class BaseOpenAILLM:
|
|||
client_initialization_params=client_initialization_params,
|
||||
client_type=client_type,
|
||||
)
|
||||
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key)
|
||||
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(
|
||||
_cache_key, loop_bound=_is_loop_bound(client_initialization_params)
|
||||
)
|
||||
return _cached_client
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -330,6 +336,7 @@ class BaseOpenAILLM:
|
|||
value=openai_client,
|
||||
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
litellm_owned_client=litellm_owned_client,
|
||||
loop_bound=_is_loop_bound(client_initialization_params),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -5,18 +5,30 @@ This module provides guardrail translation support for OpenAI's text completion
|
|||
The handler processes the 'prompt' parameter for guardrails.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
StreamingScanKey,
|
||||
StreamTransformSink,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_chat_stream_usage,
|
||||
stream_item_field,
|
||||
stream_item_items,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.types.utils import TextCompletionResponse
|
||||
|
||||
|
||||
class OpenAITextCompletionHandler(BaseTranslation):
|
||||
|
|
@ -32,6 +44,9 @@ class OpenAITextCompletionHandler(BaseTranslation):
|
|||
- A list of strings (for batch completions)
|
||||
"""
|
||||
|
||||
delivers_ended_stream_rewrites = True
|
||||
assembles_streamed_response = True
|
||||
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -125,7 +140,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict | None = None,
|
||||
request_data: dict[str, object] | None = None,
|
||||
) -> "TextCompletionResponse":
|
||||
"""
|
||||
Process output response by applying guardrails to completion text.
|
||||
|
|
@ -162,11 +177,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
|
|||
if "response" not in request_data:
|
||||
request_data["response"] = response
|
||||
|
||||
# Add user API key metadata with prefixed keys
|
||||
if "litellm_metadata" not in request_data:
|
||||
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
|
||||
if user_metadata:
|
||||
request_data["litellm_metadata"] = user_metadata
|
||||
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
|
||||
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
# Include model information from the response if available
|
||||
|
|
@ -195,3 +206,168 @@ class OpenAITextCompletionHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
return response
|
||||
|
||||
async def process_output_streaming_response(
|
||||
self,
|
||||
responses_so_far: list[object], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict[str, object] | None = None, # mutable-ok: the hooks share and mutate the request payload
|
||||
stream_transform_sink: StreamTransformSink | None = None,
|
||||
deliver_ended_stream_rewrites: bool = False,
|
||||
) -> Sequence[object]:
|
||||
chunks: Final = _text_completion_chunks(responses_so_far)
|
||||
if not self._check_streaming_has_ended(chunks):
|
||||
return responses_so_far
|
||||
assembled: Final = self._assemble_ended_stream(chunks)
|
||||
pre_guardrail_texts: Final = _choice_texts(assembled)
|
||||
await self.process_output_response(
|
||||
response=assembled,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
)
|
||||
if not deliver_ended_stream_rewrites:
|
||||
return responses_so_far
|
||||
rewrites: Final = tuple(
|
||||
after for before, after in zip(pre_guardrail_texts, _choice_texts(assembled)) if after != before
|
||||
)
|
||||
for choice_index, text in rewrites:
|
||||
self._write_choice_text(chunks, choice_index, text)
|
||||
return responses_so_far
|
||||
|
||||
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
|
||||
chunks: Final = _text_completion_chunks(responses_so_far)
|
||||
if not self._check_streaming_has_ended(chunks):
|
||||
return None
|
||||
assembled: Final = self._assemble_ended_stream(chunks)
|
||||
return StreamingScanKey(texts=tuple(text for _, text in _choice_texts(assembled)), stream_ended=True)
|
||||
|
||||
def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]:
|
||||
chunks: Final = _text_completion_chunks(responses_so_far)
|
||||
streamed_choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks))
|
||||
finished_indices: Final = frozenset(
|
||||
_choice_index(choice) for choice in streamed_choices if choice.finish_reason
|
||||
)
|
||||
unfinished_indices: Final = tuple(
|
||||
dict.fromkeys(
|
||||
_choice_index(choice) for choice in streamed_choices if _choice_index(choice) not in finished_indices
|
||||
)
|
||||
)
|
||||
if not unfinished_indices:
|
||||
return tuple(responses_so_far)
|
||||
terminator: Final = TextCompletionResponse(
|
||||
id=chunks[0].id,
|
||||
model=chunks[0].model,
|
||||
object="text_completion",
|
||||
choices=[TextChoices(index=index, text="", finish_reason="length") for index in unfinished_indices],
|
||||
)
|
||||
return (*responses_so_far, terminator)
|
||||
|
||||
def build_block_sse_chunks(
|
||||
self,
|
||||
exc: "ModifyResponseException",
|
||||
stream_started: bool = False,
|
||||
responses_so_far: Sequence[object] | None = None,
|
||||
) -> Sequence[bytes]:
|
||||
streamed: Final = _text_completion_chunks(responses_so_far or ())
|
||||
source: Final = streamed[0] if streamed else TextCompletionResponse(model=exc.model)
|
||||
message_chunk: Final = TextCompletionResponse(
|
||||
id=source.id,
|
||||
created=source.created,
|
||||
model=source.model or exc.model,
|
||||
object="text_completion",
|
||||
choices=[TextChoices(index=0, text=exc.message)],
|
||||
).model_copy(update={"usage": None})
|
||||
final_chunk: Final = TextCompletionResponse(
|
||||
id=source.id,
|
||||
created=source.created,
|
||||
model=source.model or exc.model,
|
||||
object="text_completion",
|
||||
choices=[TextChoices(index=0, text="", finish_reason="content_filter")],
|
||||
).model_copy(update={"usage": _blocked_stream_usage(exc.original_response)})
|
||||
return tuple(
|
||||
f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk)
|
||||
)
|
||||
|
||||
def build_stream_error_items(
|
||||
self,
|
||||
exc: "HTTPException",
|
||||
responses_so_far: Sequence[object] | None = None,
|
||||
) -> Sequence[bytes] | None:
|
||||
from litellm.proxy.common_request_processing import sse_error_payload
|
||||
|
||||
_, error_obj = sse_error_payload(exc)
|
||||
return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),)
|
||||
|
||||
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
|
||||
streamed_choices: Final = chain.from_iterable(stream_item_items(item, "choices") for item in responses_so_far)
|
||||
return any(stream_item_field(choice, "finish_reason") is not None for choice in streamed_choices)
|
||||
|
||||
@staticmethod
|
||||
def _assemble_ended_stream(responses_so_far: Sequence[TextCompletionResponse]) -> TextCompletionResponse:
|
||||
streamed_choices: Final = tuple(chain.from_iterable(item.choices for item in responses_so_far))
|
||||
choice_indices: Final = tuple(dict.fromkeys(_choice_index(choice) for choice in streamed_choices))
|
||||
assembled_choices: Final = [
|
||||
TextChoices(
|
||||
index=index,
|
||||
text="".join(_choice_text(choice) for choice in streamed_choices if _choice_index(choice) == index),
|
||||
finish_reason=next(
|
||||
(
|
||||
choice.finish_reason
|
||||
for choice in streamed_choices
|
||||
if _choice_index(choice) == index and choice.finish_reason
|
||||
),
|
||||
None,
|
||||
),
|
||||
)
|
||||
for index in choice_indices
|
||||
]
|
||||
first: Final = responses_so_far[0]
|
||||
return TextCompletionResponse(
|
||||
id=first.id,
|
||||
created=first.created,
|
||||
model=first.model,
|
||||
object="text_completion",
|
||||
choices=assembled_choices,
|
||||
usage=next((item.usage for item in reversed(responses_so_far) if item.usage is not None), None),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _write_choice_text(responses_so_far: Sequence[TextCompletionResponse], choice_index: int, text: str) -> None:
|
||||
streamed_choices: Final = chain.from_iterable(item.choices for item in responses_so_far)
|
||||
head, *tail = (choice for choice in streamed_choices if _choice_index(choice) == choice_index)
|
||||
head.text = text
|
||||
for choice in tail:
|
||||
if isinstance(stream_item_field(choice, "text"), str):
|
||||
choice.text = ""
|
||||
|
||||
|
||||
def _blocked_stream_usage(original_response: object) -> Usage | None:
|
||||
prompt_tokens, completion_tokens = blocked_chat_stream_usage(original_response)
|
||||
if isinstance(original_response, list) and prompt_tokens == completion_tokens == 0:
|
||||
return None
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _text_completion_chunks(responses_so_far: Sequence[object]) -> tuple[TextCompletionResponse, ...]:
|
||||
return tuple(item for item in responses_so_far if isinstance(item, TextCompletionResponse))
|
||||
|
||||
|
||||
def _choice_text(choice: TextChoices) -> str:
|
||||
text: Final = stream_item_field(choice, "text")
|
||||
return text if isinstance(text, str) else ""
|
||||
|
||||
|
||||
def _choice_index(choice: TextChoices) -> int:
|
||||
return choice.index or 0
|
||||
|
||||
|
||||
def _choice_texts(response: TextCompletionResponse) -> tuple[tuple[int, str], ...]:
|
||||
return tuple((_choice_index(choice), _choice_text(choice)) for choice in response.choices)
|
||||
|
|
|
|||
0
litellm/llms/opensearch/__init__.py
Normal file
0
litellm/llms/opensearch/__init__.py
Normal file
0
litellm/llms/opensearch/vector_stores/__init__.py
Normal file
0
litellm/llms/opensearch/vector_stores/__init__.py
Normal file
415
litellm/llms/opensearch/vector_stores/transformation.py
Normal file
415
litellm/llms/opensearch/vector_stores/transformation.py
Normal file
|
|
@ -0,0 +1,415 @@
|
|||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Final, Literal, NoReturn
|
||||
from urllib.parse import quote, urlsplit
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import AuthenticationError, BadRequestError
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.vector_store.transformation import (
|
||||
BaseQueryEmbeddingVectorStoreConfig,
|
||||
VectorStoreEmbeddingExecutor,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_json_post
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.llms.bedrock import AwsAuthParams
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
from litellm.types.vector_stores import (
|
||||
BaseVectorStoreAuthCredentials,
|
||||
VectorStoreCreateOptionalRequestParams,
|
||||
VectorStoreIndexEndpoints,
|
||||
VectorStoreResultContent,
|
||||
VectorStoreSearchOptionalRequestParams,
|
||||
VectorStoreSearchResponse,
|
||||
VectorStoreSearchResult,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
DEFAULT_EMBEDDING_FIELD: Final = "embedding"
|
||||
DEFAULT_TEXT_FIELD: Final = "text"
|
||||
DEFAULT_MAX_NUM_RESULTS: Final = 10
|
||||
MAX_MAX_NUM_RESULTS: Final = 50
|
||||
_PARAM_PREFIX: Final = "opensearch_"
|
||||
_AWS_HOSTS: Final = (
|
||||
re.compile(r"\.(?P<region>[a-z0-9-]+)\.(?P<service>aoss|aoss-fips|es)\.amazonaws\.com(?:\.cn)?\Z"),
|
||||
re.compile(r"\.(?P<service>aoss|aos)\.(?P<region>[a-z0-9-]+)\.on\.aws\Z"),
|
||||
)
|
||||
_ACCOUNT_ID_LABEL: Final = re.compile(r"\A\d{12}\.")
|
||||
_QUERY_VECTOR: Final = TypeAdapter(tuple[float, ...])
|
||||
_ATTRIBUTE_VALUE_TYPES: Final = (str, int, float, bool)
|
||||
_SEARCH_ONLY_MESSAGE: Final = (
|
||||
"The OpenSearch vector store is search-only. Create the k-NN index in OpenSearch, "
|
||||
"then register it here with its name as opensearch_index_name or vector_store_id."
|
||||
)
|
||||
_PER_ACCOUNT_MESSAGE: Final = (
|
||||
"api_base is an OpenSearch Serverless per-account endpoint, which needs a collection header the vector store "
|
||||
"does not send. Use the collection endpoint instead, for example https://<collection-id>.<region>.aoss.amazonaws.com."
|
||||
)
|
||||
_SIGV4_ONLY_MESSAGE: Final = (
|
||||
"The OpenSearch vector store signs every request with AWS SigV4 and does not support api_key. "
|
||||
"Remove api_key and give the vector store AWS credentials instead."
|
||||
)
|
||||
_CREDENTIALS_HINT: Final = (
|
||||
"Set aws_access_key_id and aws_secret_access_key, aws_role_name, or aws_profile_name on the vector store, "
|
||||
"or run LiteLLM with an AWS identity."
|
||||
)
|
||||
|
||||
|
||||
def _bad_request(message: str) -> BadRequestError:
|
||||
return BadRequestError(message=message, model=None, llm_provider="opensearch")
|
||||
|
||||
|
||||
def _unauthenticated(reason: str) -> AuthenticationError:
|
||||
return AuthenticationError(
|
||||
message=f"No usable AWS credentials to sign the OpenSearch request ({reason}). {_CREDENTIALS_HINT}",
|
||||
llm_provider="opensearch",
|
||||
model=None,
|
||||
)
|
||||
|
||||
|
||||
def _usable_credentials(credentials: object) -> "Credentials":
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
if isinstance(credentials, Credentials):
|
||||
return credentials
|
||||
raise _unauthenticated("none were found")
|
||||
|
||||
|
||||
class _OpenSearchParams(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
litellm_embedding_model: str | None = None
|
||||
opensearch_embedding_field: str | None = None
|
||||
opensearch_index_name: str | None = None
|
||||
opensearch_text_field: str | None = None
|
||||
|
||||
@property
|
||||
def embedding_field(self) -> str:
|
||||
return self.opensearch_embedding_field or DEFAULT_EMBEDDING_FIELD
|
||||
|
||||
@property
|
||||
def text_field(self) -> str:
|
||||
return self.opensearch_text_field or DEFAULT_TEXT_FIELD
|
||||
|
||||
def index_name(self, vector_store_id: str) -> str:
|
||||
return self.opensearch_index_name or vector_store_id
|
||||
|
||||
def require_embedding_model(self) -> str:
|
||||
if not self.litellm_embedding_model:
|
||||
raise _bad_request(
|
||||
"litellm_embedding_model is required for the OpenSearch vector store. It must be the model that "
|
||||
f"produced the vectors in '{self.embedding_field}', for example openai/text-embedding-3-small."
|
||||
)
|
||||
return self.litellm_embedding_model
|
||||
|
||||
|
||||
class _SigningRegion(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
aws_region_name: str | None = None
|
||||
|
||||
|
||||
class _SearchOptions(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
max_num_results: int | None = None
|
||||
filters: Mapping[str, object] | None = None
|
||||
ranking_options: Mapping[str, object] | None = None
|
||||
rewrite_query: bool | None = None
|
||||
|
||||
@property
|
||||
def unsupported(self) -> tuple[str, ...]:
|
||||
requested: Final = (
|
||||
("filters", self.filters),
|
||||
("ranking_options", self.ranking_options),
|
||||
("rewrite_query", self.rewrite_query),
|
||||
)
|
||||
return tuple(name for name, value in requested if value)
|
||||
|
||||
@property
|
||||
def k(self) -> int:
|
||||
if self.max_num_results is None:
|
||||
return DEFAULT_MAX_NUM_RESULTS
|
||||
if not 1 <= self.max_num_results <= MAX_MAX_NUM_RESULTS:
|
||||
raise _bad_request(
|
||||
f"max_num_results must be between 1 and {MAX_MAX_NUM_RESULTS}, got {self.max_num_results}"
|
||||
)
|
||||
return self.max_num_results
|
||||
|
||||
|
||||
_KNOWN_PARAMS: Final = frozenset(name for name in _OpenSearchParams.model_fields if name.startswith(_PARAM_PREFIX))
|
||||
|
||||
|
||||
class _LoggedSearch(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
query: str | tuple[str, ...] = ""
|
||||
litellm_params: _OpenSearchParams = _OpenSearchParams()
|
||||
|
||||
@property
|
||||
def search_query(self) -> str:
|
||||
return BaseQueryEmbeddingVectorStoreConfig.query_text(self.query)
|
||||
|
||||
|
||||
class _Hit(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str = Field(alias="_id")
|
||||
score: float | None = Field(default=None, alias="_score")
|
||||
source: Mapping[str, object] = Field(default_factory=dict, alias="_source")
|
||||
|
||||
|
||||
class _Hits(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
hits: Sequence[_Hit]
|
||||
|
||||
|
||||
class _SearchResponse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
hits: _Hits
|
||||
|
||||
|
||||
def aws_signing_target(url: str) -> tuple[Literal["aoss", "es"], str | None]:
|
||||
hostname: Final = urlsplit(url).hostname or ""
|
||||
match: Final = next(filter(None, (pattern.search(hostname) for pattern in _AWS_HOSTS)), None)
|
||||
if match is None:
|
||||
return "es", None
|
||||
return ("aoss" if match["service"].startswith("aoss") else "es"), match["region"]
|
||||
|
||||
|
||||
def _search_result(hit: _Hit, params: _OpenSearchParams) -> VectorStoreSearchResult:
|
||||
text: Final = hit.source.get(params.text_field)
|
||||
filename: Final = hit.source.get("filename")
|
||||
return VectorStoreSearchResult(
|
||||
score=hit.score,
|
||||
content=[VectorStoreResultContent(type="text", text=text)] if isinstance(text, str) else [],
|
||||
file_id=hit.id,
|
||||
filename=filename if isinstance(filename, str) else None,
|
||||
attributes={
|
||||
key: value
|
||||
for key, value in hit.source.items()
|
||||
if key not in (params.text_field, params.embedding_field) and isinstance(value, _ATTRIBUTE_VALUE_TYPES)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequestsWithAWS):
|
||||
def __init__(self, aws: BaseAWSLLM | None = None) -> None:
|
||||
self._aws: Final = aws if aws is not None else BaseAWSLLM()
|
||||
|
||||
def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials:
|
||||
return BaseVectorStoreAuthCredentials()
|
||||
|
||||
def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints:
|
||||
return VectorStoreIndexEndpoints(read=[], write=[])
|
||||
|
||||
def validate_environment(
|
||||
self, headers: Mapping[str, object], litellm_params: GenericLiteLLMParams | None
|
||||
) -> dict[str, object]: # mutable-ok: the shared HTTP handler requires writable headers
|
||||
params: Final = dict(litellm_params or {})
|
||||
self._signing_credentials(params, self._signing_region(params, str(params.get("api_base") or "")))
|
||||
return {**headers, "Content-Type": "application/json"}
|
||||
|
||||
def _signing_region(self, params: Mapping[str, object], api_base: str) -> str:
|
||||
host_region: Final = aws_signing_target(api_base)[1]
|
||||
return self._aws.get_aws_region_name_for_non_llm_api_calls(
|
||||
host_region or _SigningRegion.model_validate(params).aws_region_name
|
||||
)
|
||||
|
||||
def _signing_credentials(self, params: Mapping[str, object], region: str) -> "Credentials":
|
||||
if params.get("api_key"):
|
||||
raise _bad_request(_SIGV4_ONLY_MESSAGE)
|
||||
auth: Final = AwsAuthParams.model_validate(params)
|
||||
if bool(auth.aws_access_key_id) != bool(auth.aws_secret_access_key):
|
||||
raise _unauthenticated("aws_access_key_id and aws_secret_access_key must be set together")
|
||||
ensure_optional_import("botocore")
|
||||
from botocore.exceptions import (
|
||||
CredentialRetrievalError,
|
||||
NoCredentialsError,
|
||||
PartialCredentialsError,
|
||||
ProfileNotFound,
|
||||
)
|
||||
|
||||
try:
|
||||
return _usable_credentials(self._aws.resolve_credentials(auth, region))
|
||||
except CredentialRetrievalError as error:
|
||||
verbose_logger.warning("OpenSearch vector store AWS credential source failed: %s", error)
|
||||
raise _unauthenticated("the AWS credential source failed, see the LiteLLM logs") from None
|
||||
except (NoCredentialsError, PartialCredentialsError, ProfileNotFound) as error:
|
||||
raise _unauthenticated(str(error)) from None
|
||||
|
||||
def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
|
||||
parsed: Final = urlsplit(api_base or "")
|
||||
if parsed.scheme not in ("http", "https") or not parsed.hostname:
|
||||
raise _bad_request(
|
||||
"api_base must be the OpenSearch endpoint URL, for example "
|
||||
"https://<collection-id>.<region>.aoss.amazonaws.com or https://<domain-endpoint>."
|
||||
)
|
||||
if _ACCOUNT_ID_LABEL.match(parsed.hostname) and aws_signing_target(api_base or "")[0] == "aoss":
|
||||
raise _bad_request(_PER_ACCOUNT_MESSAGE)
|
||||
return (api_base or "").rstrip("/")
|
||||
|
||||
def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> NoReturn:
|
||||
raise _bad_request(_SEARCH_ONLY_MESSAGE)
|
||||
|
||||
def get_vector_store_list_url(self, api_base: str) -> NoReturn:
|
||||
raise _bad_request(_SEARCH_ONLY_MESSAGE)
|
||||
|
||||
@staticmethod
|
||||
def _params(
|
||||
litellm_params: Mapping[str, object],
|
||||
optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
extra_body: Mapping[str, object] | None,
|
||||
) -> tuple[_OpenSearchParams, int]:
|
||||
unknown: Final = sorted(
|
||||
key for key in litellm_params if key.startswith(_PARAM_PREFIX) and key not in _KNOWN_PARAMS
|
||||
)
|
||||
if unknown:
|
||||
raise _bad_request(
|
||||
f"Unknown OpenSearch vector store parameter(s): {', '.join(unknown)}. "
|
||||
f"Supported: {', '.join(sorted(_KNOWN_PARAMS))}."
|
||||
)
|
||||
if extra_body:
|
||||
raise _bad_request("The OpenSearch vector store does not support extra_body.")
|
||||
try:
|
||||
options: Final = _SearchOptions.model_validate(optional_params)
|
||||
except ValidationError as error:
|
||||
raise _bad_request(f"Invalid vector store search parameters: {error}") from None
|
||||
if options.unsupported and not drops_unsupported_params(litellm_params.get("drop_params")):
|
||||
raise _bad_request(
|
||||
f"The OpenSearch vector store does not support {', '.join(options.unsupported)} yet. "
|
||||
"Remove it from the request, or set drop_params to drop it."
|
||||
)
|
||||
try:
|
||||
return _OpenSearchParams.model_validate(litellm_params), options.k
|
||||
except ValidationError:
|
||||
raise _bad_request(
|
||||
"opensearch_embedding_field, opensearch_index_name, opensearch_text_field, and litellm_embedding_model "
|
||||
"must be strings."
|
||||
) from None
|
||||
|
||||
@staticmethod
|
||||
def _search_request(
|
||||
vector_store_id: str,
|
||||
embedding: EmbeddingResponse,
|
||||
k: int,
|
||||
params: _OpenSearchParams,
|
||||
api_base: str,
|
||||
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
|
||||
vector: Final = _QUERY_VECTOR.validate_python(embedding.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here
|
||||
return (
|
||||
f"{api_base}/{quote(params.index_name(vector_store_id), safe='')}/_search",
|
||||
{
|
||||
"size": k,
|
||||
"query": {"knn": {params.embedding_field: {"vector": vector, "k": k}}},
|
||||
"_source": {"excludes": [params.embedding_field]},
|
||||
},
|
||||
)
|
||||
|
||||
def transform_search_vector_store_request(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
query: str | Sequence[str],
|
||||
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
api_base: str,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
litellm_params: Mapping[str, object],
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
|
||||
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
|
||||
params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body)
|
||||
embedding: Final = self.query_embedding_executor(embedding_executor, None).embed(
|
||||
params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params)
|
||||
)
|
||||
return self._search_request(vector_store_id, embedding, k, params, api_base)
|
||||
|
||||
async def atransform_search_vector_store_request(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
query: str | Sequence[str],
|
||||
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
api_base: str,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
litellm_params: Mapping[str, object],
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
|
||||
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
|
||||
params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body)
|
||||
embedding: Final = await self.query_embedding_executor(embedding_executor, None).aembed(
|
||||
params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params)
|
||||
)
|
||||
return self._search_request(vector_store_id, embedding, k, params, api_base)
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
optional_params: Mapping[str, object],
|
||||
request_data: Mapping[str, object],
|
||||
api_base: str,
|
||||
api_key: str | None = None,
|
||||
) -> tuple[dict[str, str], bytes]: # mutable-ok: the shared HTTP handler requires writable headers
|
||||
region: Final = self._signing_region(optional_params, api_base)
|
||||
body: Final = json.dumps(request_data)
|
||||
signed: Final = sign_aws_json_post(
|
||||
get_credentials=partial(self._signing_credentials, optional_params, region),
|
||||
service_name=aws_signing_target(api_base)[0],
|
||||
aws_region_name=region,
|
||||
url=api_base,
|
||||
body=body,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"x-amz-content-sha256": hashlib.sha256(body.encode()).hexdigest(),
|
||||
},
|
||||
)
|
||||
signed_headers: Final = dict(signed.headers.items())
|
||||
signed_names: Final = frozenset(name.lower() for name in signed_headers)
|
||||
return {
|
||||
**{name: value for name, value in headers.items() if name.lower() not in signed_names},
|
||||
**signed_headers,
|
||||
}, body.encode()
|
||||
|
||||
def transform_search_vector_store_response(
|
||||
self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj"
|
||||
) -> VectorStoreSearchResponse:
|
||||
if not response.is_success:
|
||||
raise BaseLLMException(status_code=response.status_code, message=response.text, headers=response.headers)
|
||||
try:
|
||||
hits: Final = _SearchResponse.model_validate_json(response.content).hits.hits
|
||||
except ValidationError:
|
||||
raise BaseLLMException(
|
||||
status_code=502,
|
||||
message="OpenSearch answered the search without a hits.hits list.",
|
||||
headers=response.headers,
|
||||
) from None
|
||||
logged: Final = _LoggedSearch.model_validate(litellm_logging_obj.model_call_details)
|
||||
return VectorStoreSearchResponse(
|
||||
object="vector_store.search_results.page",
|
||||
search_query=logged.search_query,
|
||||
data=[_search_result(hit, logged.litellm_params) for hit in hits],
|
||||
)
|
||||
|
||||
def validate_create_vector_store(self) -> NoReturn:
|
||||
raise _bad_request(_SEARCH_ONLY_MESSAGE)
|
||||
|
||||
def transform_create_vector_store_request(
|
||||
self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str
|
||||
) -> NoReturn:
|
||||
raise _bad_request(_SEARCH_ONLY_MESSAGE)
|
||||
|
||||
def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn:
|
||||
raise _bad_request(_SEARCH_ONLY_MESSAGE)
|
||||
|
|
@ -11,6 +11,7 @@ import httpx
|
|||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import message_with_file_parts_as_image_urls
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
|
@ -104,7 +105,7 @@ def _message_role_mapping(role: str) -> type[ChatMessage]:
|
|||
def _messages_to_sap_template(messages: list[AllMessageValues]) -> list:
|
||||
template: Final = []
|
||||
for message in messages:
|
||||
folded_message = _fold_message_cache_control(message)
|
||||
folded_message = _fold_message_cache_control(message_with_file_parts_as_image_urls(message))
|
||||
message_role_class = _message_role_mapping(str(folded_message["role"]))
|
||||
validated_message = validate_dict(folded_message, message_role_class)
|
||||
template.append(validated_message)
|
||||
|
|
|
|||
|
|
@ -112,7 +112,7 @@ def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
|
|||
pdf/text data URIs become document blocks rather than malformed image blocks.
|
||||
"""
|
||||
image_url: Final = block.get("image_url")
|
||||
url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url")
|
||||
url: Final = image_url if isinstance(image_url, str) else _str_field(image_url, "url")
|
||||
if not url:
|
||||
return block
|
||||
|
||||
|
|
@ -121,18 +121,31 @@ def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
|
|||
if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document"
|
||||
else create_anthropic_image_param(
|
||||
image_url if isinstance(image_url, dict) else url,
|
||||
format=_image_url_field(image_url, "format"),
|
||||
format=_str_field(image_url, "format"),
|
||||
is_bedrock_invoke=True,
|
||||
)
|
||||
)
|
||||
return _with_cache_control(converted, block)
|
||||
|
||||
|
||||
def _convert_file_to_anthropic(block: Mapping[str, object]) -> object:
|
||||
file_data: Final = _str_field(block.get("file"), "file_data")
|
||||
if not file_data:
|
||||
return block
|
||||
return _with_cache_control(
|
||||
anthropic_process_openai_file_message({"type": "file", "file": {"file_data": file_data}}), block
|
||||
)
|
||||
|
||||
|
||||
def _with_cache_control(converted: Mapping[str, object], block: Mapping[str, object]) -> Mapping[str, object]:
|
||||
cache_control: Final = block.get("cache_control")
|
||||
if cache_control is None:
|
||||
return converted
|
||||
return {**converted, "cache_control": cache_control}
|
||||
|
||||
|
||||
def _image_url_field(image_url: object, key: str) -> str | None:
|
||||
value: Final = image_url.get(key) if isinstance(image_url, dict) else None
|
||||
def _str_field(container: object, key: str) -> str | None:
|
||||
value: Final = container.get(key) if isinstance(container, dict) else None
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
|
|
@ -141,15 +154,23 @@ def _data_uri_media_type(url: str) -> str:
|
|||
return match.group(1) if match else ""
|
||||
|
||||
|
||||
def _convert_image_url_blocks_to_anthropic(content: object) -> object:
|
||||
def _convert_media_block_to_anthropic(block: object) -> object:
|
||||
mapping: Final = as_str_mapping(block)
|
||||
if mapping is None:
|
||||
return block
|
||||
match mapping.get("type"):
|
||||
case "image_url":
|
||||
return _convert_image_url_to_anthropic(mapping)
|
||||
case "file":
|
||||
return _convert_file_to_anthropic(mapping)
|
||||
case _:
|
||||
return block
|
||||
|
||||
|
||||
def _convert_media_blocks_to_anthropic(content: object) -> object:
|
||||
if not isinstance(content, list):
|
||||
return content
|
||||
return [
|
||||
_convert_image_url_to_anthropic(block)
|
||||
if isinstance(block, Mapping) and block.get("type") == "image_url"
|
||||
else block
|
||||
for block in content
|
||||
]
|
||||
return [_convert_media_block_to_anthropic(block) for block in content]
|
||||
|
||||
|
||||
def _convert_tool_result_to_anthropic(
|
||||
|
|
@ -391,7 +412,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
|
|||
conversation.append(
|
||||
{
|
||||
"role": role,
|
||||
"content": _convert_image_url_blocks_to_anthropic(content),
|
||||
"content": _convert_media_blocks_to_anthropic(content),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue