mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_registry_audit_2026_09_01
This commit is contained in:
commit
23977fc290
223 changed files with 13464 additions and 1871 deletions
|
|
@ -1483,7 +1483,7 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
|
||||
installing_litellm_on_python_3_13:
|
||||
docker:
|
||||
|
|
@ -1507,9 +1507,9 @@ jobs:
|
|||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
|
||||
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
|
||||
|
||||
installing_litellm_on_python_legacy_migration_resolver:
|
||||
installing_litellm_on_python_v2_migration_resolver:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
|
|
@ -1536,10 +1536,10 @@ jobs:
|
|||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Run legacy migration resolver proxy smoke test
|
||||
name: Run v2 migration resolver proxy smoke test
|
||||
command: |
|
||||
uv run --no-sync python -m pytest -vv \
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver
|
||||
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver
|
||||
|
||||
helm_chart_testing:
|
||||
machine:
|
||||
|
|
@ -2879,8 +2879,7 @@ jobs:
|
|||
command: |
|
||||
if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \
|
||||
(grep -q "Database setup failed after multiple retries" docker_output.log || \
|
||||
grep -q "ERROR: Application startup failed. Exiting." docker_output.log || \
|
||||
grep -q "Database migration cannot proceed" docker_output.log); then
|
||||
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
|
||||
echo "Expected error found. Test passed."
|
||||
else
|
||||
echo "Expected error not found. Test failed."
|
||||
|
|
@ -3012,7 +3011,7 @@ workflows:
|
|||
filters: *main_branches
|
||||
- installing_litellm_on_python_3_13:
|
||||
filters: *main_branches
|
||||
- installing_litellm_on_python_legacy_migration_resolver:
|
||||
- installing_litellm_on_python_v2_migration_resolver:
|
||||
filters: *main_branches
|
||||
- helm_chart_testing:
|
||||
requires:
|
||||
|
|
|
|||
6
.github/workflows/image-scan.yml
vendored
6
.github/workflows/image-scan.yml
vendored
|
|
@ -80,7 +80,7 @@ jobs:
|
|||
LITELLM_IMAGE: litellm-image-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
# Scans the whole shipped artifact: OS/apk plus every language package
|
||||
# baked into the image, including ones no lockfile declares (e.g. prisma's
|
||||
|
|
@ -124,7 +124,7 @@ jobs:
|
|||
LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
migrations-image:
|
||||
name: migrations-image
|
||||
|
|
@ -185,7 +185,7 @@ jobs:
|
|||
LITELLM_COMPONENT_PORT: "4000"
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
ui-image:
|
||||
name: ui-image
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
|
|||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
|
||||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
@ -101,6 +103,12 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
# The base image only configures Chainguard's authenticated apk repo, which
|
||||
# requires an enterprise subscription. Add the public Wolfi repo so `apk add`
|
||||
# also works for anyone installing extra packages into a running container.
|
||||
# https://github.com/BerriAI/litellm/issues/33518
|
||||
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
|
||||
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
|
||||
|
||||
|
|
|
|||
|
|
@ -117,7 +117,7 @@
|
|||
"limit": 111
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 695
|
||||
"limit": 692
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -85,6 +86,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -97,6 +98,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13 \
|
||||
--no-sources-package litellm-proxy-extras; \
|
||||
else \
|
||||
|
|
@ -106,6 +108,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13; \
|
||||
fi
|
||||
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
|
|||
| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` |
|
||||
| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A |
|
||||
| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A |
|
||||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
|
||||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A |
|
||||
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
|
||||
|
|
@ -212,6 +212,8 @@ service, the **Proxy Endpoint** should be set to `http://<RELEASE>-litellm:4000`
|
|||
The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey`
|
||||
was not provided to the helm command line, the `masterkey` is a randomly
|
||||
generated string in the `sk-...` format stored in the `<RELEASE>-litellm-masterkey` Kubernetes Secret.
|
||||
The key is generated once on the first install; later `helm upgrade` runs reuse the
|
||||
value already in that Secret, so upgrading never rotates the master key.
|
||||
|
||||
```bash
|
||||
kubectl -n litellm get secret <RELEASE>-litellm-masterkey -o jsonpath="{.data.masterkey}"
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
{{- if not .Values.masterkeySecretName }}
|
||||
{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }}
|
||||
{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }}
|
||||
{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }}
|
||||
{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }}
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
name: {{ include "litellm.fullname" . }}-masterkey
|
||||
name: {{ $secretName }}
|
||||
data:
|
||||
masterkey: {{ $masterkey | b64enc }}
|
||||
type: Opaque
|
||||
|
|
|
|||
|
|
@ -15,6 +15,53 @@ tests:
|
|||
# Note: The masterkey is generated as "sk-<18-random-chars>" in plain text,
|
||||
# but stored as base64 encoded in Kubernetes secret (requirement).
|
||||
# "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern.
|
||||
- it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
masterkeySecretName: ""
|
||||
kubernetesProvider:
|
||||
scheme:
|
||||
"v1/Secret":
|
||||
gvr:
|
||||
version: "v1"
|
||||
resource: "secrets"
|
||||
namespaced: true
|
||||
objects:
|
||||
- kind: Secret
|
||||
apiVersion: v1
|
||||
metadata:
|
||||
name: RELEASE-NAME-litellm-masterkey
|
||||
namespace: NAMESPACE
|
||||
data:
|
||||
masterkey: c2stZXhpc3Rpbmcta2V5
|
||||
asserts:
|
||||
- equal:
|
||||
path: data.masterkey
|
||||
value: c2stZXhpc3Rpbmcta2V5
|
||||
- it: should let an explicit masterkey value override the one already stored in the cluster
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
masterkeySecretName: ""
|
||||
masterkey: sk-explicit
|
||||
kubernetesProvider:
|
||||
scheme:
|
||||
"v1/Secret":
|
||||
gvr:
|
||||
version: "v1"
|
||||
resource: "secrets"
|
||||
namespaced: true
|
||||
objects:
|
||||
- kind: Secret
|
||||
apiVersion: v1
|
||||
metadata:
|
||||
name: RELEASE-NAME-litellm-masterkey
|
||||
namespace: NAMESPACE
|
||||
data:
|
||||
masterkey: c2stZXhpc3Rpbmcta2V5
|
||||
asserts:
|
||||
- equal:
|
||||
path: data.masterkey
|
||||
value: c2stZXhwbGljaXQ=
|
||||
- it: should not create a secret if masterkeySecretName is set
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
|
|
|
|||
|
|
@ -7,8 +7,7 @@ import subprocess
|
|||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Optional
|
||||
from typing import Optional
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.replica_identity import (
|
||||
|
|
@ -51,17 +50,6 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
|
|||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
_PRISMA_ATTEMPTS: Final = 4
|
||||
|
||||
_TRANSIENT_PRISMA_FAILURES: Final = MappingProxyType(
|
||||
{
|
||||
"deadlock detected": "a deadlock on the migration advisory lock (a concurrent migrate deploy)",
|
||||
"P1001": "an unreachable database server",
|
||||
"P1002": "a database server that timed out",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
|
||||
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
|
||||
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
|
||||
|
|
@ -286,23 +274,6 @@ class ProxyExtrasDBManager:
|
|||
env=prisma_env,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _transient_prisma_failure(stderr: str) -> str | None:
|
||||
"""Why a failed prisma command is worth retrying, or None.
|
||||
|
||||
v1 retried every failure, so it absorbed a database that was not up yet
|
||||
or another instance holding the migration lock. v2 fails fast, which is
|
||||
right for a broken migration and wrong for these.
|
||||
"""
|
||||
return next(
|
||||
(
|
||||
reason
|
||||
for marker, reason in _TRANSIENT_PRISMA_FAILURES.items()
|
||||
if marker in stderr
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_permission_error(error_message: str) -> bool:
|
||||
"""
|
||||
|
|
@ -684,7 +655,7 @@ class ProxyExtrasDBManager:
|
|||
@staticmethod
|
||||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
"""
|
||||
v2 migration resolver (what the proxy CLI selects by default).
|
||||
v2 migration resolver (opt-in via --use_v2_migration_resolver).
|
||||
|
||||
Runs `prisma migrate deploy` and handles standard recovery paths
|
||||
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
|
||||
|
|
@ -705,46 +676,20 @@ class ProxyExtrasDBManager:
|
|||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
for attempt in range(_PRISMA_ATTEMPTS):
|
||||
try:
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.info(
|
||||
"prisma db push attempt %s timed out, retrying",
|
||||
attempt + 1,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
except subprocess.CalledProcessError as e:
|
||||
stderr = e.stderr or ""
|
||||
transient = ProxyExtrasDBManager._transient_prisma_failure(
|
||||
stderr
|
||||
)
|
||||
# Re-raise as RuntimeError so proxy_cli.py's
|
||||
# `except RuntimeError` catches it and exits cleanly.
|
||||
if transient is None or attempt == _PRISMA_ATTEMPTS - 1:
|
||||
raise RuntimeError(
|
||||
f"prisma db push failed.\n\nDetail: {e}"
|
||||
f"\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
logger.info(
|
||||
"prisma db push attempt %s failed on %s, retrying. "
|
||||
"Prisma error:\n%s",
|
||||
attempt + 1,
|
||||
transient,
|
||||
stderr,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
raise RuntimeError(
|
||||
f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts."
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
return True
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
subprocess.TimeoutExpired,
|
||||
) as e:
|
||||
# Re-raise as RuntimeError so proxy_cli.py's
|
||||
# `except RuntimeError` catches it and exits cleanly.
|
||||
raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
||||
|
|
@ -754,7 +699,7 @@ class ProxyExtrasDBManager:
|
|||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
for attempt in range(_PRISMA_ATTEMPTS):
|
||||
for attempt in range(4):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
|
|
@ -869,36 +814,16 @@ class ProxyExtrasDBManager:
|
|||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
transient = ProxyExtrasDBManager._transient_prisma_failure(stderr)
|
||||
if transient is None:
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
if attempt == _PRISMA_ATTEMPTS - 1:
|
||||
raise RuntimeError(
|
||||
f"Database migration failed after "
|
||||
f"{_PRISMA_ATTEMPTS} attempts on {transient}. "
|
||||
"Check database connectivity and load."
|
||||
f"\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
logger.info(
|
||||
"prisma migrate deploy attempt %s failed on %s, retrying. "
|
||||
"Prisma error:\n%s",
|
||||
attempt + 1,
|
||||
transient,
|
||||
stderr,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
raise RuntimeError(
|
||||
f"Database migration failed after {_PRISMA_ATTEMPTS} "
|
||||
"attempts (retry loop exhausted by timeouts or repeated "
|
||||
"idempotent-recovery continues). Check database connectivity, "
|
||||
"load, and _prisma_migrations ledger state."
|
||||
"Database migration failed after 4 attempts (retry loop "
|
||||
"exhausted by timeouts or repeated idempotent-recovery "
|
||||
"continues). Check database connectivity, load, and "
|
||||
"_prisma_migrations ledger state."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
@ -946,11 +871,10 @@ class ProxyExtrasDBManager:
|
|||
|
||||
Args:
|
||||
use_migrate: Whether to use prisma migrate instead of db push
|
||||
use_v2_resolver: Run the v2 migration resolver (safer during
|
||||
use_v2_resolver: Opt into the v2 migration resolver (safer during
|
||||
rolling deploys; does not run the diff-and-force recovery
|
||||
that causes schema thrashing). Defaults to False here so
|
||||
direct callers keep the old behavior; the proxy CLI passes
|
||||
True, so the proxy's runtime default is v2.
|
||||
that causes schema thrashing). Defaults to False for
|
||||
backwards compatibility.
|
||||
|
||||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
|
|
@ -968,7 +892,7 @@ class ProxyExtrasDBManager:
|
|||
@staticmethod
|
||||
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
if use_v2_resolver:
|
||||
logger.info("Using v2 migration resolver")
|
||||
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
|
||||
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)
|
||||
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
|
|
|
|||
0
litellm-proxy-extras/tests/__init__.py
Normal file
0
litellm-proxy-extras/tests/__init__.py
Normal file
242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
Normal file
242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
"""Regression tests for ProxyExtrasDBManager v2 migration resolver.
|
||||
|
||||
The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
|
||||
`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1
|
||||
(default) behavior is unchanged from pre-fix.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm_proxy_extras.utils import (
|
||||
ProxyExtrasDBManager,
|
||||
_max_migration_timestamp,
|
||||
_migration_timestamp,
|
||||
)
|
||||
|
||||
|
||||
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
|
||||
def _run(*args, **kwargs):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=returncode,
|
||||
cmd=args[0],
|
||||
stderr=stderr,
|
||||
output="",
|
||||
)
|
||||
|
||||
return _run
|
||||
|
||||
|
||||
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a permission failure during migrate deploy raises RuntimeError."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3018\nMigration name: 20250326162113_baseline\n"
|
||||
"Database error code: 42501\npermission denied for schema public"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="permission"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
|
||||
'Reason: syntax error at or near "BRKN" LINE 42'
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_removes_connection_limit():
|
||||
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
|
||||
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
|
||||
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
assert "connection_limit" not in stripped
|
||||
assert "pool_timeout" not in stripped
|
||||
assert "sslmode=require" in stripped
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_passthrough_no_query():
|
||||
"""URLs without query strings are returned unchanged."""
|
||||
url = "postgresql://u:p@h:5432/db"
|
||||
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
|
||||
|
||||
|
||||
def test_migration_timestamp_extracts_leading_digits():
|
||||
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
|
||||
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
|
||||
|
||||
|
||||
def test_migration_timestamp_returns_zero_on_malformed():
|
||||
assert _migration_timestamp("0_init") == 0
|
||||
assert _migration_timestamp("not_a_migration") == 0
|
||||
|
||||
|
||||
def test_max_migration_timestamp():
|
||||
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
|
||||
assert _max_migration_timestamp(names) == 20260415000000
|
||||
|
||||
|
||||
def test_max_migration_timestamp_empty_set():
|
||||
assert _max_migration_timestamp(set()) == 0
|
||||
|
||||
|
||||
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
|
||||
|
||||
This is the existing buggy behavior — we're not fixing it in v1, only
|
||||
offering v2 as opt-in. This test pins the default so that a future
|
||||
inadvertent default flip is caught.
|
||||
"""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
# Stub `prisma migrate deploy` to claim success with pending migrations
|
||||
# applied, which is the code path that triggers the legacy post-migration
|
||||
# sanity check (a call to _resolve_all_migrations).
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
def fake_run(cmd, *args, **kwargs):
|
||||
return FakeResult()
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
|
||||
def fake_resolve(*args, **kwargs):
|
||||
resolve_called["n"] += 1
|
||||
|
||||
monkeypatch.setattr("subprocess.run", fake_run)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
|
||||
|
||||
|
||||
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
|
||||
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = "db push error"
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="prisma db push failed"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
||||
"""_warn_if_db_ahead_of_head must never raise — it's informational.
|
||||
|
||||
Non-connection DB errors (e.g. InsufficientPrivilege from a user
|
||||
without SELECT on _prisma_migrations) must be caught, not propagated.
|
||||
"""
|
||||
import psycopg
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class _FakeConn:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def execute(self, *a, **kw):
|
||||
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
|
||||
raise psycopg.errors.InsufficientPrivilege("permission denied")
|
||||
|
||||
def _fake_connect(*a, **kw):
|
||||
return _FakeConn()
|
||||
|
||||
monkeypatch.setattr("psycopg.connect", _fake_connect)
|
||||
|
||||
# Must not raise.
|
||||
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
|
||||
|
||||
|
||||
def test_v2_resolve_specific_migration_failure_raises_runtime_error(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""If marking a migration as applied fails inside P3009 idempotent
|
||||
recovery, the subprocess error must be re-raised as RuntimeError so
|
||||
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
|
||||
)
|
||||
|
||||
# First call: migrate deploy -> P3009 idempotent error.
|
||||
# Recovery path tries _resolve_specific_migration; that also raises.
|
||||
def _failing_resolve(*a, **kw):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd="prisma migrate resolve --applied",
|
||||
stderr="resolve failed",
|
||||
output="",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
|
||||
)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
|
||||
"relation already exists"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(
|
||||
RuntimeError, match="Failed to mark migration .* as applied"
|
||||
):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_all_migrations",
|
||||
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
|
||||
)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
|
||||
|
|
@ -13,6 +13,12 @@ DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
|
|||
DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
|
||||
DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512))
|
||||
# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
|
||||
MAX_S3_OBJECT_KEY_BYTES: Final = 1024
|
||||
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
|
||||
S3_PREFIX_DIGEST_CHARS: Final = 16
|
||||
# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
|
||||
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
|
||||
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
|
||||
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
|
||||
|
|
@ -130,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
|
|||
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
|
||||
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
|
||||
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
|
||||
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
|
||||
|
||||
# Allowlist of commands permitted for MCP stdio transport.
|
||||
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
|
||||
|
|
@ -1583,6 +1590,7 @@ KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
|
|||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
|
||||
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
|
||||
MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job"
|
||||
USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job"
|
||||
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
|
||||
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
|
||||
SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning"
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import base64
|
|||
import os
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from datetime import timedelta
|
||||
from functools import partial
|
||||
from importlib import metadata
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
|
|
@ -47,7 +48,8 @@ from mcp.types import Tool as MCPTool
|
|||
from pydantic import AnyUrl
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
|
||||
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
||||
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
from litellm.types.llms.custom_http import VerifyTypes
|
||||
from litellm.types.mcp import (
|
||||
|
|
@ -603,17 +605,19 @@ class MCPClient:
|
|||
"""
|
||||
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
|
||||
|
||||
async def _list_tools_operation(session: ClientSession):
|
||||
return await session.list_tools()
|
||||
|
||||
try:
|
||||
result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
|
||||
tool_count: Final = len(result.tools)
|
||||
tool_names: Final = [tool.name for tool in result.tools]
|
||||
# A per-server timeout above the global default extends the whole-walk deadline
|
||||
listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)
|
||||
tools: Final = await self.run_with_session(
|
||||
partial(list_tools_with_pagination, listing_deadline=listing_deadline),
|
||||
quiet_on_error=raise_on_error,
|
||||
)
|
||||
tool_count: Final = len(tools)
|
||||
tool_names: Final = tuple(tool.name for tool in tools)
|
||||
verbose_logger.info(
|
||||
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
|
||||
)
|
||||
return result.tools
|
||||
return tools
|
||||
except asyncio.CancelledError:
|
||||
verbose_logger.warning("MCP client list_tools was cancelled")
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -1,14 +1,22 @@
|
|||
import json
|
||||
from typing import Final, Literal
|
||||
|
||||
import anyio
|
||||
from mcp import ClientSession
|
||||
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
|
||||
from mcp.types import CallToolResult as MCPCallToolResult
|
||||
from mcp.types import PaginatedRequestParams
|
||||
from mcp.types import Tool as MCPTool
|
||||
from openai.types.chat import ChatCompletionToolParam
|
||||
from openai.types.responses.function_tool_param import FunctionToolParam
|
||||
from openai.types.shared_params.function_definition import FunctionDefinition
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
MCP_CLIENT_TIMEOUT,
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesTool
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
|
||||
|
|
@ -90,6 +98,64 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
|
|||
)
|
||||
|
||||
|
||||
async def list_tools_with_pagination(
|
||||
session: ClientSession, listing_deadline: float | None = None
|
||||
) -> list[MCPTool]: # mutable-ok: list return contract
|
||||
"""Collect tools from every tools/list page by following nextCursor.
|
||||
|
||||
Stops and returns the tools collected so far when the upstream repeats a
|
||||
cursor, the page cap is reached, or the whole-walk deadline expires, so a
|
||||
buggy or slow upstream yields a partial catalog instead of an error.
|
||||
listing_deadline overrides the default whole-walk deadline; callers with a
|
||||
per-server timeout above the global default pass it through here.
|
||||
"""
|
||||
tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools
|
||||
seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops
|
||||
cursor: str | None = None # rebind-ok: advances to each page's nextCursor
|
||||
# The per-request session read timeout restarts on every page, so a multi-page
|
||||
# walk needs its own overall deadline. max() keeps the pre-pagination guarantee
|
||||
# that a single page slower than the listing timeout but within the client
|
||||
# timeout still succeeds.
|
||||
effective_deadline: Final = (
|
||||
listing_deadline if listing_deadline is not None else max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT)
|
||||
)
|
||||
|
||||
with anyio.move_on_after(effective_deadline):
|
||||
for _ in range(MCP_TOOL_LISTING_MAX_PAGES):
|
||||
result = (
|
||||
await session.list_tools()
|
||||
if cursor is None
|
||||
else await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
|
||||
)
|
||||
tools.extend(result.tools)
|
||||
|
||||
next_cursor = getattr(result, "nextCursor", None)
|
||||
if not isinstance(next_cursor, str) or not next_cursor:
|
||||
return tools
|
||||
if next_cursor in seen_cursors:
|
||||
verbose_logger.warning(
|
||||
"MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far",
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
seen_cursors.add(next_cursor)
|
||||
cursor = next_cursor
|
||||
|
||||
verbose_logger.warning(
|
||||
"MCP server tools/list pagination exceeded the maximum of %s pages; returning %s tools collected so far",
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
|
||||
verbose_logger.warning(
|
||||
"MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far",
|
||||
effective_deadline,
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
|
||||
|
||||
async def load_mcp_tools(
|
||||
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
|
||||
) -> list[MCPTool] | list[ChatCompletionToolParam]:
|
||||
|
|
@ -103,10 +169,12 @@ async def load_mcp_tools(
|
|||
|
||||
If format is set to "openai", the tools are converted to OpenAI API compatible tools.
|
||||
"""
|
||||
tools: Final = await session.list_tools()
|
||||
tools: Final = await list_tools_with_pagination(session)
|
||||
if format == "openai":
|
||||
return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools]
|
||||
return tools.tools
|
||||
return [ # mutable-ok: public API returns a list
|
||||
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools
|
||||
]
|
||||
return tools
|
||||
|
||||
|
||||
########################################################
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ from .utils import process_slack_alerting_variables
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.router import Router as _Router
|
||||
|
||||
Router = _Router
|
||||
|
|
@ -1944,6 +1945,69 @@ Model Info:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error sending weekly spend report %s", e)
|
||||
|
||||
async def send_user_spend_alerts(self, prisma_client: "PrismaClient | None" = None) -> None:
|
||||
"""Check per-user daily/monthly spend thresholds and spend anomalies, alerting once per user per period."""
|
||||
if self.alerting is None or "slack" not in self.alerting:
|
||||
return
|
||||
|
||||
thresholds_enabled: Final = AlertType.user_spend_thresholds in self.alert_types
|
||||
anomalies_enabled: Final = AlertType.user_spend_anomalies in self.alert_types
|
||||
if not thresholds_enabled and not anomalies_enabled:
|
||||
return
|
||||
|
||||
if prisma_client is None:
|
||||
from litellm.proxy.proxy_server import prisma_client as global_prisma_client
|
||||
|
||||
prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client
|
||||
if prisma_client is None:
|
||||
return
|
||||
|
||||
from litellm.integrations.SlackAlerting.user_spend_alerts import (
|
||||
evaluate_user_spend,
|
||||
fetch_user_spend_rows,
|
||||
)
|
||||
|
||||
try:
|
||||
today: Final = datetime.datetime.now(datetime.timezone.utc).date()
|
||||
rows: Final = await fetch_user_spend_rows(
|
||||
prisma_client=prisma_client,
|
||||
today=today,
|
||||
baseline_days=self.alerting_args.spend_anomaly_baseline_days,
|
||||
)
|
||||
all_events: Final = tuple(
|
||||
event
|
||||
for row in rows
|
||||
for event in evaluate_user_spend(
|
||||
row=row,
|
||||
args=self.alerting_args,
|
||||
today=today,
|
||||
thresholds_enabled=thresholds_enabled,
|
||||
anomalies_enabled=anomalies_enabled,
|
||||
)
|
||||
)
|
||||
cached_flags: Final = await asyncio.gather(
|
||||
*(self.internal_usage_cache.async_get_cache(key=event.cache_key) for event in all_events)
|
||||
)
|
||||
new_events: Final = tuple(event for event, cached in zip(all_events, cached_flags) if not cached)
|
||||
for alert_type in (AlertType.user_spend_thresholds, AlertType.user_spend_anomalies):
|
||||
typed_events = tuple(event for event in new_events if event.alert_type == alert_type)
|
||||
if not typed_events:
|
||||
continue
|
||||
await self.send_alert(
|
||||
message="\n\n".join(event.message for event in typed_events),
|
||||
level="High",
|
||||
alert_type=alert_type,
|
||||
alerting_metadata={}, # mutable-ok: send_alert takes a dict payload
|
||||
)
|
||||
for event in typed_events:
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=event.cache_key,
|
||||
value="SENT",
|
||||
ttl=event.cache_ttl,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # background job must not crash the scheduler
|
||||
verbose_proxy_logger.exception("Error sending user spend alerts: %s", e)
|
||||
|
||||
async def send_fallback_stats_from_prometheus(self):
|
||||
"""
|
||||
Helper to send fallback statistics from prometheus server -> to slack
|
||||
|
|
|
|||
139
litellm/integrations/SlackAlerting/user_spend_alerts.py
Normal file
139
litellm/integrations/SlackAlerting/user_spend_alerts.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
"""Per-user daily/monthly spend threshold alerts and spend anomaly detection."""
|
||||
|
||||
import datetime
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.constants import HOURS_IN_A_DAY
|
||||
from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
DAY_SECONDS: Final = HOURS_IN_A_DAY * 60 * 60
|
||||
MONTHLY_ALERT_TTL_SECONDS: Final = 32 * DAY_SECONDS
|
||||
|
||||
USER_SPEND_QUERY: Final = """
|
||||
SELECT
|
||||
user_id,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date = $1), 0)::float AS daily_spend,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0)::float AS monthly_spend,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date >= $3 AND date < $1), 0)::float AS baseline_spend
|
||||
FROM "LiteLLM_DailyUserSpend"
|
||||
WHERE date >= LEAST($2, $3) AND user_id IS NOT NULL
|
||||
GROUP BY user_id
|
||||
HAVING COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0) > 0
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserSpendRow:
|
||||
user_id: str
|
||||
daily_spend: float
|
||||
monthly_spend: float
|
||||
baseline_spend: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserSpendAlertEvent:
|
||||
kind: Literal["daily_threshold", "monthly_threshold", "anomaly"]
|
||||
alert_type: AlertType
|
||||
message: str
|
||||
cache_key: str
|
||||
cache_ttl: int
|
||||
|
||||
|
||||
USER_SPEND_ROWS_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...])
|
||||
|
||||
|
||||
async def fetch_user_spend_rows(
|
||||
prisma_client: "PrismaClient",
|
||||
today: datetime.date,
|
||||
baseline_days: int,
|
||||
) -> tuple[UserSpendRow, ...]:
|
||||
today_str: Final = today.strftime("%Y-%m-%d")
|
||||
month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d")
|
||||
baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d")
|
||||
raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
|
||||
return USER_SPEND_ROWS_ADAPTER.validate_python(raw)
|
||||
|
||||
|
||||
def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
|
||||
threshold: Final = args.daily_spend_per_user_threshold
|
||||
if threshold is None or row.daily_spend < threshold:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="daily_threshold",
|
||||
alert_type=AlertType.user_spend_thresholds,
|
||||
message=(
|
||||
f"User Daily Spend Threshold Crossed:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend Today: `${row.daily_spend:.2f}`\n"
|
||||
f"Daily Threshold: `${threshold:.2f}`"
|
||||
),
|
||||
cache_key=f"user_spend_alert_daily_{row.user_id}_{today_str}",
|
||||
cache_ttl=DAY_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _monthly_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, month_str: str) -> UserSpendAlertEvent | None:
|
||||
threshold: Final = args.monthly_spend_per_user_threshold
|
||||
if threshold is None or row.monthly_spend < threshold:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="monthly_threshold",
|
||||
alert_type=AlertType.user_spend_thresholds,
|
||||
message=(
|
||||
f"User Monthly Spend Threshold Crossed:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend This Month: `${row.monthly_spend:.2f}`\n"
|
||||
f"Monthly Threshold: `${threshold:.2f}`"
|
||||
),
|
||||
cache_key=f"user_spend_alert_monthly_{row.user_id}_{month_str}",
|
||||
cache_ttl=MONTHLY_ALERT_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _anomaly_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
|
||||
if row.daily_spend < args.spend_anomaly_min_spend:
|
||||
return None
|
||||
baseline_daily_avg: Final = row.baseline_spend / args.spend_anomaly_baseline_days
|
||||
if row.baseline_spend > 0 and row.daily_spend <= args.spend_anomaly_multiplier * baseline_daily_avg:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="anomaly",
|
||||
alert_type=AlertType.user_spend_anomalies,
|
||||
message=(
|
||||
f"User Spend Anomaly Detected:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend Today: `${row.daily_spend:.2f}`\n"
|
||||
f"Daily Average (last {args.spend_anomaly_baseline_days} days): `${baseline_daily_avg:.2f}`\n"
|
||||
f"Trigger: spend above `{args.spend_anomaly_multiplier}x` the daily average "
|
||||
f"(minimum `${args.spend_anomaly_min_spend:.2f}`)"
|
||||
),
|
||||
cache_key=f"user_spend_alert_anomaly_{row.user_id}_{today_str}",
|
||||
cache_ttl=DAY_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def evaluate_user_spend(
|
||||
row: UserSpendRow,
|
||||
args: SlackAlertingArgs,
|
||||
today: datetime.date,
|
||||
thresholds_enabled: bool,
|
||||
anomalies_enabled: bool,
|
||||
) -> tuple[UserSpendAlertEvent, ...]:
|
||||
today_str: Final = today.strftime("%Y-%m-%d")
|
||||
month_str: Final = today.strftime("%Y-%m")
|
||||
threshold_events: Final = (
|
||||
(
|
||||
_daily_threshold_event(row=row, args=args, today_str=today_str),
|
||||
_monthly_threshold_event(row=row, args=args, month_str=month_str),
|
||||
)
|
||||
if thresholds_enabled
|
||||
else ()
|
||||
)
|
||||
anomaly_events: Final = (_anomaly_event(row=row, args=args, today_str=today_str),) if anomalies_enabled else ()
|
||||
return tuple(event for event in (*threshold_events, *anomaly_events) if event is not None)
|
||||
|
|
@ -11,6 +11,7 @@ import json
|
|||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
import httpx
|
||||
|
|
@ -30,12 +31,16 @@ from litellm.integrations.datadog.datadog_mock_client import (
|
|||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
handle_any_messages_to_chat_completion_str_messages_conversion,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
|
||||
from litellm.types.integrations.datadog_llm_obs import *
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
|
|
@ -44,6 +49,189 @@ from litellm.types.utils import (
|
|||
StandardLoggingPayloadErrorInformation,
|
||||
)
|
||||
|
||||
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""}
|
||||
_MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024
|
||||
|
||||
|
||||
def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]:
|
||||
"""The value at `key` when it is a mapping, else an empty one."""
|
||||
value: Final = source.get(key)
|
||||
return value if isinstance(value, dict) else _EMPTY_MAPPING
|
||||
|
||||
|
||||
def _content_blocks(message: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
return ()
|
||||
return tuple(block for block in content if isinstance(block, dict))
|
||||
|
||||
|
||||
def _to_dd_arguments(raw_arguments: object) -> dict[str, Any] | str:
|
||||
"""
|
||||
Arguments as the object LLM Obs types them as, or the raw string when they are not one.
|
||||
|
||||
Strings past the size bound ship unparsed: decoding multiplies memory on hostile compact
|
||||
JSON, and the raw string is what the intake receives either way.
|
||||
"""
|
||||
if not isinstance(raw_arguments, str):
|
||||
return raw_arguments if isinstance(raw_arguments, dict) else str(raw_arguments)
|
||||
if len(raw_arguments) > _MAX_PARSED_TOOL_ARGUMENT_CHARS:
|
||||
return raw_arguments
|
||||
parsed: Final = safe_json_loads(raw_arguments)
|
||||
return parsed if isinstance(parsed, dict) else raw_arguments
|
||||
|
||||
|
||||
def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]:
|
||||
"""
|
||||
The tool calls a message carries, in LLM Obs' ToolCall schema, from either dialect.
|
||||
|
||||
OpenAI puts them in `tool_calls` with the callee nested under `function` and `arguments`
|
||||
serialized; Anthropic puts them in `content` as `tool_use` blocks with `input` already an
|
||||
object. LLM Obs reads `name` / `arguments` / `tool_id` either way.
|
||||
"""
|
||||
raw_tool_calls: Final = message.get("tool_calls")
|
||||
openai_calls: Final = tuple(
|
||||
ToolCall(
|
||||
name=function.get("name", ""),
|
||||
arguments=_to_dd_arguments(function.get("arguments", "")),
|
||||
tool_id=tool_call.get("id", ""),
|
||||
type=tool_call.get("type", "function"),
|
||||
)
|
||||
for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ())
|
||||
if isinstance(tool_call, dict)
|
||||
for function in [_mapping_field(tool_call, "function")]
|
||||
)
|
||||
anthropic_calls: Final = tuple(
|
||||
ToolCall(
|
||||
name=block.get("name", ""),
|
||||
arguments=_to_dd_arguments(block.get("input") or {}),
|
||||
tool_id=block.get("id", ""),
|
||||
type="tool_use",
|
||||
)
|
||||
for block in _content_blocks(message)
|
||||
if block.get("type") == "tool_use"
|
||||
)
|
||||
return openai_calls + anthropic_calls
|
||||
|
||||
|
||||
def _to_dd_tool_results(message: Mapping[str, Any], tool_call_names: Mapping[str, str]) -> tuple[ToolResult, ...]:
|
||||
"""
|
||||
The tool results a message carries, linked back to the call each answers.
|
||||
|
||||
OpenAI models a result as a whole `role: "tool"` message keyed by `tool_call_id`;
|
||||
Anthropic nests `tool_result` blocks inside a user message, keyed by `tool_use_id`.
|
||||
"""
|
||||
|
||||
def to_result(tool_id: str, result: object) -> ToolResult:
|
||||
return ToolResult(
|
||||
name=tool_call_names.get(tool_id, ""),
|
||||
result=result if isinstance(result, str) else safe_dumps(result),
|
||||
tool_id=tool_id,
|
||||
type="function",
|
||||
)
|
||||
|
||||
if message.get("role") == "tool":
|
||||
return (to_result(str(message.get("tool_call_id", "")), message.get("content") or ""),)
|
||||
return tuple(
|
||||
to_result(str(block.get("tool_use_id", "")), block.get("content") or "")
|
||||
for block in _content_blocks(message)
|
||||
if block.get("type") == "tool_result"
|
||||
)
|
||||
|
||||
|
||||
def _tool_call_names_by_id(messages: Sequence[object]) -> Mapping[str, str]:
|
||||
"""Ids to tool names for result linking; reads names structurally and parses nothing."""
|
||||
openai_pairs: Final = tuple(
|
||||
(tool_call.get("id"), function.get("name", ""))
|
||||
for message in messages
|
||||
if isinstance(message, dict) and isinstance(message.get("tool_calls"), list)
|
||||
for tool_call in message["tool_calls"]
|
||||
if isinstance(tool_call, dict)
|
||||
for function in [_mapping_field(tool_call, "function")]
|
||||
)
|
||||
anthropic_pairs: Final = tuple(
|
||||
(block.get("id"), block.get("name", ""))
|
||||
for message in messages
|
||||
if isinstance(message, dict)
|
||||
for block in _content_blocks(message)
|
||||
if block.get("type") == "tool_use"
|
||||
)
|
||||
return MappingProxyType({str(tool_id): str(name) for tool_id, name in openai_pairs + anthropic_pairs if tool_id})
|
||||
|
||||
|
||||
def _to_dd_message(message: object, tool_call_names: Mapping[str, str]) -> Message:
|
||||
"""
|
||||
Map one chat message onto LLM Obs' Message schema, adding fields and never destroying content.
|
||||
|
||||
Content collapses to its text only when it has text; a content list with none (tool blocks,
|
||||
images) rides along unchanged so nothing the caller logged is lost. Tool calls and results
|
||||
move into the fields the LLM Obs Tools panel reads, from both the OpenAI and Anthropic shapes.
|
||||
"""
|
||||
if not isinstance(message, dict):
|
||||
converted: Final = handle_any_messages_to_chat_completion_str_messages_conversion(message)
|
||||
return converted[0] if converted else _EMPTY_MESSAGE
|
||||
|
||||
text: Final = convert_content_list_to_str(message) # pyright: ignore[reportArgumentType] # caller-supplied dict
|
||||
original_content: Final = message.get("content")
|
||||
content: Final = (
|
||||
text if text or not isinstance(original_content, list) or not original_content else original_content
|
||||
)
|
||||
reasoning: Final = message.get("reasoning_content")
|
||||
tool_calls: Final = _to_dd_tool_calls(message)
|
||||
tool_results: Final = _to_dd_tool_results(message, tool_call_names)
|
||||
dd_message: Final[Message] = {
|
||||
"role": message.get("role", ""),
|
||||
"content": content,
|
||||
**({"reasoning_content": reasoning} if reasoning is not None else {}),
|
||||
**({"tool_calls": tool_calls} if tool_calls else {}),
|
||||
**({"tool_results": tool_results} if tool_results else {}),
|
||||
}
|
||||
return dd_message
|
||||
|
||||
|
||||
def _to_dd_messages(messages: object) -> tuple[Message, ...]:
|
||||
"""Map a whole conversation, resolving each tool result against the calls that precede it."""
|
||||
if messages is None:
|
||||
return ()
|
||||
if not isinstance(messages, list):
|
||||
return tuple(handle_any_messages_to_chat_completion_str_messages_conversion(messages))
|
||||
tool_call_names: Final = _tool_call_names_by_id(messages)
|
||||
return tuple(_to_dd_message(message, tool_call_names) for message in messages)
|
||||
|
||||
|
||||
def _to_dd_tool_definition(entry: Mapping[str, Any]) -> ToolDefinition | None:
|
||||
function: Final = entry.get("function")
|
||||
declared: Final[Mapping[str, Any]] = function if isinstance(function, dict) else entry
|
||||
name: Final = declared.get("name")
|
||||
if not name:
|
||||
return None
|
||||
schema: Final = declared.get("parameters") or declared.get("input_schema")
|
||||
description: Final = declared.get("description", "")
|
||||
if not isinstance(schema, dict):
|
||||
return ToolDefinition(name=name, description=description)
|
||||
return ToolDefinition(name=name, description=description, schema=schema)
|
||||
|
||||
|
||||
def _to_dd_tool_definitions(model_parameters: object) -> tuple[ToolDefinition, ...]:
|
||||
"""
|
||||
Map the request's declared tools onto LLM Obs' ToolDefinition schema.
|
||||
|
||||
Handles the wrapped chat-completions shape and the bare shape the Anthropic and
|
||||
Responses surfaces use, since both reach this logger through `model_parameters`.
|
||||
"""
|
||||
if not isinstance(model_parameters, dict):
|
||||
return ()
|
||||
raw_tools: Final = model_parameters.get("tools") or model_parameters.get("functions")
|
||||
if not isinstance(raw_tools, list):
|
||||
return ()
|
||||
return tuple(
|
||||
definition
|
||||
for entry in raw_tools
|
||||
if isinstance(entry, dict)
|
||||
if (definition := _to_dd_tool_definition(entry)) is not None
|
||||
)
|
||||
|
||||
|
||||
class DataDogLLMObsLogger(CustomBatchLogger):
|
||||
def __init__(self, **kwargs):
|
||||
|
|
@ -222,12 +410,9 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
if standard_logging_payload is None:
|
||||
raise Exception("DataDogLLMObs: standard_logging_object is not set")
|
||||
|
||||
messages = standard_logging_payload["messages"]
|
||||
messages = self._ensure_string_content(messages=messages)
|
||||
|
||||
metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {})
|
||||
|
||||
input_meta: Final = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages))
|
||||
input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"]))
|
||||
output_meta: Final = OutputMeta(
|
||||
messages=self._get_response_messages(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
|
|
@ -241,22 +426,20 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
if isinstance(metadata, dict):
|
||||
metadata_parent_id = metadata.get("parent_id")
|
||||
|
||||
meta: Final = Meta(
|
||||
kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id),
|
||||
input=input_meta,
|
||||
output=output_meta,
|
||||
metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload),
|
||||
error=error_info,
|
||||
)
|
||||
tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters"))
|
||||
span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id)
|
||||
payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload)
|
||||
|
||||
# Calculate metrics (you may need to adjust these based on available data)
|
||||
metrics: Final = LLMMetrics(
|
||||
input_tokens=float(standard_logging_payload.get("prompt_tokens", 0)),
|
||||
output_tokens=float(standard_logging_payload.get("completion_tokens", 0)),
|
||||
total_tokens=float(standard_logging_payload.get("total_tokens", 0)),
|
||||
total_cost=float(standard_logging_payload.get("response_cost", 0)),
|
||||
time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload),
|
||||
)
|
||||
meta: Final[Meta] = {
|
||||
"kind": span_kind,
|
||||
"input": input_meta,
|
||||
"output": output_meta,
|
||||
"metadata": payload_metadata,
|
||||
"error": error_info,
|
||||
**({"tool_definitions": tool_definitions} if tool_definitions else {}),
|
||||
}
|
||||
|
||||
metrics: Final = self._assemble_metrics(standard_logging_payload)
|
||||
|
||||
payload: Final[LLMObsPayload] = LLMObsPayload(
|
||||
parent_id=metadata_parent_id if metadata_parent_id else "undefined",
|
||||
|
|
@ -314,6 +497,45 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
)
|
||||
return error_info
|
||||
|
||||
def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics:
|
||||
"""
|
||||
Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from.
|
||||
|
||||
Cache counts resolve through the same owners the savings dashboard uses, so every provider
|
||||
spelling is covered, and `non_cached_input_tokens` subtracts BOTH cache categories because
|
||||
litellm's normalized prompt count includes both (the invariant the cost calculator's custom
|
||||
pricing helper documents). A zero residual on a fully cached request is real data and is
|
||||
emitted; a zero read or write count is absence and is not.
|
||||
"""
|
||||
prompt_tokens: Final = float(standard_logging_payload.get("prompt_tokens", 0))
|
||||
completion_tokens: Final = float(standard_logging_payload.get("completion_tokens", 0))
|
||||
total_tokens: Final = float(standard_logging_payload.get("total_tokens", 0))
|
||||
total_cost: Final = float(standard_logging_payload.get("response_cost", 0))
|
||||
time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload)
|
||||
|
||||
raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object")
|
||||
usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None
|
||||
cache_read: Final = float(extract_cache_read_tokens(usage_object))
|
||||
cache_write: Final = float(extract_cache_creation_tokens(usage_object))
|
||||
|
||||
metrics: Final[LLMMetrics] = {
|
||||
"input_tokens": prompt_tokens,
|
||||
"output_tokens": completion_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
"total_cost": total_cost,
|
||||
"time_to_first_token": time_to_first_token,
|
||||
**(
|
||||
{
|
||||
**({"cache_read_input_tokens": cache_read} if cache_read else {}),
|
||||
**({"cache_write_input_tokens": cache_write} if cache_write else {}),
|
||||
"non_cached_input_tokens": max(prompt_tokens - cache_read - cache_write, 0.0),
|
||||
}
|
||||
if cache_read or cache_write
|
||||
else {}
|
||||
),
|
||||
}
|
||||
return metrics
|
||||
|
||||
def _get_time_to_first_token_seconds(self, standard_logging_payload: StandardLoggingPayload) -> float:
|
||||
"""
|
||||
Get the time to first token in seconds
|
||||
|
|
@ -335,7 +557,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
|
||||
def _get_response_messages(
|
||||
self, standard_logging_payload: StandardLoggingPayload, call_type: str | None
|
||||
) -> list[object]:
|
||||
) -> tuple[Message, ...]:
|
||||
"""
|
||||
Get the messages from the response object
|
||||
|
||||
|
|
@ -344,7 +566,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
|
||||
response_obj = standard_logging_payload.get("response")
|
||||
if response_obj is None:
|
||||
return []
|
||||
return ()
|
||||
|
||||
# edge case: handle response_obj is a string representation of a dict
|
||||
if isinstance(response_obj, str):
|
||||
|
|
@ -357,7 +579,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
# fallback to json parsing
|
||||
response_obj = json.loads(str(response_obj))
|
||||
except json.JSONDecodeError:
|
||||
return []
|
||||
return ()
|
||||
|
||||
if call_type in [
|
||||
CallTypes.completion.value,
|
||||
|
|
@ -375,12 +597,12 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
if isinstance(response_obj, dict) and "choices" in response_obj:
|
||||
choices: Final = response_obj["choices"]
|
||||
if choices and len(choices) > 0 and "message" in choices[0]:
|
||||
return [choices[0]["message"]]
|
||||
return []
|
||||
return _to_dd_messages([choices[0]["message"]])
|
||||
return ()
|
||||
except (KeyError, IndexError, TypeError):
|
||||
# In case of any error accessing the response structure, return empty list
|
||||
return []
|
||||
return []
|
||||
return ()
|
||||
return ()
|
||||
|
||||
def _get_datadog_span_kind(
|
||||
self, call_type: str | None, parent_id: str | None = None
|
||||
|
|
@ -485,17 +707,6 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
# Default fallback for unknown or passthrough operations
|
||||
return "llm"
|
||||
|
||||
def _ensure_string_content(self, messages: str | Sequence[object] | Mapping[object, object] | None) -> list[object]:
|
||||
if messages is None:
|
||||
return []
|
||||
if isinstance(messages, str):
|
||||
return [messages]
|
||||
elif isinstance(messages, list):
|
||||
return [message for message in messages]
|
||||
elif isinstance(messages, dict):
|
||||
return [str(messages.get("content", ""))]
|
||||
return []
|
||||
|
||||
def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
|
||||
"""
|
||||
Fields to track in DD LLM Observability metadata from litellm standard logging payload
|
||||
|
|
@ -524,10 +735,6 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
spend_metrics: Final = self._get_spend_metrics(standard_logging_payload)
|
||||
_metadata.update({"spend_metrics": dict(spend_metrics)})
|
||||
|
||||
## extract tool calls and add to metadata
|
||||
tool_call_metadata: Final = self._extract_tool_call_metadata(standard_logging_payload)
|
||||
_metadata.update(tool_call_metadata)
|
||||
|
||||
_standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {}
|
||||
_metadata.update(_standard_logging_metadata)
|
||||
return _metadata
|
||||
|
|
@ -647,107 +854,3 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
verbose_logger.debug("Original value: %s", user_api_key_budget_reset_at)
|
||||
|
||||
return spend_metrics
|
||||
|
||||
def _process_input_messages_preserving_tool_calls(self, messages: Sequence[object]) -> list[dict[str, object]]:
|
||||
"""
|
||||
Process input messages while preserving tool_calls and tool message types.
|
||||
|
||||
This bypasses the lossy string conversion when tool calls are present,
|
||||
allowing complex nested tool_calls objects to be preserved for Datadog.
|
||||
"""
|
||||
processed: Final = []
|
||||
for msg in messages:
|
||||
if isinstance(msg, dict):
|
||||
# Preserve messages with tool_calls or tool role as-is
|
||||
if "tool_calls" in msg or msg.get("role") == "tool":
|
||||
processed.append(msg)
|
||||
else:
|
||||
# For regular messages, still apply string conversion
|
||||
converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg])
|
||||
processed.extend(converted)
|
||||
else:
|
||||
# For non-dict messages, apply string conversion
|
||||
converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg])
|
||||
processed.extend(converted)
|
||||
return processed
|
||||
|
||||
@staticmethod
|
||||
def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, object]:
|
||||
"""
|
||||
Extract tool call information into key-value pairs for Datadog metadata.
|
||||
|
||||
Similar to OpenTelemetry's implementation but adapted for Datadog's format.
|
||||
"""
|
||||
kv_pairs: Final[dict[str, object]] = {}
|
||||
for idx, tool_call in enumerate(tool_calls):
|
||||
try:
|
||||
# Extract tool call ID
|
||||
tool_id = tool_call.get("id")
|
||||
if tool_id:
|
||||
kv_pairs[f"tool_calls.{idx}.id"] = tool_id
|
||||
|
||||
# Extract tool call type
|
||||
tool_type = tool_call.get("type")
|
||||
if tool_type:
|
||||
kv_pairs[f"tool_calls.{idx}.type"] = tool_type
|
||||
|
||||
# Extract function information
|
||||
function = tool_call.get("function")
|
||||
if function:
|
||||
function_name = function.get("name")
|
||||
if function_name:
|
||||
kv_pairs[f"tool_calls.{idx}.function.name"] = function_name
|
||||
|
||||
function_arguments = function.get("arguments")
|
||||
if function_arguments:
|
||||
# Store arguments as JSON string for Datadog
|
||||
if isinstance(function_arguments, str):
|
||||
kv_pairs[f"tool_calls.{idx}.function.arguments"] = function_arguments
|
||||
else:
|
||||
import json
|
||||
|
||||
kv_pairs[f"tool_calls.{idx}.function.arguments"] = json.dumps(function_arguments)
|
||||
except (KeyError, TypeError, ValueError) as e:
|
||||
verbose_logger.debug("DataDogLLMObs: Error processing tool call %s: %s", idx, e)
|
||||
continue
|
||||
|
||||
return kv_pairs
|
||||
|
||||
def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
|
||||
"""
|
||||
Extract tool call information from both input messages and response for Datadog metadata.
|
||||
"""
|
||||
tool_call_metadata: Final[dict[str, object]] = {}
|
||||
|
||||
try:
|
||||
# Extract tool calls from input messages
|
||||
messages: Final = standard_logging_payload.get("messages", [])
|
||||
if messages and isinstance(messages, list):
|
||||
for message in messages:
|
||||
if isinstance(message, dict) and "tool_calls" in message:
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
input_tool_calls_kv = self._tool_calls_kv_pair(tool_calls)
|
||||
# Prefix with "input_" to distinguish from response tool calls
|
||||
for key, value in input_tool_calls_kv.items():
|
||||
tool_call_metadata[f"input_{key}"] = value
|
||||
|
||||
# Extract tool calls from response
|
||||
response_obj: Final = standard_logging_payload.get("response")
|
||||
if response_obj and isinstance(response_obj, dict):
|
||||
choices: Final = response_obj.get("choices", [])
|
||||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
message = choice.get("message")
|
||||
if message and isinstance(message, dict):
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
response_tool_calls_kv = self._tool_calls_kv_pair(tool_calls)
|
||||
# Prefix with "output_" to distinguish from input tool calls
|
||||
for key, value in response_tool_calls_kv.items():
|
||||
tool_call_metadata[f"output_{key}"] = value
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.debug("DataDogLLMObs: Error extracting tool call metadata: %s", e)
|
||||
|
||||
return tool_call_metadata
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import json
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, cast
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
|
|
@ -62,6 +63,31 @@ if TYPE_CHECKING:
|
|||
# --- typed sub-structures ---------------------------------------------------- #
|
||||
|
||||
|
||||
def _cache_token_value(*values: object) -> int | None:
|
||||
explicit_zero = False
|
||||
invalid_before_zero = False
|
||||
for raw_value in values:
|
||||
if raw_value is None:
|
||||
continue
|
||||
if isinstance(raw_value, bool):
|
||||
parsed = None
|
||||
else:
|
||||
try:
|
||||
parsed = as_int(raw_value)
|
||||
except (OverflowError, ValueError):
|
||||
parsed = None
|
||||
if parsed is None:
|
||||
if not explicit_zero:
|
||||
invalid_before_zero = True
|
||||
elif parsed > 0:
|
||||
return parsed
|
||||
elif parsed == 0:
|
||||
explicit_zero = True
|
||||
elif not explicit_zero:
|
||||
invalid_before_zero = True
|
||||
return 0 if explicit_zero and not invalid_before_zero else None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LLMRequestParams:
|
||||
temperature: float | None = None
|
||||
|
|
@ -104,12 +130,25 @@ class LLMUsage:
|
|||
metadata: Final[Mapping[str, object]] = payload.get("metadata") or {}
|
||||
raw_usage: Final = metadata.get("usage_object")
|
||||
usage_object: Final[Mapping[str, object]] = raw_usage if isinstance(raw_usage, Mapping) else {}
|
||||
raw_details: Final = usage_object.get("prompt_tokens_details")
|
||||
prompt_details: Final[Mapping[str, object]] = (
|
||||
raw_details if isinstance(raw_details, Mapping) else MappingProxyType({})
|
||||
)
|
||||
return cls(
|
||||
input_tokens=as_int(payload.get("prompt_tokens")),
|
||||
output_tokens=as_int(payload.get("completion_tokens")),
|
||||
total_tokens=as_int(payload.get("total_tokens")),
|
||||
cache_creation_input_tokens=as_int(usage_object.get("cache_creation_input_tokens")),
|
||||
cache_read_input_tokens=as_int(usage_object.get("cache_read_input_tokens")),
|
||||
cache_creation_input_tokens=_cache_token_value(
|
||||
usage_object.get("cache_creation_input_tokens"),
|
||||
prompt_details.get("cache_write_tokens"),
|
||||
prompt_details.get("cache_creation_tokens"),
|
||||
prompt_details.get("cache_creation_input_tokens"),
|
||||
),
|
||||
cache_read_input_tokens=_cache_token_value(
|
||||
usage_object.get("cache_read_input_tokens"),
|
||||
prompt_details.get("cached_tokens"),
|
||||
usage_object.get("prompt_cache_hit_tokens"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import math
|
|||
import os
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import replace
|
||||
from datetime import datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast
|
||||
|
||||
|
|
@ -58,7 +59,10 @@ from litellm.types.utils import (
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from prometheus_client import Gauge
|
||||
from prometheus_client.metrics import MetricWrapperBase
|
||||
|
||||
from litellm.router import Router
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
|
|
@ -67,6 +71,8 @@ _TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
|
|||
|
||||
_DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0
|
||||
|
||||
UNRECOGNIZED_REQUESTED_MODEL_LABEL: Final = "other"
|
||||
|
||||
_NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
|
||||
(
|
||||
"guardrail_name",
|
||||
|
|
@ -154,6 +160,44 @@ def _get_budget_metrics_per_request_timeout() -> float:
|
|||
return parsed
|
||||
|
||||
|
||||
def _get_proxy_llm_router() -> Router | None:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
except Exception:
|
||||
return None
|
||||
return llm_router
|
||||
|
||||
|
||||
def _bounded_requested_model_label(requested_model: str | None, router_originated: bool = False) -> str | None:
|
||||
"""
|
||||
Bound ``requested_model`` label cardinality: names the router recognizes
|
||||
(model names, deployment ids, aliases, routing groups, team public model
|
||||
names) or matches via a global or team wildcard/pattern route keep their
|
||||
own label value; any other client-supplied string collapses into the
|
||||
single ``other`` bucket. With no proxy router to vouch for the string,
|
||||
client-supplied values collapse to ``other`` while ``router_originated``
|
||||
values (emitted by an SDK ``Router``'s own deployment failure and
|
||||
fallback events, where the proxy router never exists) pass through.
|
||||
"""
|
||||
if not requested_model:
|
||||
return requested_model
|
||||
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):
|
||||
return requested_model
|
||||
if requested_model in llm_router.team_public_model_names:
|
||||
return requested_model
|
||||
if llm_router.pattern_router.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()
|
||||
):
|
||||
return requested_model
|
||||
return UNRECOGNIZED_REQUESTED_MODEL_LABEL
|
||||
|
||||
|
||||
class PrometheusLogger(CustomLogger):
|
||||
# Class variables or attributes
|
||||
|
||||
|
|
@ -434,6 +478,30 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"),
|
||||
)
|
||||
|
||||
self.litellm_api_key_rate_limit_allowed_metric = self._gauge_factory(
|
||||
"litellm_api_key_rate_limit_allowed_metric",
|
||||
"Configured rate limit for the API Key in the current window (rpm_limit / tpm_limit), by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_allowed_metric"),
|
||||
)
|
||||
|
||||
self.litellm_api_key_rate_limit_used_metric = self._gauge_factory(
|
||||
"litellm_api_key_rate_limit_used_metric",
|
||||
"Requests or tokens the API Key has consumed in the current rate limit window, by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_used_metric"),
|
||||
)
|
||||
|
||||
self.litellm_team_rate_limit_allowed_metric = self._gauge_factory(
|
||||
"litellm_team_rate_limit_allowed_metric",
|
||||
"Configured rate limit for the Team in the current window (team rpm_limit / tpm_limit), by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_team_rate_limit_allowed_metric"),
|
||||
)
|
||||
|
||||
self.litellm_team_rate_limit_used_metric = self._gauge_factory(
|
||||
"litellm_team_rate_limit_used_metric",
|
||||
"Requests or tokens the Team has consumed in the current rate limit window, by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_team_rate_limit_used_metric"),
|
||||
)
|
||||
|
||||
########################################
|
||||
# LLM API Deployment Metrics / analytics
|
||||
########################################
|
||||
|
|
@ -1433,6 +1501,11 @@ class PrometheusLogger(CustomLogger):
|
|||
model_id=enum_values.model_id,
|
||||
)
|
||||
|
||||
self._set_key_and_team_rate_limit_metrics(
|
||||
standard_logging_payload=standard_logging_payload, # pyright: ignore[reportArgumentType] # isinstance(dict) above narrows the TypedDict to dict[Unknown, Unknown]
|
||||
enum_values=enum_values,
|
||||
)
|
||||
|
||||
# set latency metrics
|
||||
self._set_latency_metrics(
|
||||
kwargs=kwargs,
|
||||
|
|
@ -1960,17 +2033,102 @@ class PrometheusLogger(CustomLogger):
|
|||
"""
|
||||
if standard_logging_payload is None:
|
||||
return None
|
||||
return PrometheusLogger._get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
header_name=f"x-ratelimit-model_per_key-remaining-{rate_limit_type}",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
header_name: str,
|
||||
) -> int | None:
|
||||
hidden_params: Final = standard_logging_payload.get("hidden_params")
|
||||
if hidden_params is None:
|
||||
return None
|
||||
additional_headers: Final = hidden_params.get("additional_headers")
|
||||
additional_headers: Final[Mapping[str, object] | None] = hidden_params.get("additional_headers")
|
||||
if additional_headers is None:
|
||||
return None
|
||||
value: Final = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}")
|
||||
value: Final = additional_headers.get(header_name)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return None
|
||||
return value
|
||||
|
||||
def _set_key_and_team_rate_limit_metrics(
|
||||
self,
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
) -> None:
|
||||
"""
|
||||
Export the key-level and team-level RPM / TPM limit and current window
|
||||
usage from the ``x-ratelimit-{api_key,team}-{limit,remaining}-*``
|
||||
headers the v3 rate limiter mirrors into the logging payload. The
|
||||
limiter already read these counters (from Redis when configured) on
|
||||
the request path, so no extra store lookup happens here. Descriptors
|
||||
without a configured limit emit no header, so their series is removed
|
||||
rather than left at the value from before the limit was dropped.
|
||||
"""
|
||||
descriptor_gauges: Final[
|
||||
tuple[tuple[Literal["api_key", "team"], DEFINED_PROMETHEUS_METRICS, Gauge, Gauge], ...]
|
||||
] = (
|
||||
(
|
||||
"api_key",
|
||||
"litellm_api_key_rate_limit_allowed_metric",
|
||||
self.litellm_api_key_rate_limit_allowed_metric,
|
||||
self.litellm_api_key_rate_limit_used_metric,
|
||||
),
|
||||
(
|
||||
"team",
|
||||
"litellm_team_rate_limit_allowed_metric",
|
||||
self.litellm_team_rate_limit_allowed_metric,
|
||||
self.litellm_team_rate_limit_used_metric,
|
||||
),
|
||||
)
|
||||
for descriptor_key, metric_name, allowed_gauge, used_gauge in descriptor_gauges:
|
||||
for rate_limit_type in ("requests", "tokens"):
|
||||
self._set_rate_limit_allowed_and_used_gauges(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
enum_values=enum_values,
|
||||
descriptor_key=descriptor_key,
|
||||
metric_name=metric_name,
|
||||
allowed_gauge=allowed_gauge,
|
||||
used_gauge=used_gauge,
|
||||
rate_limit_type=rate_limit_type,
|
||||
)
|
||||
|
||||
def _set_rate_limit_allowed_and_used_gauges(
|
||||
self,
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
descriptor_key: Literal["api_key", "team"],
|
||||
metric_name: DEFINED_PROMETHEUS_METRICS,
|
||||
allowed_gauge: Gauge,
|
||||
used_gauge: Gauge,
|
||||
rate_limit_type: Literal["requests", "tokens"],
|
||||
) -> None:
|
||||
limit: Final = self._get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
header_name=f"x-ratelimit-{descriptor_key}-limit-{rate_limit_type}",
|
||||
)
|
||||
remaining: Final = self._get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
header_name=f"x-ratelimit-{descriptor_key}-remaining-{rate_limit_type}",
|
||||
)
|
||||
labelled_values: Final = replace(enum_values, rate_limit_type=rate_limit_type)
|
||||
labelnames: Final = self.get_labels_for_metric(metric_name)
|
||||
labels: Final = prometheus_label_factory(
|
||||
supported_enum_labels=labelnames,
|
||||
enum_values=labelled_values,
|
||||
label_context=PrometheusLabelFactoryContext(labelled_values),
|
||||
)
|
||||
if limit is None or remaining is None:
|
||||
label_values: Final = tuple(labels.get(label) for label in labelnames)
|
||||
self._bounded_prometheus_series_tracker.remove_series(allowed_gauge, label_values)
|
||||
self._bounded_prometheus_series_tracker.remove_series(used_gauge, label_values)
|
||||
return
|
||||
allowed_gauge.labels(**labels).set(limit)
|
||||
used_gauge.labels(**labels).set(limit - remaining)
|
||||
|
||||
def _set_virtual_key_rate_limit_metrics(
|
||||
self,
|
||||
user_api_key: str | None,
|
||||
|
|
@ -2407,7 +2565,7 @@ class PrometheusLogger(CustomLogger):
|
|||
team_alias=user_api_key_dict.team_alias,
|
||||
org_id=user_api_key_dict.org_id,
|
||||
org_alias=user_api_key_dict.organization_alias,
|
||||
requested_model=request_data.get("model", ""),
|
||||
requested_model=_bounded_requested_model_label(request_data.get("model", "")),
|
||||
status_code=str(status_code),
|
||||
exception_status=str(status_code),
|
||||
exception_class=self._get_exception_class_name(original_exception),
|
||||
|
|
@ -2627,7 +2785,9 @@ class PrometheusLogger(CustomLogger):
|
|||
label_model_id = ""
|
||||
label_api_base = ""
|
||||
label_api_provider = ""
|
||||
label_requested_model = litellm_model_name or model_group or ""
|
||||
label_requested_model = (
|
||||
_bounded_requested_model_label(litellm_model_name or model_group, router_originated=True) or ""
|
||||
)
|
||||
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
litellm_model_name=label_litellm_model_name,
|
||||
|
|
@ -3186,7 +3346,7 @@ class PrometheusLogger(CustomLogger):
|
|||
_tags: Final = cast(list[str], kwargs.get("tags") or [])
|
||||
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
requested_model=original_model_group,
|
||||
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
|
||||
fallback_model=_new_model,
|
||||
hashed_api_key=standard_metadata["user_api_key_hash"],
|
||||
api_key_alias=standard_metadata["user_api_key_alias"],
|
||||
|
|
@ -3227,7 +3387,7 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
requested_model=original_model_group,
|
||||
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
|
||||
fallback_model=_new_model,
|
||||
hashed_api_key=standard_metadata["user_api_key_hash"],
|
||||
api_key_alias=standard_metadata["user_api_key_alias"],
|
||||
|
|
|
|||
|
|
@ -60,6 +60,10 @@ class BoundedPrometheusSeriesTracker:
|
|||
break
|
||||
del series[tracked_label_values]
|
||||
|
||||
def remove_series(self, metric: object, label_values: tuple[str | None, ...]) -> bool:
|
||||
"""Drop one child series, True when it is gone (removed or never existed)."""
|
||||
return self._remove_metric_child(metric, label_values)
|
||||
|
||||
def _should_run_ttl_cleanup(
|
||||
self,
|
||||
metric_name: str,
|
||||
|
|
|
|||
|
|
@ -1,11 +1,18 @@
|
|||
#### What this does ####
|
||||
# On success + failure, log events to Supabase
|
||||
|
||||
import hashlib
|
||||
from datetime import datetime
|
||||
from typing import Final, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import (
|
||||
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES,
|
||||
MAX_S3_OBJECT_KEY_BYTES,
|
||||
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES,
|
||||
S3_PREFIX_DIGEST_CHARS,
|
||||
)
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
|
|
@ -133,9 +140,7 @@ class S3Logger:
|
|||
s3_file_name,
|
||||
)
|
||||
|
||||
s3_object_download_filename: Final = (
|
||||
"time-" + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") + "_" + payload["id"] + ".json"
|
||||
)
|
||||
s3_object_download_filename: Final = get_s3_object_download_filename(start_time, payload["id"])
|
||||
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
|
|
@ -198,6 +203,47 @@ def resolve_sse_params(
|
|||
return algorithm, valid_key_id
|
||||
|
||||
|
||||
S3_MIN_BOUNDED_FILE_NAME_BYTES: Final = 64
|
||||
|
||||
|
||||
def _truncate_to_utf8_bytes(value: str, max_bytes: int) -> str:
|
||||
"""Trim `value` so its UTF-8 encoding fits `max_bytes`, never splitting a character."""
|
||||
if max_bytes <= 0:
|
||||
return ""
|
||||
encoded: Final = value.encode("utf-8")
|
||||
if len(encoded) <= max_bytes:
|
||||
return value
|
||||
return encoded[:max_bytes].decode("utf-8", errors="ignore")
|
||||
|
||||
|
||||
def get_s3_object_download_filename(start_time: datetime, response_id: str) -> str:
|
||||
"""Content-Disposition filename for the uploaded object, bounded to the metadata header cap."""
|
||||
sanitized_response_id: Final = response_id.replace("/", "_").replace('"', "_")
|
||||
file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{response_id}"
|
||||
sanitized_file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{sanitized_response_id}"
|
||||
budget: Final = MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES - len(b".json")
|
||||
if len(sanitized_file_name.encode("utf-8")) <= budget:
|
||||
return sanitized_file_name + ".json"
|
||||
return _bounded_s3_file_name(file_name, sanitized_file_name, budget) + ".json"
|
||||
|
||||
|
||||
def _bounded_s3_file_name(s3_file_name: str, sanitized_s3_file_name: str, max_bytes: int) -> str:
|
||||
"""As much of the file name as `max_bytes` allows, then the sha256 of the whole name."""
|
||||
digest: Final = hashlib.sha256(s3_file_name.encode("utf-8")).hexdigest()
|
||||
head_budget: Final = min(S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, max_bytes - len(digest) - 1)
|
||||
head: Final = _truncate_to_utf8_bytes(sanitized_s3_file_name, head_budget)
|
||||
return f"{head}_{digest}" if head else digest
|
||||
|
||||
|
||||
def _bounded_s3_prefix(configured_prefix: str, max_bytes: int) -> str:
|
||||
"""As much of the configured prefix as fits, then a digest segment naming the full prefix."""
|
||||
digest_segment: Final = hashlib.sha256(configured_prefix.encode("utf-8")).hexdigest()[:S3_PREFIX_DIGEST_CHARS] + "/"
|
||||
if max_bytes < len(digest_segment):
|
||||
return ""
|
||||
head: Final = _truncate_to_utf8_bytes(configured_prefix, max_bytes - len(digest_segment) - 1).rstrip("/")
|
||||
return f"{head}/{digest_segment}" if head else digest_segment
|
||||
|
||||
|
||||
def get_s3_object_key(
|
||||
s3_path: str,
|
||||
prefix: str,
|
||||
|
|
@ -205,12 +251,23 @@ def get_s3_object_key(
|
|||
s3_file_name: str,
|
||||
) -> str:
|
||||
sanitized_s3_file_name: Final = s3_file_name.replace("/", "_")
|
||||
s3_object_key = (
|
||||
(s3_path.rstrip("/") + "/" if s3_path else "")
|
||||
+ prefix
|
||||
+ start_time.strftime("%Y-%m-%d")
|
||||
+ "/"
|
||||
+ sanitized_s3_file_name
|
||||
) # we need the s3 key to include the time, so we log cache hits too
|
||||
s3_object_key += ".json"
|
||||
return s3_object_key
|
||||
configured_prefix: Final = (s3_path.rstrip("/") + "/" if s3_path else "") + prefix
|
||||
date_segment: Final = start_time.strftime("%Y-%m-%d") + "/"
|
||||
# we need the s3 key to include the time, so we log cache hits too
|
||||
s3_object_key: Final = configured_prefix + date_segment + sanitized_s3_file_name + ".json"
|
||||
if len(s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES:
|
||||
return s3_object_key
|
||||
|
||||
# shorten the response id first and only trim the configured prefix if that is what does not
|
||||
# fit, so prefix scoped IAM policies and lifecycle rules keep matching
|
||||
budget: Final = MAX_S3_OBJECT_KEY_BYTES - len(date_segment.encode("utf-8")) - len(b".json")
|
||||
prefix_bytes: Final = len(configured_prefix.encode("utf-8"))
|
||||
if prefix_bytes + S3_MIN_BOUNDED_FILE_NAME_BYTES <= budget:
|
||||
bounded_file_name: Final = _bounded_s3_file_name(s3_file_name, sanitized_s3_file_name, budget - prefix_bytes)
|
||||
return configured_prefix + date_segment + bounded_file_name + ".json"
|
||||
|
||||
shortest_file_name: Final = _bounded_s3_file_name(
|
||||
s3_file_name, sanitized_s3_file_name, S3_MIN_BOUNDED_FILE_NAME_BYTES
|
||||
)
|
||||
bounded_prefix: Final = _bounded_s3_prefix(configured_prefix, budget - len(shortest_file_name.encode("utf-8")))
|
||||
return bounded_prefix + date_segment + shortest_file_name + ".json"
|
||||
|
|
|
|||
|
|
@ -16,7 +16,11 @@ from urllib.parse import quote
|
|||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
|
||||
from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params
|
||||
from litellm.integrations.s3 import (
|
||||
get_s3_object_download_filename,
|
||||
get_s3_object_key,
|
||||
resolve_sse_params,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
|
|
@ -259,11 +263,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
now: Final = datetime.now(timezone.utc)
|
||||
audit_log_id: Final = audit_log.get("id", "unknown")
|
||||
|
||||
s3_path = cast(str | None, self.s3_path) or ""
|
||||
s3_path = s3_path.rstrip("/") + "/" if s3_path else ""
|
||||
|
||||
s3_object_key: Final = (
|
||||
f"{s3_path}audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log_id}.json"
|
||||
s3_object_key: Final = get_s3_object_key(
|
||||
cast(str | None, self.s3_path) or "",
|
||||
"audit_logs/",
|
||||
now,
|
||||
f"{now.strftime('%H-%M-%S')}_{audit_log_id}",
|
||||
)
|
||||
|
||||
element: Final = s3BatchLoggingElement(
|
||||
|
|
@ -463,9 +467,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
)
|
||||
verbose_logger.debug("s3_object_key=%s", s3_object_key)
|
||||
|
||||
s3_object_download_filename: Final = (
|
||||
f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json"
|
||||
)
|
||||
s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"])
|
||||
|
||||
return s3BatchLoggingElement(
|
||||
payload=dict(standard_logging_payload),
|
||||
|
|
|
|||
|
|
@ -127,7 +127,7 @@ def handle_anthropic_text_model_custom_llm_provider(
|
|||
return model, custom_llm_provider
|
||||
|
||||
|
||||
def declared_authenticating_provider(model: str, custom_llm_provider: str | None = None) -> str | None:
|
||||
def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
|
||||
"""The authenticating provider this pair already names, or None.
|
||||
|
||||
get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, because their
|
||||
|
|
@ -135,7 +135,7 @@ def declared_authenticating_provider(model: str, custom_llm_provider: str | None
|
|||
and for a declared pair the resolver's answer is the declaration itself, so metadata callers
|
||||
adopt the declaration instead of resolving.
|
||||
"""
|
||||
declared: Final = custom_llm_provider or model.split("/", 1)[0]
|
||||
declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
|
||||
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import asyncio
|
|||
import json
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -154,18 +155,6 @@ class GetModelCostMap:
|
|||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def fetch_remote_model_cost_map(url: str, timeout: int = 5) -> dict:
|
||||
"""
|
||||
Fetch the model cost map from a remote URL.
|
||||
|
||||
Returns the parsed JSON dict. Raises on network/parse errors
|
||||
(caller is expected to handle).
|
||||
"""
|
||||
response: Final = httpx.get(url, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
|
||||
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
|
||||
|
|
@ -212,6 +201,13 @@ class _AsyncGetClient(Protocol):
|
|||
def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ...
|
||||
|
||||
|
||||
class _SyncGetClient(Protocol):
|
||||
def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ...
|
||||
|
||||
|
||||
_FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable
|
||||
|
||||
|
||||
def _default_reload_client() -> _AsyncGetClient:
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
|
@ -219,13 +215,30 @@ def _default_reload_client() -> _AsyncGetClient:
|
|||
return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap)
|
||||
|
||||
|
||||
async def _attempt_fetch(
|
||||
client: _AsyncGetClient, url: str, timeout: int
|
||||
) -> ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable:
|
||||
def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome:
|
||||
reason: Final = f"{type(error).__name__} fetching {url}: {error}"
|
||||
if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)):
|
||||
return ModelCostMapReloadUnavailable(reason=reason)
|
||||
return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None)
|
||||
|
||||
|
||||
async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
|
||||
try:
|
||||
response: Final = await client.get(url, timeout=timeout)
|
||||
except httpx.HTTPError as e:
|
||||
return _FetchAttemptRetryable(reason=f"{type(e).__name__} fetching {url}: {e}", retry_after_seconds=None)
|
||||
except (httpx.HTTPError, httpx.InvalidURL) as e:
|
||||
return _classify_fetch_error(e, url)
|
||||
return _classify_fetch_response(response, url)
|
||||
|
||||
|
||||
def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
|
||||
try:
|
||||
response: Final = client.get(url, timeout=timeout)
|
||||
except (httpx.HTTPError, httpx.InvalidURL) as e:
|
||||
return _classify_fetch_error(e, url)
|
||||
return _classify_fetch_response(response, url)
|
||||
|
||||
|
||||
def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemptOutcome:
|
||||
if response.status_code in RETRYABLE_FETCH_STATUS_CODES:
|
||||
return _FetchAttemptRetryable(
|
||||
reason=f"HTTP {response.status_code} from {url}",
|
||||
|
|
@ -242,6 +255,22 @@ async def _attempt_fetch(
|
|||
return ModelCostMapReloaded(model_cost_map=parsed)
|
||||
|
||||
|
||||
def _next_retry_wait(
|
||||
outcome: _FetchAttemptRetryable, attempt: int, max_attempts: int, rng: random.Random
|
||||
) -> float | ModelCostMapReloadUnavailable:
|
||||
if attempt == max_attempts:
|
||||
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
|
||||
wait_seconds: Final = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
|
||||
attempt,
|
||||
max_attempts,
|
||||
outcome.reason,
|
||||
wait_seconds,
|
||||
)
|
||||
return wait_seconds
|
||||
|
||||
|
||||
async def _fetch_remote_model_cost_map_with_retry(
|
||||
url: str,
|
||||
timeout: int,
|
||||
|
|
@ -254,20 +283,32 @@ async def _fetch_remote_model_cost_map_with_retry(
|
|||
outcome = await _attempt_fetch(client=client, url=url, timeout=timeout)
|
||||
if not isinstance(outcome, _FetchAttemptRetryable):
|
||||
return outcome
|
||||
if attempt == max_attempts:
|
||||
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
|
||||
wait_seconds = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
|
||||
attempt,
|
||||
max_attempts,
|
||||
outcome.reason,
|
||||
wait_seconds,
|
||||
)
|
||||
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
|
||||
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
|
||||
return wait_seconds
|
||||
await sleep(wait_seconds)
|
||||
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
|
||||
|
||||
|
||||
def _fetch_remote_model_cost_map_with_retry_sync(
|
||||
url: str,
|
||||
timeout: int,
|
||||
max_attempts: int,
|
||||
sleep: Callable[[float], None],
|
||||
rng: random.Random,
|
||||
client: _SyncGetClient,
|
||||
) -> ModelCostMapReloadResult:
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
|
||||
if not isinstance(outcome, _FetchAttemptRetryable):
|
||||
return outcome
|
||||
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
|
||||
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
|
||||
return wait_seconds
|
||||
sleep(wait_seconds)
|
||||
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
|
||||
|
||||
|
||||
async def refetch_model_cost_map(
|
||||
url: str,
|
||||
timeout: int = 5,
|
||||
|
|
@ -423,13 +464,21 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
|
|||
return _expand_model_aliases(model_cost)
|
||||
|
||||
|
||||
def get_model_cost_map(url: str) -> dict:
|
||||
def get_model_cost_map(
|
||||
url: str,
|
||||
timeout: int = 5,
|
||||
max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
rng: random.Random | None = None,
|
||||
client: "_SyncGetClient | None" = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Public entry point — returns the model cost map dict.
|
||||
|
||||
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
|
||||
2. Otherwise fetches from ``url``, validates integrity, and falls back
|
||||
to the local backup on any failure.
|
||||
2. Otherwise fetches from ``url``, retrying transient HTTP errors
|
||||
(429/5xx/transport) with Retry-After-aware backoff, validates
|
||||
integrity, and falls back to the local backup on any failure.
|
||||
|
||||
Only the backup model count is cached (a single int) for validation.
|
||||
The full backup dict is only parsed when it must be *returned* as a
|
||||
|
|
@ -448,17 +497,24 @@ def get_model_cost_map(url: str) -> dict:
|
|||
_cost_map_source_info.url = url
|
||||
_cost_map_source_info.is_env_forced = False
|
||||
|
||||
try:
|
||||
content: Final = GetModelCostMap.fetch_remote_model_cost_map(url)
|
||||
except Exception as e:
|
||||
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
|
||||
url=url,
|
||||
timeout=timeout,
|
||||
max_attempts=max_attempts,
|
||||
sleep=sleep,
|
||||
rng=rng if rng is not None else random.Random(),
|
||||
client=client if client is not None else httpx,
|
||||
)
|
||||
if isinstance(result, ModelCostMapReloadUnavailable):
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
|
||||
url,
|
||||
str(e),
|
||||
result.reason,
|
||||
)
|
||||
_cost_map_source_info.source = "local"
|
||||
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {e}"
|
||||
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
|
||||
return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
|
||||
content: Final = result.model_cost_map
|
||||
|
||||
# Validate using cached count (cheap int comparison, no file I/O)
|
||||
if not GetModelCostMap.validate_model_cost_map(
|
||||
|
|
|
|||
|
|
@ -111,6 +111,7 @@ from litellm.types.mcp import MCPPostCallResponseObject
|
|||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import (
|
||||
DEPLOYMENT_SCOPED_PRICING_FIELDS,
|
||||
CachingDetails,
|
||||
CallTypes,
|
||||
CostBreakdown,
|
||||
|
|
@ -255,6 +256,7 @@ _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggi
|
|||
|
||||
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
|
||||
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
|
||||
|
||||
sentry_sdk_instance = None
|
||||
capture_exception = None
|
||||
|
|
@ -2957,13 +2959,25 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"Model=%s not found in completion cost map. Setting 'response_cost' to None", self.model
|
||||
)
|
||||
self.model_call_details["response_cost"] = None
|
||||
except Exception: # noqa: BLE001 # cost calculation must never block later callbacks (slot release)
|
||||
verbose_logger.exception(
|
||||
"Error calculating streaming response cost for model=%s. Setting 'response_cost' to None",
|
||||
self.model,
|
||||
)
|
||||
self.model_call_details["response_cost"] = None
|
||||
|
||||
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
|
||||
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
|
||||
complete_streaming_response, start_time, end_time
|
||||
)
|
||||
try:
|
||||
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
|
||||
complete_streaming_response, start_time, end_time
|
||||
)
|
||||
except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release)
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception building the standard logging payload "
|
||||
"for a streaming response; callbacks still run without it"
|
||||
)
|
||||
|
||||
# print standard logging payload
|
||||
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
|
||||
|
|
@ -3003,32 +3017,39 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
## LOGGING HOOK ##
|
||||
|
||||
for callback in callbacks:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
try:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=self.model_call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=self.model_call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
continue
|
||||
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
continue
|
||||
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
elif isinstance(callback, CustomLogger):
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
result=result, litellm_logging_obj=self, custom_logger=callback
|
||||
)
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
elif isinstance(callback, CustomLogger):
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
result=result, litellm_logging_obj=self, custom_logger=callback
|
||||
)
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release)
|
||||
verbose_logger.error(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s",
|
||||
traceback.format_exc(),
|
||||
)
|
||||
self._handle_callback_failure(callback=callback)
|
||||
|
||||
self.has_run_logging(event_type="async_success")
|
||||
|
||||
|
|
@ -5033,7 +5054,9 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool:
|
|||
"""
|
||||
Check if the model uses custom pricing
|
||||
|
||||
Returns True if any of `SPECIAL_MODEL_INFO_PARAMS` are present in `litellm_params` or `model_info`
|
||||
Returns True if any custom pricing field is present in `litellm_params`, or if
|
||||
any custom pricing or deployment-scoped pricing field (such as
|
||||
``off_peak_pricing``) is present in the metadata ``model_info``
|
||||
"""
|
||||
if litellm_params is None:
|
||||
return False
|
||||
|
|
@ -5051,7 +5074,7 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool:
|
|||
model_info: dict = metadata.get("model_info", {}) or {}
|
||||
|
||||
if model_info:
|
||||
matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys()
|
||||
matching_keys = _MODEL_INFO_CUSTOM_PRICING_KEYS & model_info.keys()
|
||||
for key in matching_keys:
|
||||
if model_info.get(key) is not None:
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@
|
|||
## Helper utilities for cost_per_token()
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone, tzinfo
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal, TypedDict, cast
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -290,10 +292,187 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float,
|
|||
)
|
||||
|
||||
|
||||
def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_time: datetime | None = None) -> bool:
|
||||
"""Return True if current_time (UTC, defaulting to now) falls inside any off-peak window.
|
||||
|
||||
off_peak_hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of such strings for providers
|
||||
with multiple daily windows (e.g. ["16:30-00:30", "04:00-06:00"]). A window may wrap past
|
||||
midnight, and a window whose start equals its end covers the whole day. The start is
|
||||
inclusive and the end is exclusive; malformed windows are ignored.
|
||||
|
||||
An aware current_time is converted to UTC. A naive one is taken to already be UTC rather
|
||||
than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(),
|
||||
or every window shifts by the host's offset.
|
||||
"""
|
||||
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
|
||||
now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time()
|
||||
windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc
|
||||
for window in windows:
|
||||
try:
|
||||
start_str, end_str = window.split("-")
|
||||
start = datetime.strptime(start_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time()
|
||||
end = datetime.strptime(end_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time()
|
||||
except (ValueError, AttributeError):
|
||||
continue
|
||||
if start < end:
|
||||
if start <= now < end:
|
||||
return True
|
||||
elif now >= start or now < end:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
_WEEKDAY_NUMBERS: Final = MappingProxyType(
|
||||
{
|
||||
"mon": 1,
|
||||
"monday": 1,
|
||||
"tue": 2,
|
||||
"tues": 2,
|
||||
"tuesday": 2,
|
||||
"wed": 3,
|
||||
"wednesday": 3,
|
||||
"thu": 4,
|
||||
"thur": 4,
|
||||
"thurs": 4,
|
||||
"thursday": 4,
|
||||
"fri": 5,
|
||||
"friday": 5,
|
||||
"sat": 6,
|
||||
"saturday": 6,
|
||||
"sun": 7,
|
||||
"sunday": 7,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _normalize_weekday(value: object) -> int | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value if 1 <= value <= 7 else None
|
||||
if isinstance(value, str):
|
||||
return _WEEKDAY_NUMBERS.get(value.strip().lower())
|
||||
return None
|
||||
|
||||
|
||||
def _weekday_calendar(weekday_timezone: object) -> tzinfo:
|
||||
if isinstance(weekday_timezone, str) and weekday_timezone.strip():
|
||||
try:
|
||||
return ZoneInfo(weekday_timezone.strip())
|
||||
except (ValueError, ZoneInfoNotFoundError):
|
||||
return timezone.utc
|
||||
return timezone.utc
|
||||
|
||||
|
||||
def _matches_weekdays(reference_utc: datetime, weekdays: object, weekday_timezone: object) -> bool:
|
||||
"""Return True when reference_utc falls on one of the rule's weekdays, read on the calendar
|
||||
named by weekday_timezone (default UTC). An absent weekdays means every day. The calendar
|
||||
matters even when UTC and vendor-local weekdays agree at every currently priced hour: a
|
||||
window past 16:00 UTC is where an Asia/Shanghai weekday diverges from the UTC one.
|
||||
"""
|
||||
if weekdays is None:
|
||||
return True
|
||||
if isinstance(weekdays, str) or not isinstance(weekdays, Sequence):
|
||||
return False
|
||||
allowed: Final = frozenset(day for day in map(_normalize_weekday, weekdays) if day is not None)
|
||||
return reference_utc.astimezone(_weekday_calendar(weekday_timezone)).isoweekday() in allowed
|
||||
|
||||
|
||||
def _as_window_strings(value: object) -> tuple[str, ...]:
|
||||
if isinstance(value, str):
|
||||
return (value,)
|
||||
if isinstance(value, Sequence):
|
||||
return tuple(entry for entry in value if isinstance(entry, str))
|
||||
return ()
|
||||
|
||||
|
||||
def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = None) -> bool:
|
||||
"""Return True when current_time (UTC, defaulting to now) is off-peak under the block's
|
||||
rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose
|
||||
hours apply only on its weekdays.
|
||||
"""
|
||||
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
|
||||
reference_utc: Final = (
|
||||
reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc)
|
||||
)
|
||||
flat_windows: Final = _as_window_strings(off_peak.get("hours_utc"))
|
||||
if flat_windows and _is_within_off_peak_window(flat_windows, reference_utc):
|
||||
return True
|
||||
windows: Final = off_peak.get("windows")
|
||||
if isinstance(windows, str) or not isinstance(windows, Sequence):
|
||||
return False
|
||||
weekday_timezone: Final = off_peak.get("weekday_timezone")
|
||||
for rule in windows:
|
||||
if not isinstance(rule, Mapping):
|
||||
continue
|
||||
rule_windows = _as_window_strings(rule.get("hours_utc"))
|
||||
if not rule_windows:
|
||||
continue
|
||||
if not _matches_weekdays(reference_utc, rule.get("weekdays"), weekday_timezone):
|
||||
continue
|
||||
if _is_within_off_peak_window(rule_windows, reference_utc):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _coerce_off_peak_rate(value: object, default: float) -> float:
|
||||
if isinstance(value, bool):
|
||||
return default
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
return default
|
||||
return default
|
||||
|
||||
|
||||
def _apply_off_peak_pricing(
|
||||
model_info: ModelInfo,
|
||||
current_time: datetime | None,
|
||||
prompt_base_cost: float,
|
||||
completion_base_cost: float,
|
||||
cache_read_cost: float,
|
||||
) -> tuple[float, float, float]:
|
||||
"""Swap in off-peak per-token rates when the current UTC time is inside one of the model's
|
||||
off_peak_pricing rules, the every-day hours_utc windows or a day-of-week-qualified entry in
|
||||
windows. An off-peak rate replaces the rate that would otherwise apply rather than
|
||||
discounting it, so a model that also has tiered or above-threshold pricing bills the flat
|
||||
off-peak rate for the whole request while the window is open. Any rate left unset in
|
||||
off_peak_pricing falls back to the standard rate.
|
||||
"""
|
||||
off_peak: Final = model_info.get("off_peak_pricing")
|
||||
if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time):
|
||||
return prompt_base_cost, completion_base_cost, cache_read_cost
|
||||
return (
|
||||
_coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost),
|
||||
_coerce_off_peak_rate(off_peak.get("output_cost_per_token"), completion_base_cost),
|
||||
_coerce_off_peak_rate(off_peak.get("cache_read_input_token_cost"), cache_read_cost),
|
||||
)
|
||||
|
||||
|
||||
def _apply_off_peak_to_base_costs(
|
||||
model_info: ModelInfo,
|
||||
current_time: datetime | None,
|
||||
base_costs: tuple[float, float, float, float, float],
|
||||
) -> tuple[float, float, float, float, float]:
|
||||
"""Apply off-peak rates to an already-resolved set of base costs, whichever pricing path
|
||||
produced them. Cache-creation rates are passed through untouched, since off_peak_pricing
|
||||
has no field for them.
|
||||
"""
|
||||
prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
|
||||
off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing(
|
||||
model_info, current_time, prompt, completion, cache_read
|
||||
)
|
||||
return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read)
|
||||
|
||||
|
||||
def _get_token_base_cost(
|
||||
model_info: ModelInfo,
|
||||
usage: Usage,
|
||||
service_tier: str | None = None,
|
||||
current_time: datetime | None = None,
|
||||
*,
|
||||
threshold_is_inclusive: bool = False,
|
||||
) -> tuple[float, float, float, float, float]:
|
||||
|
|
@ -311,7 +490,7 @@ def _get_token_base_cost(
|
|||
"""
|
||||
tiered_base_costs: Final = _get_tiered_base_costs(model_info=model_info, usage=usage)
|
||||
if tiered_base_costs is not None:
|
||||
return tiered_base_costs
|
||||
return _apply_off_peak_to_base_costs(model_info, current_time, tiered_base_costs)
|
||||
|
||||
# Get service tier aware cost keys
|
||||
input_cost_key: Final = _get_service_tier_cost_key("input_cost_per_token", service_tier)
|
||||
|
|
@ -345,12 +524,16 @@ def _get_token_base_cost(
|
|||
k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES)
|
||||
]
|
||||
if not threshold_keys:
|
||||
return (
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
return _apply_off_peak_to_base_costs(
|
||||
model_info,
|
||||
current_time,
|
||||
(
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
),
|
||||
)
|
||||
|
||||
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
|
||||
|
|
@ -451,12 +634,16 @@ def _get_token_base_cost(
|
|||
except Exception:
|
||||
continue
|
||||
|
||||
return (
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
return _apply_off_peak_to_base_costs(
|
||||
model_info,
|
||||
current_time,
|
||||
(
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1281,6 +1281,19 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin
|
|||
return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo
|
||||
|
||||
|
||||
def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]:
|
||||
function: Final = tool.get("function")
|
||||
if not isinstance(function, dict):
|
||||
return tool
|
||||
parameters: Final = function.get("parameters")
|
||||
if not isinstance(parameters, dict):
|
||||
return tool
|
||||
flattened: Final = flatten_top_level_schema_combinators(parameters)
|
||||
if flattened is parameters:
|
||||
return tool
|
||||
return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts
|
||||
|
||||
|
||||
def _get_image_mime_type_from_url(url: str) -> str | None:
|
||||
"""
|
||||
Get mime type for common image URLs
|
||||
|
|
|
|||
|
|
@ -1694,6 +1694,18 @@ def convert_function_to_anthropic_tool_invoke(
|
|||
raise e
|
||||
|
||||
|
||||
def _find_server_tool_result(
|
||||
tool_id: str,
|
||||
web_search_results: Sequence[object] | None,
|
||||
tool_results: Sequence[object] | None,
|
||||
) -> dict[str, object] | None:
|
||||
candidates: Final = (*(web_search_results or ()), *(tool_results or ()))
|
||||
return next(
|
||||
(result for result in candidates if isinstance(result, dict) and result.get("tool_use_id") == tool_id),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def convert_to_anthropic_tool_invoke(
|
||||
tool_calls: list[ChatCompletionAssistantToolCall],
|
||||
web_search_results: list[Any] | None = None,
|
||||
|
|
@ -1758,32 +1770,22 @@ def convert_to_anthropic_tool_invoke(
|
|||
context="Anthropic tool invoke",
|
||||
)
|
||||
|
||||
# Check if this is a server-side tool (web_search, tool_search, etc.)
|
||||
# Server tool IDs start with "srvtoolu_"
|
||||
if tool_id.startswith("srvtoolu_"):
|
||||
# Create server_tool_use block instead of tool_use
|
||||
_anthropic_server_tool_use: dict[str, object] = {
|
||||
"type": "server_tool_use",
|
||||
"id": tool_id,
|
||||
"name": tool_name,
|
||||
"input": tool_input,
|
||||
}
|
||||
anthropic_tool_invoke.append(_anthropic_server_tool_use)
|
||||
|
||||
# Add corresponding tool result if available.
|
||||
# Check both web_search_results (web_search_tool_result / web_fetch_tool_result)
|
||||
# and tool_results (bash_code_execution_tool_result, etc.)
|
||||
_all_tool_results: list[Any] = []
|
||||
if web_search_results:
|
||||
_all_tool_results.extend(web_search_results)
|
||||
if tool_results:
|
||||
_all_tool_results.extend(tool_results)
|
||||
for result in _all_tool_results:
|
||||
if result.get("tool_use_id") == tool_id:
|
||||
anthropic_tool_invoke.append(result)
|
||||
break
|
||||
server_tool_result = (
|
||||
_find_server_tool_result(tool_id, web_search_results, tool_results)
|
||||
if tool_id.startswith("srvtoolu_")
|
||||
else None
|
||||
)
|
||||
if server_tool_result is not None:
|
||||
anthropic_tool_invoke.append(
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": tool_id,
|
||||
"name": tool_name,
|
||||
"input": tool_input,
|
||||
}
|
||||
)
|
||||
anthropic_tool_invoke.append(server_tool_result)
|
||||
else:
|
||||
# Regular tool_use
|
||||
sanitized_tool_id = _sanitize_anthropic_tool_use_id(tool_id)
|
||||
_anthropic_tool_use_param = AnthropicMessagesToolUseParam(
|
||||
type="tool_use",
|
||||
|
|
@ -4955,10 +4957,13 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str:
|
|||
|
||||
|
||||
def add_cache_point_tool_block(tool: dict, model: str | None = None) -> BedrockToolBlock | None:
|
||||
from litellm.llms.bedrock.common_utils import is_claude_4_5_on_bedrock
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
bedrock_model_accepts_cache_points,
|
||||
is_claude_4_5_on_bedrock,
|
||||
)
|
||||
|
||||
cache_control: Final = tool.get("cache_control", None)
|
||||
if cache_control is not None:
|
||||
if cache_control is not None and bedrock_model_accepts_cache_points(model):
|
||||
cache_point: Final = cache_control.get("type", "ephemeral")
|
||||
if cache_point == "ephemeral":
|
||||
cache_point_block: Final[CachePointBlock] = {"type": "default"}
|
||||
|
|
|
|||
|
|
@ -1516,7 +1516,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
_tool = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled)
|
||||
if _tool is None:
|
||||
continue
|
||||
if not is_thinking_enabled:
|
||||
if not is_thinking_enabled and not AnthropicModelInfo.forced_tool_use_unsupported(model):
|
||||
_tool_choice = {
|
||||
"name": RESPONSE_FORMAT_TOOL_NAME,
|
||||
"type": "tool",
|
||||
|
|
|
|||
|
|
@ -325,13 +325,17 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
status_code=400,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def forced_tool_use_unsupported(model: str) -> bool:
|
||||
return AnthropicModelInfo._get_model_capability(model, "supports_forced_tool_use") is False
|
||||
|
||||
@staticmethod
|
||||
def forced_tool_use_downgraded(model: str, drop_params: bool) -> bool:
|
||||
"""True when the model map flags the model with
|
||||
``supports_forced_tool_use: false`` (Fable 5.1 / Mythos 5.1 400 on
|
||||
``any``/``tool``) and ``drop_params`` asks for the ``auto`` downgrade;
|
||||
raises a clean client-side 400 for such models without ``drop_params``."""
|
||||
if AnthropicModelInfo._get_model_capability(model, "supports_forced_tool_use") is not False:
|
||||
if not AnthropicModelInfo.forced_tool_use_unsupported(model):
|
||||
return False
|
||||
if not (litellm.drop_params or drop_params):
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
|
|
|
|||
|
|
@ -86,22 +86,41 @@ def _decoded_sse_data_line(line: bytes) -> object | None:
|
|||
return None
|
||||
|
||||
|
||||
def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
|
||||
def _anthropic_event_payload(chunk: object, event_type: str) -> Mapping[str, object] | None:
|
||||
if isinstance(chunk, dict):
|
||||
return chunk if chunk.get("type") == "error" else None
|
||||
return chunk if chunk.get("type") == event_type else None
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines())
|
||||
return next(
|
||||
(
|
||||
candidate
|
||||
for candidate in decoded_lines
|
||||
if isinstance(candidate, dict) and candidate.get("type") == "error"
|
||||
if isinstance(candidate, dict) and candidate.get("type") == event_type
|
||||
),
|
||||
None,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
|
||||
return _anthropic_event_payload(chunk, "error")
|
||||
|
||||
|
||||
def parse_anthropic_refusal_stop_details(chunk: object) -> Mapping[str, object] | None:
|
||||
"""
|
||||
Return the ``stop_details`` object of an Anthropic SSE ``message_delta``
|
||||
chunk whose delta carries ``stop_reason: "refusal"`` (a safeguard refusal:
|
||||
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback),
|
||||
or None for any other chunk, a plain refusal without ``stop_details`` included.
|
||||
"""
|
||||
payload: Final = _anthropic_event_payload(chunk, "message_delta")
|
||||
delta: Final = payload.get("delta") if payload is not None else None
|
||||
if not isinstance(delta, dict) or delta.get("stop_reason") != "refusal":
|
||||
return None
|
||||
stop_details: Final = delta.get("stop_details")
|
||||
return stop_details if isinstance(stop_details, dict) else None
|
||||
|
||||
|
||||
def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None:
|
||||
"""Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None."""
|
||||
payload: Final = _anthropic_error_event_payload(chunk)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,40 @@
|
|||
from collections.abc import Mapping
|
||||
from functools import lru_cache
|
||||
from typing import Any, Final, cast, get_type_hints
|
||||
from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints
|
||||
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.exceptions import ContentPolicyViolationError
|
||||
|
||||
|
||||
def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None:
|
||||
"""
|
||||
Return the ``stop_details`` of an Anthropic Messages response refused by a
|
||||
safeguard (``stop_reason: "refusal"`` carrying ``stop_details``:
|
||||
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback),
|
||||
or None for any other response, a plain refusal without ``stop_details`` included.
|
||||
"""
|
||||
if not isinstance(response, dict) or response.get("stop_reason") != "refusal":
|
||||
return None
|
||||
stop_details: Final = response.get("stop_details")
|
||||
return stop_details if isinstance(stop_details, dict) else None
|
||||
|
||||
|
||||
def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError":
|
||||
"""The exception a safeguard-refused Anthropic response converts into so the
|
||||
content-policy fallback chain can re-dispatch it."""
|
||||
from litellm.exceptions import ContentPolicyViolationError
|
||||
|
||||
return ContentPolicyViolationError(
|
||||
message=f"Anthropic safeguard refusal (category: {stop_details.get('category')}).",
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _anthropic_messages_optional_param_keys() -> frozenset[str]:
|
||||
|
|
@ -100,14 +129,12 @@ def mock_response(
|
|||
model=model,
|
||||
)
|
||||
return AnthropicMessagesResponse(
|
||||
**{
|
||||
"content": [{"text": mock_response, "type": "text"}],
|
||||
"id": "msg_013Zva2CMHLNnXjNJJKqJ2EF",
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"role": "assistant",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"type": "message",
|
||||
"usage": {"input_tokens": 2095, "output_tokens": 503},
|
||||
}
|
||||
content=[{"text": mock_response, "type": "text"}],
|
||||
id="msg_013Zva2CMHLNnXjNJJKqJ2EF",
|
||||
model="claude-sonnet-4-20250514",
|
||||
role="assistant",
|
||||
stop_reason="end_turn",
|
||||
stop_sequence=None,
|
||||
type="message",
|
||||
usage={"input_tokens": 2095, "output_tokens": 503},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from httpx._models import Headers, Response
|
||||
|
|
@ -6,6 +8,7 @@ import litellm
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
drop_tool_reference_parts_from_tool_messages,
|
||||
hoist_images_from_tool_messages,
|
||||
tool_with_flattened_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_azure_openai_messages,
|
||||
|
|
@ -32,6 +35,19 @@ else:
|
|||
LoggingClass = Any
|
||||
|
||||
|
||||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
tools: Final = optional_params.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return _NO_TOOLS_UPDATE
|
||||
flattened: Final = [ # mutable-ok: request tools are a JSON list
|
||||
tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools
|
||||
]
|
||||
return MappingProxyType({"tools": flattened})
|
||||
|
||||
|
||||
class AzureOpenAIConfig(BaseConfig):
|
||||
"""
|
||||
Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions
|
||||
|
|
@ -261,6 +277,7 @@ class AzureOpenAIConfig(BaseConfig):
|
|||
"model": model,
|
||||
"messages": azure_messages,
|
||||
**optional_params,
|
||||
**flattened_tools_update(optional_params),
|
||||
}
|
||||
|
||||
def transform_response(
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.utils import get_model_info, supports_reasoning
|
||||
|
||||
from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig
|
||||
from .gpt_transformation import flattened_tools_update
|
||||
|
||||
|
||||
class AzureOpenAIO1Config(OpenAIOSeriesConfig):
|
||||
|
|
@ -108,4 +109,8 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig):
|
|||
headers: dict,
|
||||
) -> dict:
|
||||
model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name
|
||||
return super().transform_request(model, messages, optional_params, litellm_params, headers)
|
||||
flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict
|
||||
**optional_params,
|
||||
**flattened_tools_update(optional_params),
|
||||
}
|
||||
return super().transform_request(model, messages, flattened_params, litellm_params, headers)
|
||||
|
|
|
|||
|
|
@ -155,8 +155,8 @@ class BaseTranslation(ABC):
|
|||
self,
|
||||
exc: "ModifyResponseException",
|
||||
stream_started: bool = False,
|
||||
responses_so_far: list[Any] | None = None,
|
||||
) -> list[bytes] | None:
|
||||
responses_so_far: Sequence[Any] | None = None,
|
||||
) -> Sequence[bytes] | None:
|
||||
"""
|
||||
Build the streaming chunks that deliver a guardrail block message and
|
||||
cleanly terminate the stream in this provider's wire format.
|
||||
|
|
|
|||
|
|
@ -124,6 +124,61 @@ def blocked_responses_api_usage(original_response: object) -> ResponseAPIUsage:
|
|||
)
|
||||
|
||||
|
||||
def stream_item_field(item: object, field: str) -> object | None:
|
||||
if isinstance(item, dict):
|
||||
return item.get(field)
|
||||
return getattr(item, field, None)
|
||||
|
||||
|
||||
def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]:
|
||||
"""
|
||||
``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked
|
||||
chat completions stream.
|
||||
|
||||
A mid-stream block carries the chunks received so far as a list; real usage
|
||||
rides on the final chunk when the upstream sent one
|
||||
(``stream_options.include_usage``). Non-list originals defer to
|
||||
``blocked_response_usage``.
|
||||
"""
|
||||
if not isinstance(original_response, list):
|
||||
usage: Final = blocked_response_usage(original_response)
|
||||
return usage.get("input_tokens", 0), usage.get("output_tokens", 0)
|
||||
usage_obj: Final = next(
|
||||
(
|
||||
chunk_usage
|
||||
for item in reversed(original_response)
|
||||
if (chunk_usage := stream_item_field(item, "usage")) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
return (
|
||||
_usage_tokens(usage_obj, "prompt_tokens", "input_tokens"),
|
||||
_usage_tokens(usage_obj, "completion_tokens", "output_tokens"),
|
||||
)
|
||||
|
||||
|
||||
def blocked_responses_stream_usage(original_response: object) -> ResponseAPIUsage:
|
||||
"""
|
||||
``ResponseAPIUsage`` for a synthetic guardrail-blocked /v1/responses stream.
|
||||
|
||||
A mid-stream block carries the events received so far as a list; real usage
|
||||
rides on the ``response.completed`` event's response when the upstream sent
|
||||
one. Non-list originals defer to ``blocked_responses_api_usage``.
|
||||
"""
|
||||
if not isinstance(original_response, list):
|
||||
return blocked_responses_api_usage(original_response)
|
||||
completed: Final = next(
|
||||
(
|
||||
response
|
||||
for item in reversed(original_response)
|
||||
if stream_item_field(item, "type") == "response.completed"
|
||||
and (response := stream_item_field(item, "response")) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
return blocked_responses_api_usage(completed)
|
||||
|
||||
|
||||
def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool:
|
||||
per: Final = getattr(guardrail_to_apply, "skip_system_message_in_guardrail", None)
|
||||
if per is not None:
|
||||
|
|
|
|||
|
|
@ -1442,7 +1442,7 @@ class BaseAWSLLM:
|
|||
@tracer.wrap()
|
||||
def get_request_headers(
|
||||
self,
|
||||
credentials: Credentials,
|
||||
credentials: Credentials | None,
|
||||
aws_region_name: str,
|
||||
extra_headers: dict | None,
|
||||
endpoint_url: str,
|
||||
|
|
@ -1469,9 +1469,13 @@ class BaseAWSLLM:
|
|||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.exceptions import NoCredentialsError
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
|
||||
if credentials is None:
|
||||
raise NoCredentialsError()
|
||||
|
||||
# Filter headers for AWS signature calculation
|
||||
# AWS SigV4 only includes specific headers in signature calculation
|
||||
aws_signature_headers: Final = self._filter_headers_for_aws_signature(headers)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -24,6 +26,22 @@ from ..common_utils import BedrockError, _get_all_bedrock_regions
|
|||
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
|
||||
|
||||
|
||||
def _sigv4_principal(credentials: Credentials | None) -> Mapping[str, str]:
|
||||
if credentials is None:
|
||||
return MappingProxyType({})
|
||||
return MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in (
|
||||
("aws_access_key_id", credentials.access_key),
|
||||
("aws_secret_access_key", credentials.secret_key),
|
||||
("aws_session_token", credentials.token),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def make_sync_call(
|
||||
client: HTTPHandler | None,
|
||||
api_base: str,
|
||||
|
|
@ -95,7 +113,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
stream,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
credentials: Credentials,
|
||||
credentials: Credentials | None,
|
||||
logger_fn=None,
|
||||
headers={},
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
|
|
@ -167,7 +185,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
stream,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
credentials: Credentials,
|
||||
credentials: Credentials | None,
|
||||
logger_fn=None,
|
||||
headers: dict = {},
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
|
|
@ -331,7 +349,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
|
||||
litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls
|
||||
|
||||
credentials: Final[Credentials] = self.get_credentials(
|
||||
credentials: Final[Credentials | None] = self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
|
|
@ -368,19 +386,13 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
# The Rust core owns the whole call for the subset it accepts. Ask
|
||||
# before transforming so whichever path runs emits pre_call once, and
|
||||
# hand down the credentials, region and endpoint this handler already
|
||||
# resolved so both paths sign as the same principal.
|
||||
# resolved so both paths sign as the same principal. Bearer-token auth
|
||||
# resolves no SigV4 principal at all, and each path reads that token
|
||||
# itself.
|
||||
rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy
|
||||
**optional_params,
|
||||
**{ # mutable-ok: merged into its mutable parent above
|
||||
key: value
|
||||
for key, value in (
|
||||
("aws_access_key_id", credentials.access_key),
|
||||
("aws_secret_access_key", credentials.secret_key),
|
||||
("aws_session_token", credentials.token),
|
||||
("aws_region_name", aws_region_name),
|
||||
)
|
||||
if value is not None
|
||||
},
|
||||
**_sigv4_principal(credentials),
|
||||
"aws_region_name": aws_region_name,
|
||||
}
|
||||
serves_via_rust: Final = rust_chat_completions_accepts(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -87,6 +87,7 @@ from ..common_utils import (
|
|||
BedrockError,
|
||||
BedrockModelInfo,
|
||||
bedrock_converse_supports_parallel_tool_use_config,
|
||||
bedrock_model_accepts_cache_points,
|
||||
get_anthropic_beta_from_headers,
|
||||
get_bedrock_tool_name,
|
||||
is_bedrock_application_inference_profile_arn,
|
||||
|
|
@ -1073,6 +1074,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
if (
|
||||
litellm.utils.supports_tool_choice(model=model, custom_llm_provider=self.custom_llm_provider)
|
||||
and not is_thinking_enabled
|
||||
and not AnthropicModelInfo.forced_tool_use_unsupported(model)
|
||||
):
|
||||
optional_params["tool_choice"] = ToolChoiceValuesBlock(
|
||||
tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME)
|
||||
|
|
@ -1148,7 +1150,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
model: str | None = None,
|
||||
) -> SystemContentBlock | ContentBlock | None:
|
||||
cache_control: Final = message_block.get("cache_control", None)
|
||||
if cache_control is None:
|
||||
if cache_control is None or not bedrock_model_accepts_cache_points(model):
|
||||
return None
|
||||
|
||||
cache_point: Final = self._build_cache_point_block(cache_control, model)
|
||||
|
|
@ -1612,7 +1614,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
# Append cachePoint to tools if cache_control_injection_points has tool_config
|
||||
cache_injection_points: Final = additional_request_params.pop("cache_control_injection_points", None)
|
||||
if cache_injection_points and len(bedrock_tools) > 0:
|
||||
if cache_injection_points and len(bedrock_tools) > 0 and bedrock_model_accepts_cache_points(model):
|
||||
for point in cache_injection_points:
|
||||
if point.get("location") == "tool_config":
|
||||
cache_point = self._build_cache_point_block(point.get("control"), model)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final
|
|||
import httpx
|
||||
|
||||
from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_anthropic_image_obj,
|
||||
)
|
||||
|
|
@ -11,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
|||
convert_url_to_base64,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
|
||||
AmazonInvokeConfig,
|
||||
)
|
||||
|
|
@ -74,10 +76,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
drop_params: bool,
|
||||
) -> dict:
|
||||
# Force tool-based structured outputs for Bedrock Invoke
|
||||
# (similar to VertexAI fix in #19201)
|
||||
# Bedrock Invoke doesn't support output_format parameter
|
||||
# (similar to VertexAI fix in #19201) unless the model map advertises
|
||||
# native structured output
|
||||
from litellm.utils import supports_native_structured_output
|
||||
|
||||
original_model: Final = model
|
||||
if "response_format" in non_default_params:
|
||||
if "response_format" in non_default_params and not supports_native_structured_output(
|
||||
model=model, custom_llm_provider="bedrock"
|
||||
):
|
||||
# Use a model name that forces tool-based approach
|
||||
model = "claude-3-sonnet-20240229"
|
||||
|
||||
|
|
@ -101,6 +107,16 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
# Restore original model name
|
||||
model = original_model
|
||||
|
||||
# The stub model hides the original model from the parent's forced-tool-use backstop
|
||||
response_format_tool_choice: Final = optional_params.get("tool_choice")
|
||||
if (
|
||||
"response_format" in non_default_params
|
||||
and isinstance(response_format_tool_choice, dict)
|
||||
and response_format_tool_choice.get("name") == RESPONSE_FORMAT_TOOL_NAME
|
||||
and AnthropicModelInfo.forced_tool_use_unsupported(original_model)
|
||||
):
|
||||
optional_params.pop("tool_choice")
|
||||
|
||||
return optional_params
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -816,6 +816,30 @@ def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def bedrock_model_accepts_cache_points(model: str | None) -> bool:
|
||||
"""
|
||||
Whether Converse ``cachePoint`` blocks may be sent to this model.
|
||||
|
||||
Bedrock rejects requests carrying cachePoint blocks for models without prompt
|
||||
caching support ("You invoked an unsupported model or your request did not allow
|
||||
prompt caching"), so a model whose cost-map entry does not declare
|
||||
``supports_prompt_caching`` must not receive them. A model absent from the map
|
||||
(an application inference profile ARN, a model newer than the map) keeps emitting
|
||||
so existing caching setups never silently degrade. ``litellm.utils.supports_prompt_caching``
|
||||
is not reusable here: it returns False for unmapped models, the opposite polarity.
|
||||
"""
|
||||
if model is None:
|
||||
return True
|
||||
entries: Final = tuple(
|
||||
entry
|
||||
for candidate in (model, get_bedrock_base_model(model))
|
||||
if (entry := litellm.model_cost.get(candidate)) is not None
|
||||
)
|
||||
if not entries:
|
||||
return True
|
||||
return any(entry.get("supports_prompt_caching") is True for entry in entries)
|
||||
|
||||
|
||||
def is_claude_4_5_on_bedrock(model: str) -> bool:
|
||||
"""
|
||||
Check if the model supports Bedrock prompt caching with an extended '1h' TTL
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import concurrent.futures
|
|||
import contextlib
|
||||
import os
|
||||
import ssl
|
||||
import sys
|
||||
import typing
|
||||
import urllib.request
|
||||
from collections.abc import Callable, Generator
|
||||
|
|
@ -75,10 +76,22 @@ except ImportError:
|
|||
pass
|
||||
|
||||
|
||||
def _current_task_is_cancelling() -> bool:
|
||||
task: Final = asyncio.current_task()
|
||||
if task is None or sys.version_info < (3, 11):
|
||||
return True
|
||||
return task.cancelling() > 0
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def map_aiohttp_exceptions() -> Generator[None, None, None]:
|
||||
try:
|
||||
yield
|
||||
except asyncio.CancelledError as exc:
|
||||
# a closing connector cancels its shielded DNS task; that surfaces here without the request task being cancelled
|
||||
if _current_task_is_cancelling():
|
||||
raise
|
||||
raise httpx.ConnectError("aiohttp transport cancelled the request internally") from exc
|
||||
except Exception as exc:
|
||||
mapped_exc: type[Exception] | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ Support for gpt model family
|
|||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -21,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
drop_tool_reference_parts_from_tool_messages,
|
||||
get_tool_call_names,
|
||||
hoist_images_from_tool_messages,
|
||||
tool_with_flattened_parameters,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
async_convert_url_to_base64,
|
||||
|
|
@ -65,6 +67,9 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
||||
"""
|
||||
Reference: https://platform.openai.com/docs/api-reference/chat/create
|
||||
|
|
@ -397,6 +402,21 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
)
|
||||
return messages, tools
|
||||
|
||||
def _targets_openai_hosted_endpoint(
|
||||
self,
|
||||
custom_llm_provider: str | None,
|
||||
api_base: str | None,
|
||||
) -> bool:
|
||||
if custom_llm_provider != "openai":
|
||||
return False
|
||||
resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE")
|
||||
if not resolved_api_base:
|
||||
return True
|
||||
hostname: Final = urlparse(resolved_api_base).hostname
|
||||
if hostname is None:
|
||||
return True
|
||||
return hostname == "openai.com" or hostname.endswith(".openai.com")
|
||||
|
||||
def _should_preserve_cache_control_for_endpoint(
|
||||
self,
|
||||
custom_llm_provider: str | None,
|
||||
|
|
@ -408,15 +428,34 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
api_base. Those can understand cache_control, so it must survive there.
|
||||
Real OpenAI cannot, so it is still stripped for an openai.com host.
|
||||
"""
|
||||
if custom_llm_provider != "openai":
|
||||
return False
|
||||
resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE")
|
||||
if not resolved_api_base:
|
||||
return False
|
||||
hostname: Final = urlparse(resolved_api_base).hostname
|
||||
if hostname is None:
|
||||
return False
|
||||
return hostname != "openai.com" and not hostname.endswith(".openai.com")
|
||||
return custom_llm_provider == "openai" and not self._targets_openai_hosted_endpoint(
|
||||
custom_llm_provider, api_base
|
||||
)
|
||||
|
||||
def _flattened_tools_update_for_openai(
|
||||
self,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
OpenAI's chat completions validator rejects tool `parameters` carrying
|
||||
'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every
|
||||
model family, unlike the Responses API, where GPT-5+ accepts them.
|
||||
"""
|
||||
tools: Final = optional_params.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return _NO_TOOLS_UPDATE
|
||||
provider: Final = litellm_params.get("custom_llm_provider")
|
||||
raw_api_base: Final = litellm_params.get("api_base")
|
||||
if not self._targets_openai_hosted_endpoint(
|
||||
provider if isinstance(provider, str) else None,
|
||||
raw_api_base if isinstance(raw_api_base, str) else None,
|
||||
):
|
||||
return _NO_TOOLS_UPDATE
|
||||
flattened: Final = [ # mutable-ok: request tools are a JSON list
|
||||
tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools
|
||||
]
|
||||
return MappingProxyType({"tools": flattened})
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
|
|
@ -443,11 +482,14 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
optional_params["tools"] = tools
|
||||
|
||||
optional_params.pop("max_retries", None)
|
||||
if not optional_params.get("tools") and not optional_params.get("functions"):
|
||||
optional_params.pop("tool_choice", None)
|
||||
|
||||
return {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
**optional_params,
|
||||
**self._flattened_tools_update_for_openai(optional_params, litellm_params),
|
||||
}
|
||||
|
||||
async def async_transform_request(
|
||||
|
|
@ -473,10 +515,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
if tools is not None and len(tools) > 0:
|
||||
optional_params["tools"] = tools
|
||||
if self.__class__._is_base_class:
|
||||
if not optional_params.get("tools") and not optional_params.get("functions"):
|
||||
optional_params.pop("tool_choice", None)
|
||||
return {
|
||||
"model": model,
|
||||
"messages": transformed_messages,
|
||||
**optional_params,
|
||||
**self._flattened_tools_update_for_openai(optional_params, litellm_params),
|
||||
}
|
||||
else:
|
||||
## allow for any object specific behaviour to be handled
|
||||
|
|
|
|||
|
|
@ -14,9 +14,14 @@ Pattern Overview:
|
|||
This pattern can be replicated for other message formats (e.g., Anthropic).
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
|
|
@ -24,6 +29,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
|||
StreamTransformSink,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_chat_stream_usage,
|
||||
effective_scan_only_tool_results_for_guardrail,
|
||||
effective_skip_system_message_for_guardrail,
|
||||
effective_skip_tool_message_for_guardrail,
|
||||
|
|
@ -32,6 +38,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
openai_tool_name,
|
||||
role_out_of_guardrail_scope,
|
||||
scoped_structured_message_indices,
|
||||
stream_item_field,
|
||||
)
|
||||
from litellm.main import stream_chunk_builder
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
|
|
@ -49,7 +56,10 @@ from litellm.types.utils import (
|
|||
if TYPE_CHECKING:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
ModifyResponseException,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
|
@ -1005,3 +1015,129 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
else:
|
||||
# Subsequent chunks - clear the text
|
||||
content_item["text"] = ""
|
||||
|
||||
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
|
||||
"""
|
||||
True once any relayed chunk carries a non-null ``finish_reason``.
|
||||
|
||||
The unified guardrail's ``end_of_stream_only`` streaming path probes
|
||||
this via ``hasattr`` to withhold the terminal chunks until
|
||||
end-of-stream moderation runs, so a block can replace the finish
|
||||
instead of trailing after a ``finish_reason`` the client already saw.
|
||||
"""
|
||||
return any(
|
||||
stream_item_field(choice, "finish_reason") is not None
|
||||
for item in responses_so_far
|
||||
for choice in _stream_chunk_choices(item)
|
||||
)
|
||||
|
||||
def build_block_sse_chunks(
|
||||
self,
|
||||
exc: "ModifyResponseException",
|
||||
stream_started: bool = False,
|
||||
responses_so_far: Sequence[object] | None = None,
|
||||
) -> Sequence[bytes]:
|
||||
"""
|
||||
Build OpenAI chat-completions SSE chunks that deliver the guardrail
|
||||
block message and terminate the stream cleanly, mirroring the
|
||||
non-streaming block response: ``finish_reason`` ``content_filter`` plus
|
||||
the real usage the upstream call consumed.
|
||||
|
||||
- ``stream_started`` False (buffered / pre-stream): nothing has been
|
||||
sent, so open a standalone completion with a ``role`` delta.
|
||||
- ``stream_started`` True (sampling / mid-stream): chunks already
|
||||
reached the client, so continue the in-progress completion (reuse its
|
||||
id/created/model, content-only delta).
|
||||
|
||||
The proxy's data generator appends ``data: [DONE]`` itself.
|
||||
"""
|
||||
chunk_id, created, model = _blocked_stream_identity(exc, responses_so_far or ())
|
||||
prompt_tokens, completion_tokens = blocked_chat_stream_usage(exc.original_response)
|
||||
continuation_delta: Final[_BlockedChunkDelta] = {"content": exc.message}
|
||||
standalone_delta: Final[_BlockedChunkDelta] = {"role": "assistant", "content": exc.message}
|
||||
message_chunk: Final[_BlockedChunk] = {
|
||||
"id": chunk_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model,
|
||||
"choices": (
|
||||
{
|
||||
"index": 0,
|
||||
"delta": continuation_delta if stream_started else standalone_delta,
|
||||
"finish_reason": None,
|
||||
},
|
||||
),
|
||||
}
|
||||
final_chunk: Final[_BlockedChunk] = {
|
||||
"id": chunk_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model,
|
||||
"choices": ({"index": 0, "delta": {}, "finish_reason": "content_filter"},),
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
},
|
||||
}
|
||||
return _chat_sse_chunk(message_chunk), _chat_sse_chunk(final_chunk)
|
||||
|
||||
|
||||
class _BlockedChunkDelta(TypedDict, total=False):
|
||||
role: ReadOnly[str]
|
||||
content: ReadOnly[str]
|
||||
|
||||
|
||||
class _BlockedChunkChoice(TypedDict):
|
||||
index: ReadOnly[int]
|
||||
delta: ReadOnly[_BlockedChunkDelta]
|
||||
finish_reason: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _BlockedChunkUsage(TypedDict):
|
||||
prompt_tokens: ReadOnly[int]
|
||||
completion_tokens: ReadOnly[int]
|
||||
total_tokens: ReadOnly[int]
|
||||
|
||||
|
||||
class _BlockedChunk(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
object: ReadOnly[str]
|
||||
created: ReadOnly[int]
|
||||
model: ReadOnly[str]
|
||||
choices: ReadOnly[tuple[_BlockedChunkChoice, ...]]
|
||||
usage: NotRequired[ReadOnly[_BlockedChunkUsage]]
|
||||
|
||||
|
||||
def _chat_sse_chunk(payload: _BlockedChunk) -> bytes:
|
||||
return f"data: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def _stream_chunk_choices(item: object) -> Sequence[object]:
|
||||
choices: Final = stream_item_field(item, "choices")
|
||||
if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes)):
|
||||
return choices
|
||||
return ()
|
||||
|
||||
|
||||
def _blocked_stream_identity(
|
||||
exc: "ModifyResponseException", responses_so_far: Sequence[object]
|
||||
) -> tuple[str, int, str]:
|
||||
identified: Final = next(
|
||||
(
|
||||
(chunk_id, item)
|
||||
for item in responses_so_far
|
||||
if isinstance(chunk_id := stream_item_field(item, "id"), str) and chunk_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
if identified is None:
|
||||
return f"chatcmpl-{uuid.uuid4()}", int(time.time()), exc.model
|
||||
chunk_id, source = identified
|
||||
created: Final = stream_item_field(source, "created")
|
||||
model: Final = stream_item_field(source, "model")
|
||||
return (
|
||||
chunk_id,
|
||||
created if isinstance(created, int) else int(time.time()),
|
||||
model if isinstance(model, str) and model else exc.model,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -28,12 +28,16 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
|
|||
- text: str
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast
|
||||
|
||||
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
|
||||
from openai.types.responses.tool_param import FunctionToolParam
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -41,17 +45,33 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
|
|||
OpenAiResponsesToChatCompletionStreamIterator,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_responses_stream_usage,
|
||||
stream_item_field,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
BaseLiteLLMOpenAIResponseObject,
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolParam,
|
||||
ContentPartAddedEvent,
|
||||
ContentPartDoneEvent,
|
||||
ContentPartDonePartOutputText,
|
||||
ErrorEvent,
|
||||
ErrorEventError,
|
||||
OpenAIMcpServerTool,
|
||||
OutputItemAddedEvent,
|
||||
OutputItemDoneEvent,
|
||||
OutputTextDeltaEvent,
|
||||
OutputTextDoneEvent,
|
||||
ResponseAPIUsage,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
ResponsesAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.responses.main import (
|
||||
GenericResponseOutputItem,
|
||||
|
|
@ -63,11 +83,13 @@ from litellm.types.utils import GenericGuardrailAPIInputs
|
|||
if TYPE_CHECKING:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
ModifyResponseException,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.llms.openai import ResponseInputParam
|
||||
from litellm.types.utils import ResponsesAPIResponse
|
||||
|
||||
|
||||
class ResponseOutputEnvelope(TypedDict, total=False):
|
||||
|
|
@ -82,6 +104,10 @@ class ResponsesStreamChunk(TypedDict, total=False):
|
|||
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
delta: ReadOnly[str]
|
||||
item_id: ReadOnly[str]
|
||||
output_index: ReadOnly[int]
|
||||
content_index: ReadOnly[int]
|
||||
|
||||
|
||||
def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
|
||||
|
|
@ -658,8 +684,32 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
|
||||
"""
|
||||
Get the string so far from the responses so far.
|
||||
|
||||
``response.output_text.done`` events carry the whole part in ``text``, while
|
||||
``response.output_text.delta`` events carry fragments in ``delta``. A stream
|
||||
that dies before its done event (``response.failed`` / ``response.incomplete``)
|
||||
has text only in deltas, so per content part the done text wins when present
|
||||
and the joined deltas fill in otherwise, never both.
|
||||
"""
|
||||
return "".join([response.get("text", "") for response in responses_so_far])
|
||||
keyed_events: Final = tuple(
|
||||
(
|
||||
(event.get("item_id"), event.get("output_index"), event.get("content_index")),
|
||||
event.get("text"),
|
||||
event.get("delta"),
|
||||
)
|
||||
for event in responses_so_far
|
||||
if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str)
|
||||
)
|
||||
|
||||
def part_text(part_key: tuple[object, object, object]) -> str:
|
||||
done_texts: Final = tuple(
|
||||
text for key, text, _ in keyed_events if key == part_key and isinstance(text, str)
|
||||
)
|
||||
if done_texts:
|
||||
return done_texts[-1]
|
||||
return "".join(delta for key, _, delta in keyed_events if key == part_key and isinstance(delta, str))
|
||||
|
||||
return "".join(part_text(key) for key in dict.fromkeys(key for key, _, _ in keyed_events))
|
||||
|
||||
def _has_text_content(self, response: "ResponsesAPIResponse") -> bool:
|
||||
"""
|
||||
|
|
@ -837,3 +887,331 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
content[content_idx]["text"] = guardrail_response
|
||||
elif hasattr(content[content_idx], "text"):
|
||||
content[content_idx].text = guardrail_response
|
||||
|
||||
def build_block_sse_chunks(
|
||||
self,
|
||||
exc: "ModifyResponseException",
|
||||
stream_started: bool = False,
|
||||
responses_so_far: Sequence[object] | None = None,
|
||||
) -> Sequence[bytes]:
|
||||
"""
|
||||
Build Responses API SSE events that deliver the guardrail block message
|
||||
and terminate the stream cleanly, mirroring the non-streaming block
|
||||
response: a completed response whose only output is the violation text,
|
||||
with the real usage the upstream call consumed.
|
||||
|
||||
- ``stream_started`` False (buffered / pre-stream): nothing has been
|
||||
sent, so emit the full synthetic sequence (``response.created``
|
||||
through ``response.completed``).
|
||||
- ``stream_started`` True (sampling / mid-stream): events already
|
||||
reached the client, so continue the in-progress response: close the
|
||||
output item still open on the wire, deliver the block message as a
|
||||
new output item under the same response id, and close with a
|
||||
``response.completed`` carrying only the replacement item.
|
||||
|
||||
The proxy's data generator appends ``data: [DONE]`` itself.
|
||||
"""
|
||||
events: Final = (
|
||||
self._block_continuation_events(exc, responses_so_far or ())
|
||||
if stream_started
|
||||
else self._standalone_block_events(exc)
|
||||
)
|
||||
return tuple(
|
||||
f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True, serialize_as_any=True)}\n\n".encode()
|
||||
for event in events
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _standalone_block_events(exc: "ModifyResponseException") -> Sequence[ResponsesAPIStreamingResponse]:
|
||||
from litellm.responses.streaming_iterator import build_synthetic_response_events
|
||||
|
||||
return build_synthetic_response_events(
|
||||
transformed=_blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model),
|
||||
logging_obj=None,
|
||||
chunk_size=max(len(exc.message), 1),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _block_continuation_events(
|
||||
exc: "ModifyResponseException", responses_so_far: Sequence[object]
|
||||
) -> Sequence[ResponsesAPIStreamingResponse]:
|
||||
response_id, model, output_index = _continuation_identity(exc, responses_so_far)
|
||||
item: Final = _blocked_output_item(exc)
|
||||
item_id: Final = item.id
|
||||
part: Final[_BlockedContentPart] = {"type": "output_text", "text": exc.message, "annotations": ()}
|
||||
done_part: Final[_BlockedDoneContentPart] = {
|
||||
"type": "output_text",
|
||||
"text": exc.message,
|
||||
"annotations": (),
|
||||
"logprobs": None,
|
||||
}
|
||||
return (
|
||||
*_open_item_closing_events(responses_so_far),
|
||||
OutputItemAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
|
||||
output_index=output_index,
|
||||
item=item,
|
||||
),
|
||||
ContentPartAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=0,
|
||||
part=BaseLiteLLMOpenAIResponseObject.model_validate(part),
|
||||
),
|
||||
OutputTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=0,
|
||||
delta=exc.message,
|
||||
),
|
||||
OutputTextDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=0,
|
||||
text=exc.message,
|
||||
),
|
||||
ContentPartDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
|
||||
item_id=item_id,
|
||||
output_index=output_index,
|
||||
content_index=0,
|
||||
part=ContentPartDonePartOutputText.model_validate(done_part),
|
||||
),
|
||||
OutputItemDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=output_index,
|
||||
item=item,
|
||||
),
|
||||
ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=_blocked_response(exc, response_id=response_id, model=model, output_item=item),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _BlockedContentPart(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
annotations: ReadOnly[tuple[object, ...]]
|
||||
|
||||
|
||||
class _BlockedDoneContentPart(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
annotations: ReadOnly[tuple[object, ...]]
|
||||
logprobs: ReadOnly[None]
|
||||
|
||||
|
||||
class _BlockedItemPayload(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
id: ReadOnly[str]
|
||||
status: ReadOnly[str]
|
||||
role: ReadOnly[str]
|
||||
content: ReadOnly[tuple[_BlockedContentPart, ...]]
|
||||
|
||||
|
||||
class _BlockedResponsePayload(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
object: ReadOnly[str]
|
||||
created_at: ReadOnly[int]
|
||||
model: ReadOnly[str]
|
||||
output: ReadOnly[tuple[GenericResponseOutputItem, ...]]
|
||||
status: ReadOnly[str]
|
||||
usage: ReadOnly[ResponseAPIUsage]
|
||||
|
||||
|
||||
def _blocked_output_item(exc: "ModifyResponseException") -> GenericResponseOutputItem:
|
||||
payload: Final[_BlockedItemPayload] = {
|
||||
"type": "message",
|
||||
"id": f"msg_{uuid.uuid4()}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": ({"type": "output_text", "text": exc.message, "annotations": ()},),
|
||||
}
|
||||
return GenericResponseOutputItem.model_validate(payload)
|
||||
|
||||
|
||||
def _blocked_response(
|
||||
exc: "ModifyResponseException",
|
||||
response_id: str,
|
||||
model: str,
|
||||
output_item: GenericResponseOutputItem | None = None,
|
||||
) -> ResponsesAPIResponse:
|
||||
payload: Final[_BlockedResponsePayload] = {
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"created_at": int(time.time()),
|
||||
"model": model,
|
||||
"output": (output_item if output_item is not None else _blocked_output_item(exc),),
|
||||
"status": "completed",
|
||||
"usage": blocked_responses_stream_usage(exc.original_response),
|
||||
}
|
||||
return ResponsesAPIResponse.model_validate(payload)
|
||||
|
||||
|
||||
def _continuation_identity(exc: "ModifyResponseException", responses_so_far: Sequence[object]) -> tuple[str, str, int]:
|
||||
responses: Final = tuple(
|
||||
response for item in responses_so_far if (response := stream_item_field(item, "response")) is not None
|
||||
)
|
||||
response_id: Final = next(
|
||||
(rid for response in responses if isinstance(rid := stream_item_field(response, "id"), str) and rid),
|
||||
f"resp_{uuid.uuid4()}",
|
||||
)
|
||||
model: Final = next(
|
||||
(m for response in responses if isinstance(m := stream_item_field(response, "model"), str) and m),
|
||||
exc.model,
|
||||
)
|
||||
indices: Final = tuple(
|
||||
index for item in responses_so_far if isinstance(index := stream_item_field(item, "output_index"), int)
|
||||
)
|
||||
return response_id, model, max(indices) + 1 if indices else 0
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OpenItemState:
|
||||
item_id: str
|
||||
item_type: str
|
||||
role: str
|
||||
output_index: int
|
||||
content_index: int
|
||||
text: str
|
||||
part_open: bool
|
||||
payload: object
|
||||
|
||||
|
||||
def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | None:
|
||||
typed: Final = tuple((stream_item_field(event, "type"), event) for event in responses_so_far)
|
||||
added: Final = tuple(
|
||||
(added_index, stream_item_field(event, "item"))
|
||||
for event_type, event in typed
|
||||
if event_type == "response.output_item.added"
|
||||
and isinstance(added_index := stream_item_field(event, "output_index"), int)
|
||||
)
|
||||
done_indices: Final = frozenset(
|
||||
done_index
|
||||
for event_type, event in typed
|
||||
if event_type == "response.output_item.done"
|
||||
and isinstance(done_index := stream_item_field(event, "output_index"), int)
|
||||
)
|
||||
open_added: Final = tuple((index, payload) for index, payload in added if index not in done_indices)
|
||||
if not open_added:
|
||||
return None
|
||||
output_index, item_payload = open_added[-1]
|
||||
if item_payload is None:
|
||||
return None
|
||||
item_id: Final = stream_item_field(item_payload, "id")
|
||||
if not isinstance(item_id, str) or not item_id:
|
||||
return None
|
||||
raw_type: Final = stream_item_field(item_payload, "type")
|
||||
raw_role: Final = stream_item_field(item_payload, "role")
|
||||
part_added: Final = tuple(
|
||||
part_index
|
||||
for event_type, event in typed
|
||||
if event_type == "response.content_part.added"
|
||||
and stream_item_field(event, "item_id") == item_id
|
||||
and isinstance(part_index := stream_item_field(event, "content_index"), int)
|
||||
)
|
||||
part_done: Final = frozenset(
|
||||
part_done_index
|
||||
for event_type, event in typed
|
||||
if event_type == "response.content_part.done"
|
||||
and stream_item_field(event, "item_id") == item_id
|
||||
and isinstance(part_done_index := stream_item_field(event, "content_index"), int)
|
||||
)
|
||||
open_parts: Final = tuple(index for index in part_added if index not in part_done)
|
||||
text: Final = "".join(
|
||||
delta
|
||||
for event_type, event in typed
|
||||
if event_type == "response.output_text.delta"
|
||||
and stream_item_field(event, "item_id") == item_id
|
||||
and isinstance(delta := stream_item_field(event, "delta"), str)
|
||||
)
|
||||
return _OpenItemState(
|
||||
item_id=item_id,
|
||||
item_type=raw_type if isinstance(raw_type, str) and raw_type else "message",
|
||||
role=raw_role if isinstance(raw_role, str) and raw_role else "assistant",
|
||||
output_index=output_index,
|
||||
content_index=open_parts[-1] if open_parts else 0,
|
||||
text=text,
|
||||
part_open=bool(open_parts),
|
||||
payload=item_payload,
|
||||
)
|
||||
|
||||
|
||||
_item_fields_adapter: Final = TypeAdapter(Mapping[str, object])
|
||||
_no_item_fields: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _incomplete_item_fields(payload: object) -> Mapping[str, object]:
|
||||
raw: Final = payload.model_dump() if isinstance(payload, BaseModel) else payload
|
||||
if not isinstance(raw, dict):
|
||||
return _no_item_fields
|
||||
return _item_fields_adapter.validate_python(raw)
|
||||
|
||||
|
||||
def _open_item_closing_events(responses_so_far: Sequence[object]) -> Sequence[ResponsesAPIStreamingResponse]:
|
||||
"""Close the output item still in progress on the relayed stream before the
|
||||
block item is appended: strict Responses clients reject a
|
||||
``response.completed`` that arrives while an earlier ``output_item.added``
|
||||
was never closed. A message item closes ``completed`` with exactly the text
|
||||
the client has received so far; any other item type (a function call the
|
||||
guardrail rejected, for instance) closes ``incomplete`` so the synthetic
|
||||
done event can never authorize acting on it."""
|
||||
open_item: Final = _open_item_state(responses_so_far)
|
||||
if open_item is None:
|
||||
return ()
|
||||
if open_item.item_type != "message":
|
||||
return (
|
||||
OutputItemDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=open_item.output_index,
|
||||
item=BaseLiteLLMOpenAIResponseObject.model_validate(
|
||||
MappingProxyType({**_incomplete_item_fields(open_item.payload), "status": "incomplete"})
|
||||
),
|
||||
),
|
||||
)
|
||||
partial_part: Final[_BlockedContentPart] = {
|
||||
"type": "output_text",
|
||||
"text": open_item.text,
|
||||
"annotations": (),
|
||||
}
|
||||
closed_payload: Final[_BlockedItemPayload] = {
|
||||
"type": open_item.item_type,
|
||||
"id": open_item.item_id,
|
||||
"status": "completed",
|
||||
"role": open_item.role,
|
||||
"content": (partial_part,),
|
||||
}
|
||||
item_done: Final = OutputItemDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=open_item.output_index,
|
||||
item=GenericResponseOutputItem.model_validate(closed_payload),
|
||||
)
|
||||
if not open_item.part_open:
|
||||
return (item_done,)
|
||||
partial_done_part: Final[_BlockedDoneContentPart] = {
|
||||
"type": "output_text",
|
||||
"text": open_item.text,
|
||||
"annotations": (),
|
||||
"logprobs": None,
|
||||
}
|
||||
return (
|
||||
OutputTextDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
|
||||
item_id=open_item.item_id,
|
||||
output_index=open_item.output_index,
|
||||
content_index=open_item.content_index,
|
||||
text=open_item.text,
|
||||
),
|
||||
ContentPartDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
|
||||
item_id=open_item.item_id,
|
||||
output_index=open_item.output_index,
|
||||
content_index=open_item.content_index,
|
||||
part=ContentPartDonePartOutputText.model_validate(partial_done_part),
|
||||
),
|
||||
item_done,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
|
|||
)
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import *
|
||||
from litellm.types.responses.main import *
|
||||
|
|
@ -36,6 +37,7 @@ else:
|
|||
_NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4")
|
||||
_PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
|
||||
_PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
|
||||
|
||||
|
||||
class _DeleteResponseBody(TypedDict):
|
||||
|
|
@ -210,8 +212,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
)
|
||||
if sanitized_tools is not None:
|
||||
response_api_optional_request_params["tools"] = sanitized_tools
|
||||
replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input)
|
||||
final_request_params: Final = dict(
|
||||
ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params)
|
||||
ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params)
|
||||
)
|
||||
|
||||
return final_request_params
|
||||
|
|
@ -248,6 +251,23 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
|
||||
return input, tools
|
||||
|
||||
def _drop_foreign_tool_call_item_ids(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
|
||||
if self.custom_llm_provider not in _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS or not isinstance(input, list):
|
||||
return input
|
||||
sanitized_items: Final = [self._without_foreign_tool_call_item_id(item) for item in input]
|
||||
return cast("ResponseInputParam", sanitized_items) # cast-ok: items keep their shape, minus a rejected id
|
||||
|
||||
@staticmethod
|
||||
def _without_foreign_tool_call_item_id(item: object) -> object:
|
||||
if not isinstance(item, dict):
|
||||
return item
|
||||
item_type: Final = item.get("type")
|
||||
item_id: Final = item.get("id")
|
||||
genuine_prefix: Final = TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE.get(item_type) if isinstance(item_type, str) else None
|
||||
if genuine_prefix is None or not isinstance(item_id, str) or item_id.startswith(genuine_prefix):
|
||||
return item
|
||||
return {key: value for key, value in item.items() if key != "id"} # mutable-ok: outgoing JSON request item
|
||||
|
||||
def _flatten_tool_schema_combinators_for_openai(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -773,7 +793,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
)
|
||||
if sanitized_tools is not None:
|
||||
response_api_optional_request_params["tools"] = sanitized_tools
|
||||
data: Final = dict(ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params))
|
||||
replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input)
|
||||
data: Final = dict(
|
||||
ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params)
|
||||
)
|
||||
|
||||
return url, data
|
||||
|
||||
|
|
|
|||
90
litellm/llms/parallel_ai/search/cost_calculator.py
Normal file
90
litellm/llms/parallel_ai/search/cost_calculator.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
PARALLEL_AI_DEFAULT_RESULTS: Final = 10
|
||||
PARALLEL_AI_ADDITIONAL_RESULT_COST: Final = 0.001
|
||||
PARALLEL_AI_USAGE_PARAM: Final = "_parallel_ai_usage"
|
||||
PARALLEL_AI_STANDARD_SEARCH_MODEL: Final = "parallel_ai/search"
|
||||
PARALLEL_AI_FAST_SEARCH_MODEL: Final = "parallel_ai/search-fast"
|
||||
PARALLEL_AI_TURBO_SEARCH_MODEL: Final = "parallel_ai/search-turbo"
|
||||
PARALLEL_AI_PRICING_MODEL_BY_MODE: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"fast": PARALLEL_AI_FAST_SEARCH_MODEL,
|
||||
"turbo": PARALLEL_AI_TURBO_SEARCH_MODEL,
|
||||
}
|
||||
)
|
||||
ADVANCED_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _non_negative_int(value: object) -> int | None:
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def _usage_count(usage: Sequence[Mapping[str, object]], sku: str) -> int | None:
|
||||
counts: Final = tuple(
|
||||
count
|
||||
for item in usage
|
||||
if item.get("name") == sku
|
||||
if (count := _non_negative_int(item.get("count"))) is not None
|
||||
)
|
||||
return sum(counts) if counts else None
|
||||
|
||||
|
||||
def _effective_mode(optional_params: Mapping[str, object]) -> str:
|
||||
mode: Final = optional_params.get("mode")
|
||||
if isinstance(mode, str):
|
||||
return mode
|
||||
|
||||
processor: Final = optional_params.get("processor")
|
||||
if processor == "pro":
|
||||
return "advanced"
|
||||
return "basic"
|
||||
|
||||
|
||||
def _effective_max_results(optional_params: Mapping[str, object]) -> int:
|
||||
try:
|
||||
advanced_settings: Final = ADVANCED_SETTINGS_ADAPTER.validate_python(optional_params.get("advanced_settings"))
|
||||
advanced_max_results: Final = _non_negative_int(advanced_settings.get("max_results"))
|
||||
if advanced_max_results is not None:
|
||||
return advanced_max_results
|
||||
except ValidationError:
|
||||
pass
|
||||
|
||||
max_results: Final = _non_negative_int(optional_params.get("max_results"))
|
||||
return max_results if max_results is not None else PARALLEL_AI_DEFAULT_RESULTS
|
||||
|
||||
|
||||
def _request_cost(mode: str) -> float:
|
||||
pricing_model: Final = PARALLEL_AI_PRICING_MODEL_BY_MODE.get(mode, PARALLEL_AI_STANDARD_SEARCH_MODEL)
|
||||
model_info: Final = get_model_info(model=pricing_model, custom_llm_provider="parallel_ai")
|
||||
return float(model_info.get("input_cost_per_query") or 0.0)
|
||||
|
||||
|
||||
def _additional_results(
|
||||
optional_params: Mapping[str, object],
|
||||
usage: Sequence[Mapping[str, object]] | None,
|
||||
) -> int:
|
||||
usage_count: Final = _usage_count(usage, "sku_search_additional_results") if usage is not None else None
|
||||
if usage_count is not None:
|
||||
return usage_count
|
||||
if usage is not None:
|
||||
return 0
|
||||
return max(_effective_max_results(optional_params) - PARALLEL_AI_DEFAULT_RESULTS, 0)
|
||||
|
||||
|
||||
def parallel_ai_search_cost(
|
||||
optional_params: Mapping[str, object],
|
||||
usage: Sequence[Mapping[str, object]] | None,
|
||||
) -> float:
|
||||
request_cost: Final = _request_cost(_effective_mode(optional_params))
|
||||
request_count_from_usage: Final = _usage_count(usage, "sku_search") if usage is not None else None
|
||||
request_count: Final = request_count_from_usage if request_count_from_usage is not None else 1
|
||||
additional_results: Final = _additional_results(optional_params, usage)
|
||||
return request_count * request_cost + additional_results * PARALLEL_AI_ADDITIONAL_RESULT_COST
|
||||
|
|
@ -4,9 +4,13 @@ Calls Parallel AI's /v1/search endpoint to search the web.
|
|||
Parallel AI API Reference: https://docs.parallel.ai/api-reference/search/search
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypedDict
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.search.transformation import (
|
||||
|
|
@ -14,9 +18,29 @@ from litellm.llms.base_llm.search.transformation import (
|
|||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from litellm.llms.parallel_ai.search.cost_calculator import PARALLEL_AI_USAGE_PARAM
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
|
||||
class _ParallelAIV1SearchResult(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
url: str | None = None
|
||||
title: str | None = None
|
||||
publish_date: str | None = None
|
||||
excerpts: Sequence[str] | None = None
|
||||
|
||||
|
||||
class _ParallelAIV1SearchResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
search_id: str | None = None
|
||||
session_id: str | None = None
|
||||
results: Sequence[_ParallelAIV1SearchResult] = ()
|
||||
usage: Sequence[Mapping[str, object]] | None = None
|
||||
warnings: Sequence[Mapping[str, object]] | None = None
|
||||
|
||||
|
||||
class _ParallelAISourcePolicy(TypedDict, total=False):
|
||||
include_domains: list[str]
|
||||
exclude_domains: list[str]
|
||||
|
|
@ -27,10 +51,16 @@ class _ParallelAIExcerptSettings(TypedDict, total=False):
|
|||
max_chars_per_result: int
|
||||
|
||||
|
||||
class _ParallelAIFetchPolicy(TypedDict, total=False):
|
||||
max_age_seconds: ReadOnly[int]
|
||||
timeout_seconds: ReadOnly[float]
|
||||
disable_cache_fallback: ReadOnly[bool]
|
||||
|
||||
|
||||
class _ParallelAIAdvancedSettings(TypedDict, total=False):
|
||||
source_policy: _ParallelAISourcePolicy
|
||||
excerpt_settings: _ParallelAIExcerptSettings
|
||||
fetch_policy: dict
|
||||
fetch_policy: _ParallelAIFetchPolicy
|
||||
location: str
|
||||
max_results: int
|
||||
|
||||
|
|
@ -43,14 +73,14 @@ class ParallelAISearchRequest(TypedDict, total=False):
|
|||
|
||||
search_queries: list[str] # Required - at least one keyword search query
|
||||
objective: str # Optional - natural-language description of search goal
|
||||
mode: str # Optional - 'turbo', 'basic', or 'advanced' (default 'advanced')
|
||||
mode: str # Optional - 'turbo', 'fast', 'basic', or 'advanced' (default 'advanced')
|
||||
max_chars_total: int # Optional - upper bound on total excerpt characters
|
||||
session_id: str # Optional - tracks calls across search/extract requests
|
||||
client_model: str # Optional - model consuming the results
|
||||
advanced_settings: _ParallelAIAdvancedSettings
|
||||
|
||||
|
||||
LEGACY_PROCESSOR_TO_MODE: Final = {"base": "basic", "pro": "advanced"}
|
||||
LEGACY_PROCESSOR_TO_MODE: Final = MappingProxyType({"base": "basic", "pro": "advanced"})
|
||||
|
||||
|
||||
class ParallelAISearchConfig(BaseSearchConfig):
|
||||
|
|
@ -67,16 +97,16 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
api_base: str | None = None,
|
||||
**kwargs,
|
||||
) -> dict:
|
||||
api_key = self.resolve_server_api_key(
|
||||
resolved_api_key: Final = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"),
|
||||
base_env_var="PARALLEL_AI_API_BASE",
|
||||
default_api_base=self.PARALLEL_AI_API_BASE,
|
||||
)
|
||||
if not api_key:
|
||||
if not resolved_api_key:
|
||||
raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.")
|
||||
headers["x-api-key"] = api_key
|
||||
headers["x-api-key"] = resolved_api_key
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
|
||||
|
|
@ -87,13 +117,12 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
data: dict | list[dict] | None = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
api_base = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE
|
||||
resolved_api_base: Final = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE
|
||||
|
||||
api_base = api_base.rstrip("/")
|
||||
if not api_base.endswith("/v1/search"):
|
||||
api_base = f"{api_base.removesuffix('/v1')}/v1/search"
|
||||
|
||||
return api_base
|
||||
trimmed: Final = resolved_api_base.rstrip("/")
|
||||
if trimmed.endswith("/v1/search"):
|
||||
return trimmed
|
||||
return f"{trimmed.removesuffix('/v1')}/v1/search"
|
||||
|
||||
def transform_search_request(
|
||||
self,
|
||||
|
|
@ -109,14 +138,17 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
- If string: maps to `search_queries` (single item) and `objective`
|
||||
- If list: maps to `search_queries` (keyword queries)
|
||||
optional_params: Optional parameters for the request
|
||||
- mode: Search mode ('turbo', 'basic', 'advanced'); defaults to 'basic'
|
||||
- mode: Search mode ('turbo', 'fast', 'basic', 'advanced'); defaults to 'basic'
|
||||
- processor: Legacy v1beta param; 'base' maps to mode 'basic', 'pro' to 'advanced'
|
||||
- max_results: Maximum number of search results -> `advanced_settings.max_results`
|
||||
- search_domain_filter: Domains to include -> `advanced_settings.source_policy.include_domains`
|
||||
- search_domain_filter / include_domains: Domains to include -> `advanced_settings.source_policy.include_domains`
|
||||
- exclude_domains: Domains to exclude -> `advanced_settings.source_policy.exclude_domains`
|
||||
- country: ISO 3166-1 alpha-2 code -> `advanced_settings.location`
|
||||
- after_date: RFC 3339 date (YYYY-MM-DD) -> `advanced_settings.source_policy.after_date`
|
||||
- country / location: ISO 3166-1 alpha-2 code -> `advanced_settings.location`
|
||||
- max_chars_per_result: -> `advanced_settings.excerpt_settings.max_chars_per_result`
|
||||
- Any other params are passed through to the request body as-is
|
||||
- fetch_policy: Cache vs live-fetch policy -> `advanced_settings.fetch_policy`
|
||||
- Any other params (objective, max_chars_total, session_id, client_model, ...)
|
||||
are passed through to the request body as-is
|
||||
|
||||
Returns:
|
||||
Dict with request data following the v1 search request spec
|
||||
|
|
@ -137,7 +169,7 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
mode = LEGACY_PROCESSOR_TO_MODE.get(processor, processor)
|
||||
# the v1 API defaults to 'advanced' when mode is omitted; default to 'basic'
|
||||
# instead to keep v1beta's default tier (processor 'base') and litellm's
|
||||
# $0.004/query cost map entry for `parallel_ai/search` accurate
|
||||
# cost map entry for `parallel_ai/search` accurate
|
||||
request_data["mode"] = mode or "basic"
|
||||
|
||||
advanced_settings: Final[_ParallelAIAdvancedSettings] = {}
|
||||
|
|
@ -148,17 +180,29 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
if "country" in params:
|
||||
advanced_settings["location"] = params.pop("country")
|
||||
|
||||
if "location" in params:
|
||||
advanced_settings["location"] = params.pop("location")
|
||||
|
||||
if "max_chars_per_result" in params:
|
||||
advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")}
|
||||
|
||||
if "fetch_policy" in params:
|
||||
advanced_settings["fetch_policy"] = params.pop("fetch_policy")
|
||||
|
||||
source_policy: Final[_ParallelAISourcePolicy] = {}
|
||||
|
||||
if "search_domain_filter" in params:
|
||||
source_policy["include_domains"] = params.pop("search_domain_filter")
|
||||
|
||||
if "include_domains" in params:
|
||||
source_policy["include_domains"] = params.pop("include_domains")
|
||||
|
||||
if "exclude_domains" in params:
|
||||
source_policy["exclude_domains"] = params.pop("exclude_domains")
|
||||
|
||||
if "after_date" in params:
|
||||
source_policy["after_date"] = params.pop("after_date")
|
||||
|
||||
if source_policy:
|
||||
advanced_settings["source_policy"] = source_policy
|
||||
|
||||
|
|
@ -170,9 +214,11 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
# unified-spec param with no v1 equivalent
|
||||
params.pop("max_tokens_per_page", None)
|
||||
|
||||
result_data: Final[dict] = dict(request_data)
|
||||
result_data.update(params)
|
||||
return result_data
|
||||
# reserved for the provider's own reported usage, which prices the request;
|
||||
# a caller-supplied value would otherwise set its own cost
|
||||
params.pop(PARALLEL_AI_USAGE_PARAM, None)
|
||||
|
||||
return {**request_data, **params}
|
||||
|
||||
def transform_search_response(
|
||||
self,
|
||||
|
|
@ -186,26 +232,49 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
Parallel AI -> LiteLLM mappings:
|
||||
- results[].title -> SearchResult.title
|
||||
- results[].url -> SearchResult.url
|
||||
- results[].excerpts (array) -> SearchResult.snippet (joined string)
|
||||
- results[].excerpts (array) -> SearchResult.snippet (joined string); the raw
|
||||
array is preserved as an extra `excerpts` field on each result
|
||||
- results[].publish_date -> SearchResult.date
|
||||
- search_id / session_id / warnings are preserved as extra fields on the
|
||||
response; usage is preserved as `parallel_usage` (the `usage` name is
|
||||
reserved for LiteLLM's token-usage object)
|
||||
"""
|
||||
response_json: Final = raw_response.json()
|
||||
parsed: Final = _ParallelAIV1SearchResponse.model_validate(raw_response.json())
|
||||
|
||||
results: Final = []
|
||||
for result in response_json.get("results", []):
|
||||
excerpts = result.get("excerpts") or []
|
||||
snippet = " ... ".join(excerpts) if excerpts else ""
|
||||
# written unconditionally: leaving a caller-supplied value in place when the
|
||||
# provider reports no usage would let the caller price its own request
|
||||
logging_obj.optional_params = {
|
||||
**logging_obj.optional_params,
|
||||
PARALLEL_AI_USAGE_PARAM: parsed.usage,
|
||||
}
|
||||
|
||||
search_result = SearchResult(
|
||||
title=result.get("title") or "",
|
||||
url=result.get("url") or "",
|
||||
snippet=snippet,
|
||||
date=result.get("publish_date"),
|
||||
last_updated=None,
|
||||
results: Final = tuple(
|
||||
SearchResult.model_validate(
|
||||
MappingProxyType(
|
||||
{
|
||||
"title": result.title or "",
|
||||
"url": result.url or "",
|
||||
"snippet": " ... ".join(result.excerpts or ()),
|
||||
"date": result.publish_date,
|
||||
"last_updated": None,
|
||||
"excerpts": result.excerpts or (),
|
||||
}
|
||||
)
|
||||
)
|
||||
results.append(search_result)
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
object="search",
|
||||
for result in parsed.results
|
||||
)
|
||||
|
||||
extra_fields: Final = MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in (
|
||||
("search_id", parsed.search_id),
|
||||
("session_id", parsed.session_id),
|
||||
("parallel_usage", parsed.usage),
|
||||
("warnings", parsed.warnings),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
|
||||
return SearchResponse.model_validate(MappingProxyType({"results": results, "object": "search", **extra_fields}))
|
||||
|
|
|
|||
|
|
@ -153,14 +153,17 @@ class VertexAIAnthropicConfig(AnthropicConfig):
|
|||
drop_params: bool,
|
||||
) -> dict:
|
||||
"""
|
||||
Override parent method to ensure VertexAI always uses tool-based structured outputs.
|
||||
VertexAI doesn't support the output_format parameter, so we force all models
|
||||
to use the tool-based approach for structured outputs.
|
||||
Override parent method so VertexAI uses tool-based structured outputs
|
||||
unless the vertex map entry advertises native structured output
|
||||
(``output_format``, which Vertex AI Claude forwards for those models).
|
||||
"""
|
||||
# Temporarily override model name to force tool-based approach
|
||||
# This ensures Claude Sonnet 4.5 uses tools instead of output_format
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
|
||||
original_model: Final = model
|
||||
if "response_format" in non_default_params:
|
||||
native_structured_output: Final = AnthropicModelInfo._get_provider_resolved_capability(
|
||||
model, "supports_native_structured_output", "vertex_ai"
|
||||
)
|
||||
if "response_format" in non_default_params and native_structured_output is not True:
|
||||
model = "claude-3-sonnet-20240229" # Use a model that will use tool-based approach
|
||||
|
||||
# Call parent method with potentially modified model name
|
||||
|
|
|
|||
|
|
@ -1482,7 +1482,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -1557,7 +1557,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -1632,7 +1632,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -1707,7 +1707,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -3254,6 +3254,7 @@
|
|||
"supports_computer_use": true,
|
||||
"supports_forced_tool_use": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
|
|
@ -38565,12 +38566,22 @@
|
|||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
|
||||
},
|
||||
"parallel_ai/search": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-fast": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-pro": {
|
||||
"input_cost_per_query": 0.009,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-turbo": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
|
|
@ -44142,6 +44153,7 @@
|
|||
"supports_computer_use": true,
|
||||
"supports_forced_tool_use": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
|
|
@ -44212,6 +44224,7 @@
|
|||
"supports_computer_use": true,
|
||||
"supports_forced_tool_use": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ Canonical definition for ``litellm_usertable``. Re-exported from
|
|||
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import ConfigDict, Field, model_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
|
||||
from litellm.models.organization_membership import (
|
||||
|
|
@ -67,3 +67,11 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase):
|
|||
if not self.models:
|
||||
return True
|
||||
return model_name in self.models
|
||||
|
||||
|
||||
class SCIMPlaceholder(BaseModel):
|
||||
"""A user row keyed by a value that names another account by SSO identity or email."""
|
||||
|
||||
placeholder_user_id: str
|
||||
resolved_user_ids: tuple[str, ...]
|
||||
team_ids: tuple[str, ...]
|
||||
|
|
|
|||
|
|
@ -1,13 +1,16 @@
|
|||
import asyncio
|
||||
import importlib
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT
|
||||
from litellm.exceptions import (
|
||||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
|
|
@ -18,8 +21,11 @@ from litellm.proxy._experimental.mcp_server.exceptions import (
|
|||
MCPUpstreamAuthError,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
|
||||
ServerListOk,
|
||||
ServerOutcome,
|
||||
classify_list_exception,
|
||||
list_fault_http_status,
|
||||
outcome_wire_value,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
|
||||
acting_user_auth,
|
||||
|
|
@ -86,8 +92,6 @@ def _connection_error_message(exc: BaseException) -> str:
|
|||
|
||||
|
||||
if MCP_AVAILABLE:
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
|
||||
global_mcp_server_manager,
|
||||
|
|
@ -99,6 +103,7 @@ if MCP_AVAILABLE:
|
|||
ListMCPToolsRestAPIResponseObject,
|
||||
MCPInfo,
|
||||
MCPServer,
|
||||
_aggregate_server_key, # pyright: ignore[reportPrivateUsage] # same per-server key as the tools/list _meta outcomes
|
||||
_apply_toolset_scope,
|
||||
_fire_mcp_tool_call_logging,
|
||||
execute_mcp_tool,
|
||||
|
|
@ -803,9 +808,6 @@ if MCP_AVAILABLE:
|
|||
list(allowed_server_ids_set), _rest_client_ip
|
||||
)
|
||||
|
||||
list_tools_result: Final = []
|
||||
error_message = None
|
||||
|
||||
# If server_id is specified, only query that specific server
|
||||
if server_id:
|
||||
return await _list_tools_for_single_server(
|
||||
|
|
@ -849,22 +851,19 @@ if MCP_AVAILABLE:
|
|||
else {}
|
||||
)
|
||||
|
||||
# Query all servers the user has access to
|
||||
errors: Final = []
|
||||
for allowed_server_id in allowed_server_ids:
|
||||
server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id)
|
||||
if server is None:
|
||||
continue
|
||||
|
||||
server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header)
|
||||
user_oauth_extra_headers = await _get_user_oauth_extra_headers(
|
||||
async def list_server(
|
||||
server: MCPServer,
|
||||
) -> tuple[Sequence[ListMCPToolsRestAPIResponseObject], ServerOutcome]:
|
||||
server_auth_header: Final = _get_server_auth_header(
|
||||
server, mcp_server_auth_headers, mcp_auth_header
|
||||
)
|
||||
user_oauth_extra_headers: Final = await _get_user_oauth_extra_headers(
|
||||
server,
|
||||
user_api_key_dict,
|
||||
prefetched_creds=prefetched_oauth_creds,
|
||||
)
|
||||
|
||||
try:
|
||||
tools_result = await _get_tools_for_single_server(
|
||||
tools_result: Final = await _get_tools_for_single_server(
|
||||
server,
|
||||
server_auth_header,
|
||||
raw_headers_from_request,
|
||||
|
|
@ -872,24 +871,36 @@ if MCP_AVAILABLE:
|
|||
extra_headers=user_oauth_extra_headers,
|
||||
apply_tool_filters=apply_tool_filters,
|
||||
)
|
||||
list_tools_result.extend(tools_result)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error getting tools from %s: %s", server.name, e)
|
||||
errors.append(
|
||||
f"{get_server_prefix(server)}: {classify_list_exception(e).tag}"
|
||||
if isinstance(e, (MCPServerListError, MCPUpstreamAuthError))
|
||||
else f"{get_server_prefix(server)}: {e}"
|
||||
)
|
||||
continue
|
||||
return (), classify_list_exception(e)
|
||||
return tools_result, ServerListOk(tool_count=len(tools_result))
|
||||
|
||||
if errors and not list_tools_result:
|
||||
error_message = "Failed to get tools from servers: " + "; ".join(errors)
|
||||
|
||||
return {
|
||||
"tools": list_tools_result,
|
||||
"error": "partial_failure" if error_message else None,
|
||||
"message": (error_message if error_message else "Successfully retrieved tools"),
|
||||
}
|
||||
# Query all servers the user has access to
|
||||
queried_servers: Final = tuple(
|
||||
server
|
||||
for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids)
|
||||
if server is not None
|
||||
)
|
||||
listings: Final = tuple([await list_server(server) for server in queried_servers])
|
||||
list_tools_result: Final = [tool for tools, _ in listings for tool in tools]
|
||||
server_outcomes: Final = MappingProxyType(
|
||||
{_aggregate_server_key(server): outcome for server, (_, outcome) in zip(queried_servers, listings)}
|
||||
)
|
||||
errors: Final = tuple(
|
||||
f"{key}: {outcome.tag}" for key, outcome in server_outcomes.items() if outcome.tag != "ok"
|
||||
)
|
||||
error_message: Final = (
|
||||
"Failed to get tools from servers: " + "; ".join(errors)
|
||||
if errors and not list_tools_result
|
||||
else None
|
||||
)
|
||||
return {
|
||||
"tools": list_tools_result,
|
||||
"error": "partial_failure" if error_message else None,
|
||||
"message": (error_message if error_message else "Successfully retrieved tools"),
|
||||
"server_outcomes": {key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items()},
|
||||
}
|
||||
|
||||
except MCPUpstreamAuthError as e:
|
||||
# Surface upstream pass-through 401/403 challenges to the client so
|
||||
|
|
@ -1173,6 +1184,7 @@ if MCP_AVAILABLE:
|
|||
transport=request.transport,
|
||||
auth_type=request.auth_type,
|
||||
mcp_info=request.mcp_info,
|
||||
timeout=request.timeout,
|
||||
command=request.command,
|
||||
args=request.args,
|
||||
env=request.env,
|
||||
|
|
@ -1402,11 +1414,28 @@ if MCP_AVAILABLE:
|
|||
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
|
||||
|
||||
async def _list_tools_operation(client):
|
||||
async def _list_tools_session_operation(session):
|
||||
return await session.list_tools()
|
||||
|
||||
list_tools_response: Final = await client.run_with_session(_list_tools_session_operation)
|
||||
list_tools_result: Final[list[MCPTool]] = list_tools_response.tools
|
||||
# Bound the whole pagination walk: without this the preview is limited only by the
|
||||
# per-request timeout times the page cap. max() keeps the pre-pagination guarantee
|
||||
# that a single slow page within the client timeout still succeeds, and a
|
||||
# per-server timeout above the global default extends the deadline with it.
|
||||
listing_deadline: Final = max(
|
||||
getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT,
|
||||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
list_tools_result = None # rebind-ok: set inside the timeout scope below
|
||||
with anyio.move_on_after(listing_deadline):
|
||||
list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above
|
||||
if list_tools_result is None:
|
||||
verbose_logger.warning(
|
||||
"MCP tools/list preview timed out after %s seconds while paginating upstream tools",
|
||||
listing_deadline,
|
||||
)
|
||||
return { # mutable-ok: error response payload
|
||||
"status": "error",
|
||||
"error": True,
|
||||
"message": f"Timed out listing tools after {listing_deadline} seconds. "
|
||||
"The MCP server may be responding slowly or paginating excessively.",
|
||||
}
|
||||
model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result]
|
||||
return {
|
||||
"tools": model_dumped_tools,
|
||||
|
|
|
|||
|
|
@ -32002,6 +32002,62 @@
|
|||
"title": "SCIMPatchOperation",
|
||||
"type": "object"
|
||||
},
|
||||
"SCIMPlaceholder": {
|
||||
"description": "A user row keyed by a value that names another account by SSO identity or email.",
|
||||
"properties": {
|
||||
"placeholder_user_id": {
|
||||
"title": "Placeholder User Id",
|
||||
"type": "string"
|
||||
},
|
||||
"resolved_user_ids": {
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Resolved User Ids",
|
||||
"type": "array"
|
||||
},
|
||||
"team_ids": {
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Team Ids",
|
||||
"type": "array"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"placeholder_user_id",
|
||||
"resolved_user_ids",
|
||||
"team_ids"
|
||||
],
|
||||
"title": "SCIMPlaceholder",
|
||||
"type": "object"
|
||||
},
|
||||
"SCIMPlaceholderMergeResult": {
|
||||
"properties": {
|
||||
"merged_into_user_id": {
|
||||
"title": "Merged Into User Id",
|
||||
"type": "string"
|
||||
},
|
||||
"placeholder_user_id": {
|
||||
"title": "Placeholder User Id",
|
||||
"type": "string"
|
||||
},
|
||||
"team_ids": {
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Team Ids",
|
||||
"type": "array"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"placeholder_user_id",
|
||||
"merged_into_user_id",
|
||||
"team_ids"
|
||||
],
|
||||
"title": "SCIMPlaceholderMergeResult",
|
||||
"type": "object"
|
||||
},
|
||||
"SCIMServiceProviderConfig": {
|
||||
"properties": {
|
||||
"authenticationSchemes": {
|
||||
|
|
@ -33641,6 +33697,129 @@
|
|||
"scim"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/scim/v2/placeholders": {
|
||||
"get": {
|
||||
"description": "List user rows whose id is another account's SSO identity or email.\n\nAn earlier release provisioned a group member it could not match as a user keyed\nby the raw member value, and that row now shadows the account the value really\nnames, so every push of that member is refused. This lists those rows so an\noperator can fold each one into the account it shadows with\n``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of\nits own or owns virtual keys is left out: someone uses that account.",
|
||||
"operationId": "list_placeholders_scim_v2_placeholders_get",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "query",
|
||||
"name": "feature",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Feature"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/SCIMPlaceholder"
|
||||
},
|
||||
"title": "Response List Placeholders Scim V2 Placeholders Get",
|
||||
"type": "array"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "List Placeholders",
|
||||
"tags": [
|
||||
"scim"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/scim/v2/placeholders/{user_id}/merge": {
|
||||
"post": {
|
||||
"description": "Fold a placeholder user into the one account its id names by SSO identity or email.\n\nThe account is added to every team the placeholder is on, then the placeholder is\ndeleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group\npush resolves the member value to the real account. Refused with 409 when the row\nhas an SSO identity of its own, owns virtual keys, or names no account or several.",
|
||||
"operationId": "merge_placeholder_scim_v2_placeholders__user_id__merge_post",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "user_id",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "User ID",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"in": "query",
|
||||
"name": "feature",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Feature"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/SCIMPlaceholderMergeResult"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Merge Placeholder",
|
||||
"tags": [
|
||||
"scim"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1216,6 +1216,13 @@ class GenerateKeyRequest(KeyRequestBase):
|
|||
organization_id: str | None = None
|
||||
project_id: str | None = None
|
||||
|
||||
@field_validator("team_id", mode="before")
|
||||
@classmethod
|
||||
def treat_cleared_team_id_as_unset(cls, v: object) -> object:
|
||||
if v == "":
|
||||
return None
|
||||
return v
|
||||
|
||||
|
||||
class GenerateKeyResponse(KeyRequestBase):
|
||||
key: str
|
||||
|
|
@ -2436,9 +2443,22 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
database_socket_timeout: float | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"Prisma `socket_timeout` URL param (seconds). When set, an idle/slow "
|
||||
"connection that has not produced data within this window is closed. "
|
||||
"This is the main knob for capping idle DB connections from LiteLLM."
|
||||
"Prisma `socket_timeout` URL param (seconds). When set, an in-flight "
|
||||
"operation that has not produced data within this window is aborted. "
|
||||
"For capping how long idle pooled connections are kept, see "
|
||||
"`database_max_idle_connection_lifetime`."
|
||||
),
|
||||
)
|
||||
database_max_idle_connection_lifetime: float | None = Field(
|
||||
60,
|
||||
description=(
|
||||
"Prisma `max_idle_connection_lifetime` URL param (seconds). A pooled "
|
||||
"connection idle longer than this is closed and replaced instead of "
|
||||
"being handed to the next request. Defaults to 60 so connections are "
|
||||
"recycled before common infra idle timeouts (AWS NLB / RDS Proxy "
|
||||
"~350s, many LBs 60-350s) silently drop them and requests fail with "
|
||||
"`Error { kind: Closed }`. A value pinned on the DATABASE_URL or set "
|
||||
"via `database_extra_connection_params` takes precedence."
|
||||
),
|
||||
)
|
||||
database_extra_connection_params: dict[str, Any] | None = Field(
|
||||
|
|
|
|||
174
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
Normal file
174
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
"""
|
||||
Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only
|
||||
stream event carrying a model, so streamed responses report the requested model like
|
||||
non-streaming ones do.
|
||||
|
||||
Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the
|
||||
provider passthrough path) or as event dicts (fake-stream and agentic paths).
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
_MESSAGE_START_EVENT: Final = "message_start"
|
||||
_MESSAGE_START_MARKER: Final = b"message_start"
|
||||
_SSE_DATA_FIELD: Final = "data:"
|
||||
_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n")
|
||||
_MAX_HELD_BYTES: Final = 65536
|
||||
_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"')
|
||||
|
||||
_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _restamped_event(event: Mapping[str, object], requested_model: str) -> Mapping[str, object] | None:
|
||||
message: Final = event.get("message")
|
||||
if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict):
|
||||
return None
|
||||
if message.get("model") == requested_model:
|
||||
return None
|
||||
return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is
|
||||
|
||||
|
||||
def _restamped_data_line(line: str, requested_model: str) -> str | None:
|
||||
stripped: Final = line.strip()
|
||||
if not stripped.startswith(_SSE_DATA_FIELD):
|
||||
return None
|
||||
payload: Final = stripped[len(_SSE_DATA_FIELD) :].strip()
|
||||
if not payload or payload == "[DONE]":
|
||||
return None
|
||||
try:
|
||||
event: Final = _EVENT_ADAPTER.validate_json(payload)
|
||||
except ValidationError:
|
||||
return None
|
||||
restamped: Final = _restamped_event(event, requested_model)
|
||||
if restamped is None:
|
||||
return None
|
||||
terminator: Final = line[len(line.rstrip("\r\n")) :]
|
||||
return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}"
|
||||
|
||||
|
||||
def _restamped_frame(frame: str, requested_model: str) -> str | None:
|
||||
lines: Final = frame.splitlines(keepends=True)
|
||||
restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines)
|
||||
if all(line is None for line in restamped):
|
||||
return None
|
||||
return "".join(new if new is not None else old for new, old in zip(restamped, lines))
|
||||
|
||||
|
||||
def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object:
|
||||
"""
|
||||
Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``.
|
||||
|
||||
Chunks that carry no model are returned unchanged.
|
||||
"""
|
||||
if isinstance(chunk, dict):
|
||||
try:
|
||||
event: Final = _EVENT_ADAPTER.validate_python(chunk)
|
||||
except ValidationError:
|
||||
return chunk
|
||||
return _restamped_event(event, requested_model) or chunk
|
||||
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
if _MESSAGE_START_EVENT.encode() not in chunk:
|
||||
return chunk
|
||||
restamped_bytes: Final = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model)
|
||||
return chunk if restamped_bytes is None else restamped_bytes.encode("utf-8")
|
||||
|
||||
if isinstance(chunk, str):
|
||||
if _MESSAGE_START_EVENT not in chunk:
|
||||
return chunk
|
||||
restamped_text: Final = _restamped_frame(chunk, requested_model)
|
||||
return chunk if restamped_text is None else restamped_text
|
||||
|
||||
return chunk
|
||||
|
||||
|
||||
def _is_ping_frame(frame: bytes) -> bool:
|
||||
return any(marker in frame for marker in _PING_MARKERS)
|
||||
|
||||
|
||||
class AnthropicStreamModelRestamper:
|
||||
"""
|
||||
Per-stream restamper for the encoded passthrough path, where chunks are raw
|
||||
transport reads: the ``message_start`` SSE frame can arrive split across
|
||||
chunks or coalesced with later frames. Complete frames (``\\n\\n``,
|
||||
``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator
|
||||
closes them and an incomplete tail is held until it completes, so the
|
||||
restamp never misses a torn frame; ``flush`` returns whatever is still held
|
||||
when the stream ends so no bytes are swallowed. Once ``message_start`` has
|
||||
been handled, or the first real event proves the stream carries none, every
|
||||
later chunk passes through untouched.
|
||||
"""
|
||||
|
||||
def __init__(self, requested_model: str) -> None:
|
||||
self._requested_model: Final = requested_model
|
||||
self._held = b""
|
||||
self._armed = True
|
||||
|
||||
def process(self, chunk: object) -> object:
|
||||
if not self._armed:
|
||||
return chunk
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
return self._process_encoded(bytes(chunk))
|
||||
if isinstance(chunk, str):
|
||||
return self._process_encoded(chunk.encode("utf-8"))
|
||||
restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model)
|
||||
if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"):
|
||||
self._armed = False
|
||||
return restamped
|
||||
|
||||
def flush(self) -> bytes:
|
||||
held: Final = self._held
|
||||
self._held = b""
|
||||
self._armed = False
|
||||
if not held:
|
||||
return b""
|
||||
restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model)
|
||||
return restamped if isinstance(restamped, bytes) else held
|
||||
|
||||
def _process_encoded(self, data: bytes) -> bytes:
|
||||
combined: Final = self._held + data
|
||||
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined))
|
||||
if not boundaries:
|
||||
if len(combined) > _MAX_HELD_BYTES:
|
||||
self._held = b""
|
||||
self._armed = False
|
||||
return combined
|
||||
self._held = combined
|
||||
return b""
|
||||
emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]])
|
||||
tail: Final = combined[boundaries[-1] :]
|
||||
if not self._armed:
|
||||
self._held = b""
|
||||
return emitted + tail
|
||||
self._held = tail
|
||||
return emitted
|
||||
|
||||
def _restamped_closed_block(self, closed: bytes) -> bytes:
|
||||
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed))
|
||||
frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries))
|
||||
decider: Final = next(
|
||||
(
|
||||
index
|
||||
for index, frame in enumerate(frames)
|
||||
if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame))
|
||||
),
|
||||
None,
|
||||
)
|
||||
if decider is None:
|
||||
return closed
|
||||
self._armed = False
|
||||
if _MESSAGE_START_MARKER not in frames[decider]:
|
||||
return closed
|
||||
restamped_text: Final = _restamped_frame(
|
||||
frames[decider].decode("utf-8", errors="ignore"), self._requested_model
|
||||
)
|
||||
if restamped_text is None:
|
||||
return closed
|
||||
return b"".join(
|
||||
restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames)
|
||||
)
|
||||
|
|
@ -9,6 +9,7 @@ import click
|
|||
import requests
|
||||
|
||||
from .auth import context_secret_vault, get_stored_api_key, login
|
||||
from .cmd_quoting import quote_for_cmd
|
||||
|
||||
ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN"
|
||||
|
|
@ -151,31 +152,9 @@ def verify_proxy_key(
|
|||
|
||||
|
||||
_WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"})
|
||||
_CMD_PERCENT_GUARD: Final = "%%cd:~,%"
|
||||
_CMD_LINE_BREAKS: Final = ("\r", "\n")
|
||||
|
||||
|
||||
def _double_trailing_backslashes(segment: str) -> str:
|
||||
bare: Final = segment.rstrip("\\")
|
||||
return bare + "\\" * 2 * (len(segment) - len(bare))
|
||||
|
||||
|
||||
def _quote_for_cmd(token: str) -> str:
|
||||
"""Quote one token so both parsers that read it see the original text.
|
||||
|
||||
Follows the algorithm the Rust standard library settled on for batch files
|
||||
after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a
|
||||
quoted string on a lone `"` and so wants an embedded one doubled, and the
|
||||
shim's own interpreter, which re-splits `%*` under C runtime rules where a
|
||||
backslash escapes the quote that follows it, so every backslash run standing
|
||||
before a quote is doubled. Quoting cannot stop cmd expanding `%VAR%`, so each
|
||||
`%` is prefixed with `%%cd:~,`: the zero-length substring of the always
|
||||
defined `cd` expands to nothing and leaves no `%` pair for cmd to match.
|
||||
"""
|
||||
escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"'))
|
||||
return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"'
|
||||
|
||||
|
||||
def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]:
|
||||
"""Build what CreateProcess runs, routing batch shims through cmd.exe.
|
||||
|
||||
|
|
@ -202,7 +181,7 @@ def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]:
|
|||
f"Cannot pass an argument containing a line break to `{os.path.basename(path)}` on "
|
||||
"Windows: cmd.exe ends the command line there, so the agent would silently lose it."
|
||||
)
|
||||
inner: Final = " ".join(_quote_for_cmd(token) for token in (path, *rest))
|
||||
inner: Final = " ".join(quote_for_cmd(token) for token in (path, *rest))
|
||||
return f'cmd.exe /d /e:on /v:off /s /c "{inner}"'
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ live here rather than in either command module.
|
|||
|
||||
import shlex
|
||||
import shutil
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
|
@ -17,6 +18,8 @@ from pydantic import JsonValue, TypeAdapter, ValidationError
|
|||
|
||||
from litellm.litellm_core_utils.private_json import write_private_json
|
||||
|
||||
from .cmd_quoting import quote_for_cmd
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
|
|
@ -87,9 +90,12 @@ def merge_claude_settings(
|
|||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Claude Code hands the string to the system shell, `sh` on POSIX and cmd.exe
|
||||
on Windows, so every token is quoted for the shell that will read it.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
|
|
@ -106,7 +112,8 @@ def resolve_api_key_helper(base_url: str) -> str:
|
|||
raise ClaudeSettingsError(
|
||||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} --base-url {shlex.quote(base_url)} auth print-token"
|
||||
quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote
|
||||
return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token"))
|
||||
|
||||
|
||||
def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
|
||||
|
|
|
|||
26
litellm/proxy/client/cli/commands/cmd_quoting.py
Normal file
26
litellm/proxy/client/cli/commands/cmd_quoting.py
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
"""Quoting for command lines that cmd.exe reads before handing them to a program."""
|
||||
|
||||
from typing import Final
|
||||
|
||||
_CMD_PERCENT_GUARD: Final = "%%cd:~,%"
|
||||
|
||||
|
||||
def _double_trailing_backslashes(segment: str) -> str:
|
||||
bare: Final = segment.rstrip("\\")
|
||||
return bare + "\\" * 2 * (len(segment) - len(bare))
|
||||
|
||||
|
||||
def quote_for_cmd(token: str) -> str:
|
||||
"""Quote one token so both parsers that read it see the original text.
|
||||
|
||||
Follows the algorithm the Rust standard library settled on for batch files
|
||||
after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a
|
||||
quoted string on a lone `"` and so wants an embedded one doubled, and the
|
||||
program's own C runtime argv split, where a backslash escapes the quote that
|
||||
follows it, so every backslash run standing before a quote is doubled.
|
||||
Quoting cannot stop cmd expanding `%VAR%`, so each `%` is prefixed with
|
||||
`%%cd:~,`: the zero-length substring of the always defined `cd` expands to
|
||||
nothing and leaves no `%` pair for cmd to match.
|
||||
"""
|
||||
escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"'))
|
||||
return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"'
|
||||
|
|
@ -176,6 +176,9 @@ if TYPE_CHECKING:
|
|||
ProxyConfig = _ProxyConfig
|
||||
else:
|
||||
ProxyConfig = Any
|
||||
from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
|
||||
AnthropicStreamModelRestamper,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
add_litellm_data_to_request,
|
||||
refresh_proxy_server_request_body_snapshot,
|
||||
|
|
@ -2490,6 +2493,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data=self.data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
request=request,
|
||||
restamp_model=(
|
||||
None if _should_return_raw_model_name(self.data) else requested_model_from_client
|
||||
),
|
||||
)
|
||||
return await create_response(
|
||||
generator=wrap_sse_stream_with_keepalive_pings(
|
||||
|
|
@ -3442,6 +3448,16 @@ class ProxyBaseLLMRequestProcessing:
|
|||
else:
|
||||
return chunk
|
||||
|
||||
@staticmethod
|
||||
def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer:
|
||||
if restamper is None:
|
||||
return ProxyBaseLLMRequestProcessing.return_sse_chunk
|
||||
|
||||
def serialize(chunk: object) -> str:
|
||||
return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk))
|
||||
|
||||
return serialize
|
||||
|
||||
@staticmethod
|
||||
async def _finalize_streaming_generator_cleanup(
|
||||
request: Request | None,
|
||||
|
|
@ -3502,11 +3518,16 @@ class ProxyBaseLLMRequestProcessing:
|
|||
serialize_chunk: StreamChunkSerializer,
|
||||
serialize_error: StreamErrorSerializer,
|
||||
request: Request | None = None,
|
||||
flush_tail: Callable[[], bytes] | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
|
||||
cost injection, then yields chunks via serialize_chunk; on exception runs
|
||||
failure hook and yields via serialize_error. Use for SSE or NDJSON.
|
||||
|
||||
``flush_tail`` runs once after the upstream iterator completes cleanly and
|
||||
its non-empty result is yielded, so a serializer that buffers bytes across
|
||||
chunks can emit anything still held at end of stream.
|
||||
"""
|
||||
verbose_proxy_logger.debug("inside generator")
|
||||
# Resolve per-stream (not per-chunk) whether the heavy per-chunk path
|
||||
|
|
@ -3569,6 +3590,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# so it must not suppress that refund.
|
||||
delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
yield serialize_chunk(chunk)
|
||||
held_tail: Final = flush_tail() if flush_tail is not None else b""
|
||||
if held_tail:
|
||||
yield serialize_chunk(held_tail)
|
||||
stream_completed = True
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
# Client disconnected mid-stream. CancelledError / GeneratorExit
|
||||
|
|
@ -3579,8 +3603,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# billing and release exactly once. This is the outermost generator
|
||||
# Starlette closes on disconnect, so the nested iterator hook (which
|
||||
# only sees GeneratorExit on GC) cannot own the refund.
|
||||
if not stream_completed:
|
||||
client_disconnected = True
|
||||
client_disconnected = not stream_completed
|
||||
if not delivered_chunk and not _withheld_provider_output(response):
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
release_budget_reservation_on_cancel,
|
||||
|
|
@ -3634,6 +3657,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data: dict,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
request: Request | None = None,
|
||||
restamp_model: str | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Anthropic /messages and Google /generateContent streaming data generator require SSE events.
|
||||
|
|
@ -3642,17 +3666,23 @@ class ProxyBaseLLMRequestProcessing:
|
|||
SSE serializers directly (rather than re-wrapping it in another
|
||||
``async for: yield`` trampoline), so a streamed chunk traverses one
|
||||
fewer async-generator layer / coroutine resume on the hot path.
|
||||
|
||||
``restamp_model`` publishes that name on the Anthropic ``message_start``
|
||||
event in place of the provider's model, matching what the non-streaming
|
||||
response reports.
|
||||
"""
|
||||
restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None
|
||||
return ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper),
|
||||
serialize_error=lambda proxy_exc: (
|
||||
f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n"
|
||||
),
|
||||
request=request,
|
||||
flush_tail=None if restamper is None else restamper.flush,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
|
|||
|
|
@ -82,10 +82,30 @@ CONNECTION_PARAM_KEYS: Final[frozenset[str]] = frozenset(
|
|||
"pool_timeout",
|
||||
"connect_timeout",
|
||||
"socket_timeout",
|
||||
"max_idle_connection_lifetime",
|
||||
"pgbouncer",
|
||||
}
|
||||
)
|
||||
|
||||
# Quaint never tests pooled connections on checkout and keeps them idle for
|
||||
# 300s by default, past many infra idle timeouts, so dead sockets surface as
|
||||
# `Error { kind: Closed }`. 60s recycles them first; explicit values win.
|
||||
DEFAULT_MAX_IDLE_CONNECTION_LIFETIME: Final = 60
|
||||
IDLE_LIFETIME_DEFAULT_PARAMS: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{"max_idle_connection_lifetime": DEFAULT_MAX_IDLE_CONNECTION_LIFETIME}
|
||||
)
|
||||
|
||||
|
||||
def idle_lifetime_params(configured: float | None) -> Mapping[str, str | int | float]:
|
||||
"""The `max_idle_connection_lifetime` to add to URLs that do not pin one.
|
||||
|
||||
Applied via ``add_missing_query_params`` so a URL-pinned value always wins,
|
||||
whether the operator configured `database_max_idle_connection_lifetime` or not.
|
||||
"""
|
||||
if configured is None:
|
||||
return IDLE_LIFETIME_DEFAULT_PARAMS
|
||||
return MappingProxyType({"max_idle_connection_lifetime": configured})
|
||||
|
||||
|
||||
def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) -> str:
|
||||
"""Return ``url`` with the ``params`` it does not already carry appended.
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.caching import DualCache
|
|||
from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name
|
||||
|
|
@ -215,6 +216,16 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]:
|
|||
return redacted if isinstance(redacted, list) else assessments
|
||||
|
||||
|
||||
_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
|
||||
|
||||
|
||||
def _is_responses_api_route(request_route: str | None) -> bool:
|
||||
if request_route is None:
|
||||
return False
|
||||
call_types: Final = get_call_types_for_route(request_route)
|
||||
return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types)
|
||||
|
||||
|
||||
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
||||
# During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
|
||||
# OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
|
||||
|
|
@ -2709,6 +2720,24 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
yield streamed_chunk
|
||||
return
|
||||
|
||||
# Responses-API events are neither chat-completions chunks nor raw
|
||||
# Anthropic SSE, so the assembly below cannot scan them; the unified
|
||||
# guardrail's translation layer can, with buffering semantics kept.
|
||||
if _is_responses_api_route(user_api_key_dict.request_route):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
async for translated_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
guardrail_to_apply=self,
|
||||
buffer_until_moderated_default=True,
|
||||
):
|
||||
yield translated_chunk
|
||||
return
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
|
||||
from litellm.main import stream_chunk_builder
|
||||
|
|
|
|||
|
|
@ -54,7 +54,10 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
|
|||
coordination_redis_cache,
|
||||
publish_config_change,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
|
|
@ -701,6 +704,12 @@ async def patch_model(
|
|||
param="blocked",
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
litellm_params=patch_data.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
existing_litellm_params=db_model.litellm_params,
|
||||
)
|
||||
|
||||
_raise_on_strategy_router_write_violation(
|
||||
incoming_params=patch_data.litellm_params,
|
||||
existing_params=db_model.litellm_params,
|
||||
|
|
@ -1464,6 +1473,32 @@ class ModelManagementAuthChecks:
|
|||
)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def can_user_attach_credential(
|
||||
litellm_params: GenericLiteLLMParams | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
existing_litellm_params: GenericLiteLLMParams | None = None,
|
||||
) -> Literal[True]:
|
||||
if litellm_params is None or litellm_params.litellm_credential_name is None:
|
||||
return True
|
||||
if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None:
|
||||
existing_credential_name: Final = decrypt_value_helper(
|
||||
value=existing_litellm_params.litellm_credential_name,
|
||||
key="litellm_credential_name",
|
||||
exception_type="debug",
|
||||
return_original_value=True,
|
||||
)
|
||||
if litellm_params.litellm_credential_name == existing_credential_name:
|
||||
return True
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return True
|
||||
raise ProxyException(
|
||||
message=f"Only a proxy admin can attach a stored credential (litellm_credential_name) to a model. Your role={user_api_key_dict.user_role}.",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
param="litellm_credential_name",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def allow_team_model_action(
|
||||
model_params: Deployment | updateDeployment,
|
||||
|
|
@ -1786,6 +1821,11 @@ async def add_new_model(
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
litellm_params=model_params.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
_raise_on_strategy_router_write_violation(
|
||||
incoming_params=model_params.litellm_params,
|
||||
existing_params=None,
|
||||
|
|
@ -1958,6 +1998,12 @@ async def update_model(
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
litellm_params=model_params.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
existing_litellm_params=deployment.litellm_params,
|
||||
)
|
||||
|
||||
_raise_on_strategy_router_write_violation(
|
||||
incoming_params=model_params.litellm_params,
|
||||
existing_params=deployment.litellm_params,
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.models.user import SCIMPlaceholder
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
|
|
@ -585,6 +586,37 @@ async def _users_named_by_member_value(
|
|||
return tuple(dict.fromkeys(row.user_id for row in rows))
|
||||
|
||||
|
||||
async def _accounts_named_by_member_value(value: str, prisma_client: PrismaClient) -> tuple[str, ...]:
|
||||
"""Every user id this member value names, by user id, SSO identity or email.
|
||||
|
||||
Classification needs to know whether the value is one account's ``user_id`` and
|
||||
whether it names any other account, so all three fields are read in one pass. The
|
||||
id is compared exactly and unstripped, as a primary key lookup would; the
|
||||
identities compare as ``_users_named_by_member_value`` describes. Two rows are
|
||||
enough to tell one account from several, so the read stops there. Only a full
|
||||
read that lacks the row keyed by the value leaves that row's existence open, and
|
||||
only then is the id read on its own.
|
||||
"""
|
||||
subject: Final = value.strip()
|
||||
email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"}
|
||||
users: Final = _table(UserRepository(prisma_client))
|
||||
rows: Final = await users.find_many(
|
||||
where={ # mutable-ok: Prisma filter
|
||||
"OR": [ # mutable-ok: Prisma filter
|
||||
{"user_id": value}, # mutable-ok: Prisma filter
|
||||
{"sso_user_id": subject}, # mutable-ok: Prisma filter
|
||||
{"user_email": email}, # mutable-ok: Prisma filter
|
||||
],
|
||||
},
|
||||
take=2,
|
||||
)
|
||||
named: Final = tuple(dict.fromkeys(row.user_id for row in rows))
|
||||
if len(named) < 2 or value in named:
|
||||
return named
|
||||
keyed: Final = await users.find_unique(where={"user_id": value})
|
||||
return named if keyed is None else (value, *named)
|
||||
|
||||
|
||||
async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember:
|
||||
"""
|
||||
Decide what a single SCIM group member refers to.
|
||||
|
|
@ -627,11 +659,9 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient
|
|||
if member_type == "group":
|
||||
return _SkippedGroupMember(value=value, reason="nested_group")
|
||||
|
||||
user: Final = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": value})
|
||||
if user is not None:
|
||||
shared_with: Final = tuple(
|
||||
other for other in await _users_named_by_member_value(value, prisma_client) if other != value
|
||||
)
|
||||
named: Final = await _accounts_named_by_member_value(value, prisma_client)
|
||||
if value in named:
|
||||
shared_with: Final = tuple(other for other in named if other != value)
|
||||
if shared_with:
|
||||
verbose_proxy_logger.warning(
|
||||
"SCIM: group member '%s' is one account's user id and is also account '%s' by SSO identity or email, "
|
||||
|
|
@ -651,7 +681,6 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient
|
|||
if team is not None and _team_metadata_has_scim_provenance(team.metadata):
|
||||
return _SkippedGroupMember(value=value, reason="existing_team")
|
||||
|
||||
named: Final = await _users_named_by_member_value(value, prisma_client)
|
||||
if len(named) == 1:
|
||||
verbose_proxy_logger.info(
|
||||
"SCIM: group member '%s' matched user_id '%s' by SSO identity or email",
|
||||
|
|
@ -1834,6 +1863,89 @@ async def delete_user(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
@scim_router.get(
|
||||
"/placeholders",
|
||||
response_model=tuple[SCIMPlaceholder, ...],
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
)
|
||||
async def list_placeholders() -> tuple[SCIMPlaceholder, ...]:
|
||||
"""
|
||||
List user rows whose id is another account's SSO identity or email.
|
||||
|
||||
An earlier release provisioned a group member it could not match as a user keyed
|
||||
by the raw member value, and that row now shadows the account the value really
|
||||
names, so every push of that member is refused. This lists those rows so an
|
||||
operator can fold each one into the account it shadows with
|
||||
``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of
|
||||
its own or owns virtual keys is left out: someone uses that account.
|
||||
"""
|
||||
try:
|
||||
prisma_client: Final = await _get_prisma_client_or_raise_exception()
|
||||
async with prisma_client.tx() as tx:
|
||||
return await UserRepository(prisma_client).find_shadowing_placeholders(tx)
|
||||
except Exception as e:
|
||||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
def _placeholder_rejection(placeholder: LiteLLM_UserTable, resolved: tuple[str, ...], key_count: int) -> str | None:
|
||||
if placeholder.sso_user_id is not None:
|
||||
return f"User '{placeholder.user_id}' has an SSO identity of its own, so it is an account someone signs in to"
|
||||
if key_count:
|
||||
return f"User '{placeholder.user_id}' owns {key_count} virtual keys. Move or delete them before merging it"
|
||||
if not resolved:
|
||||
return f"User '{placeholder.user_id}' shadows no account: no other user has that id as SSO identity or email"
|
||||
if len(resolved) > 1:
|
||||
return (
|
||||
f"User '{placeholder.user_id}' names {len(resolved)} accounts ({', '.join(resolved)}). Resolve that first"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
@scim_router.post(
|
||||
"/placeholders/{user_id}/merge",
|
||||
response_model=SCIMPlaceholderMergeResult,
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
)
|
||||
async def merge_placeholder(
|
||||
user_id: str = Path(..., title="User ID"),
|
||||
) -> SCIMPlaceholderMergeResult:
|
||||
"""
|
||||
Fold a placeholder user into the one account its id names by SSO identity or email.
|
||||
|
||||
The account is added to every team the placeholder is on, then the placeholder is
|
||||
deleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group
|
||||
push resolves the member value to the real account. Refused with 409 when the row
|
||||
has an SSO identity of its own, owns virtual keys, or names no account or several.
|
||||
"""
|
||||
try:
|
||||
prisma_client: Final = await _get_prisma_client_or_raise_exception()
|
||||
placeholder: Final = await _check_user_exists(user_id)
|
||||
resolved: Final = tuple(
|
||||
other for other in await _users_named_by_member_value(user_id, prisma_client, take=None) if other != user_id
|
||||
)
|
||||
owned_keys: Final[_UserIdWhere] = {"user_id": user_id}
|
||||
keys: Final = await _table(VerificationTokenRepository(prisma_client)).find_many(where=owned_keys)
|
||||
rejection: Final = _placeholder_rejection(placeholder, resolved, len(keys))
|
||||
if rejection is not None:
|
||||
detail: Final[_ScimErrorDetail] = {"error": rejection}
|
||||
raise HTTPException(status_code=409, detail=detail)
|
||||
|
||||
target_user_id: Final = resolved[0]
|
||||
team_ids: Final = tuple(placeholder.teams)
|
||||
for team_id in team_ids:
|
||||
await _add_user_to_team(user_id=target_user_id, team_id=team_id)
|
||||
await delete_user(user_id=user_id)
|
||||
await _recompute_scim_member_roles(prisma_client, (target_user_id,))
|
||||
verbose_proxy_logger.info(
|
||||
"SCIM: merged placeholder user '%s' into '%s', moving teams %s", user_id, target_user_id, team_ids
|
||||
)
|
||||
return SCIMPlaceholderMergeResult(
|
||||
placeholder_user_id=user_id, merged_into_user_id=target_user_id, team_ids=team_ids
|
||||
)
|
||||
except Exception as e:
|
||||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
def _parse_member_entry(entry: object) -> SCIMMember | None:
|
||||
"""Parse one entry of a SCIM patch value, or None when it carries no id."""
|
||||
if isinstance(entry, str):
|
||||
|
|
|
|||
|
|
@ -812,6 +812,8 @@ def _resolve_team_callback_wiring(
|
|||
else { # mutable-ok: Logging arg
|
||||
**callback_vars,
|
||||
TRUSTED_CALLBACK_VARS_FIELD: callback_vars,
|
||||
"metadata": {}, # mutable-ok: Logging arg
|
||||
"model_info": {}, # mutable-ok: Logging arg
|
||||
}
|
||||
)
|
||||
return _TeamCallbackWiring(
|
||||
|
|
|
|||
|
|
@ -913,14 +913,13 @@ class ProxyInitializationHelpers:
|
|||
envvar="ENFORCE_PRISMA_MIGRATION_CHECK",
|
||||
)
|
||||
@click.option(
|
||||
"--use_v2_migration_resolver/--use_legacy_migration_resolver",
|
||||
default=True,
|
||||
"--use_v2_migration_resolver",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help=(
|
||||
"Which database migration resolver to run at startup. The default v2 "
|
||||
"resolver avoids the diff-and-force recovery path that can cause schema "
|
||||
"thrashing during rolling deploys where two LiteLLM versions contend for "
|
||||
"the same DB. Pass --use_legacy_migration_resolver, or set "
|
||||
"USE_V2_MIGRATION_RESOLVER=false, to fall back to v1."
|
||||
"Opt into the v2 migration resolver. Avoids the diff-and-force recovery "
|
||||
"path that can cause schema thrashing during rolling deploys where two "
|
||||
"LiteLLM versions contend for the same DB. Default is the v1 resolver."
|
||||
),
|
||||
envvar="USE_V2_MIGRATION_RESOLVER",
|
||||
)
|
||||
|
|
@ -1226,6 +1225,7 @@ def run_server(
|
|||
if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None:
|
||||
from litellm.proxy.db.db_url_settings import (
|
||||
add_missing_query_params,
|
||||
idle_lifetime_params,
|
||||
reader_shareable_params,
|
||||
unsupported_db_scheme,
|
||||
unsupported_db_scheme_message,
|
||||
|
|
@ -1254,6 +1254,9 @@ def run_server(
|
|||
disable_prepared_statements=db_disable_prepared_statements,
|
||||
extra_params=db_extra_connection_params,
|
||||
)
|
||||
lifetime_params: Final = idle_lifetime_params(
|
||||
general_settings.get("database_max_idle_connection_lifetime")
|
||||
)
|
||||
if os.getenv("DATABASE_URL", None) is not None:
|
||||
database_url = get_secret("DATABASE_URL", default_value=None)
|
||||
resolved_url: Final[str | None] = str(database_url) if database_url else None
|
||||
|
|
@ -1271,11 +1274,11 @@ def run_server(
|
|||
writer_url,
|
||||
connection_url_params,
|
||||
)
|
||||
os.environ["DATABASE_URL"] = modified_url
|
||||
os.environ["DATABASE_URL"] = add_missing_query_params(modified_url, lifetime_params)
|
||||
if os.getenv("DIRECT_URL", None) is not None:
|
||||
database_url = os.getenv("DIRECT_URL")
|
||||
modified_url = append_query_params(database_url, connection_url_params)
|
||||
os.environ["DIRECT_URL"] = modified_url
|
||||
os.environ["DIRECT_URL"] = add_missing_query_params(modified_url, lifetime_params)
|
||||
# The reader pool is a real pool against the same configured cap, so it
|
||||
# gets the allowlisted pool params. Schema-affecting ones, including any
|
||||
# the operator smuggled in through database_extra_connection_params, stay
|
||||
|
|
@ -1289,10 +1292,13 @@ def run_server(
|
|||
db_lock_timeout,
|
||||
)
|
||||
os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params(
|
||||
_with_query_value(read_replica_url, "options", reader_options)
|
||||
if reader_options
|
||||
else read_replica_url,
|
||||
reader_shareable_params(connection_url_params),
|
||||
add_missing_query_params(
|
||||
_with_query_value(read_replica_url, "options", reader_options)
|
||||
if reader_options
|
||||
else read_replica_url,
|
||||
reader_shareable_params(connection_url_params),
|
||||
),
|
||||
lifetime_params,
|
||||
)
|
||||
subprocess.run(["prisma"], capture_output=True)
|
||||
is_prisma_runnable = True
|
||||
|
|
@ -1311,11 +1317,10 @@ def run_server(
|
|||
else:
|
||||
if not use_v2_migration_resolver:
|
||||
print(
|
||||
"\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration resolver. "
|
||||
"The default v2 resolver is safer: it avoids the diff-and-force "
|
||||
"recovery that caused schema thrashing during rolling deploys. "
|
||||
"Remove --use_legacy_migration_resolver / "
|
||||
"USE_V2_MIGRATION_RESOLVER=false to switch back to it.\033[0m"
|
||||
"\033[1;33mLiteLLM Proxy: Using default (v1) migration resolver. "
|
||||
"If your deployment has seen schema thrashing during rolling "
|
||||
"deploys, try --use_v2_migration_resolver (safer: avoids the "
|
||||
"diff-and-force recovery that caused the thrash).\033[0m"
|
||||
)
|
||||
try:
|
||||
setup_ok: Final = PrismaManager.setup_database(
|
||||
|
|
@ -1323,10 +1328,10 @@ def run_server(
|
|||
use_v2_resolver=use_v2_migration_resolver,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
# Raised on unrecoverable migration errors: permission
|
||||
# failures from either resolver, the v2 resolver's
|
||||
# non-idempotent failures, and any `prisma db push`
|
||||
# against a partitioned LiteLLM_SpendLogs.
|
||||
# Raised on unrecoverable migration errors: the v2
|
||||
# resolver's non-idempotent failures and permission
|
||||
# issues, and any `prisma db push` against a
|
||||
# partitioned LiteLLM_SpendLogs.
|
||||
print(
|
||||
f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m",
|
||||
file=sys.stderr,
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ from typing import (
|
|||
import anyio
|
||||
import websockets
|
||||
import websockets.exceptions
|
||||
from pydantic import BaseModel, Json, JsonValue
|
||||
from pydantic import BaseModel, Json, JsonValue, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, assert_never
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -253,6 +253,7 @@ from litellm.constants import (
|
|||
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
|
||||
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
|
||||
USER_SPEND_ALERTS_JOB_ID,
|
||||
WEEKLY_SPEND_REPORT_JOB_ID,
|
||||
)
|
||||
from litellm.exceptions import RejectedRequestError
|
||||
|
|
@ -9866,6 +9867,35 @@ class ProxyStartupEvent:
|
|||
replace_existing=True,
|
||||
)
|
||||
|
||||
slack_alerting_args: Final = proxy_logging_obj.slack_alerting_instance.alerting_args
|
||||
user_spend_check_interval: Final = (
|
||||
slack_alerting_args.user_spend_check_interval
|
||||
if isinstance(slack_alerting_args, SlackAlertingArgs) # pyright: ignore[reportUnnecessaryIsInstance] # tests inject a mock slack_alerting_instance
|
||||
else SlackAlertingArgs().user_spend_check_interval
|
||||
)
|
||||
|
||||
async def _scheduled_user_spend_alerts() -> None:
|
||||
if (
|
||||
await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=USER_SPEND_ALERTS_JOB_ID,
|
||||
ttl=max(user_spend_check_interval - 60, 60),
|
||||
allow_reentrant=False,
|
||||
)
|
||||
is False
|
||||
):
|
||||
return
|
||||
await proxy_logging_obj.slack_alerting_instance.send_user_spend_alerts()
|
||||
|
||||
scheduler.add_job(
|
||||
_scheduled_user_spend_alerts,
|
||||
"interval",
|
||||
seconds=user_spend_check_interval,
|
||||
next_run_time=datetime.now(timezone.utc) + timedelta(seconds=10 + random.randint(0, 60)),
|
||||
id=USER_SPEND_ALERTS_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
if os.getenv("PROMETHEUS_URL"):
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
|
|
@ -12481,11 +12511,21 @@ async def supported_openai_params(model: str):
|
|||
--header 'Authorization: Bearer sk-1234'
|
||||
```
|
||||
"""
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
|
||||
global llm_router
|
||||
try:
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
|
||||
resolved_models: Final = llm_router.resolved_litellm_models(model) if llm_router is not None else ()
|
||||
target_model: Final = resolved_models[0] if resolved_models else model
|
||||
declared_provider: Final = declared_authenticating_provider(target_model)
|
||||
litellm_model, custom_llm_provider = (
|
||||
(target_model.removeprefix(f"{declared_provider}/"), declared_provider)
|
||||
if declared_provider is not None
|
||||
else litellm.get_llm_provider(model=target_model)[:2]
|
||||
)
|
||||
return {
|
||||
"supported_openai_params": litellm.get_supported_openai_params(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
model=litellm_model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
}
|
||||
except Exception:
|
||||
|
|
@ -14962,17 +15002,25 @@ async def alerting_settings(
|
|||
alerting_args_dict = {}
|
||||
alerting_values = None
|
||||
|
||||
allowed_args: Final = {
|
||||
"slack_alerting": {"type": "Boolean"},
|
||||
"daily_report_frequency": {"type": "Integer"},
|
||||
"report_check_interval": {"type": "Integer"},
|
||||
"budget_alert_ttl": {"type": "Integer"},
|
||||
"outage_alert_ttl": {"type": "Integer"},
|
||||
"region_outage_alert_ttl": {"type": "Integer"},
|
||||
"minor_outage_alert_threshold": {"type": "Integer"},
|
||||
"major_outage_alert_threshold": {"type": "Integer"},
|
||||
"max_outage_alert_list_size": {"type": "Integer"},
|
||||
}
|
||||
allowed_args: Final = MappingProxyType(
|
||||
{
|
||||
"slack_alerting": "Boolean",
|
||||
"daily_report_frequency": "Integer",
|
||||
"report_check_interval": "Integer",
|
||||
"budget_alert_ttl": "Integer",
|
||||
"outage_alert_ttl": "Integer",
|
||||
"region_outage_alert_ttl": "Integer",
|
||||
"minor_outage_alert_threshold": "Integer",
|
||||
"major_outage_alert_threshold": "Integer",
|
||||
"max_outage_alert_list_size": "Integer",
|
||||
"daily_spend_per_user_threshold": "Float",
|
||||
"monthly_spend_per_user_threshold": "Float",
|
||||
"spend_anomaly_multiplier": "Float",
|
||||
"spend_anomaly_baseline_days": "Integer",
|
||||
"spend_anomaly_min_spend": "Float",
|
||||
"user_spend_check_interval": "Integer",
|
||||
}
|
||||
)
|
||||
|
||||
_slack_alerting: Final[SlackAlerting] = proxy_logging_obj.slack_alerting_instance
|
||||
_slack_alerting_args_dict: Final = _slack_alerting.alerting_args.model_dump()
|
||||
|
|
@ -14987,7 +15035,7 @@ async def alerting_settings(
|
|||
|
||||
_response_obj = ConfigList(
|
||||
field_name="slack_alerting",
|
||||
field_type=allowed_args["slack_alerting"]["type"],
|
||||
field_type=allowed_args["slack_alerting"],
|
||||
field_description="Enable slack alerting for monitoring proxy in production: llm outages, budgets, spend tracking failures.",
|
||||
field_value=is_slack_enabled,
|
||||
stored_in_db=True if alerting_values is not None else False,
|
||||
|
|
@ -15006,7 +15054,7 @@ async def alerting_settings(
|
|||
|
||||
_response_obj = ConfigList(
|
||||
field_name=field_name,
|
||||
field_type=allowed_args[field_name]["type"],
|
||||
field_type=allowed_args[field_name],
|
||||
field_description=field_info.description or "",
|
||||
field_value=_slack_alerting_args_dict.get(field_name, None),
|
||||
stored_in_db=_stored_in_db,
|
||||
|
|
@ -16434,6 +16482,16 @@ async def update_config_general_settings(
|
|||
detail={"error": f"Invalid type of field value={type(data.field_value)} passed in."},
|
||||
)
|
||||
|
||||
if data.field_name == "alerting_args":
|
||||
try:
|
||||
SlackAlertingArgs.model_validate(data.field_value)
|
||||
except ValidationError as e:
|
||||
errors: Final = "; ".join(f"{'.'.join(str(loc) for loc in err['loc'])}: {err['msg']}" for err in e.errors())
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Invalid alerting_args: {errors}"},
|
||||
)
|
||||
|
||||
## get general settings from db
|
||||
db_general_settings: Final = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
|
|
|
|||
|
|
@ -120,13 +120,15 @@ async def _apply_over_budget_reservation_policy(
|
|||
applied_entries: list[dict[str, float | str]],
|
||||
reservation_cost: float,
|
||||
current_spend: float,
|
||||
fail_closed_budget_enforcement: bool = False,
|
||||
) -> float:
|
||||
"""
|
||||
Decide what to do when a counter is over budget, and return the reservation
|
||||
cost to carry into the next counter. Three outcomes: an over-budget key that
|
||||
opted into throttling releases its own reservation (the rate limiter slows
|
||||
it) and keeps the cost; a partially-remaining budget resizes the reservation
|
||||
down to what is left; anything else hard-blocks by raising.
|
||||
down to what is left, unless strict enforcement is on, because the known
|
||||
estimate already does not fit; anything else hard-blocks by raising.
|
||||
"""
|
||||
if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token):
|
||||
await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost)
|
||||
|
|
@ -134,21 +136,36 @@ async def _apply_over_budget_reservation_policy(
|
|||
return reservation_cost
|
||||
|
||||
remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost)
|
||||
if remaining_before_reservation > 1e-12:
|
||||
await _resize_applied_reservation(
|
||||
entries=applied_entries,
|
||||
current_reserved_cost=reservation_cost,
|
||||
new_reserved_cost=remaining_before_reservation,
|
||||
if remaining_before_reservation <= 1e-12:
|
||||
_raise_counter_budget_exceeded(counter=counter, current_cost=current_spend)
|
||||
if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12:
|
||||
_raise_counter_budget_exceeded(
|
||||
counter=counter,
|
||||
current_cost=current_spend - reservation_cost,
|
||||
estimated_cost=reservation_cost,
|
||||
)
|
||||
return remaining_before_reservation
|
||||
await _resize_applied_reservation(
|
||||
entries=applied_entries,
|
||||
current_reserved_cost=reservation_cost,
|
||||
new_reserved_cost=remaining_before_reservation,
|
||||
)
|
||||
return remaining_before_reservation
|
||||
|
||||
|
||||
def _raise_counter_budget_exceeded(
|
||||
counter: _BudgetCounter,
|
||||
current_cost: float,
|
||||
estimated_cost: float | None = None,
|
||||
) -> NoReturn:
|
||||
estimate_detail: Final = "" if estimated_cost is None else f"Estimated request cost: {estimated_cost}, "
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=current_spend,
|
||||
current_cost=current_cost,
|
||||
max_budget=counter.max_budget,
|
||||
message=(
|
||||
"Budget has been exceeded! "
|
||||
f"{counter.entity_type}={counter.entity_id} "
|
||||
f"Current cost: {current_spend}, "
|
||||
f"Current cost: {current_cost}, "
|
||||
f"{estimate_detail}"
|
||||
f"Max budget: {counter.max_budget}"
|
||||
),
|
||||
entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type),
|
||||
|
|
@ -258,6 +275,7 @@ async def reserve_budget_for_request(
|
|||
applied_entries=applied_entries,
|
||||
reservation_cost=reservation_cost,
|
||||
current_spend=current_spend,
|
||||
fail_closed_budget_enforcement=fail_closed_budget_enforcement,
|
||||
)
|
||||
continue
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -6,15 +6,34 @@ import json
|
|||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.models.user import LiteLLM_UserTable
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.models.user import LiteLLM_UserTable, SCIMPlaceholder
|
||||
from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
from prisma import models as prisma_models
|
||||
|
||||
_JSON_ENCODED_COLUMNS: Final = frozenset({"metadata", "model_spend", "model_max_budget"})
|
||||
|
||||
_SHADOWING_PLACEHOLDERS_SQL: Final = """
|
||||
SELECT p.user_id AS placeholder_user_id,
|
||||
array_agg(r.user_id ORDER BY r.user_id) AS resolved_user_ids,
|
||||
p.teams AS team_ids
|
||||
FROM "LiteLLM_UserTable" p
|
||||
JOIN "LiteLLM_UserTable" r
|
||||
ON r.user_id <> p.user_id
|
||||
AND (r.sso_user_id = p.user_id OR LOWER(r.user_email) = LOWER(p.user_id))
|
||||
WHERE p.sso_user_id IS NULL
|
||||
AND NOT EXISTS (SELECT 1 FROM "LiteLLM_VerificationToken" k WHERE k.user_id = p.user_id)
|
||||
GROUP BY p.user_id, p.teams
|
||||
ORDER BY p.user_id
|
||||
"""
|
||||
|
||||
_PLACEHOLDER_ROWS_ADAPTER: Final = TypeAdapter(tuple[SCIMPlaceholder, ...])
|
||||
|
||||
|
||||
class UserRepository(BaseRepository[LiteLLM_UserTable]):
|
||||
"""Repository for user database operations."""
|
||||
|
|
@ -59,6 +78,11 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
|
|||
"""Find all users in a team."""
|
||||
return await self.find_many(where={"teams": {"has": team_id}})
|
||||
|
||||
async def find_shadowing_placeholders(self, tx: "Prisma") -> tuple[SCIMPlaceholder, ...]:
|
||||
"""Users with no SSO id and no virtual keys whose id is another user's SSO id or email."""
|
||||
rows: Final = await tx.query_raw(_SHADOWING_PLACEHOLDERS_SQL)
|
||||
return _PLACEHOLDER_ROWS_ADAPTER.validate_python(rows)
|
||||
|
||||
async def count_billable_users(self) -> int:
|
||||
"""Number of users that count toward the license seat limit.
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ logic.
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
|
@ -28,6 +29,15 @@ from litellm.types.llms.openai import (
|
|||
|
||||
_MAX_ARGUMENTS_LEN: Final = 1_000_000
|
||||
|
||||
TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE: Final = MappingProxyType({"function_call": "fc", "custom_tool_call": "ctc"})
|
||||
|
||||
|
||||
def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str:
|
||||
prefix: Final = TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE.get(item_type)
|
||||
if prefix is None or not tool_id or tool_id.startswith(prefix):
|
||||
return tool_id
|
||||
return f"{prefix}_{tool_id}"
|
||||
|
||||
|
||||
def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]:
|
||||
"""Extract names of tools originally defined as ``type: "custom"``."""
|
||||
|
|
@ -103,7 +113,7 @@ def build_tool_call_item_kwargs(
|
|||
item_type: Final = "custom_tool_call" if custom else "function_call"
|
||||
kwargs: Final[dict[str, str]] = {
|
||||
"type": item_type,
|
||||
"id": call_id,
|
||||
"id": openai_shaped_tool_call_item_id(item_type, call_id),
|
||||
"call_id": call_id,
|
||||
"name": name,
|
||||
"status": status,
|
||||
|
|
|
|||
|
|
@ -114,6 +114,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = []
|
||||
self._tool_output_index_by_call_id: dict[str, int] = {}
|
||||
self._tool_args_by_call_id: dict[str, str] = {}
|
||||
self._tool_item_id_by_call_id: dict[str, str] = {} # mutable-ok: filled per call id as tool call events stream
|
||||
self._tool_call_id_by_index: dict[int, str] = {}
|
||||
self._ambiguous_tool_call_indexes: set[int] = set()
|
||||
self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item
|
||||
|
|
@ -227,6 +228,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
|
||||
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
event = OutputItemAddedEvent(
|
||||
|
|
@ -248,7 +250,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._sequence_number += 1
|
||||
delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
|
||||
item_id=call_id,
|
||||
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
|
||||
output_index=output_index,
|
||||
delta=delta_chunk,
|
||||
)
|
||||
|
|
@ -300,6 +302,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
|
||||
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
event = OutputItemAddedEvent(
|
||||
|
|
@ -325,7 +328,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._sequence_number += 1
|
||||
delta_event = FunctionCallArgumentsDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
|
||||
item_id=call_id,
|
||||
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
|
||||
output_index=output_index,
|
||||
delta=delta_chunk,
|
||||
)
|
||||
|
|
@ -335,7 +338,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._sequence_number += 1
|
||||
done_event = FunctionCallArgumentsDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
|
||||
item_id=call_id,
|
||||
item_id=self._tool_item_id_by_call_id.get(call_id, call_id),
|
||||
output_index=output_index,
|
||||
arguments=final_args,
|
||||
)
|
||||
|
|
@ -345,6 +348,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._sequence_number += 1
|
||||
names = self._custom_tool_names
|
||||
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names)
|
||||
item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"])
|
||||
if tool_namespace:
|
||||
item_kwargs["namespace"] = tool_namespace
|
||||
item_done_event = OutputItemDoneEvent(
|
||||
|
|
|
|||
|
|
@ -93,6 +93,7 @@ from .custom_tools import (
|
|||
convert_custom_tool_to_function_tool,
|
||||
extract_custom_tool_names,
|
||||
is_custom_tool_call,
|
||||
openai_shaped_tool_call_item_id,
|
||||
serialize_tool_call_arguments,
|
||||
unwrap_custom_tool_arguments,
|
||||
validated_allowed_callers,
|
||||
|
|
@ -2034,7 +2035,7 @@ class LiteLLMCompletionResponsesConfig:
|
|||
custom_item = CustomToolCallOutputItem(
|
||||
type="custom_tool_call",
|
||||
call_id=tool_id,
|
||||
id=tool_id,
|
||||
id=openai_shaped_tool_call_item_id("custom_tool_call", tool_id),
|
||||
name=tool_name,
|
||||
input=input_str,
|
||||
status=function_definition.get("status") or "completed",
|
||||
|
|
@ -2065,7 +2066,7 @@ class LiteLLMCompletionResponsesConfig:
|
|||
name=tool_name,
|
||||
arguments=tool_arguments,
|
||||
call_id=tool_id,
|
||||
id=tool_id,
|
||||
id=openai_shaped_tool_call_item_id("function_call", tool_id),
|
||||
type="function_call",
|
||||
status=function_definition.get("status") or "completed",
|
||||
)
|
||||
|
|
@ -2502,8 +2503,7 @@ class LiteLLMCompletionResponsesConfig:
|
|||
choice=choice,
|
||||
)
|
||||
message_output_items.extend(image_generation_items)
|
||||
else:
|
||||
# Regular message output
|
||||
elif choice.message.content is not None:
|
||||
message_output_items.append(
|
||||
GenericResponseOutputItem(
|
||||
type="message",
|
||||
|
|
|
|||
|
|
@ -1023,7 +1023,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
transformed: ResponsesAPIResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> None:
|
||||
self._events: list[ResponsesAPIStreamingResponse] = _build_synthetic_response_events(
|
||||
self._events: Sequence[ResponsesAPIStreamingResponse] = build_synthetic_response_events(
|
||||
transformed=transformed,
|
||||
logging_obj=logging_obj,
|
||||
chunk_size=self.CHUNK_SIZE,
|
||||
|
|
@ -1090,7 +1090,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
transformed: ResponsesAPIResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> None:
|
||||
self._events = _build_synthetic_response_events(
|
||||
self._events = build_synthetic_response_events(
|
||||
transformed=transformed,
|
||||
logging_obj=logging_obj,
|
||||
chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE,
|
||||
|
|
@ -1274,10 +1274,10 @@ def _add_text_like_part_events(
|
|||
)
|
||||
|
||||
|
||||
def _build_synthetic_response_events(
|
||||
def build_synthetic_response_events(
|
||||
*,
|
||||
transformed: ResponsesAPIResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
logging_obj: LiteLLMLoggingObj | None,
|
||||
chunk_size: int,
|
||||
) -> list[ResponsesAPIStreamingResponse]:
|
||||
openai_types: Final = _get_openai_response_types()
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ import traceback
|
|||
import weakref
|
||||
from collections import defaultdict
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence
|
||||
from functools import lru_cache
|
||||
from functools import lru_cache, partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
|
||||
|
||||
|
|
@ -143,7 +143,11 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
from litellm.router_utils.fallback_event_handlers import (
|
||||
AttemptedFallbackTargets,
|
||||
_check_non_standard_fallback_format,
|
||||
get_fallback_model_group,
|
||||
clear_pre_routing_selection,
|
||||
fallback_lookup_groups,
|
||||
get_fallback_model_group_for_lookup_groups,
|
||||
get_pre_routing_selection,
|
||||
record_pre_routing_selection,
|
||||
run_async_fallback,
|
||||
)
|
||||
from litellm.router_utils.get_retry_from_policy import (
|
||||
|
|
@ -821,6 +825,7 @@ class Router:
|
|||
self._zero_cost_cache: dict[str, bool] = {}
|
||||
self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None
|
||||
self._init_routing_groups(None)
|
||||
self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = ()
|
||||
|
||||
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
|
||||
self.model_group_affinity_config = model_group_affinity_config
|
||||
|
|
@ -4918,6 +4923,19 @@ class Router:
|
|||
)
|
||||
response = await response
|
||||
|
||||
if self._should_raise_anthropic_refusal_error(
|
||||
model=model,
|
||||
original_generic_function=original_generic_function,
|
||||
response=response,
|
||||
kwargs=kwargs,
|
||||
):
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
|
||||
safeguard_refusal_error,
|
||||
)
|
||||
|
||||
refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape
|
||||
raise safeguard_refusal_error(model=model, stop_details=refusal_details)
|
||||
|
||||
self.success_calls[model_name] += 1
|
||||
verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
|
||||
|
||||
|
|
@ -4964,6 +4982,11 @@ class Router:
|
|||
# fallback to the original reference for any non-picklable value.
|
||||
# The original_generic_function is preserved so the per-attempt
|
||||
# helper knows which underlying API to call on fallback.
|
||||
# The pre-routing hook stamps its tier selection into this bucket during the primary
|
||||
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
|
||||
# bucket, so the post-call carry-over below always has somewhere to read and write.
|
||||
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
|
||||
|
||||
fallback_kwargs: Final[dict[str, object]] = kwargs.copy()
|
||||
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
|
||||
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
|
||||
|
|
@ -4973,6 +4996,14 @@ class Router:
|
|||
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
|
||||
|
||||
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
|
||||
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
|
||||
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
|
||||
clear_pre_routing_selection(fallback_kwargs)
|
||||
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
|
||||
if live_pre_routing_selection is not None:
|
||||
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
|
||||
|
||||
if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator):
|
||||
return await self._aresponses_streaming_iterator(
|
||||
response=response,
|
||||
|
|
@ -5030,6 +5061,10 @@ class Router:
|
|||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
aclose_if_supported,
|
||||
parse_anthropic_error_event,
|
||||
parse_anthropic_refusal_stop_details,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
|
||||
safeguard_refusal_error,
|
||||
)
|
||||
|
||||
source_iterator: Final = response
|
||||
|
|
@ -5068,13 +5103,35 @@ class Router:
|
|||
continue
|
||||
if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
|
||||
has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit
|
||||
error_event = parse_anthropic_error_event(chunk)
|
||||
# A transport can split one SSE data line across byte chunks, so pre-content
|
||||
# detection parses the accumulated buffer plus the current chunk, never the
|
||||
# chunk alone; the buffer is already capped, which bounds this window too.
|
||||
parse_window = ( # rebind-ok: freshly computed each iteration, never carried over
|
||||
b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime
|
||||
if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime
|
||||
else chunk
|
||||
)
|
||||
error_event = parse_anthropic_error_event(parse_window)
|
||||
retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over
|
||||
not has_generated_content
|
||||
and error_event is not None
|
||||
and _is_retriable_anthropic_status(error_event[2])
|
||||
and not _anthropic_stream_error_is_gateway_verdict(chunk)
|
||||
)
|
||||
refusal_stop_details = ( # rebind-ok: freshly computed each iteration, never carried over
|
||||
parse_anthropic_refusal_stop_details(parse_window)
|
||||
if not has_generated_content and error_event is None
|
||||
else None
|
||||
)
|
||||
if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs):
|
||||
refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details)
|
||||
raise MidStreamFallbackError(
|
||||
message=refusal_error.message,
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
original_exception=refusal_error,
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
if not has_generated_content and not retriable_pending_error and error_event is None:
|
||||
buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk)
|
||||
continue
|
||||
|
|
@ -5186,8 +5243,13 @@ class Router:
|
|||
kwargs=initial_kwargs,
|
||||
metadata_variable_name="litellm_metadata",
|
||||
)
|
||||
# The content-policy dispatch branch matches on the trigger's own type, so a refusal's
|
||||
# MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted.
|
||||
fallback_trigger: Final[Exception] = (
|
||||
e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e
|
||||
)
|
||||
fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success
|
||||
e=e,
|
||||
e=fallback_trigger,
|
||||
disable_fallbacks=False,
|
||||
fallbacks=fallbacks,
|
||||
context_window_fallbacks=context_window_fallbacks,
|
||||
|
|
@ -5243,6 +5305,11 @@ class Router:
|
|||
# share, leaking primary-deployment metadata into the mid-stream
|
||||
# fallback request. safe_deep_copy avoids deep-copying the full
|
||||
# kwargs (which can hold non-deepcopyable logging handles/clients).
|
||||
# The pre-routing hook stamps its tier selection into this bucket during the primary
|
||||
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
|
||||
# bucket, so the post-call carry-over below always has somewhere to read and write.
|
||||
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
|
||||
|
||||
fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry
|
||||
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
|
||||
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
|
||||
|
|
@ -5252,6 +5319,14 @@ class Router:
|
|||
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
|
||||
|
||||
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
|
||||
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
|
||||
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
|
||||
clear_pre_routing_selection(fallback_kwargs)
|
||||
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
|
||||
if live_pre_routing_selection is not None:
|
||||
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
|
||||
|
||||
if kwargs.get("stream") and hasattr(response, "__aiter__"):
|
||||
return await self._aanthropic_messages_streaming_iterator(
|
||||
response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator
|
||||
|
|
@ -6807,6 +6882,9 @@ class Router:
|
|||
original_exception: Final = e
|
||||
fallback_model_group = None
|
||||
original_model_group: Final[str | None] = kwargs.get("model")
|
||||
# A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier
|
||||
# behind the router name, and fallbacks are configured per tier, not per router.
|
||||
lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group)
|
||||
fallback_failure_exception_str = ""
|
||||
|
||||
if disable_fallbacks is True or original_model_group is None:
|
||||
|
|
@ -6851,15 +6929,15 @@ class Router:
|
|||
]
|
||||
# Get external fallbacks — handle both standard and non-standard formats
|
||||
external_fallback_group: list | None = None
|
||||
if fallbacks is not None and model_group is not None:
|
||||
if fallbacks is not None and lookup_groups:
|
||||
if _check_non_standard_fallback_format(fallbacks=fallbacks):
|
||||
# Non-standard formats (e.g. ["claude-3-haiku"] or
|
||||
# [{"model": "...", "messages": [...]}]) are passed through directly
|
||||
external_fallback_group = fallbacks
|
||||
else:
|
||||
external_fallback_group, generic_idx = get_fallback_model_group(
|
||||
external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=fallbacks,
|
||||
model_group=cast(str, model_group),
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
if external_fallback_group is None and generic_idx is not None:
|
||||
external_fallback_group = fallbacks[generic_idx]["*"]
|
||||
|
|
@ -6917,9 +6995,9 @@ class Router:
|
|||
if isinstance(e, litellm.ContextWindowExceededError):
|
||||
if context_window_fallbacks is not None:
|
||||
context_window_fallback_model_group: Final[list[str] | None] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
self._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=context_window_fallbacks,
|
||||
model_group=model_group,
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
)
|
||||
if context_window_fallback_model_group is None:
|
||||
|
|
@ -6950,9 +7028,9 @@ class Router:
|
|||
elif isinstance(e, litellm.ContentPolicyViolationError):
|
||||
if content_policy_fallbacks is not None:
|
||||
content_policy_fallback_model_group: Final[list[str] | None] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
self._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
model_group=model_group,
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
)
|
||||
if content_policy_fallback_model_group is None:
|
||||
|
|
@ -6979,14 +7057,14 @@ class Router:
|
|||
|
||||
if litellm.expose_router_debug_in_errors:
|
||||
e.message += f"\n{error_message}"
|
||||
if fallbacks is not None and model_group is not None:
|
||||
if fallbacks is not None and lookup_groups:
|
||||
verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks))
|
||||
(
|
||||
fallback_model_group,
|
||||
generic_fallback_idx,
|
||||
) = get_fallback_model_group(
|
||||
) = get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}]
|
||||
model_group=cast(str, model_group),
|
||||
lookup_groups=lookup_groups,
|
||||
)
|
||||
## if none, check for generic fallback
|
||||
if fallback_model_group is None and generic_fallback_idx is not None:
|
||||
|
|
@ -6995,12 +7073,12 @@ class Router:
|
|||
if fallback_model_group is None:
|
||||
masked_fallbacks: Final = mask_sensitive_structure(fallbacks)
|
||||
verbose_router_logger.info(
|
||||
"No fallback model group found for original model_group=%s. Fallbacks=%s",
|
||||
model_group,
|
||||
"No fallback model group found for lookup_groups=%s. Fallbacks=%s",
|
||||
" -> ".join(lookup_groups),
|
||||
masked_fallbacks,
|
||||
)
|
||||
if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors:
|
||||
original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={masked_fallbacks}"
|
||||
original_exception.message += f"No fallback model group found for lookup_groups={' -> '.join(lookup_groups)}. Fallbacks={masked_fallbacks}"
|
||||
raise original_exception
|
||||
|
||||
input_kwargs.update(
|
||||
|
|
@ -7046,6 +7124,7 @@ class Router:
|
|||
If it fails after num_retries, fall back to another model group
|
||||
"""
|
||||
model_group: Final[str | None] = kwargs.get("model")
|
||||
clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary
|
||||
if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets):
|
||||
_fallback_metadata_key: Final = _get_router_metadata_variable_name(
|
||||
function_name=getattr(kwargs.get("original_function"), "__name__", None)
|
||||
|
|
@ -7471,6 +7550,24 @@ class Router:
|
|||
break
|
||||
return fallback_model_group
|
||||
|
||||
def _get_fallback_model_group_for_lookup_groups(
|
||||
self,
|
||||
fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the sibling resolver's contract
|
||||
lookup_groups: tuple[str, ...],
|
||||
) -> list[str] | None: # mutable-ok: mirrors the sibling resolver's contract
|
||||
"""First lookup group whose exact-key chain resolves (tier first, then requested group)."""
|
||||
return next(
|
||||
(
|
||||
resolved
|
||||
for resolved in (
|
||||
self._get_fallback_model_group_from_fallbacks(fallbacks=fallbacks, model_group=group)
|
||||
for group in lookup_groups
|
||||
)
|
||||
if resolved is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
def _get_first_default_fallback(self) -> str | None:
|
||||
"""
|
||||
Returns the first model from the default_fallbacks list, if it exists.
|
||||
|
|
@ -7886,6 +7983,31 @@ class Router:
|
|||
return True
|
||||
return False
|
||||
|
||||
def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
|
||||
"""
|
||||
Whether a content-policy fallback would resolve for this request, keyed the same way
|
||||
async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook
|
||||
selected wins over the requested group. Raising without this returning True would turn
|
||||
a deliverable response into an error the fallback chain cannot recover from.
|
||||
"""
|
||||
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
|
||||
if content_policy_fallbacks is not None:
|
||||
return (
|
||||
self._get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
lookup_groups=fallback_lookup_groups(kwargs, model_group),
|
||||
)
|
||||
is not None
|
||||
)
|
||||
if self._has_default_fallbacks():
|
||||
return True
|
||||
verbose_router_logger.debug(
|
||||
"No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s",
|
||||
model_group,
|
||||
content_policy_fallbacks,
|
||||
)
|
||||
return False
|
||||
|
||||
def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool:
|
||||
"""
|
||||
Determines if a content policy error should be raised.
|
||||
|
|
@ -7898,27 +8020,26 @@ class Router:
|
|||
if response.choices[0].finish_reason != "content_filter":
|
||||
return False
|
||||
|
||||
content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
|
||||
return self._has_content_policy_fallback(model, kwargs)
|
||||
|
||||
### ONLY RAISE ERROR IF CP FALLBACK AVAILABLE ###
|
||||
if content_policy_fallbacks is not None:
|
||||
fallback_model_group = None
|
||||
for item in content_policy_fallbacks: # [{"gpt-3.5-turbo": ["gpt-4"]}]
|
||||
if list(item.keys())[0] == model:
|
||||
fallback_model_group = item[model]
|
||||
break
|
||||
|
||||
if fallback_model_group is not None:
|
||||
return True
|
||||
elif self._has_default_fallbacks(): # default fallbacks set
|
||||
return True
|
||||
|
||||
verbose_router_logger.debug(
|
||||
"Content Policy Error occurred. No available fallbacks. Returning original response. model=%s, content_policy_fallbacks=%s",
|
||||
model,
|
||||
content_policy_fallbacks,
|
||||
def _should_raise_anthropic_refusal_error(
|
||||
self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, Any]
|
||||
) -> bool:
|
||||
"""
|
||||
The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard
|
||||
refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only
|
||||
when a content-policy fallback is configured; a plain refusal without stop_details, or
|
||||
any response with nothing configured, is returned to the client unchanged.
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
|
||||
get_safeguard_refusal_stop_details,
|
||||
)
|
||||
return False
|
||||
|
||||
if getattr(original_generic_function, "__name__", "") != "anthropic_messages":
|
||||
return False
|
||||
if get_safeguard_refusal_stop_details(response) is None:
|
||||
return False
|
||||
return self._has_content_policy_fallback(model, kwargs)
|
||||
|
||||
def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None):
|
||||
_all_deployments: list = []
|
||||
|
|
@ -8165,6 +8286,52 @@ class Router:
|
|||
if backend_value is not None:
|
||||
model_info[field] = backend_value
|
||||
|
||||
@staticmethod
|
||||
def _inherit_builtin_base_rates_for_off_peak(
|
||||
model_info: dict, # mutable-ok: cost-map entry filled in place
|
||||
backend_model: str,
|
||||
custom_llm_provider: str | None,
|
||||
) -> None:
|
||||
"""Fill missing pricing fields on a deployment entry that only sets
|
||||
``off_peak_pricing``, from the backend model's built-in cost map entry.
|
||||
|
||||
Cost lookup selects the deployment-scoped entry over the shared backend
|
||||
entry only when the deployment entry carries a base pricing field, and
|
||||
``off_peak_pricing`` is deliberately kept off the shared entry, so a
|
||||
deployment spelling out only its off-peak schedule would otherwise
|
||||
never receive the discount. The backend model's entire canonical cost
|
||||
map entry is copied, field by field, so threshold, tiered,
|
||||
service-tier, cache, character, and per-second rates as well as
|
||||
companion billing fields like ``web_search_billing_unit`` and the
|
||||
regional uplift multipliers all carry over, and peak-hour billing
|
||||
through the deployment entry matches the shared backend entry exactly.
|
||||
The raw ``litellm.model_cost`` entry is the copy source rather than
|
||||
``get_model_info``'s view of it, since that view synthesizes zero flat
|
||||
token rates for backends without one and storing those would mark a
|
||||
tiered-only backend explicitly priced free. Values are deep-copied to
|
||||
keep the builtin entry isolated. User-specified fields always win;
|
||||
no-op when any base pricing field is already set or the backend model
|
||||
has no canonical entry.
|
||||
"""
|
||||
if not model_info.get("off_peak_pricing"):
|
||||
return
|
||||
if any(
|
||||
model_info.get(field) is not None
|
||||
for field in ("input_cost_per_token", "input_cost_per_second", "tiered_pricing")
|
||||
):
|
||||
return
|
||||
try:
|
||||
backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model
|
||||
return
|
||||
backend_entry: Final = litellm.model_cost.get(backend_info.get("key") or "")
|
||||
if not isinstance(backend_entry, dict):
|
||||
return
|
||||
for field, backend_value in backend_entry.items():
|
||||
if model_info.get(field) is not None or backend_value is None:
|
||||
continue
|
||||
model_info[field] = copy.deepcopy(backend_value)
|
||||
|
||||
@staticmethod
|
||||
def _inherit_builtin_tiered_output_rate(
|
||||
model_info: dict, backend_model: str, custom_llm_provider: str | None
|
||||
|
|
@ -8253,6 +8420,11 @@ class Router:
|
|||
if deployment.litellm_params.get(field) is not None:
|
||||
_model_info[field] = deployment.litellm_params[field]
|
||||
|
||||
Router._inherit_builtin_base_rates_for_off_peak(
|
||||
model_info=_model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
if _model_info.get("input_cost_per_token") is not None:
|
||||
Router._inherit_builtin_cache_pricing(
|
||||
model_info=_model_info,
|
||||
|
|
@ -8301,6 +8473,19 @@ class Router:
|
|||
return deployment
|
||||
except Exception as e:
|
||||
if self.ignore_invalid_deployments:
|
||||
if isinstance(e, litellm.BadRequestError):
|
||||
self._provider_unresolved_deployments = (
|
||||
*self._provider_unresolved_deployments,
|
||||
partial(
|
||||
self._create_deployment,
|
||||
deployment_info=deployment_info,
|
||||
_model_name=_model_name,
|
||||
_litellm_params=_litellm_params,
|
||||
_model_info=_model_info,
|
||||
declared_id=declared_id,
|
||||
duplicate_ids=duplicate_ids,
|
||||
),
|
||||
)
|
||||
verbose_router_logger.exception(
|
||||
"Error creating deployment: %s, ignoring and continuing with other deployments.", e
|
||||
)
|
||||
|
|
@ -8730,6 +8915,7 @@ class Router:
|
|||
self.quality_routers = {}
|
||||
self.complexity_routers = {}
|
||||
self.auto_routers = {}
|
||||
self._provider_unresolved_deployments = ()
|
||||
self._invalidate_model_group_info_cache()
|
||||
self._invalidate_access_groups_cache()
|
||||
# we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works
|
||||
|
|
@ -8994,6 +9180,11 @@ class Router:
|
|||
if field_value is not None:
|
||||
_model_info_dict[field] = field_value
|
||||
|
||||
Router._inherit_builtin_base_rates_for_off_peak(
|
||||
model_info=_model_info_dict,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
if _model_info_dict.get("input_cost_per_token") is not None:
|
||||
Router._inherit_builtin_cache_pricing(
|
||||
model_info=_model_info_dict,
|
||||
|
|
@ -9249,6 +9440,11 @@ class Router:
|
|||
field_value = deployment.litellm_params.get(field)
|
||||
if field_value is not None:
|
||||
model_info[field] = field_value
|
||||
Router._inherit_builtin_base_rates_for_off_peak(
|
||||
model_info=model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
if model_info.get("input_cost_per_token") is not None:
|
||||
Router._inherit_builtin_cache_pricing(
|
||||
model_info=model_info,
|
||||
|
|
@ -9342,8 +9538,12 @@ class Router:
|
|||
"""Re-assert this router's deployments onto a freshly fetched catalog.
|
||||
|
||||
Reads ``model_list`` at call time, so only deployments the router still
|
||||
serves are restored.
|
||||
serves are restored, plus any config deployment the fresh catalog now resolves.
|
||||
"""
|
||||
provider_unresolved: Final = self._provider_unresolved_deployments
|
||||
self._provider_unresolved_deployments = ()
|
||||
for create_deployment in provider_unresolved:
|
||||
create_deployment()
|
||||
for entry in tuple(self.model_list):
|
||||
try:
|
||||
deployment = entry if isinstance(entry, Deployment) else Deployment(**entry)
|
||||
|
|
@ -12026,6 +12226,7 @@ class Router:
|
|||
if pre_routing_hook_response is not None:
|
||||
model = pre_routing_hook_response.model
|
||||
messages = pre_routing_hook_response.messages
|
||||
record_pre_routing_selection(request_kwargs, model)
|
||||
if pre_routing_hook_response.litellm_params:
|
||||
accepted_tier_params: Final = self._tier_params_the_target_accepts(
|
||||
model, pre_routing_hook_response.litellm_params, request_kwargs
|
||||
|
|
@ -12141,6 +12342,7 @@ class Router:
|
|||
if pre_routing_hook_response is not None:
|
||||
model = pre_routing_hook_response.model
|
||||
messages = pre_routing_hook_response.messages
|
||||
record_pre_routing_selection(request_kwargs, model)
|
||||
if pre_routing_hook_response.litellm_params:
|
||||
accepted_tier_params: Final = self._tier_params_the_target_accepts(
|
||||
model, pre_routing_hook_response.litellm_params, request_kwargs
|
||||
|
|
|
|||
|
|
@ -214,6 +214,91 @@ def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
PRE_ROUTING_SELECTED_MODEL_KEY: Final = "pre_routing_selected_model"
|
||||
_ROUTER_METADATA_BUCKETS: Final = ("metadata", "litellm_metadata")
|
||||
|
||||
|
||||
def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selected_model: str) -> None:
|
||||
"""
|
||||
Remember which model a pre-routing hook picked, so fallback lookup can key off it.
|
||||
|
||||
Fallback resolution runs on an outer kwargs dict that ``**kwargs`` already copied, so
|
||||
writing the model there is invisible by the time routing picks a tier. The metadata
|
||||
buckets are nested dicts shared by reference across those copies, which is how the
|
||||
router already carries values back up.
|
||||
|
||||
The write goes through the proxy-internal bucket resolver, never into both buckets:
|
||||
on /v1/messages the top-level ``metadata`` dict is the provider's own request field,
|
||||
so a blanket write would forward the tier stamp upstream.
|
||||
"""
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
|
||||
if request_kwargs is None:
|
||||
return
|
||||
bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs))
|
||||
if isinstance(bucket, dict):
|
||||
bucket[PRE_ROUTING_SELECTED_MODEL_KEY] = selected_model
|
||||
|
||||
|
||||
def clear_pre_routing_selection(request_kwargs: Mapping[str, object] | None) -> None:
|
||||
"""
|
||||
Drop any selection the router did not make itself on this hop.
|
||||
|
||||
The buckets carry whatever the caller sent, so an inbound value is the caller
|
||||
choosing a fallback chain rather than the router choosing a tier. A fallback hop
|
||||
also inherits the previous hop's selection, which would key its own failure off
|
||||
the tier that already failed. Clearing at the start of every hop leaves only a
|
||||
value the pre-routing hook wrote while routing that hop.
|
||||
"""
|
||||
if request_kwargs is None:
|
||||
return
|
||||
for bucket in (request_kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS):
|
||||
if isinstance(bucket, dict) and PRE_ROUTING_SELECTED_MODEL_KEY in bucket:
|
||||
del bucket[PRE_ROUTING_SELECTED_MODEL_KEY]
|
||||
|
||||
|
||||
def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None:
|
||||
"""The model a pre-routing hook selected for this request, if one did."""
|
||||
buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS)
|
||||
selections: Final = (bucket.get(PRE_ROUTING_SELECTED_MODEL_KEY) for bucket in buckets if isinstance(bucket, dict))
|
||||
return next((selected for selected in selections if isinstance(selected, str) and selected), None)
|
||||
|
||||
|
||||
def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]:
|
||||
"""
|
||||
Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins,
|
||||
and the requested group still resolves when no tier-keyed chain exists, so configs keyed
|
||||
on the router name (the documented contract) keep working behind auto-routers.
|
||||
"""
|
||||
ordered: Final = (get_pre_routing_selection(kwargs), model_group)
|
||||
return tuple(dict.fromkeys(group for group in ordered if group))
|
||||
|
||||
|
||||
def _resolved_a_specific_chain(
|
||||
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
) -> bool:
|
||||
resolved, generic_idx = result
|
||||
if resolved is None:
|
||||
return False
|
||||
return generic_idx is None or resolved is not fallbacks[generic_idx]["*"]
|
||||
|
||||
|
||||
def get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
lookup_groups: tuple[str, ...],
|
||||
) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract
|
||||
"""
|
||||
First lookup group with a specifically-keyed chain wins; the generic "*" chain applies
|
||||
only after every group missed, so a catch-all cannot shadow a later group's own chain.
|
||||
"""
|
||||
results: Final = tuple(get_fallback_model_group(fallbacks=fallbacks, model_group=group) for group in lookup_groups)
|
||||
specific: Final = next((result for result in results if _resolved_a_specific_chain(fallbacks, result)), None)
|
||||
if specific is not None:
|
||||
return specific
|
||||
return next((result for result in results if result[0] is not None), (None, None))
|
||||
|
||||
|
||||
def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]:
|
||||
"""
|
||||
Returns:
|
||||
|
|
@ -412,6 +497,7 @@ async def run_async_fallback(
|
|||
# LOGGING
|
||||
kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception)
|
||||
verbose_router_logger.info("Falling back to model_group = %s", mask_sensitive_structure(mg))
|
||||
kwargs.pop("_target_order", None) # rebind-ok: next hop must not inherit the previous order target
|
||||
if isinstance(mg, str):
|
||||
kwargs["model"] = mg
|
||||
elif isinstance(mg, dict):
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from re import Match
|
|||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
|
||||
|
||||
|
||||
class PatternUtils:
|
||||
|
|
@ -204,7 +204,7 @@ class PatternMatchRouter:
|
|||
|
||||
return litellm_deployment_litellm_model
|
||||
|
||||
def get_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict] | None:
|
||||
def get_pattern(self, model: str | None, custom_llm_provider: str | None = None) -> list[dict] | None:
|
||||
"""
|
||||
Check if a pattern exists for the given model and custom llm provider
|
||||
|
||||
|
|
@ -215,18 +215,17 @@ class PatternMatchRouter:
|
|||
Returns:
|
||||
bool: True if pattern exists, False otherwise
|
||||
"""
|
||||
if custom_llm_provider is None:
|
||||
try:
|
||||
(
|
||||
_,
|
||||
custom_llm_provider,
|
||||
_,
|
||||
_,
|
||||
) = get_llm_provider(model=model)
|
||||
except Exception:
|
||||
# get_llm_provider raises exception when provider is unknown
|
||||
pass
|
||||
return self.route(model) or self.route(f"{custom_llm_provider}/{model}")
|
||||
provider: Final = (
|
||||
custom_llm_provider or declared_authenticating_provider(model) or self._resolved_provider(model)
|
||||
)
|
||||
return self.route(model) or self.route(f"{provider}/{model}")
|
||||
|
||||
@staticmethod
|
||||
def _resolved_provider(model: str | None) -> str | None:
|
||||
try:
|
||||
return get_llm_provider(model=model)[1] if model else None
|
||||
except Exception: # noqa: BLE001 # get_llm_provider raises when the provider is unknown; the name then routes as-is
|
||||
return None
|
||||
|
||||
def get_deployments_by_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -427,6 +427,8 @@ class DeploymentAffinityCheck(CustomLogger):
|
|||
"""
|
||||
request_kwargs = request_kwargs or {}
|
||||
typed_healthy_deployments: Final = cast(list[dict], healthy_deployments)
|
||||
if request_kwargs.get("_target_order") is not None:
|
||||
return typed_healthy_deployments
|
||||
|
||||
(
|
||||
enable_user_key,
|
||||
|
|
|
|||
|
|
@ -58,6 +58,9 @@ class PromptCachingDeploymentCheck(CustomLogger):
|
|||
request_kwargs: dict | None = None,
|
||||
parent_otel_span: Span | None = None,
|
||||
) -> list[dict]:
|
||||
if request_kwargs is not None and request_kwargs.get("_target_order") is not None:
|
||||
return healthy_deployments
|
||||
|
||||
if messages is not None and is_prompt_caching_valid_prompt(
|
||||
messages=messages,
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import random
|
|||
import traceback
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
|
|
@ -214,6 +215,15 @@ class SearchAPIRouter:
|
|||
api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials(
|
||||
tool_litellm_params=litellm_params,
|
||||
)
|
||||
protected_params: Final = frozenset(("search_provider", "api_key", "api_base"))
|
||||
search_params: Final = MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for params in (litellm_params, kwargs)
|
||||
for key, value in params.items()
|
||||
if key not in protected_params and value is not None
|
||||
}
|
||||
)
|
||||
|
||||
verbose_router_logger.debug("Selected search tool with provider: %s", search_provider)
|
||||
|
||||
|
|
@ -222,7 +232,7 @@ class SearchAPIRouter:
|
|||
search_provider=search_provider,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
**kwargs,
|
||||
**search_params,
|
||||
)
|
||||
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -2,16 +2,37 @@
|
|||
Cost calculation for search providers.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
|
||||
tuple[Mapping[str, object], ...]
|
||||
)
|
||||
EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _provider_usage(
|
||||
optional_params: Mapping[str, object] | None,
|
||||
usage_param: str,
|
||||
) -> tuple[Mapping[str, object], ...] | None:
|
||||
params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS
|
||||
raw_usage: Final[object] = params.get(usage_param)
|
||||
try:
|
||||
return PROVIDER_USAGE_ADAPTER.validate_python(raw_usage)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def search_provider_cost_per_query(
|
||||
model: str,
|
||||
custom_llm_provider: str | None = None,
|
||||
number_of_queries: int = 1,
|
||||
optional_params: dict | None = None,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculate cost for search-only providers.
|
||||
|
|
@ -28,6 +49,18 @@ def search_provider_cost_per_query(
|
|||
Returns:
|
||||
Tuple of (input_cost, output_cost) where output_cost is always 0.0
|
||||
"""
|
||||
if custom_llm_provider == "parallel_ai":
|
||||
from litellm.llms.parallel_ai.search.cost_calculator import (
|
||||
PARALLEL_AI_USAGE_PARAM,
|
||||
parallel_ai_search_cost,
|
||||
)
|
||||
|
||||
input_cost: Final = parallel_ai_search_cost(
|
||||
optional_params=optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS,
|
||||
usage=_provider_usage(optional_params, PARALLEL_AI_USAGE_PARAM),
|
||||
)
|
||||
return (input_cost, 0.0)
|
||||
|
||||
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# Check for tiered pricing (e.g., Exa AI based on max_results)
|
||||
|
|
|
|||
|
|
@ -4,21 +4,58 @@ Payloads for Datadog LLM Observability Service (LLMObs)
|
|||
API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=example#api-standards
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Literal
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
|
||||
|
||||
|
||||
class ToolCall(TypedDict, total=False):
|
||||
"""A tool call on a message, as LLM Obs names its fields."""
|
||||
|
||||
name: ReadOnly[str]
|
||||
arguments: ReadOnly[dict[str, Any] | str] # parsed object, or the raw string when it will not parse to one
|
||||
tool_id: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
|
||||
|
||||
class ToolResult(TypedDict, total=False):
|
||||
"""The result of a tool call, as LLM Obs names its fields."""
|
||||
|
||||
name: ReadOnly[str]
|
||||
result: ReadOnly[str]
|
||||
tool_id: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
|
||||
|
||||
class ToolDefinition(TypedDict, total=False):
|
||||
"""A tool the model was offered on the request."""
|
||||
|
||||
name: ReadOnly[str]
|
||||
description: ReadOnly[str]
|
||||
schema: ReadOnly[dict[str, Any]]
|
||||
|
||||
|
||||
class Message(TypedDict, total=False):
|
||||
"""A message on a span, as LLM Obs names its fields."""
|
||||
|
||||
content: ReadOnly[str]
|
||||
role: ReadOnly[str]
|
||||
reasoning_content: ReadOnly[str]
|
||||
tool_calls: ReadOnly[Sequence[ToolCall]]
|
||||
tool_results: ReadOnly[Sequence[ToolResult]]
|
||||
|
||||
|
||||
class InputMeta(TypedDict):
|
||||
messages: list[
|
||||
dict[str, Any] # changed to fit with tool calls
|
||||
messages: Sequence[
|
||||
Message | dict[str, Any] # changed to fit with tool calls
|
||||
] # Relevant Issue: https://github.com/BerriAI/litellm/issues/9494
|
||||
|
||||
|
||||
class OutputMeta(TypedDict):
|
||||
messages: list[Any]
|
||||
messages: Sequence[Any]
|
||||
|
||||
|
||||
class DDLLMObsError(TypedDict, total=False):
|
||||
|
|
@ -36,6 +73,7 @@ class Meta(TypedDict, total=False):
|
|||
output: OutputMeta # The span's output information.
|
||||
metadata: dict[str, Any]
|
||||
error: DDLLMObsError | None # Error information on the span
|
||||
tool_definitions: ReadOnly[Sequence[ToolDefinition]] # The tools offered to the model on this request
|
||||
|
||||
|
||||
class LLMMetrics(TypedDict, total=False):
|
||||
|
|
@ -45,6 +83,9 @@ class LLMMetrics(TypedDict, total=False):
|
|||
time_to_first_token: float
|
||||
time_per_output_token: float
|
||||
total_cost: float
|
||||
cache_read_input_tokens: ReadOnly[float]
|
||||
cache_write_input_tokens: ReadOnly[float]
|
||||
non_cached_input_tokens: ReadOnly[float]
|
||||
|
||||
|
||||
class LLMObsPayload(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -270,6 +270,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[
|
|||
"litellm_deployment_rpm_limit",
|
||||
"litellm_remaining_api_key_requests_for_model",
|
||||
"litellm_remaining_api_key_tokens_for_model",
|
||||
"litellm_api_key_rate_limit_allowed_metric",
|
||||
"litellm_api_key_rate_limit_used_metric",
|
||||
"litellm_team_rate_limit_allowed_metric",
|
||||
"litellm_team_rate_limit_used_metric",
|
||||
"litellm_llm_api_failed_requests_metric",
|
||||
"litellm_callback_logging_failures_metric",
|
||||
"litellm_in_flight_requests",
|
||||
|
|
@ -775,6 +779,22 @@ class PrometheusMetricLabels:
|
|||
UserAPIKeyLabelNames.MODEL_ID.value,
|
||||
]
|
||||
|
||||
litellm_api_key_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = (
|
||||
UserAPIKeyLabelNames.API_KEY_HASH.value,
|
||||
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
|
||||
UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value,
|
||||
)
|
||||
|
||||
litellm_api_key_rate_limit_used_metric = litellm_api_key_rate_limit_allowed_metric
|
||||
|
||||
litellm_team_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = (
|
||||
UserAPIKeyLabelNames.TEAM.value,
|
||||
UserAPIKeyLabelNames.TEAM_ALIAS.value,
|
||||
UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value,
|
||||
)
|
||||
|
||||
litellm_team_rate_limit_used_metric = litellm_team_rate_limit_allowed_metric
|
||||
|
||||
litellm_llm_api_failed_requests_metric = [
|
||||
UserAPIKeyLabelNames.END_USER.value,
|
||||
UserAPIKeyLabelNames.API_KEY_HASH.value,
|
||||
|
|
|
|||
|
|
@ -91,6 +91,40 @@ class SlackAlertingArgs(LiteLLMPydanticObjectBase):
|
|||
default=False,
|
||||
description="If true, the alerting payload will be printed to the console.",
|
||||
)
|
||||
daily_spend_per_user_threshold: float | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Alert when a user's spend for the current day (UTC) crosses this USD amount. Off by default.",
|
||||
)
|
||||
monthly_spend_per_user_threshold: float | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Alert when a user's spend for the current calendar month (UTC) crosses this USD amount. Off by default.",
|
||||
)
|
||||
spend_anomaly_multiplier: float = Field(
|
||||
default=3.0,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Flag a user's spend as anomalous when today's spend exceeds this multiple of their trailing daily average.",
|
||||
)
|
||||
spend_anomaly_baseline_days: int = Field(
|
||||
default=7,
|
||||
ge=1,
|
||||
description="Number of trailing days used to compute a user's daily average spend for anomaly detection.",
|
||||
)
|
||||
spend_anomaly_min_spend: float = Field(
|
||||
default=10.0,
|
||||
gt=0,
|
||||
allow_inf_nan=False,
|
||||
description="Minimum spend (USD) a user must reach today before an anomaly alert can fire. Reduces false positives.",
|
||||
)
|
||||
user_spend_check_interval: int = Field(
|
||||
default=3600,
|
||||
ge=60,
|
||||
description="How often (in seconds) to check per-user spend thresholds and anomalies. Default is hourly.",
|
||||
)
|
||||
|
||||
|
||||
class DeploymentMetrics(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -138,6 +172,8 @@ class AlertType(str, Enum):
|
|||
budget_alerts = "budget_alerts"
|
||||
spend_reports = "spend_reports"
|
||||
failed_tracking_spend = "failed_tracking_spend"
|
||||
user_spend_thresholds = "user_spend_thresholds"
|
||||
user_spend_anomalies = "user_spend_anomalies"
|
||||
|
||||
# Database alerts
|
||||
db_exceptions = "db_exceptions"
|
||||
|
|
@ -182,6 +218,7 @@ DEFAULT_ALERT_TYPES: Final[list[AlertType]] = [
|
|||
AlertType.budget_alerts,
|
||||
AlertType.spend_reports,
|
||||
AlertType.failed_tracking_spend,
|
||||
AlertType.user_spend_thresholds,
|
||||
# Database alerts
|
||||
AlertType.db_exceptions,
|
||||
# Report alerts
|
||||
|
|
|
|||
|
|
@ -78,6 +78,16 @@ class AnthropicUsage(TypedDict, total=False):
|
|||
server_tool_use: NotRequired[ReadOnly[ServerToolUsage]]
|
||||
|
||||
|
||||
class AnthropicStopDetails(TypedDict, total=False):
|
||||
"""
|
||||
Safeguard verdict accompanying a `stop_reason: "refusal"` response:
|
||||
https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback
|
||||
"""
|
||||
|
||||
category: ReadOnly[str | None]
|
||||
explanation: ReadOnly[str | None]
|
||||
|
||||
|
||||
class AnthropicMessagesResponse(TypedDict, total=False):
|
||||
"""
|
||||
Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages
|
||||
|
|
@ -90,7 +100,8 @@ class AnthropicMessagesResponse(TypedDict, total=False):
|
|||
id: str
|
||||
model: str | None # This represents the Model type from Anthropic
|
||||
role: Literal["assistant"] | None
|
||||
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None
|
||||
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] | None
|
||||
stop_details: NotRequired[ReadOnly[AnthropicStopDetails | None]]
|
||||
stop_sequence: str | None
|
||||
type: Literal["message"] | None
|
||||
usage: AnthropicUsage | None
|
||||
|
|
|
|||
|
|
@ -150,6 +150,12 @@ class SCIMGroup(SCIMResource):
|
|||
members: list[SCIMMember] | None = None
|
||||
|
||||
|
||||
class SCIMPlaceholderMergeResult(BaseModel):
|
||||
placeholder_user_id: str
|
||||
merged_into_user_id: str
|
||||
team_ids: tuple[str, ...]
|
||||
|
||||
|
||||
# SCIM List Response Models
|
||||
class SCIMListResponse(BaseModel):
|
||||
schemas: list[str] = ["urn:ietf:params:scim:api:messages:2.0:ListResponse"]
|
||||
|
|
|
|||
|
|
@ -193,6 +193,38 @@ class AgenticLoopParams(TypedDict, total=False):
|
|||
"""The LLM provider name (e.g., 'bedrock', 'anthropic')"""
|
||||
|
||||
|
||||
class OffPeakWindow(TypedDict, total=False):
|
||||
"""One off-peak rule: UTC time-of-day windows, optionally restricted to weekdays.
|
||||
|
||||
hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them; a window may wrap past
|
||||
midnight and an equal-ended window covers the whole day. weekdays is a list of days the
|
||||
rule applies on, as ISO-8601 numbers (1 = Monday .. 7 = Sunday) or English day names;
|
||||
omitted means every day. The weekday is read on the calendar named by the block's
|
||||
weekday_timezone.
|
||||
"""
|
||||
|
||||
hours_utc: ReadOnly[str | Sequence[str]]
|
||||
weekdays: ReadOnly[Sequence[int | str]]
|
||||
|
||||
|
||||
class OffPeakPricing(TypedDict, total=False):
|
||||
"""Time-windowed off-peak rates for providers that discount by time of day (e.g. DeepSeek).
|
||||
|
||||
hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them for multiple daily windows,
|
||||
applying on every day of the week; a window may wrap past midnight. windows adds
|
||||
day-of-week-qualified rules (e.g. weekend-only whole-day off-peak), matched as a union
|
||||
with hours_utc. weekday_timezone names the IANA calendar weekdays are read on, defaulting
|
||||
to UTC. Any rate left unset falls back to the standard rate.
|
||||
"""
|
||||
|
||||
hours_utc: ReadOnly[str | Sequence[str]]
|
||||
windows: ReadOnly[Sequence[OffPeakWindow]]
|
||||
weekday_timezone: ReadOnly[str]
|
||||
input_cost_per_token: ReadOnly[float]
|
||||
output_cost_per_token: ReadOnly[float]
|
||||
cache_read_input_token_cost: ReadOnly[float]
|
||||
|
||||
|
||||
class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
||||
key: Required[str] # the key in litellm.model_cost which is returned
|
||||
|
||||
|
|
@ -225,6 +257,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
# Smallest prefix this model will actually cache, whatever caching mechanism its provider uses.
|
||||
# Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT.
|
||||
prompt_cache_min_tokens: int | None
|
||||
off_peak_pricing: ReadOnly[OffPeakPricing | None] # time-windowed off-peak rates
|
||||
input_cost_per_character: float | None # only for vertex ai models
|
||||
input_cost_per_audio_token: float | None
|
||||
input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models
|
||||
|
|
@ -3486,17 +3519,22 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
return {k: v for k, v in model_info.items() if k not in cls.model_fields}
|
||||
|
||||
|
||||
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = frozenset(
|
||||
ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__
|
||||
) - frozenset(CustomPricingLiteLLMParams.model_fields)
|
||||
DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"})
|
||||
|
||||
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
|
||||
frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
|
||||
- frozenset(CustomPricingLiteLLMParams.model_fields)
|
||||
- DEPLOYMENT_SCOPED_PRICING_FIELDS
|
||||
)
|
||||
|
||||
|
||||
def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return only the fields safe to register under a shared ``{provider}/{model}``
|
||||
key in ``litellm.model_cost``: cost-map schema fields (``ModelInfoBase``) minus
|
||||
per-deployment pricing overrides. Per-deployment metadata (``id``,
|
||||
``access_via_team_ids``, arbitrary custom keys) never belongs on the shared key;
|
||||
it stays under the deployment's unique model id.
|
||||
per-deployment pricing overrides and deployment-scoped pricing blocks such as
|
||||
``off_peak_pricing``. Per-deployment metadata (``id``, ``access_via_team_ids``,
|
||||
arbitrary custom keys) never belongs on the shared key; it stays under the
|
||||
deployment's unique model id.
|
||||
"""
|
||||
return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS}
|
||||
|
||||
|
|
@ -3598,6 +3636,8 @@ all_litellm_params = (
|
|||
"client",
|
||||
"rpm",
|
||||
"tpm",
|
||||
"default_api_key_rpm_limit",
|
||||
"default_api_key_tpm_limit",
|
||||
"itpm",
|
||||
"otpm",
|
||||
"max_parallel_requests",
|
||||
|
|
|
|||
|
|
@ -2869,10 +2869,9 @@ def _update_dictionary(existing_dict: dict, new_dict: dict) -> dict:
|
|||
elif isinstance(v, dict):
|
||||
existing_nested_dict = existing_dict.get(k)
|
||||
if isinstance(existing_nested_dict, dict):
|
||||
existing_nested_dict.update(v)
|
||||
existing_dict[k] = existing_nested_dict
|
||||
existing_dict[k] = {**existing_nested_dict, **v} # mutable-ok: copy-on-write merge
|
||||
else:
|
||||
existing_dict[k] = v
|
||||
existing_dict[k] = dict(v) # mutable-ok: detached copy, never the caller's dict by reference
|
||||
else:
|
||||
existing_dict[k] = v
|
||||
|
||||
|
|
@ -4877,11 +4876,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None:
|
|||
|
||||
def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list:
|
||||
if target_order is not None:
|
||||
filtered: Final = [d for d in healthy_deployments if _get_deployment_order(d) == target_order]
|
||||
if filtered:
|
||||
return filtered
|
||||
# target_order doesn't match any deployment (e.g., external fallback model) — return all
|
||||
return healthy_deployments
|
||||
return [d for d in healthy_deployments if _get_deployment_order(d) == target_order]
|
||||
|
||||
# Default: pick min order group
|
||||
_valid_orders: Final[list[int]] = [
|
||||
|
|
@ -5860,6 +5855,7 @@ def _get_model_info_helper(
|
|||
cache_creation_input_token_cost_above_1hr=_model_info.get(
|
||||
"cache_creation_input_token_cost_above_1hr", None
|
||||
),
|
||||
off_peak_pricing=_model_info.get("off_peak_pricing", None),
|
||||
input_cost_per_character=_model_info.get("input_cost_per_character", None),
|
||||
input_cost_per_token_above_128k_tokens=_model_info.get("input_cost_per_token_above_128k_tokens", None),
|
||||
input_cost_per_token_above_200k_tokens=_model_info.get("input_cost_per_token_above_200k_tokens", None),
|
||||
|
|
|
|||
|
|
@ -1482,7 +1482,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -1557,7 +1557,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -1632,7 +1632,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -1707,7 +1707,7 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_native_structured_output": false,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_output_config": true,
|
||||
"bedrock_output_config_effort_ceiling": "xhigh",
|
||||
|
|
@ -3254,6 +3254,7 @@
|
|||
"supports_computer_use": true,
|
||||
"supports_forced_tool_use": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
|
|
@ -38565,12 +38566,22 @@
|
|||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
|
||||
},
|
||||
"parallel_ai/search": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-fast": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-pro": {
|
||||
"input_cost_per_query": 0.009,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-turbo": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
|
|
@ -44142,6 +44153,7 @@
|
|||
"supports_computer_use": true,
|
||||
"supports_forced_tool_use": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
|
|
@ -44212,6 +44224,7 @@
|
|||
"supports_computer_use": true,
|
||||
"supports_forced_tool_use": false,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
|
|
|
|||
|
|
@ -81,6 +81,7 @@ ignored_function_names = [
|
|||
"_merge_tools_from_deployment", # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name)
|
||||
"_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name)
|
||||
"has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call
|
||||
"_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name)
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,16 @@
|
|||
|
||||
The e2e suite ships results to Loki/Grafana from a standard pytest JUnit report
|
||||
(`--junitxml=e2e-report.xml`), not a bespoke log line. JUnit already records
|
||||
outcome, duration, and node id for every `<testcase>`; the only signals it cannot
|
||||
derive on its own are the normalized suite package and the coverage-registry cell
|
||||
ids a test covers. Those ride along as JUnit `<property>` entries via each item's
|
||||
`user_properties`, attached in `conftest.py::pytest_collection_modifyitems`.
|
||||
outcome, duration, and node id for every `<testcase>`; the signals it cannot
|
||||
derive on its own are the normalized suite package, the coverage-registry cell
|
||||
ids a test covers, and where the test's source lives. Those ride along as JUnit
|
||||
`<property>` entries via each item's `user_properties`, attached in
|
||||
`conftest.py::pytest_collection_modifyitems`.
|
||||
|
||||
`source` is a property rather than the `file=` / `line=` attributes pytest used
|
||||
to write, because the `xunit2` family this suite runs on drops those, and
|
||||
switching families would change the XML for every consumer of it -- the
|
||||
Buildkite Test Engine upload and the Loki pipeline included.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -14,22 +20,61 @@ from collections.abc import Iterable
|
|||
|
||||
import pytest
|
||||
|
||||
# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing
|
||||
# at runtime names this suite's place in the repo. test_junit_properties.py
|
||||
# fails from a checkout if it moves.
|
||||
SUITE_ROOT = "tests/e2e"
|
||||
|
||||
|
||||
def suite_parts(path_part: str) -> tuple[str, ...]:
|
||||
"""Path components of a suite file relative to tests/e2e, however it ran.
|
||||
|
||||
Pytest paths are rootdir-relative, and rootdir moves with the invocation: a
|
||||
repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd run (the
|
||||
runner image) gives `logging/test_x.py`. Both collapse to the same tuple.
|
||||
"""
|
||||
raw = tuple(p for p in path_part.replace("\\", "/").split("/") if p and p != ".")
|
||||
return raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw
|
||||
|
||||
|
||||
def package_from_nodeid(nodeid: str) -> str:
|
||||
"""Top-level suite package under tests/e2e/, or 'root' for top-level files.
|
||||
|
||||
Pytest nodeids are relative to the invocation cwd. Repo-root runs look like
|
||||
`tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the
|
||||
`tests/e2e` prefix so package is the suite dir either way.
|
||||
"""
|
||||
path_part = nodeid.split("::", 1)[0].replace("\\", "/")
|
||||
raw = tuple(p for p in path_part.split("/") if p and p != ".")
|
||||
parts = raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw
|
||||
"""Top-level suite package under tests/e2e/, or 'root' for top-level files."""
|
||||
parts = suite_parts(nodeid.split("::", 1)[0])
|
||||
if len(parts) <= 1:
|
||||
return "root"
|
||||
return parts[0]
|
||||
|
||||
|
||||
def source_from_location(path: str, lineno: int | None) -> str:
|
||||
"""Repo-relative `path:line` for a test, or '' when nothing is linkable.
|
||||
|
||||
`pytest.Item.location` gives a rootdir-relative path and a ZERO-based line.
|
||||
The path is re-rooted at SUITE_ROOT so consumers need not know how pytest was
|
||||
started, and the line is emitted ONE-based to match editors, tracebacks and
|
||||
code hosts. A decorated test anchors at its first decorator, which is where
|
||||
pytest reports it.
|
||||
|
||||
Empty rather than a guess for anything unlinkable: no line, a path reaching
|
||||
upward, or a path carrying a colon, which is both how an absolute Windows
|
||||
path arrives and a character `path:line` has no way to represent.
|
||||
"""
|
||||
if lineno is None:
|
||||
return ""
|
||||
normalized = path.replace("\\", "/")
|
||||
if normalized.startswith("/") or ":" in normalized or ".." in normalized.split("/"):
|
||||
return ""
|
||||
parts = suite_parts(normalized)
|
||||
if not parts:
|
||||
return ""
|
||||
return f"{'/'.join((SUITE_ROOT, *parts))}:{lineno + 1}"
|
||||
|
||||
|
||||
def source_from_item(item: pytest.Item) -> str:
|
||||
"""Read the repo-relative `path:line` off a pytest Item's reported location."""
|
||||
path, lineno, _ = item.location
|
||||
return source_from_location(path, lineno)
|
||||
|
||||
|
||||
def dedupe_covers(marker_args: Iterable[tuple[object, ...]]) -> tuple[str, ...]:
|
||||
"""Flatten @pytest.mark.covers arg lists into unique, order-preserving cell
|
||||
ids, dropping anything that is not a non-empty string."""
|
||||
|
|
@ -43,10 +88,12 @@ def covers_from_item(item: pytest.Item) -> tuple[str, ...]:
|
|||
|
||||
def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]:
|
||||
"""The custom signals a standard reporter cannot derive: the normalized suite
|
||||
package and the comma-joined coverage-registry cell ids this test covers."""
|
||||
package, the comma-joined coverage-registry cell ids this test covers, and the
|
||||
repo-relative `path:line` its source sits at."""
|
||||
return (
|
||||
("package", package_from_nodeid(item.nodeid)),
|
||||
("covers", ",".join(covers_from_item(item))),
|
||||
("source", source_from_item(item)),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -806,6 +806,7 @@ class LiteLLMParamsBody(BaseModel):
|
|||
mock_response: str | None = None
|
||||
timeout: float | None = None
|
||||
tpm: int | None = None
|
||||
weight: int | None = None
|
||||
|
||||
|
||||
ModelMode = Literal["batch", "realtime", "image_generation"]
|
||||
|
|
@ -820,6 +821,7 @@ class ModelInfoBody(BaseModel):
|
|||
mode: ModelMode | None = None
|
||||
access_groups: list[str] | None = None
|
||||
team_id: str | None = None
|
||||
allowed_fails_policy: dict[str, int] | None = None
|
||||
|
||||
|
||||
class ModelNewBody(BaseModel):
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ from models import (
|
|||
ChatMessage,
|
||||
ChatResponse,
|
||||
LiteLLMParamsBody,
|
||||
ModelInfoBody,
|
||||
ModelNewBody,
|
||||
ReliabilityChatBody,
|
||||
RouterSettingsOverride,
|
||||
)
|
||||
|
|
@ -26,6 +28,18 @@ from models import (
|
|||
REAL_MODEL = "openai/gpt-5.5"
|
||||
REAL_KEY = "os.environ/OPENAI_API_KEY"
|
||||
|
||||
# The smallest-context chat model OpenAI still serves (16385 tokens). A prompt
|
||||
# past that limit comes back as a real `context_length_exceeded` 400, which is
|
||||
# what litellm maps to ContextWindowExceededError.
|
||||
SMALL_CONTEXT_MODEL = "openai/gpt-3.5-turbo"
|
||||
SMALL_CONTEXT_LIMIT_TOKENS = 16385
|
||||
|
||||
|
||||
def oversized_prompt(marker: str) -> str:
|
||||
"""A prompt comfortably past SMALL_CONTEXT_MODEL's context limit, so the
|
||||
provider refuses it on length rather than answering a truncated version."""
|
||||
return f"{marker} " + ("token " * (SMALL_CONTEXT_LIMIT_TOKENS + 4000))
|
||||
|
||||
|
||||
def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""Register a deployment pointing at an unreachable base, so every call to it
|
||||
|
|
@ -40,6 +54,38 @@ def create_timeout_deployment(proxy: ProxyClient, name: str) -> str:
|
|||
return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001))
|
||||
|
||||
|
||||
def create_small_context_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""Register a deployment on the smallest-context model OpenAI still serves, so an
|
||||
oversized prompt earns a real context-window refusal from the provider."""
|
||||
return proxy.create_model(name, LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY))
|
||||
|
||||
|
||||
def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""The always-picked half of a retry pair: a 1ms deadline the backend always
|
||||
exceeds, all of the model group's shuffle weight, and a cooldown policy that
|
||||
benches it on its first Timeout so the retry cannot land on it again."""
|
||||
return proxy.register_model(
|
||||
ModelNewBody(
|
||||
model_name=name,
|
||||
litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001, weight=1),
|
||||
model_info=ModelInfoBody(allowed_fails_policy={"TimeoutErrorAllowedFails": 0}),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
"""The other half of a retry pair: healthy, but weight 0, so the weighted shuffle
|
||||
never opens on it. It is reachable only once its sibling is benched and the
|
||||
weighted pick falls through to a uniform one over what is left."""
|
||||
return proxy.register_model(
|
||||
ModelNewBody(
|
||||
model_name=name,
|
||||
litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, weight=0),
|
||||
model_info=ModelInfoBody(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def chat_override(
|
||||
proxy: ProxyClient,
|
||||
key: str,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,10 @@ in the x-litellm-attempted-fallbacks header. Empty content is accepted only when
|
|||
`finish_reason == "length"` and the response billed completion tokens, since
|
||||
gpt-5.5 counts reasoning against max_tokens and can consume the whole budget
|
||||
before emitting any text; a fallback that produced nothing at all still fails.
|
||||
|
||||
The context-window case is a different reroute from a plain failure: the provider
|
||||
refuses the prompt on length, and `context_window_fallbacks` is the setting that
|
||||
reroutes it, not `fallbacks`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -25,8 +29,10 @@ from reliability_support import (
|
|||
completion_tokens_of,
|
||||
content_of,
|
||||
create_bad_base_deployment,
|
||||
create_small_context_deployment,
|
||||
create_timeout_deployment,
|
||||
finish_reason_of,
|
||||
oversized_prompt,
|
||||
reasoning_tokens_of,
|
||||
)
|
||||
|
||||
|
|
@ -82,3 +88,17 @@ class TestReliabilityFallbacks:
|
|||
override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]),
|
||||
)
|
||||
_assert_served_by_fallback(resp)
|
||||
|
||||
@pytest.mark.covers("reliability.fallback.context_window.routes_to_fallback")
|
||||
def test_context_window_routes_to_fallback(
|
||||
self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
primary = f"reliability-ctxfail-{unique_marker()}"
|
||||
model_id = create_small_context_deployment(client.proxy, primary)
|
||||
resources.defer(lambda: client.proxy.delete_model(model_id))
|
||||
|
||||
resp = chat_override(
|
||||
client.proxy, scoped_key, primary, oversized_prompt(unique_marker()),
|
||||
override=RouterSettingsOverride(context_window_fallbacks=[{primary: ["gpt-5.5"]}]),
|
||||
)
|
||||
_assert_served_by_fallback(resp)
|
||||
|
|
|
|||
73
tests/e2e/router/test_reliability_retries_e2e.py
Normal file
73
tests/e2e/router/test_reliability_retries_e2e.py
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
"""Live e2e: a request that fails on its first deployment is retried inside its own
|
||||
model group and still comes back a completion.
|
||||
|
||||
The model group is a pair: an always-timing-out deployment that holds all of the
|
||||
group's shuffle weight, and a healthy backup at weight 0. The weighted pick always
|
||||
opens on the timing-out one, its first Timeout benches it (an
|
||||
`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls
|
||||
through to the only deployment left. So the customer sees a completion and the
|
||||
proxy reports that it took a retry to get there, with no random first pick in the
|
||||
middle of it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from complexity_router_client import ComplexityRouterClient
|
||||
from e2e_config import unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
from models import RouterSettingsOverride
|
||||
from reliability_support import (
|
||||
chat_override,
|
||||
completion_tokens_of,
|
||||
content_of,
|
||||
create_always_timing_out_deployment,
|
||||
create_zero_weight_backup_deployment,
|
||||
finish_reason_of,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
||||
class TestReliabilityRetries:
|
||||
@pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries")
|
||||
def test_timeout_on_first_deployment_succeeds_on_retry(
|
||||
self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
group = f"reliability-retry-{unique_marker()}"
|
||||
timing_out = create_always_timing_out_deployment(client.proxy, group)
|
||||
resources.defer(lambda: client.proxy.delete_model(timing_out))
|
||||
backup = create_zero_weight_backup_deployment(client.proxy, group)
|
||||
resources.defer(lambda: client.proxy.delete_model(backup))
|
||||
|
||||
resp = chat_override(
|
||||
client.proxy,
|
||||
scoped_key,
|
||||
group,
|
||||
f"say hi {unique_marker()}",
|
||||
override=RouterSettingsOverride(num_retries=2),
|
||||
)
|
||||
|
||||
assert resp.status_code == 200, (
|
||||
f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}"
|
||||
)
|
||||
|
||||
attempted = resp.headers.get("x-litellm-attempted-retries")
|
||||
assert attempted is not None, "response is missing the x-litellm-attempted-retries header"
|
||||
assert int(attempted) >= 1, (
|
||||
f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never "
|
||||
"opened on the timing-out deployment, so this proves nothing about retries"
|
||||
)
|
||||
|
||||
content = content_of(resp)
|
||||
finish_reason = finish_reason_of(resp)
|
||||
completion_tokens = completion_tokens_of(resp) or 0
|
||||
assert isinstance(content, str), (
|
||||
f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})"
|
||||
)
|
||||
assert content or (finish_reason == "length" and completion_tokens > 0), (
|
||||
f"the retry returned empty content with finish_reason={finish_reason!r}, "
|
||||
f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget "
|
||||
f"was spent on non-visible reasoning (body={resp.body[:300]})"
|
||||
)
|
||||
146
tests/e2e/test_junit_properties.py
Normal file
146
tests/e2e/test_junit_properties.py
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
"""Harness coverage for the custom JUnit properties.
|
||||
|
||||
No proxy and no ``e2e`` marker. Pins the two normalizations that have to agree
|
||||
about where a suite file lives -- ``package_from_nodeid`` (strip the suite root)
|
||||
and ``source_from_location`` (re-root at it) -- across both ways the suite is
|
||||
launched, plus the one-based line offset and the refusal to emit a path that
|
||||
escapes the suite. The consumers of these properties are the Loki/Grafana
|
||||
rollups and, for ``source``, the status page's per-test links to GitHub.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from junit_properties import (
|
||||
SUITE_ROOT,
|
||||
attach_result_properties,
|
||||
dedupe_covers,
|
||||
package_from_nodeid,
|
||||
result_properties,
|
||||
source_from_location,
|
||||
suite_parts,
|
||||
)
|
||||
|
||||
|
||||
class FakeMarker:
|
||||
def __init__(self, name: str, *args: object) -> None:
|
||||
self.name = name
|
||||
self.args = args
|
||||
|
||||
|
||||
class FakeItem:
|
||||
"""The three attributes junit_properties reads off a pytest Item."""
|
||||
|
||||
def __init__(
|
||||
self, nodeid: str, location: tuple[str, int | None, str], markers: tuple[FakeMarker, ...] = ()
|
||||
) -> None:
|
||||
self.nodeid = nodeid
|
||||
self.location = location
|
||||
self.user_properties: list[tuple[str, str]] = []
|
||||
self._markers = markers
|
||||
|
||||
def iter_markers(self, name: str):
|
||||
return (marker for marker in self._markers if marker.name == name)
|
||||
|
||||
|
||||
def repo_root() -> Path | None:
|
||||
"""The litellm checkout above this file, or None when there isn't one."""
|
||||
return next((p for p in Path(__file__).resolve().parents if (p / ".git").exists()), None)
|
||||
|
||||
|
||||
class TestSuiteParts:
|
||||
@pytest.mark.parametrize(
|
||||
"path",
|
||||
["logging/test_x.py", "tests/e2e/logging/test_x.py", "./logging/test_x.py", "tests\\e2e\\logging\\test_x.py"],
|
||||
)
|
||||
def test_both_invocation_shapes_collapse_to_the_same_components(self, path: str) -> None:
|
||||
"""A repo-root run and a suite-cwd run report the same file differently;
|
||||
every downstream signal has to see one spelling."""
|
||||
assert suite_parts(path) == ("logging", "test_x.py")
|
||||
|
||||
def test_top_level_suite_file_keeps_its_single_component(self) -> None:
|
||||
assert suite_parts("tests/e2e/test_fixture_mode.py") == ("test_fixture_mode.py",)
|
||||
|
||||
|
||||
class TestPackageFromNodeid:
|
||||
@pytest.mark.parametrize(
|
||||
("nodeid", "expected"),
|
||||
[
|
||||
("logging/test_x.py::TestFoo::test_bar", "logging"),
|
||||
("tests/e2e/logging/test_x.py::TestFoo::test_bar", "logging"),
|
||||
("quota_management/spend_tracking/test_x.py::test_bar", "quota_management"),
|
||||
("test_fixture_mode.py::TestParseFixtureMode::test_known_values_normalize", "root"),
|
||||
("tests/e2e/test_fixture_mode.py::test_bar", "root"),
|
||||
],
|
||||
)
|
||||
def test_package_is_the_first_dir_under_the_suite_root(self, nodeid: str, expected: str) -> None:
|
||||
assert package_from_nodeid(nodeid) == expected
|
||||
|
||||
|
||||
class TestSourceFromLocation:
|
||||
@pytest.mark.parametrize("path", ["a2a/test_a2a_agent_e2e.py", "tests/e2e/a2a/test_a2a_agent_e2e.py"])
|
||||
def test_path_is_repo_relative_however_pytest_was_started(self, path: str) -> None:
|
||||
assert source_from_location(path, 40) == "tests/e2e/a2a/test_a2a_agent_e2e.py:41"
|
||||
|
||||
def test_line_is_emitted_one_based(self) -> None:
|
||||
"""pytest.Item.location counts from 0; editors, tracebacks and GitHub's
|
||||
#L anchor all count from 1, and an off-by-one lands on the decorator."""
|
||||
assert source_from_location("a2a/test_x.py", 0) == "tests/e2e/a2a/test_x.py:1"
|
||||
|
||||
def test_top_level_suite_file_sits_directly_under_the_suite_root(self) -> None:
|
||||
assert source_from_location("test_fixture_mode.py", 39) == "tests/e2e/test_fixture_mode.py:40"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "lineno"),
|
||||
[
|
||||
("a2a/test_x.py", None),
|
||||
("/app/e2e/a2a/test_x.py", 40),
|
||||
("C:\\app\\e2e\\a2a\\test_x.py", 40),
|
||||
("../conftest.py", 40),
|
||||
("", 40),
|
||||
],
|
||||
)
|
||||
def test_nothing_linkable_yields_empty_rather_than_a_guess(self, path: str, lineno: int | None) -> None:
|
||||
"""A colon is rejected on two counts: it is how a Windows absolute path
|
||||
arrives, and `path:line` cannot represent one in the path half."""
|
||||
assert source_from_location(path, lineno) == ""
|
||||
|
||||
|
||||
class TestResultProperties:
|
||||
def test_every_test_carries_package_covers_and_source(self) -> None:
|
||||
item = FakeItem(
|
||||
"logging/test_x.py::TestFoo::test_bar",
|
||||
("logging/test_x.py", 40, "TestFoo.test_bar"),
|
||||
(FakeMarker("covers", "LOG-1", "LOG-2"),),
|
||||
)
|
||||
assert result_properties(item) == (
|
||||
("package", "logging"),
|
||||
("covers", "LOG-1,LOG-2"),
|
||||
("source", "tests/e2e/logging/test_x.py:41"),
|
||||
)
|
||||
|
||||
def test_attach_is_idempotent(self) -> None:
|
||||
"""Collection can run the hook more than once; a second pass must not
|
||||
double the <property> entries in the report."""
|
||||
item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar"))
|
||||
attach_result_properties(item)
|
||||
attach_result_properties(item)
|
||||
assert [name for name, _ in item.user_properties] == ["package", "covers", "source"]
|
||||
|
||||
|
||||
class TestSuiteRoot:
|
||||
def test_suite_root_names_this_file_s_real_home(self) -> None:
|
||||
"""SUITE_ROOT is hardcoded because the runner image has no repo to read it
|
||||
from. Where there IS a checkout, prove the constant still points at us --
|
||||
otherwise a moved tests/e2e/ ships links that 404."""
|
||||
root = repo_root()
|
||||
if root is None:
|
||||
pytest.skip("no checkout above this file (the runner image copies tests/e2e/ to /app/e2e)")
|
||||
assert (root / SUITE_ROOT / Path(__file__).name).resolve() == Path(__file__).resolve()
|
||||
|
||||
|
||||
class TestDedupeCovers:
|
||||
def test_ids_are_unique_order_preserving_and_non_empty_strings(self) -> None:
|
||||
assert dedupe_covers([("A", "B"), ("B", ""), ("C", 7)]) == ("A", "B", "C")
|
||||
|
|
@ -40,6 +40,24 @@ def prometheus_logger() -> PrometheusLogger:
|
|||
return PrometheusLogger()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def known_model_router():
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-mini",
|
||||
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "fake-key"},
|
||||
},
|
||||
{
|
||||
"model_name": "us/azure/openai/gpt-5-mini",
|
||||
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "fake-key"},
|
||||
},
|
||||
]
|
||||
)
|
||||
with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam
|
||||
yield router
|
||||
|
||||
|
||||
def create_standard_logging_payload() -> StandardLoggingPayload:
|
||||
return StandardLoggingPayload(
|
||||
id="test_id",
|
||||
|
|
@ -741,7 +759,7 @@ async def test_async_log_failure_event(prometheus_logger):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger):
|
||||
async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger, known_model_router):
|
||||
"""LiteLLM-side reject (no deployment picked) routes the requested model
|
||||
into `requested_model` and skips the partial-outage flag."""
|
||||
standard_logging_object = create_standard_logging_payload()
|
||||
|
|
@ -786,7 +804,7 @@ async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_failure_hook(prometheus_logger):
|
||||
async def test_async_post_call_failure_hook(prometheus_logger, known_model_router):
|
||||
"""
|
||||
Test for the async_post_call_failure_hook method
|
||||
|
||||
|
|
@ -1069,7 +1087,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_success_fallback_event(prometheus_logger):
|
||||
async def test_log_success_fallback_event(prometheus_logger, known_model_router):
|
||||
prometheus_logger.litellm_deployment_successful_fallbacks = MagicMock()
|
||||
|
||||
original_model_group = "gpt-5-mini"
|
||||
|
|
@ -1107,7 +1125,7 @@ async def test_log_success_fallback_event(prometheus_logger):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_failure_fallback_event(prometheus_logger):
|
||||
async def test_log_failure_fallback_event(prometheus_logger, known_model_router):
|
||||
prometheus_logger.litellm_deployment_failed_fallbacks = MagicMock()
|
||||
|
||||
original_model_group = "gpt-5-mini"
|
||||
|
|
|
|||
|
|
@ -1,571 +0,0 @@
|
|||
"""Regression tests for ProxyExtrasDBManager's v2 migration resolver.
|
||||
|
||||
v2 is the proxy CLI default; v1 stays reachable via the `use_v2_resolver`
|
||||
kwarg, which still defaults to False for direct callers.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm_proxy_extras.utils import (
|
||||
_PRISMA_ATTEMPTS,
|
||||
ProxyExtrasDBManager,
|
||||
_max_migration_timestamp,
|
||||
_migration_timestamp,
|
||||
)
|
||||
|
||||
|
||||
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
|
||||
def _run(*args, **kwargs):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=returncode,
|
||||
cmd=args[0],
|
||||
stderr=stderr,
|
||||
output="",
|
||||
)
|
||||
|
||||
return _run
|
||||
|
||||
|
||||
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a permission failure during migrate deploy raises RuntimeError."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3018\nMigration name: 20250326162113_baseline\n"
|
||||
"Database error code: 42501\npermission denied for schema public"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="permission"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
|
||||
'Reason: syntax error at or near "BRKN" LINE 42'
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_removes_connection_limit():
|
||||
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
|
||||
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
|
||||
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
assert "connection_limit" not in stripped
|
||||
assert "pool_timeout" not in stripped
|
||||
assert "sslmode=require" in stripped
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_passthrough_no_query():
|
||||
"""URLs without query strings are returned unchanged."""
|
||||
url = "postgresql://u:p@h:5432/db"
|
||||
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
|
||||
|
||||
|
||||
def test_migration_timestamp_extracts_leading_digits():
|
||||
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
|
||||
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
|
||||
|
||||
|
||||
def test_migration_timestamp_returns_zero_on_malformed():
|
||||
assert _migration_timestamp("0_init") == 0
|
||||
assert _migration_timestamp("not_a_migration") == 0
|
||||
|
||||
|
||||
def test_max_migration_timestamp():
|
||||
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
|
||||
assert _max_migration_timestamp(names) == 20260415000000
|
||||
|
||||
|
||||
def test_max_migration_timestamp_empty_set():
|
||||
assert _max_migration_timestamp(set()) == 0
|
||||
|
||||
|
||||
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
|
||||
|
||||
This is the existing buggy behavior — we're not fixing it in v1, only
|
||||
offering v2 as opt-in. This test pins the default so that a future
|
||||
inadvertent default flip is caught.
|
||||
"""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
# Stub `prisma migrate deploy` to claim success with pending migrations
|
||||
# applied, which is the code path that triggers the legacy post-migration
|
||||
# sanity check (a call to _resolve_all_migrations).
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
def fake_run(cmd, *args, **kwargs):
|
||||
return FakeResult()
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
|
||||
def fake_resolve(*args, **kwargs):
|
||||
resolve_called["n"] += 1
|
||||
|
||||
monkeypatch.setattr("subprocess.run", fake_run)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
|
||||
|
||||
|
||||
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
|
||||
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = "db push error"
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="prisma db push failed"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
||||
"""_warn_if_db_ahead_of_head must never raise — it's informational.
|
||||
|
||||
Non-connection DB errors (e.g. InsufficientPrivilege from a user
|
||||
without SELECT on _prisma_migrations) must be caught, not propagated.
|
||||
"""
|
||||
import psycopg
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class _FakeConn:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def execute(self, *a, **kw):
|
||||
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
|
||||
raise psycopg.errors.InsufficientPrivilege("permission denied")
|
||||
|
||||
connects = {"n": 0}
|
||||
|
||||
def _fake_connect(*a, **kw):
|
||||
connects["n"] += 1
|
||||
return _FakeConn()
|
||||
|
||||
monkeypatch.setattr("psycopg.connect", _fake_connect)
|
||||
|
||||
assert ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) is None
|
||||
assert connects["n"] == 1, "the failing query must actually have been reached"
|
||||
|
||||
|
||||
def test_v2_resolve_specific_migration_failure_raises_runtime_error(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""If marking a migration as applied fails inside P3009 idempotent
|
||||
recovery, the subprocess error must be re-raised as RuntimeError so
|
||||
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
|
||||
)
|
||||
|
||||
# First call: migrate deploy -> P3009 idempotent error.
|
||||
# Recovery path tries _resolve_specific_migration; that also raises.
|
||||
def _failing_resolve(*a, **kw):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd="prisma migrate resolve --applied",
|
||||
stderr="resolve failed",
|
||||
output="",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
|
||||
)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
|
||||
"relation already exists"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(
|
||||
RuntimeError, match=r"Failed to mark migration .* as applied"
|
||||
):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_all_migrations",
|
||||
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
|
||||
)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
|
||||
|
||||
|
||||
_DEADLOCK_STDERR = (
|
||||
"Error: ERROR: deadlock detected\n"
|
||||
"DETAIL: Process 277 waits for ExclusiveLock on advisory lock "
|
||||
"[17556,0,72707369,1]; blocked by process 278.\n"
|
||||
"Process 278 waits for ShareLock on virtual transaction 3/1041; "
|
||||
"blocked by process 277."
|
||||
)
|
||||
|
||||
|
||||
class _DeployApplied:
|
||||
stdout = "All migrations have been successfully applied."
|
||||
stderr = ""
|
||||
returncode = 0
|
||||
|
||||
|
||||
def _deploy_only(deploy_side_effect):
|
||||
"""subprocess.run stand-in that only intercepts `prisma migrate deploy`.
|
||||
|
||||
Scoped by argv so the Prisma toolchain check cannot consume the mock first.
|
||||
"""
|
||||
deploys = {"n": 0}
|
||||
|
||||
def _run(*args, **kwargs):
|
||||
cmd = args[0] if args else kwargs.get("args", [])
|
||||
if list(cmd)[-2:] == ["migrate", "deploy"]:
|
||||
deploys["n"] += 1
|
||||
return deploy_side_effect(deploys["n"], cmd)
|
||||
return _DeployApplied()
|
||||
|
||||
return _run, deploys
|
||||
|
||||
|
||||
def _prepare_v2_resolver(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr("time.sleep", lambda *_a, **_k: None)
|
||||
|
||||
|
||||
def test_v2_retries_transient_advisory_lock_deadlock(monkeypatch, tmp_path):
|
||||
"""v2: replicas racing `migrate deploy` deadlock on Prisma's advisory
|
||||
lock, which is transient and must be retried rather than kill the boot."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
if n == 1:
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output=""
|
||||
)
|
||||
return _DeployApplied()
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert ok is True
|
||||
assert deploys["n"] == 2, "the deadlocked deploy must be retried, not raised"
|
||||
|
||||
|
||||
def test_v2_persistent_advisory_lock_deadlock_eventually_raises(monkeypatch, tmp_path):
|
||||
"""v2: the deadlock retry is bounded, so a deadlock that never clears
|
||||
still raises instead of looping or reporting success."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=_DEADLOCK_STDERR, output=""
|
||||
)
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError, match="after 4 attempts"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert deploys["n"] == 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stderr",
|
||||
[
|
||||
"Error: P1001: Can't reach database server at `db`:`5432`",
|
||||
"Error: P1002: The database server was reached but timed out.",
|
||||
],
|
||||
)
|
||||
def test_v2_retries_transient_database_connectivity_errors(monkeypatch, tmp_path, stderr):
|
||||
"""v2: a database not accepting connections yet is retried, not fatal."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
if n == 1:
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=stderr, output=""
|
||||
)
|
||||
return _DeployApplied()
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert ok is True
|
||||
assert deploys["n"] == 2, "an unreachable database must be retried, not raised"
|
||||
|
||||
|
||||
def test_v2_unreachable_database_still_fails_after_the_retries(monkeypatch, tmp_path):
|
||||
"""v2: a genuinely unreachable database still raises once the attempts
|
||||
are spent, rather than passing as a successful migration."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd=cmd,
|
||||
stderr="Error: P1001: Can't reach database server at `db`:`5432`",
|
||||
output="",
|
||||
)
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError, match="after 4 attempts"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert deploys["n"] == 4
|
||||
|
||||
|
||||
def test_v2_exhausted_retries_report_the_prisma_error(monkeypatch, tmp_path, caplog):
|
||||
"""v2: retrying must not swallow Prisma's stderr, which is captured and is
|
||||
the only place the cause appears for an operator or a boot-log grep."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
stderr = "Error: P1001: Can't reach database server at `wrong`:`5432`"
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1, cmd=cmd, stderr=stderr, output=""
|
||||
)
|
||||
|
||||
run, _ = _deploy_only(_side_effect)
|
||||
with caplog.at_level("INFO", logger="litellm_proxy_extras"):
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
|
||||
assert "P1001" in str(exc_info.value)
|
||||
assert "P1001" in caplog.text
|
||||
|
||||
|
||||
def test_v2_db_push_retries_transient_failures(monkeypatch, tmp_path):
|
||||
"""v2: `prisma db push` retries a transient failure like v1 did.
|
||||
|
||||
Reached from the migrations Job (USE_PRISMA_DB_PUSH=true), not from the
|
||||
proxy CLI, whose --use_prisma_db_push has its own loop in prisma_client.
|
||||
"""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
pushes = {"n": 0}
|
||||
|
||||
def _run(*args, **kwargs):
|
||||
cmd = list(args[0] if args else kwargs.get("args", []))
|
||||
if cmd[-3:] != ["db", "push", "--accept-data-loss"]:
|
||||
return _DeployApplied()
|
||||
pushes["n"] += 1
|
||||
if pushes["n"] == 1:
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd=cmd,
|
||||
stderr="Error: P1001: Can't reach database server at `db`:`5432`",
|
||||
output="",
|
||||
)
|
||||
return _DeployApplied()
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_run):
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
assert ok is True
|
||||
assert pushes["n"] == 2
|
||||
|
||||
|
||||
def test_v2_db_push_retries_are_bounded_and_report_the_prisma_error(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""v2: a database that never comes back stops after _PRISMA_ATTEMPTS and
|
||||
surfaces the prisma error, rather than retrying the boot forever."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
pushes = {"n": 0}
|
||||
|
||||
def _run(*args, **kwargs):
|
||||
cmd = list(args[0] if args else kwargs.get("args", []))
|
||||
if cmd[-3:] != ["db", "push", "--accept-data-loss"]:
|
||||
return _DeployApplied()
|
||||
pushes["n"] += 1
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd=cmd,
|
||||
stderr="Error: P1001: Can't reach database server at `db`:`5432`",
|
||||
output="",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_run):
|
||||
with pytest.raises(RuntimeError) as exc:
|
||||
ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=False, use_v2_resolver=True
|
||||
)
|
||||
|
||||
assert pushes["n"] == _PRISMA_ATTEMPTS
|
||||
assert "P1001" in str(exc.value)
|
||||
|
||||
|
||||
def _db_push_only(push_side_effect):
|
||||
"""subprocess.run stand-in that only intercepts `prisma db push`."""
|
||||
pushes = {"n": 0}
|
||||
|
||||
def _run(*args, **kwargs):
|
||||
cmd = list(args[0] if args else kwargs.get("args", []))
|
||||
if cmd[-3:] != ["db", "push", "--accept-data-loss"]:
|
||||
return _DeployApplied()
|
||||
pushes["n"] += 1
|
||||
return push_side_effect(pushes["n"], cmd)
|
||||
|
||||
return _run, pushes
|
||||
|
||||
|
||||
def _timed_out_for_real():
|
||||
"""Capture what subprocess.run really puts on a TimeoutExpired.
|
||||
|
||||
Under text=True it still leaves stderr as bytes, unlike CalledProcessError,
|
||||
so hardcoding a str here would test a shape production never sees. Derived
|
||||
at import, before any test patches subprocess.run.
|
||||
"""
|
||||
try:
|
||||
subprocess.run(
|
||||
["sh", "-c", "echo 'Error: P1001 unreachable' >&2; sleep 5"],
|
||||
timeout=0.2,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
except subprocess.TimeoutExpired as e:
|
||||
return e
|
||||
raise AssertionError("the helper command was supposed to time out")
|
||||
|
||||
|
||||
_TIMEOUT_TEMPLATE = _timed_out_for_real()
|
||||
|
||||
|
||||
def _real_timeout_expired(cmd):
|
||||
return subprocess.TimeoutExpired(
|
||||
cmd=cmd,
|
||||
timeout=_TIMEOUT_TEMPLATE.timeout,
|
||||
output=_TIMEOUT_TEMPLATE.stdout,
|
||||
stderr=_TIMEOUT_TEMPLATE.stderr,
|
||||
)
|
||||
|
||||
|
||||
def test_v2_db_push_retries_a_timeout(monkeypatch, tmp_path):
|
||||
"""v2: a `prisma db push` that times out is retried, not turned into a
|
||||
TypeError by classifying its bytes stderr as if it were text."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
if n == 1:
|
||||
raise _real_timeout_expired(cmd)
|
||||
return _DeployApplied()
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False
|
||||
)
|
||||
run, pushes = _db_push_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
assert ok is True
|
||||
assert pushes["n"] == 2
|
||||
|
||||
|
||||
def test_v2_db_push_timeouts_are_bounded(monkeypatch, tmp_path):
|
||||
"""v2: a `prisma db push` that never stops timing out gives up as a
|
||||
RuntimeError, which is the only exception proxy_cli.py exits cleanly on."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise _real_timeout_expired(cmd)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False
|
||||
)
|
||||
run, pushes = _db_push_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError, match=r"prisma db push failed after \d+"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
assert pushes["n"] == _PRISMA_ATTEMPTS
|
||||
|
||||
|
||||
def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path):
|
||||
"""v2: an unrecognised deploy failure still raises on the first attempt."""
|
||||
_prepare_v2_resolver(monkeypatch, tmp_path)
|
||||
|
||||
def _side_effect(n, cmd):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd=cmd,
|
||||
stderr="Error: relation \"LiteLLM_SpendLogs\" does not exist",
|
||||
output="",
|
||||
)
|
||||
|
||||
run, deploys = _deploy_only(_side_effect)
|
||||
with patch("subprocess.run", side_effect=run):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
assert deploys["n"] == 1
|
||||
|
||||
|
||||
|
|
@ -841,7 +841,7 @@ def test_build_synthetic_response_events_covers_annotations_function_calls_and_r
|
|||
)
|
||||
|
||||
try:
|
||||
events = streaming_module._build_synthetic_response_events(
|
||||
events = streaming_module.build_synthetic_response_events(
|
||||
transformed=transformed,
|
||||
logging_obj=logging_obj,
|
||||
chunk_size=5,
|
||||
|
|
|
|||
|
|
@ -44,7 +44,11 @@ _VCR_AUTO_MARKER_SKIP_FILES = frozenset(
|
|||
{"test_vcr_redis_persister.py", "test_ws_vcr.py"}
|
||||
)
|
||||
|
||||
_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = ()
|
||||
_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = (
|
||||
"test_nvidia_nim.py::test_embedding_nvidia_nim",
|
||||
"test_litellm_proxy_provider.py::test_litellm_gateway_from_sdk_embedding[False]",
|
||||
"test_litellm_proxy_provider.py::test_litellm_gateway_from_sdk_embedding[True]",
|
||||
)
|
||||
|
||||
|
||||
_verbose_state = VerboseReporterState()
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue