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_bedrock_bearer_token_converse_crash
This commit is contained in:
commit
863199c09b
334 changed files with 22664 additions and 2695 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
|
||||
|
||||
|
|
|
|||
|
|
@ -354,6 +354,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
|
|||
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
|
||||
| [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
|
||||
| [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | |
|
||||
| [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Sagemaker Chat (`sagemaker_chat`)](https://docs.litellm.ai/docs/providers/aws_sagemaker) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 14765
|
||||
"limit": 14076
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2216
|
||||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 4493
|
||||
"limit": 4128
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5607
|
||||
"limit": 5601
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15310
|
||||
"limit": 15306
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -105,25 +105,25 @@
|
|||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38368
|
||||
"limit": 38350
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19633
|
||||
"limit": 19626
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 29908
|
||||
"limit": 29890
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 111
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 695
|
||||
"limit": 692
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 828
|
||||
"limit": 826
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 0
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -659,6 +659,8 @@ aiml_models: Set = set()
|
|||
deepgram_models: Set = set()
|
||||
elevenlabs_models: Set = set()
|
||||
dashscope_models: Set = set()
|
||||
qwencloud_models: Set = set()
|
||||
qwen_ai_platform_models: Set = set()
|
||||
moonshot_models: Set = set()
|
||||
publicai_models: Set = set()
|
||||
darkbloom_models: Set = set()
|
||||
|
|
@ -909,6 +911,10 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None:
|
|||
heroku_models.add(key)
|
||||
elif value.get("litellm_provider") == "dashscope":
|
||||
dashscope_models.add(key)
|
||||
elif value.get("litellm_provider") == "qwencloud":
|
||||
qwencloud_models.add(key)
|
||||
elif value.get("litellm_provider") == "qwen_ai_platform":
|
||||
qwen_ai_platform_models.add(key)
|
||||
elif value.get("litellm_provider") == "modelscope":
|
||||
modelscope_models.add(key)
|
||||
elif value.get("litellm_provider") == "moonshot":
|
||||
|
|
@ -1072,6 +1078,8 @@ model_list = list(
|
|||
| deepgram_models
|
||||
| elevenlabs_models
|
||||
| dashscope_models
|
||||
| qwencloud_models
|
||||
| qwen_ai_platform_models
|
||||
| moonshot_models
|
||||
| publicai_models
|
||||
| darkbloom_models
|
||||
|
|
@ -1178,6 +1186,8 @@ def _build_models_by_provider() -> dict:
|
|||
"elevenlabs": elevenlabs_models,
|
||||
"heroku": heroku_models,
|
||||
"dashscope": dashscope_models,
|
||||
"qwencloud": qwencloud_models,
|
||||
"qwen_ai_platform": qwen_ai_platform_models,
|
||||
"modelscope": modelscope_models,
|
||||
"moonshot": moonshot_models,
|
||||
"publicai": publicai_models,
|
||||
|
|
@ -2014,6 +2024,24 @@ if TYPE_CHECKING:
|
|||
from .llms.dashscope.rerank.transformation import (
|
||||
DashScopeRerankConfig as DashScopeRerankConfig,
|
||||
)
|
||||
from .llms.dashscope.qwencloud import (
|
||||
QwenCloudChatConfig as QwenCloudChatConfig,
|
||||
)
|
||||
from .llms.dashscope.qwencloud import (
|
||||
QwenCloudEmbeddingConfig as QwenCloudEmbeddingConfig,
|
||||
)
|
||||
from .llms.dashscope.qwencloud import (
|
||||
QwenCloudRerankConfig as QwenCloudRerankConfig,
|
||||
)
|
||||
from .llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformChatConfig as QwenAIPlatformChatConfig,
|
||||
)
|
||||
from .llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformEmbeddingConfig as QwenAIPlatformEmbeddingConfig,
|
||||
)
|
||||
from .llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformRerankConfig as QwenAIPlatformRerankConfig,
|
||||
)
|
||||
from .llms.modelscope.chat.transformation import (
|
||||
ModelScopeChatConfig as ModelScopeChatConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -310,6 +310,8 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"GigaChatConfig",
|
||||
"GigaChatEmbeddingConfig",
|
||||
"DashScopeChatConfig",
|
||||
"QwenCloudChatConfig",
|
||||
"QwenAIPlatformChatConfig",
|
||||
"ModelScopeChatConfig",
|
||||
"MoonshotChatConfig",
|
||||
"DockerModelRunnerChatConfig",
|
||||
|
|
@ -1172,6 +1174,14 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.dashscope.chat.transformation",
|
||||
"DashScopeChatConfig",
|
||||
),
|
||||
"QwenCloudChatConfig": (
|
||||
".llms.dashscope.qwencloud",
|
||||
"QwenCloudChatConfig",
|
||||
),
|
||||
"QwenAIPlatformChatConfig": (
|
||||
".llms.dashscope.qwen_ai_platform",
|
||||
"QwenAIPlatformChatConfig",
|
||||
),
|
||||
"GDCGeminiConfig": (
|
||||
".llms.gdc.chat.transformation",
|
||||
"GDCGeminiConfig",
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import hashlib
|
|||
import json
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from enum import Enum
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -506,7 +507,7 @@ class Cache:
|
|||
|
||||
def _get_cache_logic(
|
||||
self,
|
||||
cached_result: Any | None,
|
||||
cached_result: object | None,
|
||||
max_age: float | None,
|
||||
):
|
||||
"""
|
||||
|
|
@ -538,8 +539,8 @@ class Cache:
|
|||
return cached_result
|
||||
|
||||
@staticmethod
|
||||
def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
cache_lookup_kwargs: Final[dict[str, Any]] = {}
|
||||
def _get_safe_cache_lookup_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
cache_lookup_kwargs: Final[dict[str, object]] = {}
|
||||
for prompt_kwarg in ("messages", "input"):
|
||||
if prompt_kwarg in kwargs:
|
||||
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
|
||||
|
|
@ -552,7 +553,7 @@ class Cache:
|
|||
|
||||
@staticmethod
|
||||
def _update_metadata_from_cache_lookup_kwargs(
|
||||
original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any]
|
||||
original_kwargs: Mapping[str, object], cache_lookup_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
original_metadata: Final = original_kwargs.get("metadata")
|
||||
cache_lookup_metadata: Final = cache_lookup_kwargs.get("metadata")
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import ast
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose
|
||||
|
|
@ -39,6 +39,12 @@ if TYPE_CHECKING:
|
|||
from litellm.router import Router
|
||||
|
||||
|
||||
class _QdrantCollectionDetailsResponse(Protocol):
|
||||
"""The qdrant `/collections/{name}` response, whose body is kept as an opaque JSON object."""
|
||||
|
||||
def json(self) -> dict[str, object]: ...
|
||||
|
||||
|
||||
class QdrantSemanticCache(BaseCache):
|
||||
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
|
||||
embedding_max_input_tokens: int | None = None
|
||||
|
|
@ -115,15 +121,15 @@ class QdrantSemanticCache(BaseCache):
|
|||
raise ValueError(f"Error from qdrant checking if /collections exist {collection_exists.text}")
|
||||
|
||||
if collection_exists.json()["result"]["exists"]:
|
||||
collection_details = self.sync_client.get(
|
||||
collection_details: _QdrantCollectionDetailsResponse = self.sync_client.get(
|
||||
url=f"{self.qdrant_api_base}/collections/{self.collection_name}",
|
||||
headers=self.headers,
|
||||
)
|
||||
self.collection_info = collection_details.json()
|
||||
self.collection_info: dict[str, object] = collection_details.json()
|
||||
print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}")
|
||||
self._ensure_cache_key_payload_index()
|
||||
else:
|
||||
quantization_params: dict[str, Any]
|
||||
quantization_params: dict[str, dict[str, object]]
|
||||
if quantization_config is None or quantization_config == "binary":
|
||||
quantization_params = {
|
||||
"binary": {
|
||||
|
|
@ -214,7 +220,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router),
|
||||
)
|
||||
|
||||
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
|
||||
def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
|
||||
"""Embed via the proxy Router when it serves the model, else direct."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_model_list, llm_router
|
||||
|
|
@ -241,7 +247,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
num_retries=0,
|
||||
)
|
||||
|
||||
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
|
||||
async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_model_list, llm_router
|
||||
except ImportError:
|
||||
|
|
|
|||
|
|
@ -45,14 +45,14 @@ class ResponsesToCompletionBridgeHandler:
|
|||
return bool(stream)
|
||||
|
||||
@staticmethod
|
||||
def _is_preformatted_cached_chat_stream(result: Any) -> bool:
|
||||
def _is_preformatted_cached_chat_stream(result: object) -> bool:
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
|
||||
return isinstance(result, CustomStreamWrapper) and result.custom_llm_provider == "cached_response"
|
||||
|
||||
@staticmethod
|
||||
def _coerce_response_object(
|
||||
response_obj: Any,
|
||||
response_obj: object,
|
||||
hidden_params: dict | None,
|
||||
) -> "ResponsesAPIResponse":
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
|
|
@ -78,8 +78,8 @@ class ResponsesToCompletionBridgeHandler:
|
|||
for _ in stream_iter:
|
||||
pass
|
||||
|
||||
completed: Final = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final = getattr(completed, "response", None) if completed else None
|
||||
completed: Final[object] = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final[object] = getattr(completed, "response", None) if completed else None
|
||||
if response_obj is None:
|
||||
raise ValueError("Stream ended without a completed response")
|
||||
|
||||
|
|
@ -93,8 +93,8 @@ class ResponsesToCompletionBridgeHandler:
|
|||
async for _ in stream_iter:
|
||||
pass
|
||||
|
||||
completed: Final = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final = getattr(completed, "response", None) if completed else None
|
||||
completed: Final[object] = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final[object] = getattr(completed, "response", None) if completed else None
|
||||
if response_obj is None:
|
||||
raise ValueError("Stream ended without a completed response")
|
||||
|
||||
|
|
@ -157,7 +157,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
def completion(
|
||||
self, *args, **kwargs
|
||||
) -> Union[
|
||||
Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]],
|
||||
Coroutine[None, None, Union["ModelResponse", "CustomStreamWrapper"]],
|
||||
"ModelResponse",
|
||||
"CustomStreamWrapper",
|
||||
]:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -630,6 +637,8 @@ LITELLM_CHAT_PROVIDERS: Final = [
|
|||
"nscale",
|
||||
"nebius",
|
||||
"dashscope",
|
||||
"qwencloud",
|
||||
"qwen_ai_platform",
|
||||
"modelscope",
|
||||
"moonshot",
|
||||
"publicai",
|
||||
|
|
@ -799,6 +808,7 @@ openai_compatible_endpoints: Final[list] = [
|
|||
"inference.api.nscale.com/v1",
|
||||
"api.studio.nebius.ai/v1",
|
||||
"https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
|
||||
"https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
"https://api-inference.modelscope.cn/v1",
|
||||
"https://api.moonshot.ai/v1",
|
||||
"https://api.publicai.co/v1",
|
||||
|
|
@ -872,6 +882,8 @@ openai_compatible_providers: Final[list] = [
|
|||
"nscale",
|
||||
"nebius",
|
||||
"dashscope",
|
||||
"qwencloud",
|
||||
"qwen_ai_platform",
|
||||
"modelscope",
|
||||
"moonshot",
|
||||
"v0",
|
||||
|
|
@ -902,6 +914,8 @@ openai_text_completion_compatible_providers: Final[list] = [ # providers that s
|
|||
"featherless_ai",
|
||||
"nebius",
|
||||
"dashscope",
|
||||
"qwencloud",
|
||||
"qwen_ai_platform",
|
||||
"modelscope",
|
||||
"moonshot",
|
||||
"publicai",
|
||||
|
|
@ -1109,7 +1123,7 @@ nebius_models: Final[set] = set(
|
|||
]
|
||||
)
|
||||
|
||||
dashscope_models: Final[set] = set(
|
||||
dashscope_models: Final[frozenset] = frozenset(
|
||||
[
|
||||
"qwen-turbo",
|
||||
"qwen-plus",
|
||||
|
|
@ -1124,6 +1138,10 @@ dashscope_models: Final[set] = set(
|
|||
]
|
||||
)
|
||||
|
||||
qwencloud_models: Final[frozenset] = frozenset(dashscope_models)
|
||||
|
||||
qwen_ai_platform_models: Final[frozenset] = frozenset(dashscope_models)
|
||||
|
||||
nebius_embedding_models: Final[set] = set(
|
||||
[
|
||||
"BAAI/bge-en-icl",
|
||||
|
|
@ -1240,6 +1258,7 @@ BEDROCK_CONVERSE_MODELS: Final = [
|
|||
"openai.gpt-oss-120b-1:0",
|
||||
"anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"anthropic.claude-fable-5-1",
|
||||
"anthropic.claude-fable-5",
|
||||
"anthropic.claude-sonnet-5",
|
||||
"anthropic.claude-opus-5",
|
||||
|
|
@ -1571,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"
|
||||
|
|
|
|||
|
|
@ -641,12 +641,12 @@ def cost_per_token(
|
|||
return xai_cost_per_token(model=model, usage=usage_block)
|
||||
elif custom_llm_provider == "lemonade":
|
||||
return lemonade_cost_per_token(model=model, usage=usage_block)
|
||||
elif custom_llm_provider == "dashscope":
|
||||
elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"):
|
||||
from litellm.llms.dashscope.cost_calculator import (
|
||||
cost_per_token as dashscope_cost_per_token,
|
||||
)
|
||||
|
||||
return dashscope_cost_per_token(model=model, usage=usage_block)
|
||||
return dashscope_cost_per_token(model=model, usage=usage_block, custom_llm_provider=custom_llm_provider)
|
||||
elif custom_llm_provider == "azure_ai":
|
||||
return azure_ai_cost_per_token(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
########################################################
|
||||
|
|
|
|||
|
|
@ -52,10 +52,10 @@ class GenerateContentSetupResult(BaseModel):
|
|||
model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
model: str
|
||||
request_body: dict[str, Any]
|
||||
request_body: dict[str, object]
|
||||
custom_llm_provider: str
|
||||
generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None
|
||||
generate_content_config_dict: dict[str, Any]
|
||||
generate_content_config_dict: dict[str, object]
|
||||
native_request_fields: dict[str, object]
|
||||
litellm_params: GenericLiteLLMParams
|
||||
litellm_logging_obj: LiteLLMLoggingObj
|
||||
|
|
@ -68,7 +68,7 @@ class GenerateContentHelper:
|
|||
@staticmethod
|
||||
def mock_generate_content_response(
|
||||
mock_response: str = "This is a mock response from Google GenAI generate_content.",
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""Mock response for generate_content for testing purposes"""
|
||||
return {
|
||||
"text": mock_response,
|
||||
|
|
@ -239,9 +239,9 @@ async def agenerate_content(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -307,9 +307,9 @@ def generate_content(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -397,9 +397,9 @@ async def agenerate_content_stream(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -492,9 +492,9 @@ def generate_content_stream(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import contextvars
|
|||
import importlib
|
||||
from collections.abc import Coroutine
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
|
|
@ -151,7 +151,7 @@ def image_generation(
|
|||
*,
|
||||
aimg_generation: Literal[True],
|
||||
**kwargs,
|
||||
) -> Coroutine[Any, Any, ImageResponse]:
|
||||
) -> Coroutine[object, object, ImageResponse]:
|
||||
...
|
||||
|
||||
|
||||
|
|
@ -197,7 +197,7 @@ def image_generation(
|
|||
api_version: str | None = None,
|
||||
custom_llm_provider=None,
|
||||
**kwargs,
|
||||
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
|
||||
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
|
||||
"""
|
||||
Maps the https://api.openai.com/v1/images/generations endpoint.
|
||||
|
||||
|
|
@ -386,6 +386,8 @@ def image_generation(
|
|||
litellm.LlmProviders.VERTEX_AI,
|
||||
litellm.LlmProviders.OPENROUTER,
|
||||
litellm.LlmProviders.DASHSCOPE,
|
||||
litellm.LlmProviders.QWENCLOUD,
|
||||
litellm.LlmProviders.QWEN_AI_PLATFORM,
|
||||
):
|
||||
if image_generation_config is None:
|
||||
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
|
||||
|
|
@ -723,14 +725,14 @@ def image_edit(
|
|||
user: str | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
|
||||
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
|
||||
"""
|
||||
Maps the image edit functionality, similar to OpenAI's images/edits endpoint.
|
||||
"""
|
||||
|
|
@ -769,7 +771,7 @@ def image_edit(
|
|||
images: Final = image if isinstance(image, list) else ([image] if image is not None else [])
|
||||
|
||||
headers_from_kwargs: Final = kwargs.get("headers")
|
||||
merged_extra_headers: Final[dict[str, Any]] = {}
|
||||
merged_extra_headers: Final[dict[str, object]] = {}
|
||||
if isinstance(headers_from_kwargs, dict):
|
||||
merged_extra_headers.update(headers_from_kwargs)
|
||||
if isinstance(extra_headers, dict):
|
||||
|
|
@ -974,9 +976,9 @@ async def aimage_edit(
|
|||
user: str | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -1044,7 +1046,7 @@ async def aimage_edit(
|
|||
)
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
def __getattr__(name: str) -> type["ImageEditRequestUtils"]:
|
||||
"""Lazy import handler for images.main module"""
|
||||
if name == "ImageEditRequestUtils":
|
||||
# Lazy load ImageEditRequestUtils to avoid heavy import from images.utils at module load time
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -545,7 +546,6 @@ class SlackAlerting(CustomBatchLogger):
|
|||
# Get the appropriate budget alert type handler
|
||||
budget_alert_class: Final = get_budget_alert_type(type)
|
||||
_id: Final = budget_alert_class.get_id(user_info)
|
||||
user_info_json: Final = user_info.model_dump(exclude_none=True)
|
||||
user_info_str: Final = self._get_user_info_str(user_info)
|
||||
event_message = budget_alert_class.get_event_message()
|
||||
|
||||
|
|
@ -575,7 +575,22 @@ class SlackAlerting(CustomBatchLogger):
|
|||
webhook_event = WebhookEvent(
|
||||
event=event,
|
||||
event_message=event_message,
|
||||
**user_info_json,
|
||||
spend=user_info.spend,
|
||||
max_budget=user_info.max_budget,
|
||||
soft_budget=user_info.soft_budget,
|
||||
token=user_info.token,
|
||||
customer_id=user_info.customer_id,
|
||||
user_id=user_info.user_id,
|
||||
team_id=user_info.team_id,
|
||||
team_alias=user_info.team_alias,
|
||||
organization_id=user_info.organization_id,
|
||||
user_email=user_info.user_email,
|
||||
key_alias=user_info.key_alias,
|
||||
projected_exceeded_date=user_info.projected_exceeded_date,
|
||||
projected_spend=user_info.projected_spend,
|
||||
event_group=user_info.event_group,
|
||||
alert_emails=user_info.alert_emails,
|
||||
max_budget_alert_emails=user_info.max_budget_alert_emails,
|
||||
)
|
||||
await self.send_alert(
|
||||
message=event_message + "\n\n" + user_info_str,
|
||||
|
|
@ -657,7 +672,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
"""
|
||||
Create a standard message for a budget alert
|
||||
"""
|
||||
_all_fields_as_dict: Final = user_info.model_dump(exclude_none=True)
|
||||
_all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True)
|
||||
_all_fields_as_dict.pop("token")
|
||||
msg = ""
|
||||
for k, v in _all_fields_as_dict.items():
|
||||
|
|
@ -1006,7 +1021,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: Any):
|
||||
async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: object):
|
||||
base_model_from_user: Final = getattr(passed_model_info, "base_model", None)
|
||||
model_info = {}
|
||||
base_model = ""
|
||||
|
|
@ -1930,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
|
||||
|
|
@ -1973,7 +2051,7 @@ Model Info:
|
|||
try:
|
||||
message = f"`{event_name}`\n"
|
||||
|
||||
key_event_dict: Final = key_event.model_dump()
|
||||
key_event_dict: Final[dict[str, object]] = key_event.model_dump()
|
||||
|
||||
# Add Created by information first
|
||||
message += "*Action Done by:*\n"
|
||||
|
|
|
|||
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)
|
||||
|
|
@ -3,6 +3,7 @@ BitBucket prompt manager that integrates with LiteLLM's prompt management system
|
|||
Fetches .prompt files from BitBucket repositories and provides team-based access control.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
|
|
@ -65,7 +66,7 @@ class BitBucketTemplateManager:
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
bitbucket_config: dict[str, Any],
|
||||
bitbucket_config: Mapping[str, object],
|
||||
prompt_id: str | None = None,
|
||||
):
|
||||
self.bitbucket_config = bitbucket_config
|
||||
|
|
@ -123,7 +124,7 @@ class BitBucketTemplateManager:
|
|||
template_content = content
|
||||
|
||||
# Parse YAML frontmatter
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata: dict[str, object] = {}
|
||||
if frontmatter_str:
|
||||
try:
|
||||
import yaml
|
||||
|
|
@ -141,9 +142,9 @@ class BitBucketTemplateManager:
|
|||
metadata=metadata,
|
||||
)
|
||||
|
||||
def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]:
|
||||
def _parse_yaml_basic(self, yaml_str: str) -> dict[str, object]:
|
||||
"""Basic YAML parser for simple cases when PyYAML is not available."""
|
||||
result: Final[dict[str, Any]] = {}
|
||||
result: Final[dict[str, object]] = {}
|
||||
for line in yaml_str.split("\n"):
|
||||
line = line.strip()
|
||||
if ":" in line and not line.startswith("#"):
|
||||
|
|
@ -162,7 +163,7 @@ class BitBucketTemplateManager:
|
|||
result[key] = value.strip("\"'")
|
||||
return result
|
||||
|
||||
def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str:
|
||||
def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str:
|
||||
"""Render a template with the given variables."""
|
||||
if template_id not in self.prompts:
|
||||
raise ValueError(f"Template '{template_id}' not found")
|
||||
|
|
@ -209,7 +210,7 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
bitbucket_config: dict[str, Any],
|
||||
bitbucket_config: Mapping[str, object],
|
||||
prompt_id: str | None = None,
|
||||
):
|
||||
self.bitbucket_config = bitbucket_config
|
||||
|
|
@ -234,7 +235,7 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
def get_prompt_template(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
"""
|
||||
Get a prompt template and render it with variables.
|
||||
|
|
@ -267,12 +268,12 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
self,
|
||||
user_id: str | None,
|
||||
messages: list[AllMessageValues],
|
||||
function_call: dict[str, Any] | str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
function_call: Mapping[str, object] | str | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
prompt_id: str | None = None,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> tuple[list[AllMessageValues], dict[str, Any] | None]:
|
||||
) -> tuple[list[AllMessageValues], dict[str, object] | None]:
|
||||
"""
|
||||
Pre-call hook that processes the prompt template before making the LLM call.
|
||||
"""
|
||||
|
|
@ -316,9 +317,9 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the call
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
litellm._logging.verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e)
|
||||
verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e)
|
||||
return messages, litellm_params
|
||||
|
||||
def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]:
|
||||
|
|
@ -384,14 +385,14 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
def post_call_hook(
|
||||
self,
|
||||
user_id: str | None,
|
||||
response: Any,
|
||||
response: object,
|
||||
input_messages: list[AllMessageValues],
|
||||
function_call: dict[str, Any] | str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
function_call: Mapping[str, object] | str | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
prompt_id: str | None = None,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
"""
|
||||
Post-call hook for any post-processing after the LLM call.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -19,14 +19,29 @@
|
|||
"""Transform LiteLLM data to CloudZero AnyCost CBF format."""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Final
|
||||
from typing import Final, SupportsFloat, SupportsIndex, SupportsInt
|
||||
|
||||
import polars as pl
|
||||
from typing_extensions import Buffer
|
||||
|
||||
from ...types.integrations.cloudzero import CBFRecord
|
||||
from .cz_resource_names import CZEntityType, CZRNGenerator
|
||||
|
||||
|
||||
def _as_int(value: object) -> int:
|
||||
"""The integer form of a spend table cell, computed the way :func:`int` computes it."""
|
||||
if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex)):
|
||||
return int(value)
|
||||
raise TypeError(f"int() argument must be a string or a number, not {type(value).__name__!r}")
|
||||
|
||||
|
||||
def _as_float(value: object) -> float:
|
||||
"""The floating point form of a spend table cell, computed the way :func:`float` computes it."""
|
||||
if isinstance(value, (str, Buffer, SupportsFloat, SupportsIndex)):
|
||||
return float(value)
|
||||
raise TypeError(f"float() argument must be a string or a number, not {type(value).__name__!r}")
|
||||
|
||||
|
||||
class CBFTransformer:
|
||||
"""Transform LiteLLM usage data to CloudZero Billing Format (CBF)."""
|
||||
|
||||
|
|
@ -82,15 +97,15 @@ class CBFTransformer:
|
|||
|
||||
return pl.DataFrame(cbf_data)
|
||||
|
||||
def _create_cbf_record(self, row: dict[str, Any]) -> CBFRecord:
|
||||
def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord:
|
||||
"""Create a single CBF record from LiteLLM daily spend row."""
|
||||
|
||||
# Parse date (daily spend tables use date strings like '2025-04-19')
|
||||
usage_date: Final = self._parse_date(row.get("date"))
|
||||
|
||||
# Calculate total tokens
|
||||
prompt_tokens: Final = int(row.get("prompt_tokens", 0))
|
||||
completion_tokens: Final = int(row.get("completion_tokens", 0))
|
||||
prompt_tokens: Final = _as_int(row.get("prompt_tokens", 0))
|
||||
completion_tokens: Final = _as_int(row.get("completion_tokens", 0))
|
||||
total_tokens: Final = prompt_tokens + completion_tokens
|
||||
|
||||
# Create CloudZero Resource Name (CZRN) as resource_id
|
||||
|
|
@ -154,7 +169,7 @@ class CBFTransformer:
|
|||
"time/usage_start": (
|
||||
usage_date.isoformat() if usage_date else None
|
||||
), # Required: ISO-formatted UTC datetime
|
||||
"cost/cost": float(row.get("spend", 0.0)), # Required: billed cost
|
||||
"cost/cost": _as_float(row.get("spend", 0.0)), # Required: billed cost
|
||||
"resource/id": resource_id, # CZRN (CloudZero Resource Name)
|
||||
# Usage metrics for token consumption
|
||||
"usage/amount": total_tokens, # Numeric value of tokens consumed
|
||||
|
|
@ -187,7 +202,7 @@ class CBFTransformer:
|
|||
|
||||
return CBFRecord(cbf_record)
|
||||
|
||||
def _parse_date(self, date_str) -> datetime | None:
|
||||
def _parse_date(self, date_str: object) -> datetime | None:
|
||||
"""Parse date string from daily spend tables (e.g., '2025-04-19')."""
|
||||
if date_str is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import contextvars
|
|||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
|
||||
|
||||
|
|
@ -227,13 +228,13 @@ class CustomGuardrail(CustomLogger):
|
|||
)
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def render_violation_message(self, default: str, context: dict[str, Any] | None = None) -> str:
|
||||
def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str:
|
||||
"""Return a custom violation message if template is configured."""
|
||||
|
||||
if not self.violation_message_template:
|
||||
return default
|
||||
|
||||
format_context: Final[dict[str, Any]] = {"default_message": default}
|
||||
format_context: Final[dict[str, object]] = {"default_message": default}
|
||||
if context:
|
||||
format_context.update(context)
|
||||
try:
|
||||
|
|
@ -661,7 +662,7 @@ class CustomGuardrail(CustomLogger):
|
|||
value: Final = self._get_admin_metadata(data).get("opted_out_global_guardrails")
|
||||
return value if isinstance(value, list) else []
|
||||
|
||||
def _is_valid_response_type(self, result: Any) -> bool:
|
||||
def _is_valid_response_type(self, result: object) -> bool:
|
||||
"""
|
||||
Check if result is a valid LLMResponseTypes instance.
|
||||
|
||||
|
|
@ -722,7 +723,7 @@ class CustomGuardrail(CustomLogger):
|
|||
return None
|
||||
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
|
||||
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, Any]) -> None:
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None:
|
||||
"""
|
||||
Record that this guardrail's ``async_pre_call_hook`` already ran for this
|
||||
request, so the deployment-level hook does not run it a second time.
|
||||
|
|
@ -747,7 +748,7 @@ class CustomGuardrail(CustomLogger):
|
|||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
|
||||
def _pre_call_hook_already_ran(self, data: dict[str, Any]) -> bool:
|
||||
def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool:
|
||||
marker: Final = self._pre_call_marker()
|
||||
if marker is None:
|
||||
return False
|
||||
|
|
@ -1170,7 +1171,7 @@ class CustomGuardrail(CustomLogger):
|
|||
This gets logged on downsteam Langfuse, DataDog, etc.
|
||||
"""
|
||||
# Convert None to empty dict to satisfy type requirements
|
||||
guardrail_response: dict[str, Any] | str = {} if response is None else response
|
||||
guardrail_response: dict[str, object] | str = {} if response is None else response
|
||||
|
||||
# For apply_guardrail functions in custom_code_guardrail scenario,
|
||||
# simplify the logged response to "allow", "deny", or "mask"
|
||||
|
|
|
|||
|
|
@ -20,10 +20,11 @@ import time
|
|||
import traceback
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime as datetimeObj
|
||||
from typing import Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from httpx import Response
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -62,6 +63,18 @@ from litellm.types.utils import StandardLoggingPayload
|
|||
|
||||
from ..additional_logging_utils import AdditionalLoggingUtils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
class _DatadogLoggingKwargs(TypedDict, total=False):
|
||||
"""The subset of logging ``kwargs`` that the Datadog payload builder reads."""
|
||||
|
||||
standard_logging_object: ReadOnly[StandardLoggingPayload | None]
|
||||
|
||||
|
||||
# max number of logs DD API can accept
|
||||
|
||||
|
||||
|
|
@ -87,6 +100,11 @@ def _resolve_dd_batch_size() -> int:
|
|||
return max(1, min(value, DD_MAX_BATCH_SIZE))
|
||||
|
||||
|
||||
def _span_attribute(span: object, name: str) -> object:
|
||||
"""Read an optional attribute off whatever span object the active tracer hands back."""
|
||||
return getattr(span, name, None)
|
||||
|
||||
|
||||
class DataDogLogger(
|
||||
CustomBatchLogger,
|
||||
AdditionalLoggingUtils,
|
||||
|
|
@ -271,9 +289,9 @@ class DataDogLogger(
|
|||
self,
|
||||
request_data: dict,
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: Any,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
traceback_str: str | None = None,
|
||||
) -> Any | None:
|
||||
) -> "HTTPException | None":
|
||||
"""
|
||||
Log proxy-level failures (e.g. 401 auth, DB connection errors) to Datadog.
|
||||
|
||||
|
|
@ -297,7 +315,7 @@ class DataDogLogger(
|
|||
status_code = int(_code)
|
||||
|
||||
# Use project-standard sanitized user context when running in proxy
|
||||
user_context: dict[str, Any] = {}
|
||||
user_context: dict[str, object] = {}
|
||||
try:
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
|
|
@ -553,8 +571,8 @@ class DataDogLogger(
|
|||
|
||||
def create_datadog_logging_payload(
|
||||
self,
|
||||
kwargs: dict | Any,
|
||||
response_obj: Any,
|
||||
kwargs: _DatadogLoggingKwargs,
|
||||
response_obj: object,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
) -> DatadogPayload:
|
||||
|
|
@ -562,8 +580,8 @@ class DataDogLogger(
|
|||
Helper function to create a datadog payload for logging
|
||||
|
||||
Args:
|
||||
kwargs (Union[dict, Any]): request kwargs
|
||||
response_obj (Any): llm api response
|
||||
kwargs: request kwargs, read for its standard logging object
|
||||
response_obj: llm api response
|
||||
start_time (datetime.datetime): start time of request
|
||||
end_time (datetime.datetime): end time of request
|
||||
|
||||
|
|
@ -625,7 +643,7 @@ class DataDogLogger(
|
|||
self,
|
||||
payload: ServiceLoggerPayload,
|
||||
error: str | None = "",
|
||||
parent_otel_span: Any | None = None,
|
||||
parent_otel_span: object = None,
|
||||
start_time: datetimeObj | float | None = None,
|
||||
end_time: float | datetimeObj | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
|
|
@ -659,7 +677,7 @@ class DataDogLogger(
|
|||
self,
|
||||
payload: ServiceLoggerPayload,
|
||||
error: str | None = "",
|
||||
parent_otel_span: Any | None = None,
|
||||
parent_otel_span: object = None,
|
||||
start_time: datetimeObj | float | None = None,
|
||||
end_time: float | datetimeObj | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
|
|
@ -696,7 +714,7 @@ class DataDogLogger(
|
|||
|
||||
def _create_v0_logging_payload(
|
||||
self,
|
||||
kwargs: dict | Any,
|
||||
kwargs: dict,
|
||||
response_obj: Any,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
|
|
@ -810,11 +828,11 @@ class DataDogLogger(
|
|||
if current_span is None:
|
||||
return None
|
||||
|
||||
trace_id: Final = getattr(current_span, "trace_id", None)
|
||||
trace_id: Final = _span_attribute(current_span, "trace_id")
|
||||
if trace_id is None:
|
||||
return None
|
||||
|
||||
span_id: Final = getattr(current_span, "span_id", None)
|
||||
span_id: Final = _span_attribute(current_span, "span_id")
|
||||
trace_context: Final[dict[str, str]] = {"trace_id": str(trace_id)}
|
||||
if span_id is not None:
|
||||
trace_context["span_id"] = str(span_id)
|
||||
|
|
|
|||
|
|
@ -9,7 +9,9 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp
|
|||
import asyncio
|
||||
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
|
||||
|
|
@ -29,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,
|
||||
|
|
@ -43,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):
|
||||
|
|
@ -221,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,
|
||||
|
|
@ -240,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",
|
||||
|
|
@ -313,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
|
||||
|
|
@ -334,7 +557,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
|
||||
def _get_response_messages(
|
||||
self, standard_logging_payload: StandardLoggingPayload, call_type: str | None
|
||||
) -> list[Any]:
|
||||
) -> tuple[Message, ...]:
|
||||
"""
|
||||
Get the messages from the response object
|
||||
|
||||
|
|
@ -343,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):
|
||||
|
|
@ -356,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,
|
||||
|
|
@ -374,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
|
||||
|
|
@ -484,22 +707,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
# Default fallback for unknown or passthrough operations
|
||||
return "llm"
|
||||
|
||||
def _ensure_string_content(self, messages: str | list[Any] | dict[Any, Any] | None) -> list[Any]:
|
||||
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, Any]:
|
||||
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
|
||||
"""
|
||||
_metadata: Final[dict[str, Any]] = {
|
||||
_metadata: Final[dict[str, object]] = {
|
||||
"model_name": standard_logging_payload.get("model", "unknown"),
|
||||
"model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"),
|
||||
"id": standard_logging_payload.get("id", "unknown"),
|
||||
|
|
@ -523,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
|
||||
|
|
@ -646,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: list[Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
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, Any]:
|
||||
"""
|
||||
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, Any]] = {}
|
||||
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, Any]:
|
||||
"""
|
||||
Extract tool call information from both input messages and response for Datadog metadata.
|
||||
"""
|
||||
tool_call_metadata: Final[dict[str, Any]] = {}
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -3,12 +3,21 @@ Based on Google's GenAI Kit dotprompt implementation: https://google.github.io/d
|
|||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, Final
|
||||
|
||||
import yaml
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
|
||||
class _PromptFileJson(TypedDict):
|
||||
"""JSON form of a .prompt file: rendered template text plus its frontmatter."""
|
||||
|
||||
content: ReadOnly[NotRequired[str]]
|
||||
metadata: ReadOnly[NotRequired[dict[str, object]]]
|
||||
|
||||
|
||||
def strip_version_suffix(prompt_id: str) -> str | None:
|
||||
|
|
@ -167,7 +176,7 @@ class PromptManager:
|
|||
template_id=prompt_id,
|
||||
)
|
||||
|
||||
def _parse_frontmatter(self, content: str) -> tuple[dict[str, Any], str]:
|
||||
def _parse_frontmatter(self, content: str) -> tuple[dict[str, object], str]:
|
||||
"""Parse YAML frontmatter from prompt content."""
|
||||
# Match YAML frontmatter between --- delimiters
|
||||
frontmatter_pattern: Final = r"^---\s*\n(.*?)\n---\s*\n(.*)$"
|
||||
|
|
@ -178,7 +187,7 @@ class PromptManager:
|
|||
template_content = match.group(2)
|
||||
|
||||
try:
|
||||
frontmatter = yaml.safe_load(frontmatter_yaml) or {}
|
||||
frontmatter: dict[str, object] = yaml.safe_load(frontmatter_yaml) or {}
|
||||
except yaml.YAMLError as e:
|
||||
raise ValueError(f"Invalid YAML frontmatter: {e}")
|
||||
else:
|
||||
|
|
@ -191,7 +200,7 @@ class PromptManager:
|
|||
def render(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
version: int | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
|
|
@ -231,7 +240,7 @@ class PromptManager:
|
|||
except Exception as e:
|
||||
raise ValueError(f"Error rendering template '{prompt_id}': {e}")
|
||||
|
||||
def _validate_input(self, variables: dict[str, Any], schema: dict[str, Any]) -> None:
|
||||
def _validate_input(self, variables: Mapping[str, object], schema: Mapping[str, str]) -> None:
|
||||
"""Basic validation of input variables against schema."""
|
||||
for field_name, field_type in schema.items():
|
||||
if field_name in variables:
|
||||
|
|
@ -291,7 +300,7 @@ class PromptManager:
|
|||
"""Get a list of all available prompt IDs."""
|
||||
return list(self.prompts.keys())
|
||||
|
||||
def get_prompt_metadata(self, prompt_id: str) -> dict[str, Any] | None:
|
||||
def get_prompt_metadata(self, prompt_id: str) -> dict[str, object] | None:
|
||||
"""Get metadata for a specific prompt."""
|
||||
template: Final = self.prompts.get(prompt_id)
|
||||
return template.metadata if template else None
|
||||
|
|
@ -302,12 +311,12 @@ class PromptManager:
|
|||
if self.prompt_directory:
|
||||
self._load_prompts()
|
||||
|
||||
def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, Any] | None = None) -> None:
|
||||
def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, object] | None = None) -> None:
|
||||
"""Add a prompt template programmatically."""
|
||||
template: Final = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id)
|
||||
self.prompts[prompt_id] = template
|
||||
|
||||
def prompt_file_to_json(self, file_path: str | Path) -> dict[str, Any]:
|
||||
def prompt_file_to_json(self, file_path: str | Path) -> _PromptFileJson:
|
||||
"""Convert a .prompt file to JSON format.
|
||||
|
||||
Args:
|
||||
|
|
@ -324,7 +333,7 @@ class PromptManager:
|
|||
|
||||
return {"content": template_content.strip(), "metadata": frontmatter}
|
||||
|
||||
def json_to_prompt_file(self, prompt_data: dict[str, Any]) -> str:
|
||||
def json_to_prompt_file(self, prompt_data: _PromptFileJson) -> str:
|
||||
"""Convert JSON prompt data to .prompt file format.
|
||||
|
||||
Args:
|
||||
|
|
|
|||
|
|
@ -6,10 +6,11 @@ import re
|
|||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone, tzinfo
|
||||
from typing import Any, Final, TypedDict, cast
|
||||
from typing import Any, Final, Protocol, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -35,6 +36,34 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai"
|
|||
GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000
|
||||
|
||||
|
||||
class _GalileoLoginBody(TypedDict):
|
||||
"""Decoded body of the Galileo login response."""
|
||||
|
||||
access_token: ReadOnly[str]
|
||||
|
||||
|
||||
class _GalileoLoginResponse(Protocol):
|
||||
"""The login call's HTTP response, read for the access token it carries."""
|
||||
|
||||
def json(self) -> _GalileoLoginBody: ...
|
||||
|
||||
|
||||
class _JsonResponse(Protocol):
|
||||
"""An HTTP response read only for whatever JSON body it decodes to."""
|
||||
|
||||
def json(self) -> object: ...
|
||||
|
||||
|
||||
def _login_access_token(response: _GalileoLoginResponse) -> str:
|
||||
"""Read the bearer token out of a Galileo login response body."""
|
||||
return response.json()["access_token"]
|
||||
|
||||
|
||||
def _decoded_body(response: _JsonResponse) -> object:
|
||||
"""Decode a response body without asserting anything about its shape."""
|
||||
return response.json()
|
||||
|
||||
|
||||
class GalileoStandardLoggingFields(TypedDict, total=False):
|
||||
call_type: str
|
||||
model: str
|
||||
|
|
@ -156,7 +185,7 @@ class GalileoObserve(CustomLogger):
|
|||
},
|
||||
)
|
||||
galileo_login_response.raise_for_status()
|
||||
access_token: Final = galileo_login_response.json()["access_token"]
|
||||
access_token: Final = _login_access_token(galileo_login_response)
|
||||
self.headers = {
|
||||
"accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
|
|
@ -421,7 +450,7 @@ class GalileoObserve(CustomLogger):
|
|||
try:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error response json: %s",
|
||||
response.json(),
|
||||
_decoded_body(response),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -4,12 +4,80 @@ Now supports selecting a tag via `config["tag"]`; falls back to branch ("main").
|
|||
"""
|
||||
|
||||
import base64
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, Protocol, TypedDict
|
||||
from urllib.parse import quote
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
||||
class GitLabFilePayload(TypedDict, total=False):
|
||||
"""A repository-files API entry."""
|
||||
|
||||
content: ReadOnly[str]
|
||||
encoding: ReadOnly[str]
|
||||
|
||||
|
||||
class GitLabTreeEntry(TypedDict, total=False):
|
||||
"""A repository-tree API entry."""
|
||||
|
||||
path: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
|
||||
|
||||
class GitLabBranch(TypedDict, total=False):
|
||||
"""A repository-branches API entry."""
|
||||
|
||||
name: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
|
||||
|
||||
class GitLabFileMetadata(TypedDict):
|
||||
"""The response headers a raw file request exposes as metadata."""
|
||||
|
||||
content_type: ReadOnly[str | None]
|
||||
content_length: ReadOnly[str | None]
|
||||
last_modified: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _FileJsonResponse(Protocol):
|
||||
def json(self) -> GitLabFilePayload: ...
|
||||
|
||||
|
||||
class _TreeJsonResponse(Protocol):
|
||||
def json(self) -> Sequence[GitLabTreeEntry] | None: ...
|
||||
|
||||
|
||||
class _ProjectJsonResponse(Protocol):
|
||||
def json(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class _BranchesJsonResponse(Protocol):
|
||||
def json(self) -> Sequence[GitLabBranch] | None: ...
|
||||
|
||||
|
||||
def _file_payload(resp: _FileJsonResponse) -> GitLabFilePayload:
|
||||
"""The JSON body of a repository-files response."""
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _tree_entries(resp: _TreeJsonResponse) -> Sequence[GitLabTreeEntry]:
|
||||
"""The entries of a repository-tree response."""
|
||||
return resp.json() or []
|
||||
|
||||
|
||||
def _project_info(resp: _ProjectJsonResponse) -> Mapping[str, object]:
|
||||
"""The JSON body of a project response."""
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _branch_entries(resp: _BranchesJsonResponse) -> Sequence[GitLabBranch] | None:
|
||||
"""The JSON body of a repository-branches response."""
|
||||
return resp.json()
|
||||
|
||||
|
||||
class GitLabClient:
|
||||
"""
|
||||
Client for interacting with the GitLab API to fetch files.
|
||||
|
|
@ -42,12 +110,12 @@ class GitLabClient:
|
|||
|
||||
self.project: str | int = project
|
||||
self.access_token: str = str(access_token)
|
||||
self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth'
|
||||
self.auth_method: str = config.get("auth_method", "token") # 'token' or 'oauth'
|
||||
self.branch = config.get("branch", None)
|
||||
if not self.branch:
|
||||
self.branch = "main"
|
||||
self.tag = config.get("tag")
|
||||
self.base_url = config.get("base_url", "https://gitlab.com/api/v4")
|
||||
self.base_url: str = config.get("base_url", "https://gitlab.com/api/v4")
|
||||
|
||||
if not all([self.project, self.access_token]):
|
||||
raise ValueError("project and access_token are required")
|
||||
|
|
@ -159,7 +227,7 @@ class GitLabClient:
|
|||
if resp.status_code == 404:
|
||||
return None
|
||||
resp.raise_for_status()
|
||||
data: Final = resp.json()
|
||||
data: Final = _file_payload(resp)
|
||||
content: Final = data.get("content")
|
||||
encoding: Final = data.get("encoding", "")
|
||||
if content and encoding == "base64":
|
||||
|
|
@ -208,7 +276,7 @@ class GitLabClient:
|
|||
return []
|
||||
resp.raise_for_status()
|
||||
|
||||
data: Final = resp.json() or []
|
||||
data: Final = _tree_entries(resp)
|
||||
files: Final[list[str]] = []
|
||||
for item in data:
|
||||
if item.get("type") == "blob":
|
||||
|
|
@ -229,13 +297,13 @@ class GitLabClient:
|
|||
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
|
||||
raise Exception(f"Failed to list files in '{directory_path}': {e}")
|
||||
|
||||
def get_repository_info(self) -> dict[str, Any]:
|
||||
def get_repository_info(self) -> Mapping[str, object]:
|
||||
"""Get information about the project/repository."""
|
||||
url: Final = f"{self.base_url}/projects/{self._project_enc}"
|
||||
try:
|
||||
resp: Final = self.http_handler.get(url, headers=self.headers)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
return _project_info(resp)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to get repository info: {e}")
|
||||
|
||||
|
|
@ -247,18 +315,18 @@ class GitLabClient:
|
|||
except Exception:
|
||||
return False
|
||||
|
||||
def get_branches(self) -> list[dict[str, Any]]:
|
||||
def get_branches(self) -> list[GitLabBranch]:
|
||||
"""Get list of branches in the repository."""
|
||||
url: Final = f"{self.base_url}/projects/{self._project_enc}/repository/branches"
|
||||
try:
|
||||
resp: Final = self.http_handler.get(url, headers=self.headers)
|
||||
resp.raise_for_status()
|
||||
data: Final = resp.json()
|
||||
data: Final = _branch_entries(resp)
|
||||
return data if isinstance(data, list) else []
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to get branches: {e}")
|
||||
|
||||
def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> dict[str, Any] | None:
|
||||
def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> GitLabFileMetadata | None:
|
||||
"""
|
||||
Get minimal metadata about a file via RAW endpoint headers at a given ref.
|
||||
|
||||
|
|
|
|||
|
|
@ -89,7 +89,7 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
|
|||
|
||||
# Check prompt_tokens_details.cached_tokens (used by Gemini and other providers)
|
||||
if hasattr(usage_obj, "prompt_tokens_details"):
|
||||
prompt_tokens_details: Final = getattr(usage_obj, "prompt_tokens_details", None)
|
||||
prompt_tokens_details: Final[object] = getattr(usage_obj, "prompt_tokens_details", None)
|
||||
if prompt_tokens_details is not None and hasattr(prompt_tokens_details, "cached_tokens"):
|
||||
cached_tokens: Final = getattr(prompt_tokens_details, "cached_tokens", None)
|
||||
if cached_tokens is not None and isinstance(cached_tokens, (int, float)) and cached_tokens > 0:
|
||||
|
|
@ -623,9 +623,16 @@ class LangFuseLogger:
|
|||
)
|
||||
|
||||
# Apply custom masking function if provided
|
||||
if masking_function is not None and callable(masking_function):
|
||||
input = self._apply_masking_function(input, masking_function)
|
||||
output = self._apply_masking_function(output, masking_function)
|
||||
masked_input: Final[object] = (
|
||||
self._apply_masking_function(input, masking_function)
|
||||
if masking_function is not None and callable(masking_function)
|
||||
else input
|
||||
)
|
||||
masked_output: Final[object] = (
|
||||
self._apply_masking_function(output, masking_function)
|
||||
if masking_function is not None and callable(masking_function)
|
||||
else output
|
||||
)
|
||||
|
||||
clean_metadata = redact_user_api_key_info(metadata=clean_metadata)
|
||||
|
||||
|
|
@ -651,15 +658,15 @@ class LangFuseLogger:
|
|||
|
||||
# Special keys that are found in the function arguments and not the metadata
|
||||
if "input" in update_trace_keys:
|
||||
trace_params["input"] = input if not mask_input else "redacted-by-litellm"
|
||||
trace_params["input"] = masked_input if not mask_input else "redacted-by-litellm"
|
||||
if "output" in update_trace_keys:
|
||||
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
|
||||
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
|
||||
else: # don't overwrite an existing trace
|
||||
trace_params = {
|
||||
"id": trace_id,
|
||||
"name": trace_name,
|
||||
"session_id": session_id,
|
||||
"input": input if not mask_input else "redacted-by-litellm",
|
||||
"input": masked_input if not mask_input else "redacted-by-litellm",
|
||||
"version": clean_metadata.pop(
|
||||
"trace_version", clean_metadata.get("version", None)
|
||||
), # If provided just version, it will applied to the trace as well, if applied a trace version it will take precedence
|
||||
|
|
@ -669,9 +676,9 @@ class LangFuseLogger:
|
|||
trace_params[key.replace("trace_", "")] = clean_metadata.pop(key, None)
|
||||
|
||||
if level == "ERROR":
|
||||
trace_params["status_message"] = output
|
||||
trace_params["status_message"] = masked_output
|
||||
else:
|
||||
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
|
||||
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
|
||||
|
||||
if debug is True or (isinstance(debug, str) and debug.lower() == "true"):
|
||||
debug_metadata: Final = {
|
||||
|
|
@ -708,7 +715,7 @@ class LangFuseLogger:
|
|||
("aws_region_name", aws_region_name, bool(aws_region_name)),
|
||||
("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs),
|
||||
)
|
||||
enrichments: Final[Mapping[str, Any]] = {
|
||||
enrichments: Final[Mapping[str, object]] = {
|
||||
key: value for key, value, include in candidate_enrichments if include
|
||||
}
|
||||
|
||||
|
|
@ -802,8 +809,8 @@ class LangFuseLogger:
|
|||
"end_time": end_time,
|
||||
"model": model_name,
|
||||
"model_parameters": optional_params,
|
||||
"input": input if not mask_input else "redacted-by-litellm",
|
||||
"output": output if not mask_output else "redacted-by-litellm",
|
||||
"input": masked_input if not mask_input else "redacted-by-litellm",
|
||||
"output": masked_output if not mask_output else "redacted-by-litellm",
|
||||
"usage": usage,
|
||||
"usage_details": usage_details,
|
||||
"metadata": {
|
||||
|
|
@ -825,8 +832,8 @@ class LangFuseLogger:
|
|||
prompt_management_metadata=prompt_management_metadata,
|
||||
langfuse_client=self.Langfuse,
|
||||
)
|
||||
if output is not None and isinstance(output, str) and level == "ERROR":
|
||||
generation_params["status_message"] = output
|
||||
if masked_output is not None and isinstance(masked_output, str) and level == "ERROR":
|
||||
generation_params["status_message"] = masked_output
|
||||
|
||||
if self._supports_completion_start_time():
|
||||
generation_params["completion_start_time"] = kwargs.get("completion_start_time", None)
|
||||
|
|
@ -935,7 +942,7 @@ class LangFuseLogger:
|
|||
return Version(self.langfuse_sdk_version) >= Version("2.7.3")
|
||||
|
||||
@staticmethod
|
||||
def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any:
|
||||
def _apply_masking_function(data: object, masking_function: Callable[[object], object]) -> object:
|
||||
"""
|
||||
Apply a masking function to data, handling different data types.
|
||||
|
||||
|
|
@ -1049,7 +1056,7 @@ def _add_prompt_to_generation_params(
|
|||
generation_params: dict,
|
||||
clean_metadata: dict,
|
||||
prompt_management_metadata: StandardLoggingPromptManagementMetadata | None,
|
||||
langfuse_client: Any,
|
||||
langfuse_client: object,
|
||||
) -> dict:
|
||||
from langfuse import Langfuse
|
||||
from langfuse.model import (
|
||||
|
|
|
|||
|
|
@ -4,9 +4,12 @@ Opik Logger that logs LLM events to an Opik server
|
|||
|
||||
import asyncio
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Any, Final
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict, Unpack
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -23,7 +26,7 @@ except Exception:
|
|||
opik_client = None
|
||||
|
||||
|
||||
def _should_skip_event(kwargs: dict[str, Any]) -> bool:
|
||||
def _should_skip_event(kwargs: Mapping[str, object]) -> bool:
|
||||
"""Check if event should be skipped due to missing standard_logging_object."""
|
||||
if kwargs.get("standard_logging_object") is None:
|
||||
verbose_logger.debug("OpikLogger skipping event; no standard_logging_object found")
|
||||
|
|
@ -31,12 +34,24 @@ def _should_skip_event(kwargs: dict[str, Any]) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
class _OpikLoggerKwargs(TypedDict, total=False):
|
||||
"""Constructor options accepted by ``OpikLogger``."""
|
||||
|
||||
project_name: ReadOnly[str | None]
|
||||
url: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | None]
|
||||
workspace: ReadOnly[str | None]
|
||||
batch_size: ReadOnly[int | None]
|
||||
flush_interval: ReadOnly[int | None]
|
||||
max_queue_size: ReadOnly[int | None]
|
||||
|
||||
|
||||
class OpikLogger(CustomBatchLogger):
|
||||
"""
|
||||
Opik Logger for logging events to an Opik Server
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
def __init__(self, **kwargs: Unpack[_OpikLoggerKwargs]) -> None:
|
||||
self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
|
||||
self.sync_httpx_client = _get_httpx_client()
|
||||
|
||||
|
|
@ -95,7 +110,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
|
||||
async def async_log_success_event(
|
||||
self,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: dict[str, object],
|
||||
response_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
|
|
@ -163,7 +178,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None:
|
||||
def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None:
|
||||
try:
|
||||
response: Final = self.sync_httpx_client.post(
|
||||
url=url,
|
||||
|
|
@ -178,7 +193,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
|
||||
def log_success_event(
|
||||
self,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: dict[str, object],
|
||||
response_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
|
|
@ -247,7 +262,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None:
|
||||
async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None:
|
||||
try:
|
||||
response: Final = await self.async_httpx_client.post(
|
||||
url=url,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""Data extraction functions for Opik payload building."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm import _logging
|
||||
|
|
@ -35,8 +36,8 @@ def normalize_provider_name(provider: str | None) -> str | None:
|
|||
|
||||
|
||||
def extract_opik_metadata(
|
||||
litellm_metadata: dict[str, Any],
|
||||
standard_logging_metadata: dict[str, Any],
|
||||
litellm_metadata: Mapping[str, Any],
|
||||
standard_logging_metadata: Mapping[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Merge Opik metadata from three sources in increasing priority order:
|
||||
|
|
@ -97,7 +98,7 @@ def extract_span_identifiers(
|
|||
|
||||
|
||||
def extract_tags(
|
||||
opik_metadata: dict[str, Any],
|
||||
opik_metadata: Mapping[str, Any],
|
||||
custom_llm_provider: str | None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
|
|
@ -122,7 +123,7 @@ def apply_proxy_header_overrides(
|
|||
project_name: str,
|
||||
tags: list[str],
|
||||
thread_id: str | None,
|
||||
proxy_headers: dict[str, Any],
|
||||
proxy_headers: Mapping[str, str],
|
||||
) -> tuple[str, list[str], str | None]:
|
||||
"""
|
||||
Apply overrides from proxy request headers (opik_* prefix).
|
||||
|
|
@ -148,7 +149,7 @@ def apply_proxy_header_overrides(
|
|||
thread_id = value
|
||||
elif param_key == "tags":
|
||||
try:
|
||||
parsed_tags = json.loads(value)
|
||||
parsed_tags: object = json.loads(value)
|
||||
if isinstance(parsed_tags, list):
|
||||
tags.extend(parsed_tags)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
|
|
@ -158,11 +159,11 @@ def apply_proxy_header_overrides(
|
|||
|
||||
|
||||
def extract_and_build_metadata(
|
||||
opik_metadata: dict[str, Any],
|
||||
standard_logging_metadata: dict[str, Any],
|
||||
standard_logging_object: dict[str, Any],
|
||||
litellm_kwargs: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
opik_metadata: Mapping[str, object],
|
||||
standard_logging_metadata: Mapping[str, object],
|
||||
standard_logging_object: Mapping[str, object],
|
||||
litellm_kwargs: Mapping[str, object],
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Build the complete metadata dictionary from all available sources.
|
||||
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -11,9 +11,10 @@ identical metrics. The attribute cardinality filter is reused from v1 by import
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Final, TypeAlias
|
||||
from typing import Any, Final, Literal, Protocol, TypeAlias
|
||||
|
||||
from opentelemetry.metrics import Histogram, Meter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -151,6 +152,29 @@ METRIC_ATTRIBUTE_CEILING: Final[frozenset[str]] = frozenset(
|
|||
BOUNDED_HIDDEN_PARAM_KEYS: Final[tuple[str, ...]] = ("model_id",)
|
||||
|
||||
|
||||
class _TokenUsage(TypedDict, total=False):
|
||||
"""The token counts a response's ``usage`` carries, as the recorder reads them."""
|
||||
|
||||
prompt_tokens: ReadOnly[int]
|
||||
completion_tokens: ReadOnly[int]
|
||||
|
||||
|
||||
class _ResponseView(Protocol):
|
||||
"""The one read the recorder makes on a litellm response object."""
|
||||
|
||||
def get(self, key: Literal["usage"], /) -> _TokenUsage | None: ...
|
||||
|
||||
|
||||
class _MetricKwargs(TypedDict, total=False):
|
||||
"""The logging kwargs the recorder reads directly."""
|
||||
|
||||
call_type: ReadOnly[str | None]
|
||||
litellm_params: ReadOnly[Mapping[str, object] | None]
|
||||
response_cost: ReadOnly[float | None]
|
||||
completion_start_time: ReadOnly[datetime | float | str | None]
|
||||
api_call_start_time: ReadOnly[datetime | float | str | None]
|
||||
|
||||
|
||||
def resolve_error_type(kwargs: Mapping[str, Any]) -> str:
|
||||
"""The ``error.type`` value for a failed request.
|
||||
|
||||
|
|
@ -192,8 +216,8 @@ class GenAIMetricRecorder:
|
|||
|
||||
def record(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
response_obj: Any,
|
||||
kwargs: _MetricKwargs,
|
||||
response_obj: _ResponseView | None,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
|
|
@ -218,7 +242,7 @@ class GenAIMetricRecorder:
|
|||
|
||||
def record_failure(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
kwargs: _MetricKwargs,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
|
|
@ -342,7 +366,7 @@ class GenAIMetricRecorder:
|
|||
# Per-metric recording
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _record_token_usage(self, response_obj: Any, common_attrs: dict) -> None:
|
||||
def _record_token_usage(self, response_obj: _ResponseView | None, common_attrs: dict) -> None:
|
||||
if not response_obj:
|
||||
return
|
||||
usage: Final = response_obj.get("usage")
|
||||
|
|
@ -353,7 +377,7 @@ class GenAIMetricRecorder:
|
|||
self._metrics.token_usage.record(usage.get("prompt_tokens", 0), attributes=in_attrs)
|
||||
self._metrics.token_usage.record(usage.get("completion_tokens", 0), attributes=out_attrs)
|
||||
|
||||
def _record_time_to_first_token(self, kwargs: Mapping[str, Any], common_attrs: dict) -> None:
|
||||
def _record_time_to_first_token(self, kwargs: _MetricKwargs, common_attrs: dict) -> None:
|
||||
time_to_first_chunk: Final = time_to_first_chunk_seconds(kwargs)
|
||||
if time_to_first_chunk is None:
|
||||
return
|
||||
|
|
@ -361,15 +385,14 @@ class GenAIMetricRecorder:
|
|||
|
||||
def _record_time_per_output_token(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
response_obj: Any,
|
||||
kwargs: _MetricKwargs,
|
||||
response_obj: _ResponseView | None,
|
||||
end_time: datetime,
|
||||
duration_s: float,
|
||||
common_attrs: dict,
|
||||
) -> None:
|
||||
completion_tokens = None
|
||||
if response_obj and (usage := response_obj.get("usage")):
|
||||
completion_tokens = usage.get("completion_tokens")
|
||||
usage: Final = response_obj.get("usage") if response_obj else None
|
||||
completion_tokens: Final = usage.get("completion_tokens") if usage else None
|
||||
if completion_tokens is None or completion_tokens <= 0:
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
from litellm.types.utils import CallTypes, StandardCallbackDynamicParams
|
||||
from litellm.types.vector_stores import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
VectorStoreResultContent,
|
||||
|
|
@ -226,7 +226,7 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
self,
|
||||
request_data: dict,
|
||||
response: Any,
|
||||
call_type: Any | None,
|
||||
call_type: CallTypes | None,
|
||||
) -> Any | None:
|
||||
"""
|
||||
Add search results to the response after successful LLM call.
|
||||
|
|
@ -283,7 +283,7 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
self,
|
||||
request_data: dict,
|
||||
response_chunk: Any,
|
||||
call_type: Any | None,
|
||||
call_type: CallTypes | None,
|
||||
) -> Any | None:
|
||||
"""
|
||||
Add search results to the final streaming chunk.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -536,6 +536,14 @@ def get_llm_provider(
|
|||
)
|
||||
|
||||
|
||||
def _dashscope_family_chat_config(custom_llm_provider: str) -> "litellm.DashScopeChatConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
return litellm.QwenCloudChatConfig()
|
||||
if custom_llm_provider == "qwen_ai_platform":
|
||||
return litellm.QwenAIPlatformChatConfig()
|
||||
return litellm.DashScopeChatConfig()
|
||||
|
||||
|
||||
def _get_openai_compatible_provider_info(
|
||||
model: str,
|
||||
api_base: str | None,
|
||||
|
|
@ -785,11 +793,11 @@ def _get_openai_compatible_provider_info(
|
|||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.HerokuChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
|
||||
elif custom_llm_provider == "dashscope":
|
||||
elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"):
|
||||
(
|
||||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.DashScopeChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
|
||||
) = _dashscope_family_chat_config(custom_llm_provider)._get_openai_compatible_provider_info(api_base, api_key)
|
||||
elif custom_llm_provider == "modelscope":
|
||||
(
|
||||
api_base,
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -1500,6 +1500,6 @@ class RealTimeStreaming:
|
|||
pass
|
||||
|
||||
|
||||
def client_sent_openai_beta_realtime_header(websocket: Any) -> bool:
|
||||
def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool:
|
||||
"""True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``."""
|
||||
return RealTimeStreaming._detect_beta_header(websocket)
|
||||
|
|
|
|||
|
|
@ -73,6 +73,18 @@ class _ContentChunk(TypedDict):
|
|||
choices: Sequence[_ContentChoice]
|
||||
|
||||
|
||||
class _FunctionCallDelta(TypedDict):
|
||||
function_call: ReadOnly[FunctionCall]
|
||||
|
||||
|
||||
class _FunctionCallChoice(TypedDict):
|
||||
delta: ReadOnly[_FunctionCallDelta]
|
||||
|
||||
|
||||
class _FunctionCallChunk(TypedDict):
|
||||
choices: ReadOnly[Sequence[_FunctionCallChoice]]
|
||||
|
||||
|
||||
class _AudioDelta(TypedDict, total=False):
|
||||
audio: ChatCompletionAudioDelta | None
|
||||
|
||||
|
|
@ -588,7 +600,7 @@ class ChunkProcessor:
|
|||
|
||||
return tool_calls_list
|
||||
|
||||
def get_combined_function_call_content(self, function_call_chunks: list[dict[str, Any]]) -> FunctionCall:
|
||||
def get_combined_function_call_content(self, function_call_chunks: Sequence["_FunctionCallChunk"]) -> FunctionCall:
|
||||
argument_list: Final = []
|
||||
delta = function_call_chunks[0]["choices"][0]["delta"]
|
||||
function_call = delta.get("function_call", "")
|
||||
|
|
|
|||
|
|
@ -862,6 +862,8 @@ class CustomStreamWrapper:
|
|||
model_response: Final = ModelResponseStream(**args)
|
||||
if self.response_id is not None:
|
||||
model_response.id = self.response_id
|
||||
elif model_response.id:
|
||||
self.response_id = model_response.id
|
||||
if self.system_fingerprint is not None:
|
||||
model_response.system_fingerprint = self.system_fingerprint
|
||||
|
||||
|
|
|
|||
|
|
@ -11,8 +11,11 @@ A2A Protocol Format:
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
|
@ -23,6 +26,13 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
class _A2ATextPart(TypedDict, total=False):
|
||||
"""The subset of an A2A message part this handler reads text from."""
|
||||
|
||||
kind: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
|
||||
|
||||
class A2AGuardrailHandler(BaseTranslation):
|
||||
"""
|
||||
Handler for processing A2A Protocol messages with guardrails.
|
||||
|
|
@ -41,7 +51,7 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
data: dict,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> Any:
|
||||
) -> dict:
|
||||
"""
|
||||
Process A2A input messages by applying guardrails to text content.
|
||||
|
||||
|
|
@ -214,12 +224,12 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
|
||||
async def process_output_streaming_response(
|
||||
self,
|
||||
responses_so_far: list[Any],
|
||||
responses_so_far: list[object],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
|
||||
request_data: dict | None = None,
|
||||
) -> list[Any]:
|
||||
) -> list[object]:
|
||||
"""
|
||||
Process A2A streaming output by applying guardrails to accumulated text.
|
||||
|
||||
|
|
@ -305,11 +315,12 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
|
||||
def _parse_streaming_responses(
|
||||
self,
|
||||
responses_so_far: list[Any],
|
||||
) -> tuple[list[dict[str, Any] | None], list[tuple[int, dict[str, Any]]]]:
|
||||
responses_so_far: list[object],
|
||||
) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]:
|
||||
"""Parse JSON-RPC items, returning aligned parsed list and valid entries."""
|
||||
parsed: Final[list[dict[str, Any] | None]] = [None] * len(responses_so_far)
|
||||
parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far)
|
||||
for i, item in enumerate(responses_so_far):
|
||||
obj: dict[str, object]
|
||||
if isinstance(item, dict):
|
||||
obj = item
|
||||
elif isinstance(item, str):
|
||||
|
|
@ -326,7 +337,7 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
|
||||
def _collect_text_from_parsed_chunks(
|
||||
self,
|
||||
valid_parsed: list[tuple[int, dict[str, Any]]],
|
||||
valid_parsed: list[tuple[int, dict[str, object]]],
|
||||
) -> tuple[str, list[int]]:
|
||||
"""Collect text from parsed chunks, returning combined text and indices."""
|
||||
from litellm.llms.a2a.common_utils import extract_text_from_a2a_response
|
||||
|
|
@ -411,7 +422,7 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
|
||||
def _extract_texts_from_parts(
|
||||
self,
|
||||
parts: list[dict[str, Any]],
|
||||
parts: Sequence[_A2ATextPart],
|
||||
path: tuple[str, ...],
|
||||
texts_to_check: list[str],
|
||||
task_mappings: list[tuple[tuple[str, ...], int]],
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, cast
|
|||
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
|
|
@ -125,7 +126,25 @@ else:
|
|||
_ANTHROPIC_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]")
|
||||
_ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128
|
||||
|
||||
_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[Any], bool]]] = MappingProxyType(
|
||||
|
||||
class _AnthropicUsageIteration(TypedDict, total=False):
|
||||
"""One entry of the ``usage.iterations`` array on an Anthropic response."""
|
||||
|
||||
input_tokens: ReadOnly[int | None]
|
||||
output_tokens: ReadOnly[int | None]
|
||||
cache_creation_input_tokens: ReadOnly[int | None]
|
||||
cache_read_input_tokens: ReadOnly[int | None]
|
||||
|
||||
|
||||
class _AnthropicToolResultBlock(TypedDict, total=False):
|
||||
"""A ``*_tool_result`` content block on an Anthropic response."""
|
||||
|
||||
type: ReadOnly[str]
|
||||
tool_use_id: ReadOnly[str]
|
||||
content: ReadOnly[object]
|
||||
|
||||
|
||||
_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[object], bool]]] = MappingProxyType(
|
||||
{
|
||||
"null": lambda v: v is None,
|
||||
"boolean": lambda v: isinstance(v, bool),
|
||||
|
|
@ -440,7 +459,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
optional_params.pop("speed", None)
|
||||
|
||||
@staticmethod
|
||||
def _raise_invalid_reasoning_effort(model: str, value: Any, llm_provider: str) -> NoReturn:
|
||||
def _raise_invalid_reasoning_effort(model: str, value: object, llm_provider: str) -> NoReturn:
|
||||
"""Raise a ``BadRequestError`` for an unrecognised ``reasoning_effort``.
|
||||
|
||||
Args:
|
||||
|
|
@ -1466,7 +1485,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
)
|
||||
|
||||
if _tool_choice is not None:
|
||||
optional_params["tool_choice"] = _tool_choice
|
||||
optional_params["tool_choice"] = AnthropicConfig._apply_forced_tool_choice(
|
||||
model=model, tool_choice=_tool_choice, drop_params=drop_params
|
||||
)
|
||||
elif param == "stream" and value is True:
|
||||
optional_params["stream"] = value
|
||||
elif param == "stop" and (isinstance(value, str) or isinstance(value, list)):
|
||||
|
|
@ -1495,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",
|
||||
|
|
@ -1992,19 +2013,35 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
return data
|
||||
|
||||
def _apply_output_config(self, data: dict, model: str, optional_params: dict) -> None:
|
||||
"""Validate and apply output_config to the request data."""
|
||||
"""Validate and apply output_config to the request data.
|
||||
|
||||
The ``drop_params`` gate here is an effort gate: ``format`` is a
|
||||
structured-output field, not an effort field, so it survives the drop
|
||||
and is vetted where it is consumed (the map's
|
||||
``supports_native_structured_output`` flag on emission paths).
|
||||
"""
|
||||
if "output_config" not in optional_params:
|
||||
return
|
||||
output_config: Final = optional_params.get("output_config")
|
||||
if not output_config or not isinstance(output_config, dict):
|
||||
return
|
||||
if litellm.drop_params is True and not self._model_supports_effort_param(model, self._resolved_provider):
|
||||
if (
|
||||
litellm.drop_params is True
|
||||
and any(key != "format" for key in output_config)
|
||||
and not self._model_supports_effort_param(model, self._resolved_provider)
|
||||
):
|
||||
litellm.verbose_logger.warning(
|
||||
DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING,
|
||||
model,
|
||||
)
|
||||
optional_params.pop("output_config", None)
|
||||
data.pop("output_config", None)
|
||||
preserved_format: Final = output_config.get("format")
|
||||
if preserved_format is None:
|
||||
optional_params.pop("output_config", None)
|
||||
data.pop("output_config", None)
|
||||
return
|
||||
format_only: Final = {"format": preserved_format} # mutable-ok: json body
|
||||
optional_params["output_config"] = format_only # rebind-ok: out-param store
|
||||
data["output_config"] = format_only # rebind-ok: out-param store
|
||||
return
|
||||
effort: Final = output_config.get("effort")
|
||||
valid_efforts: Final = ["high", "medium", "low", "xhigh", "max"]
|
||||
|
|
@ -2059,22 +2096,22 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
self, completion_response: dict
|
||||
) -> tuple[
|
||||
str,
|
||||
list[Any] | None,
|
||||
list[object] | None,
|
||||
list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None,
|
||||
str | None,
|
||||
list[ChatCompletionToolCallChunk],
|
||||
list[Any] | None,
|
||||
list[Any] | None,
|
||||
list[Any] | None,
|
||||
list[object] | None,
|
||||
list[_AnthropicToolResultBlock] | None,
|
||||
list[object] | None,
|
||||
]:
|
||||
text_content = ""
|
||||
citations: list[Any] | None = None
|
||||
citations: list[object] | None = None
|
||||
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None = None
|
||||
reasoning_content: str | None = None
|
||||
tool_calls: Final[list[ChatCompletionToolCallChunk]] = []
|
||||
web_search_results: list[Any] | None = None
|
||||
tool_results: list[Any] | None = None
|
||||
compaction_blocks: list[Any] | None = None
|
||||
web_search_results: list[object] | None = None
|
||||
tool_results: list[_AnthropicToolResultBlock] | None = None
|
||||
compaction_blocks: list[object] | None = None
|
||||
for idx, content in enumerate(completion_response["content"]):
|
||||
if content["type"] == "text":
|
||||
text_content += content["text"]
|
||||
|
|
@ -2284,7 +2321,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
raw_speed: Final = _usage.get("speed")
|
||||
resolved_speed: Final = raw_speed if isinstance(raw_speed, str) else speed
|
||||
|
||||
iterations: Final[list[Any] | None] = _usage.get("iterations")
|
||||
iterations: Final[Sequence[_AnthropicUsageIteration] | None] = _usage.get("iterations")
|
||||
if iterations:
|
||||
prompt_tokens = sum(it.get("input_tokens", 0) or 0 for it in iterations)
|
||||
completion_tokens = sum(it.get("output_tokens", 0) or 0 for it in iterations)
|
||||
|
|
@ -2377,7 +2414,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
|
||||
def _build_code_interpreter_results(
|
||||
self,
|
||||
tool_results: list[Any],
|
||||
tool_results: Sequence[_AnthropicToolResultBlock],
|
||||
code_by_id: dict[str, str],
|
||||
container_id: str | None,
|
||||
) -> list[OutputCodeInterpreterCall]:
|
||||
|
|
@ -2403,11 +2440,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
def _build_provider_specific_fields(
|
||||
self,
|
||||
completion_response: dict,
|
||||
citations: list[Any] | None,
|
||||
citations: Sequence[object] | None,
|
||||
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None,
|
||||
web_search_results: list[Any] | None,
|
||||
tool_results: list[Any] | None,
|
||||
compaction_blocks: list[Any] | None,
|
||||
web_search_results: Sequence[object] | None,
|
||||
tool_results: Sequence[_AnthropicToolResultBlock] | None,
|
||||
compaction_blocks: Sequence[object] | None,
|
||||
tool_calls: list[ChatCompletionToolCallChunk],
|
||||
) -> dict[str, Any]:
|
||||
provider_specific_fields: Final[dict[str, Any]] = {
|
||||
|
|
|
|||
|
|
@ -28,10 +28,15 @@ from litellm.types.llms.anthropic import (
|
|||
ANTHROPIC_OAUTH_TOKEN_PREFIX,
|
||||
AllAnthropicToolsValues,
|
||||
AnthropicMcpServerTool,
|
||||
AnthropicMessagesToolChoice,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.proxy.model_listing import ModelInfoResponse
|
||||
|
||||
DROP_FORCED_TOOL_CHOICE_WARNING: Final = (
|
||||
"Downgrading forced tool_choice to 'auto' for model=%s (drop_params=True): this model rejects tool_choice type "
|
||||
"'any'/'tool' with a 400 because thinking is always on and a forced call would skip it."
|
||||
)
|
||||
DROP_DISABLED_THINKING_WARNING: Final = (
|
||||
"Dropping `thinking={'type': 'disabled'}` for model=%s: thinking is always on for this model and cannot be "
|
||||
"disabled (the alternative is a provider 400). The model will still think adaptively, its response can contain "
|
||||
|
|
@ -320,6 +325,45 @@ 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 not AnthropicModelInfo.forced_tool_use_unsupported(model):
|
||||
return False
|
||||
if not (litellm.drop_params or drop_params):
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} does not support forced tool use (tool_choice='required' or a named tool). "
|
||||
"Use tool_choice='auto' and tell the model in the prompt when to call the tool, or set "
|
||||
"`litellm.drop_params = True` to downgrade to 'auto' automatically."
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
litellm.verbose_logger.warning(DROP_FORCED_TOOL_CHOICE_WARNING, model)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _apply_forced_tool_choice(
|
||||
model: str,
|
||||
tool_choice: AnthropicMessagesToolChoice,
|
||||
drop_params: bool,
|
||||
) -> AnthropicMessagesToolChoice:
|
||||
if tool_choice["type"] not in ("any", "tool"):
|
||||
return tool_choice
|
||||
if not AnthropicModelInfo.forced_tool_use_downgraded(model, drop_params):
|
||||
return tool_choice
|
||||
disable_parallel: Final = tool_choice.get("disable_parallel_tool_use")
|
||||
if disable_parallel is None:
|
||||
return AnthropicMessagesToolChoice(type="auto")
|
||||
return AnthropicMessagesToolChoice(type="auto", disable_parallel_tool_use=disable_parallel)
|
||||
|
||||
@staticmethod
|
||||
def _strip_version_suffix(model: str) -> str:
|
||||
at: Final = model.rfind("@")
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import asyncio
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Coroutine
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -116,7 +116,7 @@ class AnthropicFilesHandler:
|
|||
api_key: str | None = None,
|
||||
timeout: float | httpx.Timeout = 600.0,
|
||||
max_retries: int | None = None,
|
||||
) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]:
|
||||
) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]:
|
||||
"""
|
||||
Retrieve file content from Anthropic.
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import asyncio
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Callable, Coroutine
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from openai import (
|
||||
|
|
@ -374,7 +374,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
except Exception as e:
|
||||
status_code: Final = getattr(e, "status_code", 500)
|
||||
error_headers = getattr(e, "headers", None)
|
||||
error_response: Final = getattr(e, "response", None)
|
||||
error_response: Final[object] = getattr(e, "response", None)
|
||||
error_body: Final = getattr(e, "body", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
|
|
@ -392,7 +392,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
model: str,
|
||||
api_base: str,
|
||||
data: dict,
|
||||
timeout: Any,
|
||||
timeout: float | httpx.Timeout,
|
||||
dynamic_params: bool,
|
||||
model_response: ModelResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
|
|
@ -502,7 +502,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
dynamic_params: bool,
|
||||
data: dict[str, object],
|
||||
model: str,
|
||||
timeout: Any,
|
||||
timeout: float | httpx.Timeout,
|
||||
max_retries: int,
|
||||
azure_ad_token: str | None = None,
|
||||
azure_ad_token_provider: Callable | None = None,
|
||||
|
|
@ -578,7 +578,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
dynamic_params: bool,
|
||||
data: dict,
|
||||
model: str,
|
||||
timeout: Any,
|
||||
timeout: float | httpx.Timeout,
|
||||
max_retries: int,
|
||||
azure_ad_token: str | None = None,
|
||||
azure_ad_token_provider: Callable | None = None,
|
||||
|
|
@ -634,7 +634,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
except Exception as e:
|
||||
status_code: Final = getattr(e, "status_code", 500)
|
||||
error_headers = getattr(e, "headers", None)
|
||||
error_response: Final = getattr(e, "response", None)
|
||||
error_response: Final[object] = getattr(e, "response", None)
|
||||
message: Final = getattr(e, "message", str(e))
|
||||
error_body: Final = getattr(e, "body", None)
|
||||
if error_headers is None and error_response:
|
||||
|
|
@ -754,7 +754,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
aembedding=None,
|
||||
headers: dict | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
) -> EmbeddingResponse | Coroutine[Any, Any, EmbeddingResponse]:
|
||||
) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]:
|
||||
if headers:
|
||||
optional_params["extra_headers"] = headers
|
||||
if self._client_session is None:
|
||||
|
|
@ -1268,7 +1268,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
headers["Authorization"] = f"Bearer {azure_ad_token}"
|
||||
|
||||
# init AzureOpenAI Client
|
||||
azure_client_params: Final[dict[str, Any]] = self.initialize_azure_sdk_client(
|
||||
azure_client_params: Final[dict[str, object]] = self.initialize_azure_sdk_client(
|
||||
litellm_params=litellm_params or {},
|
||||
api_key=api_key,
|
||||
model_name=model or "",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -51,15 +51,13 @@ else:
|
|||
AsyncHTTPHandler = Any
|
||||
|
||||
|
||||
class _AzureRawAnnotation(TypedDict, total=False):
|
||||
type: ReadOnly[str]
|
||||
class _AzureRawAnnotation(ChatCompletionAnnotation, total=False):
|
||||
text: ReadOnly[str]
|
||||
start_index: ReadOnly[int]
|
||||
end_index: ReadOnly[int]
|
||||
url_citation: ReadOnly[ChatCompletionAnnotationURLCitation]
|
||||
|
||||
|
||||
_TransformedAnnotation: TypeAlias = ChatCompletionAnnotation | _AzureRawAnnotation
|
||||
_TransformedAnnotation: TypeAlias = ChatCompletionAnnotation
|
||||
|
||||
|
||||
class _AzureText(TypedDict, total=False):
|
||||
|
|
@ -223,18 +221,11 @@ class AzureAIAgentsHandler:
|
|||
"""Build the ModelResponse from agent output."""
|
||||
from litellm.types.utils import Choices, Message, Usage
|
||||
|
||||
message_kwargs: Final[dict[str, Any]] = {
|
||||
"content": content,
|
||||
"role": "assistant",
|
||||
}
|
||||
if annotations:
|
||||
message_kwargs["annotations"] = annotations
|
||||
|
||||
model_response.choices = [
|
||||
Choices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
message=Message(**message_kwargs),
|
||||
message=Message(content=content, role="assistant", annotations=annotations or None),
|
||||
)
|
||||
]
|
||||
model_response.model = model
|
||||
|
|
@ -655,9 +646,6 @@ class AzureAIAgentsHandler:
|
|||
|
||||
if data_str == "[DONE]":
|
||||
# Send final chunk with finish_reason
|
||||
final_delta_kwargs: dict[str, Any] = {"content": None}
|
||||
if collected_annotations:
|
||||
final_delta_kwargs["annotations"] = collected_annotations
|
||||
final_chunk = ModelResponseStream(
|
||||
id=response_id,
|
||||
created=created,
|
||||
|
|
@ -667,7 +655,7 @@ class AzureAIAgentsHandler:
|
|||
StreamingChoices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
delta=Delta(**final_delta_kwargs),
|
||||
delta=Delta(content=None, annotations=collected_annotations or None),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -588,6 +589,10 @@ class AmazonConverseConfig(BaseConfig):
|
|||
supported_params.append("context_management")
|
||||
return supported_params
|
||||
|
||||
@staticmethod
|
||||
def _auto_tool_choice() -> ToolChoiceValuesBlock:
|
||||
return ToolChoiceValuesBlock(auto={})
|
||||
|
||||
def map_tool_choice_values(
|
||||
self, model: str, tool_choice: str | dict, drop_params: bool
|
||||
) -> ToolChoiceValuesBlock | None:
|
||||
|
|
@ -600,10 +605,14 @@ class AmazonConverseConfig(BaseConfig):
|
|||
status_code=400,
|
||||
)
|
||||
elif tool_choice == "required":
|
||||
if AnthropicModelInfo.forced_tool_use_downgraded(model, drop_params):
|
||||
return self._auto_tool_choice()
|
||||
return ToolChoiceValuesBlock(any={})
|
||||
elif tool_choice == "auto":
|
||||
return ToolChoiceValuesBlock(auto={})
|
||||
return self._auto_tool_choice()
|
||||
elif isinstance(tool_choice, dict):
|
||||
if AnthropicModelInfo.forced_tool_use_downgraded(model, drop_params):
|
||||
return self._auto_tool_choice()
|
||||
# only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
|
||||
specific_tool: Final = SpecificToolChoiceBlock(
|
||||
name=make_valid_bedrock_tool_name(tool_choice.get("function", {}).get("name", ""))
|
||||
|
|
@ -1065,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)
|
||||
|
|
@ -1140,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)
|
||||
|
|
@ -1604,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,7 +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.litellm_core_utils.litellm_logging import verbose_logger
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
convert_to_anthropic_image_obj,
|
||||
)
|
||||
|
|
@ -12,21 +12,21 @@ 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,
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
convert_bedrock_invoke_output_format_to_inline_schema,
|
||||
apply_bedrock_invoke_structured_output,
|
||||
get_anthropic_beta_from_headers,
|
||||
normalize_bedrock_opus_output_config_effort,
|
||||
normalize_custom_field_on_tools,
|
||||
normalize_tool_input_schema_types_for_bedrock_invoke,
|
||||
pop_bedrock_invoke_output_config_format,
|
||||
strip_unsupported_bedrock_invoke_output_config_keys,
|
||||
)
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import _supports_factory
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
|
@ -76,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"
|
||||
|
||||
|
|
@ -103,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
|
||||
|
|
@ -212,36 +226,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
anthropic_request.pop("model", None)
|
||||
anthropic_request.pop("stream", None)
|
||||
anthropic_request.pop("stream_chunk_size", None)
|
||||
output_format: Final = anthropic_request.pop("output_format", None)
|
||||
output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_request)
|
||||
if output_format:
|
||||
convert_bedrock_invoke_output_format_to_inline_schema(
|
||||
output_format=output_format,
|
||||
request_body=anthropic_request,
|
||||
)
|
||||
elif output_config_format:
|
||||
convert_bedrock_invoke_output_format_to_inline_schema(
|
||||
output_format=output_config_format,
|
||||
request_body=anthropic_request,
|
||||
)
|
||||
if not (
|
||||
_supports_factory(
|
||||
model=model,
|
||||
custom_llm_provider="bedrock",
|
||||
key="supports_output_config",
|
||||
)
|
||||
or AnthropicConfig._model_supports_effort_param(model, "bedrock")
|
||||
):
|
||||
if anthropic_request.pop("output_config", None) is not None:
|
||||
verbose_logger.warning(
|
||||
"Bedrock Invoke: stripping unsupported `output_config` for "
|
||||
"model=%s — neither `supports_output_config` nor any "
|
||||
"`supports_*_reasoning_effort` flag is set in "
|
||||
"model_prices_and_context_window.json. Add the capability "
|
||||
"flag to the model JSON entry if this model accepts "
|
||||
"`output_config`.",
|
||||
model,
|
||||
)
|
||||
apply_bedrock_invoke_structured_output(
|
||||
model=model,
|
||||
request_body=anthropic_request,
|
||||
)
|
||||
strip_unsupported_bedrock_invoke_output_config_keys(
|
||||
model=model,
|
||||
request_body=anthropic_request,
|
||||
)
|
||||
if "anthropic_version" not in anthropic_request:
|
||||
anthropic_request["anthropic_version"] = self.anthropic_version
|
||||
|
||||
|
|
|
|||
|
|
@ -177,6 +177,95 @@ def convert_bedrock_invoke_output_format_to_inline_schema(
|
|||
request_body["messages"] = new_messages
|
||||
|
||||
|
||||
def _bedrock_model_supports(model: str, key: str) -> bool:
|
||||
from litellm.utils import _supports_factory
|
||||
|
||||
return _supports_factory(model=model, custom_llm_provider="bedrock", key=key)
|
||||
|
||||
|
||||
def apply_bedrock_invoke_structured_output(
|
||||
model: str,
|
||||
request_body: dict[str, object], # mutable-ok: edited in place like siblings
|
||||
) -> None:
|
||||
"""
|
||||
Route Anthropic structured-output params to what the Bedrock model supports.
|
||||
|
||||
Consumes the legacy top-level ``output_format`` and the newer
|
||||
``output_config.format``, keeping the pre-existing precedence of the legacy
|
||||
field when a request carries both. Models flagged
|
||||
``supports_native_structured_output`` in the model map get the schema
|
||||
forwarded as ``output_config.format``, which Bedrock relays to the model for
|
||||
enforced structured output. For every other model the schema is inlined into
|
||||
the last user message as best-effort text, with a warning because nothing
|
||||
enforces it.
|
||||
"""
|
||||
legacy_output_format: Final = request_body.pop("output_format", None)
|
||||
output_config_format: Final = pop_bedrock_invoke_output_config_format(request_body)
|
||||
schema_format: Final = legacy_output_format if isinstance(legacy_output_format, dict) else output_config_format
|
||||
if schema_format is None:
|
||||
return
|
||||
|
||||
if _bedrock_model_supports(model, "supports_native_structured_output"):
|
||||
existing_output_config: Final = request_body.get("output_config")
|
||||
if isinstance(existing_output_config, dict):
|
||||
existing_output_config["format"] = schema_format
|
||||
else:
|
||||
request_body["output_config"] = {"format": schema_format} # rebind-ok: out-param # mutable-ok: json
|
||||
return
|
||||
|
||||
verbose_logger.warning(
|
||||
"Bedrock Invoke: model=%s does not advertise `supports_native_structured_output` "
|
||||
"in model_prices_and_context_window.json, so the JSON schema was inlined into "
|
||||
"the last user message and is NOT enforced by the model.",
|
||||
model,
|
||||
)
|
||||
convert_bedrock_invoke_output_format_to_inline_schema(
|
||||
output_format=schema_format,
|
||||
request_body=request_body,
|
||||
)
|
||||
|
||||
|
||||
def strip_unsupported_bedrock_invoke_output_config_keys(
|
||||
model: str,
|
||||
request_body: dict[str, object], # mutable-ok: edited in place like siblings
|
||||
) -> None:
|
||||
"""
|
||||
Drop ``output_config`` keys the Bedrock model does not accept.
|
||||
|
||||
``format`` survives unconditionally: it is only attached for models whose map
|
||||
entry advertises ``supports_native_structured_output``. Effort-bearing keys
|
||||
survive only when the map flags ``supports_output_config`` or a
|
||||
``supports_*_reasoning_effort`` tier; otherwise they are dropped with a
|
||||
warning so Bedrock does not reject the request.
|
||||
"""
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
output_config: Final = request_body.get("output_config")
|
||||
if not isinstance(output_config, dict):
|
||||
return
|
||||
if all(key == "format" for key in output_config):
|
||||
return
|
||||
if _bedrock_model_supports(model, "supports_output_config") or AnthropicConfig._model_supports_effort_param(
|
||||
model, "bedrock"
|
||||
):
|
||||
return
|
||||
|
||||
verbose_logger.warning(
|
||||
"Bedrock Invoke: stripping unsupported `output_config` keys for "
|
||||
"model=%s: neither `supports_output_config` nor any "
|
||||
"`supports_*_reasoning_effort` flag is set in "
|
||||
"model_prices_and_context_window.json. Add the capability "
|
||||
"flag to the model JSON entry if this model accepts "
|
||||
"`output_config`.",
|
||||
model,
|
||||
)
|
||||
preserved_format: Final = output_config.get("format")
|
||||
if preserved_format is None:
|
||||
request_body.pop("output_config", None)
|
||||
else:
|
||||
request_body["output_config"] = {"format": preserved_format} # rebind-ok: out-param # mutable-ok: json
|
||||
|
||||
|
||||
def normalize_custom_field_on_tools(request_body: dict) -> None:
|
||||
"""
|
||||
Drop the ``custom`` field from each tool, first hoisting a boolean
|
||||
|
|
@ -727,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
|
||||
|
|
|
|||
|
|
@ -29,14 +29,14 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
|
|||
AmazonInvokeConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
convert_bedrock_invoke_output_format_to_inline_schema,
|
||||
apply_bedrock_invoke_structured_output,
|
||||
ensure_bedrock_anthropic_messages_tool_names,
|
||||
get_anthropic_beta_from_headers,
|
||||
is_claude_4_5_on_bedrock,
|
||||
normalize_bedrock_opus_output_config_effort,
|
||||
normalize_custom_field_on_tools,
|
||||
normalize_tool_input_schema_types_for_bedrock_invoke,
|
||||
pop_bedrock_invoke_output_config_format,
|
||||
strip_unsupported_bedrock_invoke_output_config_keys,
|
||||
)
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
bedrock_request_metadata_headers,
|
||||
|
|
@ -51,7 +51,6 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
from litellm.utils import _supports_factory
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
|
@ -708,52 +707,25 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
# 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models)
|
||||
self._remove_ttl_from_cache_control(anthropic_messages_request=anthropic_messages_request, model=model)
|
||||
|
||||
# 5. Convert structured-output params to inline schema.
|
||||
# Bedrock Invoke doesn't support top-level `output_format`; its
|
||||
# accepted `output_config` subset is also narrower than Anthropic's, so
|
||||
# consume the newer `output_config.format` shape here instead of
|
||||
# forwarding it as an unknown nested key.
|
||||
# 5. Route structured-output params (`output_format` /
|
||||
# `output_config.format`) to native enforcement or the inline-schema
|
||||
# fallback, then strip `output_config` keys the model does not accept.
|
||||
# Ref: https://github.com/BerriAI/litellm/issues/22797
|
||||
existing_output_config: Final = anthropic_messages_request.get("output_config")
|
||||
if isinstance(existing_output_config, dict):
|
||||
anthropic_messages_request["output_config"] = dict(existing_output_config)
|
||||
output_format: Final = anthropic_messages_request.pop("output_format", None)
|
||||
output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_messages_request)
|
||||
if output_format:
|
||||
convert_bedrock_invoke_output_format_to_inline_schema(
|
||||
output_format=output_format,
|
||||
request_body=anthropic_messages_request,
|
||||
)
|
||||
elif output_config_format:
|
||||
convert_bedrock_invoke_output_format_to_inline_schema(
|
||||
output_format=output_config_format,
|
||||
request_body=anthropic_messages_request,
|
||||
)
|
||||
apply_bedrock_invoke_structured_output(
|
||||
model=model,
|
||||
request_body=anthropic_messages_request,
|
||||
)
|
||||
normalize_bedrock_opus_output_config_effort(
|
||||
model=model,
|
||||
output_config=anthropic_messages_request.get("output_config"),
|
||||
)
|
||||
|
||||
# 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models,
|
||||
# but older models do not — strip it to avoid request rejection.
|
||||
# Ref: https://github.com/BerriAI/litellm/issues/22797
|
||||
if not (
|
||||
_supports_factory(
|
||||
model=model,
|
||||
custom_llm_provider="bedrock",
|
||||
key="supports_output_config",
|
||||
)
|
||||
or AnthropicConfig._model_supports_effort_param(model, "bedrock")
|
||||
):
|
||||
if anthropic_messages_request.pop("output_config", None) is not None:
|
||||
verbose_logger.warning(
|
||||
"Bedrock Invoke: stripping unsupported `output_config` for "
|
||||
"model=%s — neither `supports_output_config` nor any "
|
||||
"`supports_*_reasoning_effort` flag is set in "
|
||||
"model_prices_and_context_window.json. Add the capability "
|
||||
"flag to the model JSON entry if this model accepts "
|
||||
"`output_config`.",
|
||||
model,
|
||||
)
|
||||
strip_unsupported_bedrock_invoke_output_config_keys(
|
||||
model=model,
|
||||
request_body=anthropic_messages_request,
|
||||
)
|
||||
|
||||
# 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it)
|
||||
# Ref: https://github.com/BerriAI/litellm/issues/22847
|
||||
|
|
@ -774,9 +746,11 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
if filtered_betas:
|
||||
anthropic_messages_request["anthropic_beta"] = filtered_betas
|
||||
|
||||
remaining_output_config: Final = anthropic_messages_request.get("output_config")
|
||||
if (
|
||||
litellm.drop_params is True
|
||||
and "output_config" in anthropic_messages_request
|
||||
and isinstance(remaining_output_config, dict)
|
||||
and any(key != "format" for key in remaining_output_config)
|
||||
and not AnthropicConfig._model_supports_effort_param(model, "bedrock")
|
||||
):
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format.
|
|||
import base64
|
||||
import json
|
||||
import uuid as uuid_lib
|
||||
from typing import Any, Final, cast
|
||||
from typing import Final, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -633,7 +633,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
List of Bedrock format messages (JSON strings)
|
||||
"""
|
||||
try:
|
||||
json_message: Final = json.loads(message)
|
||||
json_message: Final[dict[str, object]] = json.loads(message)
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.warning("Invalid JSON message: %s", message[:200])
|
||||
return []
|
||||
|
|
@ -1182,7 +1182,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
|
||||
# Create a function call arguments done event
|
||||
# This is a custom event format that matches what clients expect
|
||||
function_call_event: Final[dict[str, Any]] = {
|
||||
function_call_event: Final[dict[str, object]] = {
|
||||
"type": "response.function_call_arguments.done",
|
||||
"event_id": f"event_{uuid.uuid4()}",
|
||||
"response_id": current_response_id,
|
||||
|
|
|
|||
|
|
@ -8,9 +8,11 @@ then we poll until the result is ready.
|
|||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any, Final
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -33,6 +35,42 @@ from ..common_utils import (
|
|||
from .transformation import BlackForestLabsImageEditConfig
|
||||
|
||||
|
||||
class _BFLSubmitBody(TypedDict, total=False):
|
||||
"""Decoded body of the BFL submit response, which hands back a polling URL."""
|
||||
|
||||
errors: ReadOnly[object]
|
||||
polling_url: ReadOnly[str]
|
||||
|
||||
|
||||
class _BFLPollBody(TypedDict, total=False):
|
||||
"""Decoded body of a BFL polling response."""
|
||||
|
||||
status: ReadOnly[str]
|
||||
|
||||
|
||||
class _BFLSubmitResponse(Protocol):
|
||||
"""The submit call's HTTP response, read for its status, body text and decoded body."""
|
||||
|
||||
@property
|
||||
def status_code(self) -> int: ...
|
||||
|
||||
@property
|
||||
def text(self) -> str: ...
|
||||
|
||||
def json(self) -> _BFLSubmitBody: ...
|
||||
|
||||
|
||||
class _BFLPollResponse(Protocol):
|
||||
"""A polling call's HTTP response, read only for the task status it carries."""
|
||||
|
||||
def json(self) -> _BFLPollBody: ...
|
||||
|
||||
|
||||
def _poll_status(response: _BFLPollResponse) -> str | None:
|
||||
"""Read the task status out of a BFL polling response body."""
|
||||
return response.json().get("status")
|
||||
|
||||
|
||||
class BlackForestLabsImageEdit:
|
||||
"""
|
||||
Black Forest Labs Image Edit handler.
|
||||
|
|
@ -53,10 +91,10 @@ class BlackForestLabsImageEdit:
|
|||
litellm_params: GenericLiteLLMParams | dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_headers: Mapping[str, object] | None = None,
|
||||
client: HTTPHandler | AsyncHTTPHandler | None = None,
|
||||
aimage_edit: bool = False,
|
||||
) -> ImageResponse | Any:
|
||||
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
|
||||
"""
|
||||
Main entry point for image edit requests.
|
||||
|
||||
|
|
@ -185,7 +223,7 @@ class BlackForestLabsImageEdit:
|
|||
litellm_params: GenericLiteLLMParams | dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_headers: Mapping[str, object] | None = None,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> ImageResponse:
|
||||
"""
|
||||
|
|
@ -281,7 +319,7 @@ class BlackForestLabsImageEdit:
|
|||
|
||||
def _poll_for_result_sync(
|
||||
self,
|
||||
initial_response: httpx.Response,
|
||||
initial_response: _BFLSubmitResponse,
|
||||
headers: dict,
|
||||
sync_client: HTTPHandler,
|
||||
max_wait: float = DEFAULT_MAX_POLLING_TIME,
|
||||
|
|
@ -356,8 +394,7 @@ class BlackForestLabsImageEdit:
|
|||
message=f"Polling failed: {response.text}",
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
status = data.get("status")
|
||||
status = _poll_status(response)
|
||||
|
||||
verbose_logger.debug("BFL poll status: %s", status)
|
||||
|
||||
|
|
@ -383,7 +420,7 @@ class BlackForestLabsImageEdit:
|
|||
|
||||
async def _poll_for_result_async(
|
||||
self,
|
||||
initial_response: httpx.Response,
|
||||
initial_response: _BFLSubmitResponse,
|
||||
headers: dict,
|
||||
async_client: AsyncHTTPHandler,
|
||||
max_wait: float = DEFAULT_MAX_POLLING_TIME,
|
||||
|
|
@ -447,8 +484,7 @@ class BlackForestLabsImageEdit:
|
|||
message=f"Polling failed: {response.text}",
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
status = data.get("status")
|
||||
status = _poll_status(response)
|
||||
|
||||
verbose_logger.debug("BFL poll status: %s", status)
|
||||
|
||||
|
|
|
|||
|
|
@ -8,9 +8,11 @@ then we poll until the result is ready.
|
|||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any, Final
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from typing import Final, Protocol, TypedDict
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -33,6 +35,23 @@ from ..common_utils import (
|
|||
from .transformation import BlackForestLabsImageGenerationConfig
|
||||
|
||||
|
||||
class _BFLTaskPayload(TypedDict, total=False):
|
||||
"""The body BFL returns for a submitted or polled generation task."""
|
||||
|
||||
errors: ReadOnly[object]
|
||||
polling_url: ReadOnly[str]
|
||||
status: ReadOnly[str]
|
||||
|
||||
|
||||
class _TaskJsonResponse(Protocol):
|
||||
def json(self) -> _BFLTaskPayload: ...
|
||||
|
||||
|
||||
def _task_payload(response: _TaskJsonResponse) -> _BFLTaskPayload:
|
||||
"""The JSON body of a BFL task submission or poll response."""
|
||||
return response.json()
|
||||
|
||||
|
||||
class BlackForestLabsImageGeneration:
|
||||
"""
|
||||
Black Forest Labs Image Generation handler.
|
||||
|
|
@ -53,10 +72,10 @@ class BlackForestLabsImageGeneration:
|
|||
litellm_params: GenericLiteLLMParams | dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_headers: Mapping[str, str] | None = None,
|
||||
client: HTTPHandler | AsyncHTTPHandler | None = None,
|
||||
aimg_generation: bool = False,
|
||||
) -> ImageResponse | Any:
|
||||
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
|
||||
"""
|
||||
Main entry point for image generation requests.
|
||||
|
||||
|
|
@ -187,7 +206,7 @@ class BlackForestLabsImageGeneration:
|
|||
litellm_params: GenericLiteLLMParams | dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_headers: Mapping[str, str] | None = None,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> ImageResponse:
|
||||
"""
|
||||
|
|
@ -305,7 +324,7 @@ class BlackForestLabsImageGeneration:
|
|||
|
||||
# Parse initial response to get polling URL
|
||||
try:
|
||||
response_data: Final = initial_response.json()
|
||||
response_data: Final = _task_payload(initial_response)
|
||||
except Exception as e:
|
||||
raise BlackForestLabsError(
|
||||
status_code=initial_response.status_code,
|
||||
|
|
@ -350,7 +369,7 @@ class BlackForestLabsImageGeneration:
|
|||
message=f"Polling failed: {response.text}",
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
data = _task_payload(response)
|
||||
status = data.get("status")
|
||||
|
||||
verbose_logger.debug("BFL poll status: %s", status)
|
||||
|
|
@ -396,7 +415,7 @@ class BlackForestLabsImageGeneration:
|
|||
|
||||
# Parse initial response to get polling URL
|
||||
try:
|
||||
response_data: Final = initial_response.json()
|
||||
response_data: Final = _task_payload(initial_response)
|
||||
except Exception as e:
|
||||
raise BlackForestLabsError(
|
||||
status_code=initial_response.status_code,
|
||||
|
|
@ -441,7 +460,7 @@ class BlackForestLabsImageGeneration:
|
|||
message=f"Polling failed: {response.text}",
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
data = _task_payload(response)
|
||||
status = data.get("status")
|
||||
|
||||
verbose_logger.debug("BFL poll status: %s", status)
|
||||
|
|
|
|||
|
|
@ -4,9 +4,10 @@
|
|||
import json
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
from typing import Final
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
|
|
@ -23,6 +24,53 @@ from litellm.types.utils import TextChoices
|
|||
from litellm.utils import CustomStreamWrapper, TextCompletionResponse
|
||||
|
||||
|
||||
class _CodestralChoiceMessage(TypedDict):
|
||||
"""`choices[].message` of a Codestral FIM completion."""
|
||||
|
||||
role: ReadOnly[NotRequired[str]]
|
||||
content: ReadOnly[NotRequired[str | None]]
|
||||
|
||||
|
||||
class _CodestralChoice(TypedDict):
|
||||
"""One entry of `choices` in a Codestral FIM completion."""
|
||||
|
||||
index: ReadOnly[int]
|
||||
message: ReadOnly[NotRequired[_CodestralChoiceMessage]]
|
||||
finish_reason: ReadOnly[NotRequired[str | None]]
|
||||
logprobs: ReadOnly[NotRequired[dict[str, object] | None]]
|
||||
|
||||
|
||||
class _CodestralUsage(TypedDict):
|
||||
"""Token accounting returned alongside a Codestral FIM completion."""
|
||||
|
||||
prompt_tokens: ReadOnly[NotRequired[int]]
|
||||
completion_tokens: ReadOnly[NotRequired[int]]
|
||||
total_tokens: ReadOnly[NotRequired[int]]
|
||||
|
||||
|
||||
class _CodestralCompletionResponse(TypedDict):
|
||||
"""Body returned by the Codestral `/v1/fim/completions` endpoint."""
|
||||
|
||||
id: ReadOnly[NotRequired[str]]
|
||||
created: ReadOnly[NotRequired[int]]
|
||||
model: ReadOnly[NotRequired[str]]
|
||||
object: ReadOnly[NotRequired[str]]
|
||||
usage: ReadOnly[NotRequired[_CodestralUsage]]
|
||||
choices: ReadOnly[NotRequired[list[_CodestralChoice]]]
|
||||
|
||||
|
||||
class _CodestralHTTPResponse(Protocol):
|
||||
"""The Codestral completion response as this handler reads it."""
|
||||
|
||||
@property
|
||||
def status_code(self) -> int: ...
|
||||
|
||||
@property
|
||||
def text(self) -> str: ...
|
||||
|
||||
def json(self) -> _CodestralCompletionResponse: ...
|
||||
|
||||
|
||||
class TextCompletionCodestralError(Exception):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -115,7 +163,7 @@ class CodestralTextCompletion:
|
|||
def process_text_completion_response(
|
||||
self,
|
||||
model: str,
|
||||
response: httpx.Response,
|
||||
response: _CodestralHTTPResponse,
|
||||
model_response: TextCompletionResponse,
|
||||
stream: bool,
|
||||
logging_obj: LiteLLMLogging,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -54,6 +54,9 @@ class DashScopeChatConfig(OpenAIGPTConfig):
|
|||
dynamic_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
return api_base, dynamic_api_key
|
||||
|
||||
def _resolve_chat_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
|
|
@ -66,10 +69,7 @@ class DashScopeChatConfig(OpenAIGPTConfig):
|
|||
"""
|
||||
If api_base is not provided, use the default DashScope /chat/completions endpoint.
|
||||
"""
|
||||
if not api_base:
|
||||
api_base = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
|
||||
if not api_base.endswith("/chat/completions"):
|
||||
api_base = f"{api_base}/chat/completions"
|
||||
|
||||
return api_base
|
||||
resolved_api_base: Final = self._resolve_chat_api_base(api_base)
|
||||
if resolved_api_base.endswith("/chat/completions"):
|
||||
return resolved_api_base
|
||||
return f"{resolved_api_base}/chat/completions"
|
||||
|
|
|
|||
|
|
@ -2,9 +2,89 @@
|
|||
Common utilities for the DashScope LLM provider.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
|
||||
|
||||
|
||||
def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
from litellm.llms.dashscope.qwencloud import QwenCloudEmbeddingConfig
|
||||
|
||||
return QwenCloudEmbeddingConfig()
|
||||
if custom_llm_provider == "qwen_ai_platform":
|
||||
from litellm.llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformEmbeddingConfig,
|
||||
)
|
||||
|
||||
return QwenAIPlatformEmbeddingConfig()
|
||||
from litellm.llms.dashscope.embed.transformation import DashScopeEmbeddingConfig
|
||||
|
||||
return DashScopeEmbeddingConfig()
|
||||
|
||||
|
||||
def get_dashscope_family_rerank_config(custom_llm_provider: str) -> "BaseRerankConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
from litellm.llms.dashscope.qwencloud import QwenCloudRerankConfig
|
||||
|
||||
return QwenCloudRerankConfig()
|
||||
if custom_llm_provider == "qwen_ai_platform":
|
||||
from litellm.llms.dashscope.qwen_ai_platform import QwenAIPlatformRerankConfig
|
||||
|
||||
return QwenAIPlatformRerankConfig()
|
||||
from litellm.llms.dashscope.rerank.transformation import DashScopeRerankConfig
|
||||
|
||||
return DashScopeRerankConfig()
|
||||
|
||||
|
||||
def get_dashscope_family_image_generation_config(
|
||||
custom_llm_provider: str,
|
||||
) -> "BaseImageGenerationConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
from litellm.llms.dashscope.qwencloud import QwenCloudImageGenerationConfig
|
||||
|
||||
return QwenCloudImageGenerationConfig()
|
||||
if custom_llm_provider == "qwen_ai_platform":
|
||||
from litellm.llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformImageGenerationConfig,
|
||||
)
|
||||
|
||||
return QwenAIPlatformImageGenerationConfig()
|
||||
from litellm.llms.dashscope.image_generation.transformation import (
|
||||
DashScopeImageGenerationConfig,
|
||||
)
|
||||
|
||||
return DashScopeImageGenerationConfig()
|
||||
|
||||
|
||||
def resolve_dashscope_family_api_key(custom_llm_provider: str, api_key: str | None) -> str | None:
|
||||
if custom_llm_provider == "dashscope":
|
||||
return api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
return api_key or get_secret_str(f"{custom_llm_provider.upper()}_API_KEY") or get_secret_str("DASHSCOPE_API_KEY")
|
||||
|
||||
|
||||
def missing_dashscope_family_key_message(custom_llm_provider: str) -> str:
|
||||
if custom_llm_provider == "qwencloud":
|
||||
return (
|
||||
"Missing API key for QwenCloud. Set QWENCLOUD_API_KEY or "
|
||||
"DASHSCOPE_API_KEY environment variable or pass api_key parameter."
|
||||
)
|
||||
if custom_llm_provider == "qwen_ai_platform":
|
||||
return (
|
||||
"Missing API key for Qwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or "
|
||||
"DASHSCOPE_API_KEY environment variable or pass api_key parameter."
|
||||
)
|
||||
return "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter."
|
||||
|
||||
|
||||
class DashScopeError(BaseLLMException):
|
||||
|
|
|
|||
|
|
@ -110,7 +110,7 @@ def _calculate_completion_cost(
|
|||
return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost)
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
|
||||
def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashscope") -> tuple[float, float]:
|
||||
"""
|
||||
Calculate cost per token for Dashscope models.
|
||||
|
||||
|
|
@ -119,11 +119,12 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
|
|||
Args:
|
||||
model: Model name without provider prefix
|
||||
usage: LiteLLM Usage block
|
||||
custom_llm_provider: The provider id the request resolved to; dashscope or one of its brand aliases
|
||||
|
||||
Returns:
|
||||
Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd)
|
||||
"""
|
||||
model_info: Final = get_model_info(model=model, custom_llm_provider="dashscope")
|
||||
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
breakdown: Final = _extract_token_breakdown(usage)
|
||||
raw_tiers: Final = model_info.get("tiered_pricing")
|
||||
tiered_pricing: Final = raw_tiers if isinstance(raw_tiers, list) else None
|
||||
|
|
|
|||
|
|
@ -62,6 +62,17 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
|||
# for drop_params=False before this method is called.
|
||||
return optional_params
|
||||
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
resolved_api_key: Final = api_key if api_key is not None else get_secret_str("DASHSCOPE_API_KEY")
|
||||
if resolved_api_key is None:
|
||||
raise ValueError(
|
||||
"DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly."
|
||||
)
|
||||
return resolved_api_key
|
||||
|
||||
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
|
|
@ -72,17 +83,11 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("DASHSCOPE_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly."
|
||||
)
|
||||
default_headers: Final = {
|
||||
return {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Authorization": f"Bearer {self._resolve_api_key(api_key)}",
|
||||
**headers,
|
||||
}
|
||||
return {**default_headers, **headers}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
|
|
@ -93,8 +98,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
|||
litellm_params: dict,
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
base = api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE
|
||||
base = base.rstrip("/")
|
||||
base: Final = self._resolve_embedding_api_base(api_base).rstrip("/")
|
||||
if base.endswith("/embeddings"):
|
||||
return base
|
||||
return f"{base}/embeddings"
|
||||
|
|
|
|||
|
|
@ -91,6 +91,15 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
mapped[k] = v
|
||||
return mapped
|
||||
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
resolved_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
if not resolved_api_key:
|
||||
raise ValueError("DASHSCOPE_API_KEY is not set")
|
||||
return resolved_api_key
|
||||
|
||||
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
|
||||
return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
|
|
@ -103,7 +112,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
image_api_base: Final = (
|
||||
api_base if api_base and not api_base.rstrip("/").endswith(CHAT_COMPATIBLE_MODE_PATH) else None
|
||||
)
|
||||
return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE
|
||||
return self._resolve_image_api_base(image_api_base)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
|
|
@ -115,10 +124,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
final_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
if not final_api_key:
|
||||
raise ValueError("DASHSCOPE_API_KEY is not set")
|
||||
headers["Authorization"] = f"Bearer {final_api_key}"
|
||||
headers["Authorization"] = f"Bearer {self._resolve_api_key(api_key)}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
|
||||
|
|
|
|||
62
litellm/llms/dashscope/qwen_ai_platform.py
Normal file
62
litellm/llms/dashscope/qwen_ai_platform.py
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
from .chat.transformation import DashScopeChatConfig
|
||||
from .embed.transformation import DashScopeEmbeddingConfig
|
||||
from .image_generation.transformation import DashScopeImageGenerationConfig
|
||||
from .rerank.transformation import DashScopeRerankConfig
|
||||
|
||||
QWEN_AI_PLATFORM_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
QWEN_AI_PLATFORM_RERANK_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
QWEN_AI_PLATFORM_IMAGE_API_BASE: Final = (
|
||||
"https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
|
||||
)
|
||||
|
||||
|
||||
def _resolve_qwen_ai_platform_api_key(api_key: str | None) -> str | None:
|
||||
return api_key or get_secret_str("QWEN_AI_PLATFORM_API_KEY") or get_secret_str("DASHSCOPE_API_KEY")
|
||||
|
||||
|
||||
def _require_qwen_ai_platform_api_key(api_key: str | None) -> str:
|
||||
resolved: Final = _resolve_qwen_ai_platform_api_key(api_key)
|
||||
if resolved is None:
|
||||
raise ValueError(
|
||||
"Qwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var "
|
||||
"or pass api_key explicitly."
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
class QwenAIPlatformChatConfig(DashScopeChatConfig):
|
||||
def _get_openai_compatible_provider_info(
|
||||
self, api_base: str | None, api_key: str | None
|
||||
) -> tuple[str | None, str | None]:
|
||||
return self._resolve_chat_api_base(api_base), _resolve_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_chat_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
|
||||
|
||||
|
||||
class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig):
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
|
||||
|
||||
|
||||
class QwenAIPlatformRerankConfig(DashScopeRerankConfig):
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_rerank_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_RERANK") or QWEN_AI_PLATFORM_RERANK_API_BASE
|
||||
|
||||
|
||||
class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig):
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwen_ai_platform_api_key(api_key)
|
||||
|
||||
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
|
||||
return image_api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE") or QWEN_AI_PLATFORM_IMAGE_API_BASE
|
||||
62
litellm/llms/dashscope/qwencloud.py
Normal file
62
litellm/llms/dashscope/qwencloud.py
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
from .chat.transformation import DashScopeChatConfig
|
||||
from .embed.transformation import DashScopeEmbeddingConfig
|
||||
from .image_generation.transformation import DashScopeImageGenerationConfig
|
||||
from .rerank.transformation import DashScopeRerankConfig
|
||||
|
||||
QWENCLOUD_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
|
||||
QWENCLOUD_RERANK_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks"
|
||||
QWENCLOUD_IMAGE_API_BASE: Final = (
|
||||
"https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
|
||||
)
|
||||
|
||||
|
||||
def _resolve_qwencloud_api_key(api_key: str | None) -> str | None:
|
||||
return api_key or get_secret_str("QWENCLOUD_API_KEY") or get_secret_str("DASHSCOPE_API_KEY")
|
||||
|
||||
|
||||
def _require_qwencloud_api_key(api_key: str | None) -> str:
|
||||
resolved: Final = _resolve_qwencloud_api_key(api_key)
|
||||
if resolved is None:
|
||||
raise ValueError(
|
||||
"QwenCloud API key is required. Set 'QWENCLOUD_API_KEY' or 'DASHSCOPE_API_KEY' env var "
|
||||
"or pass api_key explicitly."
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
class QwenCloudChatConfig(DashScopeChatConfig):
|
||||
def _get_openai_compatible_provider_info(
|
||||
self, api_base: str | None, api_key: str | None
|
||||
) -> tuple[str | None, str | None]:
|
||||
return self._resolve_chat_api_base(api_base), _resolve_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_chat_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
|
||||
|
||||
|
||||
class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig):
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
|
||||
|
||||
|
||||
class QwenCloudRerankConfig(DashScopeRerankConfig):
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_rerank_api_base(self, api_base: str | None) -> str:
|
||||
return api_base or get_secret_str("QWENCLOUD_API_BASE_RERANK") or QWENCLOUD_RERANK_API_BASE
|
||||
|
||||
|
||||
class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig):
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
return _require_qwencloud_api_key(api_key)
|
||||
|
||||
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
|
||||
return image_api_base or get_secret_str("QWENCLOUD_API_BASE_IMAGE") or QWENCLOUD_IMAGE_API_BASE
|
||||
|
|
@ -58,19 +58,30 @@ class DashScopeRerankConfig(BaseRerankConfig):
|
|||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def _resolve_api_key(self, api_key: str | None) -> str:
|
||||
resolved_api_key: Final = api_key if api_key is not None else get_secret_str("DASHSCOPE_API_KEY")
|
||||
if resolved_api_key is None:
|
||||
raise ValueError(
|
||||
"DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly."
|
||||
)
|
||||
return resolved_api_key
|
||||
|
||||
def _resolve_rerank_api_base(self, api_base: str | None) -> str:
|
||||
if api_base is not None:
|
||||
return api_base
|
||||
return get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
model: str,
|
||||
optional_params: dict | None = None,
|
||||
) -> str:
|
||||
if api_base is None:
|
||||
api_base = get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL
|
||||
resolved_api_base: Final = self._resolve_rerank_api_base(api_base)
|
||||
if resolved_api_base == DEFAULT_RERANK_URL:
|
||||
return resolved_api_base
|
||||
|
||||
if api_base == DEFAULT_RERANK_URL:
|
||||
return DEFAULT_RERANK_URL
|
||||
|
||||
cleaned: Final = api_base.rstrip("/")
|
||||
cleaned: Final = resolved_api_base.rstrip("/")
|
||||
if cleaned.endswith("/reranks") or cleaned.endswith("/rerank"):
|
||||
return cleaned
|
||||
|
||||
|
|
@ -88,19 +99,12 @@ class DashScopeRerankConfig(BaseRerankConfig):
|
|||
optional_params: dict | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> dict:
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("DASHSCOPE_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly."
|
||||
)
|
||||
|
||||
default_headers: Final = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
return {
|
||||
"Authorization": f"Bearer {self._resolve_api_key(api_key)}",
|
||||
"accept": "application/json",
|
||||
"content-type": "application/json",
|
||||
**headers,
|
||||
}
|
||||
return {**default_headers, **headers}
|
||||
|
||||
def get_supported_cohere_rerank_params(self, model: str) -> list:
|
||||
return ["query", "documents", "top_n", "return_documents"]
|
||||
|
|
|
|||
|
|
@ -2,10 +2,11 @@
|
|||
Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -24,6 +25,36 @@ from litellm.types.rerank import (
|
|||
)
|
||||
|
||||
|
||||
class _DeepinfraInferenceStatus(TypedDict, total=False):
|
||||
"""The ``inference_status`` block of a DeepInfra rerank response."""
|
||||
|
||||
status: ReadOnly[str]
|
||||
runtime_ms: ReadOnly[float]
|
||||
cost: ReadOnly[float]
|
||||
tokens_generated: ReadOnly[int]
|
||||
tokens_input: ReadOnly[int]
|
||||
|
||||
|
||||
class _DeepinfraRerankResponse(TypedDict, total=False):
|
||||
"""Body of a DeepInfra ``/rerank`` response."""
|
||||
|
||||
scores: ReadOnly[Sequence[float]]
|
||||
input_tokens: ReadOnly[int]
|
||||
request_id: ReadOnly[str | None]
|
||||
inference_status: ReadOnly[_DeepinfraInferenceStatus]
|
||||
|
||||
|
||||
class _DeepinfraRerankResponseSource(Protocol):
|
||||
"""The DeepInfra ``/rerank`` HTTP response, read for the body it decodes to."""
|
||||
|
||||
def json(self) -> _DeepinfraRerankResponse: ...
|
||||
|
||||
|
||||
def _deepinfra_rerank_body(response: _DeepinfraRerankResponseSource) -> _DeepinfraRerankResponse:
|
||||
"""Decode the body of a DeepInfra ``/rerank`` response."""
|
||||
return response.json()
|
||||
|
||||
|
||||
class DeepinfraRerankConfig(BaseRerankConfig):
|
||||
"""
|
||||
Deepinfra Rerank - Follows the same Spec as Cohere Rerank
|
||||
|
|
@ -95,7 +126,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
|
|||
model: str,
|
||||
drop_params: bool,
|
||||
query: str,
|
||||
documents: list[str | dict[str, Any]],
|
||||
documents: list[str | dict[str, object]],
|
||||
custom_llm_provider: str | None = None,
|
||||
top_n: int | None = None,
|
||||
rank_fields: list[str] | None = None,
|
||||
|
|
@ -150,7 +181,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
|
|||
litellm_params: dict = {},
|
||||
) -> RerankResponse:
|
||||
try:
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final = _deepinfra_rerank_body(raw_response)
|
||||
logging_obj.post_call(original_response=raw_response.text)
|
||||
|
||||
# Extract the scores from the response
|
||||
|
|
|
|||
|
|
@ -12,9 +12,10 @@ Schema versioning:
|
|||
litellm.use_legacy_interactions_schema = True. Remove flag after June 8, 2026.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -41,6 +42,53 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
_JsonObject: TypeAlias = dict[str, object]
|
||||
|
||||
|
||||
class _InteractionPayload(TypedDict, total=False):
|
||||
"""JSON body of an Interactions API interaction, keyed as ``InteractionsAPIResponse`` fields."""
|
||||
|
||||
id: ReadOnly[str | None]
|
||||
object: ReadOnly[str | None]
|
||||
model: ReadOnly[str | None]
|
||||
agent: ReadOnly[str | None]
|
||||
status: ReadOnly[str | None]
|
||||
created: ReadOnly[str | None]
|
||||
updated: ReadOnly[str | None]
|
||||
outputs: ReadOnly[list[_JsonObject] | None]
|
||||
steps: ReadOnly[list[_JsonObject] | None]
|
||||
usage: ReadOnly[_JsonObject | None]
|
||||
|
||||
|
||||
class _CancelPayload(TypedDict, total=False):
|
||||
"""JSON body of an Interactions API cancel response."""
|
||||
|
||||
id: ReadOnly[str | None]
|
||||
status: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _InteractionPayloadSource(Protocol):
|
||||
"""An Interactions API HTTP response, read for the interaction body it decodes to."""
|
||||
|
||||
def json(self) -> _InteractionPayload: ...
|
||||
|
||||
|
||||
class _CancelPayloadSource(Protocol):
|
||||
"""An Interactions API cancel HTTP response, read for the body it decodes to."""
|
||||
|
||||
def json(self) -> _CancelPayload: ...
|
||||
|
||||
|
||||
def _interaction_body(response: _InteractionPayloadSource) -> _InteractionPayload:
|
||||
"""Decode the body of an Interactions API interaction response."""
|
||||
return response.json()
|
||||
|
||||
|
||||
def _cancel_body(response: _CancelPayloadSource) -> _CancelPayload:
|
||||
"""Decode the body of an Interactions API cancel response."""
|
||||
return response.json()
|
||||
|
||||
|
||||
class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
||||
"""
|
||||
Configuration for Google AI Studio Interactions API.
|
||||
|
|
@ -143,7 +191,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
"""
|
||||
use_legacy: Final[bool] = litellm.use_legacy_interactions_schema
|
||||
|
||||
request_body: Final[dict[str, Any]] = {}
|
||||
request_body: Final[dict[str, object]] = {}
|
||||
|
||||
# Model or Agent (one required)
|
||||
if model:
|
||||
|
|
@ -189,7 +237,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
and (not isinstance(response_format, dict) or "mime_type" not in response_format)
|
||||
):
|
||||
# Wrap the legacy schema into the new polymorphic format.
|
||||
new_rf: Final[dict[str, Any]] = {
|
||||
new_rf: Final[dict[str, object]] = {
|
||||
"type": "text",
|
||||
"mime_type": response_mime_type,
|
||||
}
|
||||
|
|
@ -215,7 +263,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
|
||||
if image_config is not None:
|
||||
# Move image_config to response_format with type=image.
|
||||
image_rf: Final[dict[str, Any]] = {"type": "image", **image_config}
|
||||
image_rf: Final[_JsonObject] = {"type": "image", **image_config}
|
||||
existing_rf: Final = request_body.get("response_format")
|
||||
if existing_rf is None:
|
||||
request_body["response_format"] = image_rf
|
||||
|
|
@ -239,7 +287,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
original_response=raw_response.text,
|
||||
additional_args={"complete_input_dict": {}},
|
||||
)
|
||||
raw_json: Final = raw_response.json()
|
||||
raw_json: Final = _interaction_body(raw_response)
|
||||
except Exception:
|
||||
raise GeminiError(
|
||||
message=raw_response.text,
|
||||
|
|
@ -290,7 +338,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> InteractionsAPIResponse:
|
||||
try:
|
||||
raw_json: Final = raw_response.json()
|
||||
raw_json: Final = _interaction_body(raw_response)
|
||||
except Exception:
|
||||
raise GeminiError(
|
||||
message=raw_response.text,
|
||||
|
|
@ -355,7 +403,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> CancelInteractionResult:
|
||||
try:
|
||||
raw_json: Final = raw_response.json()
|
||||
raw_json: Final = _cancel_body(raw_response)
|
||||
except Exception:
|
||||
raise GeminiError(
|
||||
message=raw_response.text,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import base64
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -54,8 +55,13 @@ def _convert_image_to_gemini_format(image_file) -> dict[str, str]:
|
|||
return {"bytesBase64Encoded": base64_encoded, "mimeType": mime_type}
|
||||
|
||||
|
||||
def _json_payload(raw_response: httpx.Response) -> object:
|
||||
"""Read an HTTP response body as an opaque JSON payload."""
|
||||
return raw_response.json()
|
||||
|
||||
|
||||
def _usage_video_resolution_from_parameters(
|
||||
parameters: dict[str, Any],
|
||||
parameters: Mapping[str, object],
|
||||
) -> str | None:
|
||||
"""Normalize Veo ``parameters.resolution`` for usage and cost tracking."""
|
||||
res: Final = parameters.get("resolution")
|
||||
|
|
@ -97,7 +103,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
video_create_optional_params: VideoCreateOptionalRequestParams,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Map OpenAI-style parameters to Veo format.
|
||||
|
||||
|
|
@ -111,7 +117,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
|
||||
All other params are passed through as-is to support Gemini-specific parameters.
|
||||
"""
|
||||
mapped_params: Final[dict[str, Any]] = {}
|
||||
mapped_params: Final[dict[str, object]] = {}
|
||||
|
||||
# Get supported OpenAI params (exclude "model" and "prompt" which are handled separately)
|
||||
supported_openai_params: Final = self.get_supported_openai_params(model)
|
||||
|
|
@ -312,11 +318,11 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
- status: "processing"
|
||||
- usage: includes duration_seconds and optional video_resolution for cost calculation
|
||||
"""
|
||||
response_data: Final = raw_response.json()
|
||||
response_data: Final = _json_payload(raw_response)
|
||||
|
||||
# Parse response using Pydantic model for type safety
|
||||
try:
|
||||
operation_response: Final = GeminiLongRunningOperationResponse(**response_data)
|
||||
operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to parse operation response: {e}")
|
||||
|
||||
|
|
@ -336,7 +342,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
model=model,
|
||||
)
|
||||
|
||||
usage_data: Final[dict[str, Any]] = {}
|
||||
usage_data: Final[dict[str, float | str]] = {}
|
||||
if request_data:
|
||||
parameters: Final = request_data.get("parameters", {})
|
||||
duration: Final = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
|
||||
|
|
@ -367,7 +373,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
"""
|
||||
operation_name: Final = extract_original_video_id(video_id)
|
||||
url: Final = f"{api_base.rstrip('/')}/v1beta/{operation_name}"
|
||||
params: Final[dict[str, Any]] = {}
|
||||
params: Final[dict[str, object]] = {}
|
||||
|
||||
return url, params
|
||||
|
||||
|
|
@ -403,9 +409,9 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
}
|
||||
}
|
||||
"""
|
||||
response_data: Final = raw_response.json()
|
||||
response_data: Final = _json_payload(raw_response)
|
||||
# Parse response using Pydantic model for type safety
|
||||
operation_response: Final = GeminiLongRunningOperationResponse(**response_data)
|
||||
operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data)
|
||||
|
||||
operation_name: Final = operation_response.name
|
||||
is_done: Final = operation_response.done
|
||||
|
|
@ -443,9 +449,9 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
client: Final = litellm.module_level_client
|
||||
status_response: Final = client.get(url=status_url, headers=headers)
|
||||
status_response.raise_for_status()
|
||||
response_data: Final = status_response.json()
|
||||
response_data: Final = _json_payload(status_response)
|
||||
|
||||
operation_response: Final = GeminiLongRunningOperationResponse(**response_data)
|
||||
operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data)
|
||||
|
||||
if not operation_response.done:
|
||||
raise ValueError(
|
||||
|
|
@ -458,7 +464,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
generated_samples: Final = operation_response.response.generateVideoResponse.generatedSamples
|
||||
download_url: Final = generated_samples[0].video.uri
|
||||
|
||||
params: Final[dict[str, Any]] = {}
|
||||
params: Final[dict[str, object]] = {}
|
||||
|
||||
return download_url, params
|
||||
|
||||
|
|
@ -480,7 +486,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
api_base: str,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
) -> tuple[str, dict]:
|
||||
"""
|
||||
Video remix is not supported by Veo API.
|
||||
|
|
@ -506,7 +512,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
after: str | None = None,
|
||||
limit: int | None = None,
|
||||
order: str | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_query: Mapping[str, object] | None = None,
|
||||
) -> tuple[str, dict]:
|
||||
"""
|
||||
Video list is not supported by Veo API.
|
||||
|
|
@ -547,7 +553,7 @@ class GeminiVideoConfig(BaseVideoConfig):
|
|||
"""Video delete is not supported."""
|
||||
raise NotImplementedError("Video delete is not supported by Google Veo.")
|
||||
|
||||
def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers):
|
||||
def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers):
|
||||
raise NotImplementedError("video create character is not supported for Gemini")
|
||||
|
||||
def transform_video_create_character_response(self, raw_response, logging_obj):
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -24,6 +25,8 @@ from litellm.utils import token_counter
|
|||
from ..common_utils import HuggingFaceError, hf_task_list, hf_tasks, output_parser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import tiktoken
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
LoggingClass = LiteLLMLoggingObj
|
||||
|
|
@ -31,6 +34,12 @@ else:
|
|||
LoggingClass = Any
|
||||
|
||||
|
||||
class _TokenEncoding(Protocol):
|
||||
"""Tokenizer handle the caller passes in; only `encode` is used, to count completion tokens."""
|
||||
|
||||
def encode(self, text: str, /) -> Sequence[object]: ...
|
||||
|
||||
|
||||
tgi_models_cache = None
|
||||
conv_models_cache = None
|
||||
|
||||
|
|
@ -369,7 +378,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
|
|||
model_response: ModelResponse,
|
||||
task: hf_tasks | None,
|
||||
optional_params: dict,
|
||||
encoding: Any,
|
||||
encoding: "_TokenEncoding | None",
|
||||
messages: list[AllMessageValues],
|
||||
model: str,
|
||||
):
|
||||
|
|
@ -439,9 +448,10 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
|
|||
if output_text is not None and len(output_text) > 0:
|
||||
completion_tokens = 0
|
||||
try:
|
||||
completion_tokens = len(
|
||||
encoding.encode(model_response["choices"][0]["message"].get("content", ""))
|
||||
) ##[TODO] use the llama2 tokenizer here
|
||||
if encoding is not None:
|
||||
completion_tokens = len(
|
||||
encoding.encode(model_response["choices"][0]["message"].get("content", ""))
|
||||
) ##[TODO] use the llama2 tokenizer here
|
||||
except Exception:
|
||||
# this should remain non blocking we should not block a response returning if calculating usage fails
|
||||
pass
|
||||
|
|
@ -469,7 +479,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
|
|||
messages: list[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: Any,
|
||||
encoding: "tiktoken.Encoding | None",
|
||||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ModelResponse:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -325,7 +330,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
@overload
|
||||
def _transform_messages(
|
||||
self, messages: list[AllMessageValues], model: str, is_async: Literal[True]
|
||||
) -> Coroutine[Any, Any, list[AllMessageValues]]:
|
||||
) -> Coroutine[object, object, list[AllMessageValues]]:
|
||||
...
|
||||
|
||||
@overload
|
||||
|
|
@ -341,7 +346,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
|
||||
def _transform_messages(
|
||||
self, messages: list[AllMessageValues], model: str, is_async: bool = False
|
||||
) -> list[AllMessageValues] | Coroutine[Any, Any, list[AllMessageValues]]:
|
||||
) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]:
|
||||
"""OpenAI no longer supports image_url as a string, so we need to convert it to a dict"""
|
||||
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(messages)
|
||||
hoisted_messages: Final = hoist_images_from_tool_messages(stripped_messages)
|
||||
|
|
@ -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
|
||||
|
|
@ -497,8 +542,12 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
return None
|
||||
tool_call_names: Final = get_tool_call_names(optional_params.get("tools", []))
|
||||
try:
|
||||
json_content: Final = json.loads(content)
|
||||
if json_content.get("type") == "function" and json_content.get("name") in tool_call_names:
|
||||
json_content: Final[object] = json.loads(content)
|
||||
if (
|
||||
isinstance(json_content, dict)
|
||||
and json_content.get("type") == "function"
|
||||
and json_content.get("name") in tool_call_names
|
||||
):
|
||||
return ChatCompletionMessageToolCall(
|
||||
function=Function(
|
||||
name=json_content.get("name"),
|
||||
|
|
@ -622,7 +671,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
|
|||
|
||||
## RESPONSE OBJECT
|
||||
try:
|
||||
completion_response: Final = raw_response.json()
|
||||
completion_response: Final[dict[str, object]] = raw_response.json()
|
||||
except Exception as e:
|
||||
response_headers: Final = getattr(raw_response, "headers", None)
|
||||
raise OpenAIError(
|
||||
|
|
|
|||
|
|
@ -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,8 +56,12 @@ 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
|
||||
|
||||
|
||||
class OpenAIChatCompletionsHandler(BaseTranslation):
|
||||
|
|
@ -80,7 +91,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
data: dict,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> Any:
|
||||
) -> dict:
|
||||
"""
|
||||
Process input messages by applying guardrails to text content.
|
||||
"""
|
||||
|
|
@ -329,9 +340,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
response: "ModelResponse",
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: Any | None = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict | None = None,
|
||||
) -> Any:
|
||||
) -> ModelResponse:
|
||||
"""
|
||||
Process output response by applying guardrails to text content.
|
||||
|
||||
|
|
@ -436,7 +447,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
responses_so_far: list["ModelResponseStream"],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
user_api_key_dict: Any | None = None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None" = None,
|
||||
request_data: dict | None = None,
|
||||
stream_transform_sink: StreamTransformSink | None = None,
|
||||
) -> list["ModelResponseStream"]:
|
||||
|
|
@ -486,7 +497,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
responses_so_far: list["ModelResponseStream"],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
user_api_key_dict: Any | None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None",
|
||||
request_data: dict | None,
|
||||
) -> list["ModelResponseStream"]:
|
||||
"""Block-only streaming path: run the guardrail so an in-flight BLOCK can
|
||||
|
|
@ -589,8 +600,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
def build_stream_error_items(
|
||||
self,
|
||||
exc: "HTTPException",
|
||||
responses_so_far: Sequence[Any] | None = None,
|
||||
) -> Sequence[Any] | None:
|
||||
responses_so_far: Sequence[object] | None = None,
|
||||
) -> Sequence[bytes] | None:
|
||||
import json
|
||||
|
||||
from litellm.proxy.common_request_processing import sse_error_payload
|
||||
|
|
@ -630,7 +641,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
responses_so_far: list["ModelResponseStream"],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
user_api_key_dict: Any | None,
|
||||
user_api_key_dict: "UserAPIKeyAuth | None",
|
||||
request_data: dict | None,
|
||||
sink: StreamTransformSink,
|
||||
) -> None:
|
||||
|
|
@ -794,7 +805,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
# Determine content source and tool calls based on choice type
|
||||
content = None
|
||||
tool_calls: list[Any] | None = None
|
||||
tool_calls: Sequence[object] | None = None
|
||||
if isinstance(choice, litellm.Choices):
|
||||
content = choice.message.content
|
||||
tool_calls = choice.message.tool_calls
|
||||
|
|
@ -1004,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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints
|
||||
|
||||
import httpx
|
||||
from openai.types.responses import ResponseReasoningItem
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -14,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 *
|
||||
|
|
@ -35,6 +37,37 @@ 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):
|
||||
"""Decoded body of the Responses API delete call."""
|
||||
|
||||
id: ReadOnly[str | None]
|
||||
object: ReadOnly[str | None]
|
||||
deleted: ReadOnly[bool | None]
|
||||
|
||||
|
||||
class _DeleteResponse(Protocol):
|
||||
"""The delete call's HTTP response, read for the decoded body it carries."""
|
||||
|
||||
def json(self) -> _DeleteResponseBody: ...
|
||||
|
||||
|
||||
class _JsonObjectResponse(Protocol):
|
||||
"""A Responses API HTTP response, read for the JSON object it decodes to."""
|
||||
|
||||
def json(self) -> dict[str, object]: ...
|
||||
|
||||
|
||||
def _delete_response_body(response: _DeleteResponse) -> _DeleteResponseBody:
|
||||
"""Decode a delete response body into the id, object and deleted fields it carries."""
|
||||
return response.json()
|
||||
|
||||
|
||||
def _json_object_body(response: _JsonObjectResponse) -> dict[str, object]:
|
||||
"""Decode a Responses API response body into its JSON object form."""
|
||||
return response.json()
|
||||
|
||||
|
||||
class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
||||
|
|
@ -179,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
|
||||
|
|
@ -217,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,
|
||||
|
|
@ -469,7 +520,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_event_model_class(event_type: str) -> Any:
|
||||
def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]:
|
||||
"""
|
||||
Returns the appropriate event model class based on the event type.
|
||||
|
||||
|
|
@ -583,7 +634,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
Transform the delete response API response into a DeleteResponseResult
|
||||
"""
|
||||
try:
|
||||
raw_response_json: Final = raw_response.json()
|
||||
raw_response_json: Final = _delete_response_body(raw_response)
|
||||
except Exception:
|
||||
raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code)
|
||||
return DeleteResponseResult(**raw_response_json)
|
||||
|
|
@ -618,7 +669,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
Transform the get response API response into a ResponsesAPIResponse
|
||||
"""
|
||||
try:
|
||||
raw_response_json: Final = raw_response.json()
|
||||
raw_response_json: Final = _json_object_body(raw_response)
|
||||
except Exception:
|
||||
raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code)
|
||||
raw_response_headers: Final = dict(raw_response.headers)
|
||||
|
|
@ -646,7 +697,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
) -> tuple[str, dict]:
|
||||
encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id")
|
||||
url: Final = f"{api_base}/{encoded_response_id}/input_items"
|
||||
params: Final[dict[str, Any]] = {}
|
||||
params: Final[dict[str, object]] = {}
|
||||
if after is not None:
|
||||
params["after"] = after
|
||||
if before is not None:
|
||||
|
|
@ -665,7 +716,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> dict:
|
||||
try:
|
||||
return raw_response.json()
|
||||
return _json_object_body(raw_response)
|
||||
except Exception:
|
||||
raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code)
|
||||
|
||||
|
|
@ -699,7 +750,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
Transform the cancel response API response into a ResponsesAPIResponse
|
||||
"""
|
||||
try:
|
||||
raw_response_json: Final = raw_response.json()
|
||||
raw_response_json: Final = _json_object_body(raw_response)
|
||||
except Exception:
|
||||
raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code)
|
||||
raw_response_headers: Final = dict(raw_response.headers)
|
||||
|
|
@ -742,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
|
||||
|
||||
|
|
|
|||
|
|
@ -5,10 +5,11 @@ For handling OpenAI-like chat completions, like IBM WatsonX, etc.
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Final
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Final, TypedDict
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm import LlmProviders
|
||||
|
|
@ -25,6 +26,23 @@ from ..common_utils import OpenAILikeBase, OpenAILikeError
|
|||
from .transformation import OpenAILikeChatConfig
|
||||
|
||||
|
||||
class _OpenAILikeChatCompletion(TypedDict, total=False):
|
||||
"""The chat-completion JSON body an OpenAI-like provider returns for a non-streamed call."""
|
||||
|
||||
id: ReadOnly[str]
|
||||
choices: ReadOnly[Sequence[Mapping[str, object]]]
|
||||
created: ReadOnly[int]
|
||||
model: ReadOnly[str]
|
||||
system_fingerprint: ReadOnly[str]
|
||||
usage: ReadOnly[Mapping[str, object]]
|
||||
object: ReadOnly[str]
|
||||
|
||||
|
||||
def _fake_streamed_model_response(payload: _OpenAILikeChatCompletion) -> ModelResponse:
|
||||
"""Build the single response a fake-streamed provider call replays as one chunk."""
|
||||
return ModelResponse(**payload)
|
||||
|
||||
|
||||
async def make_call(
|
||||
client: AsyncHTTPHandler | None,
|
||||
api_base: str,
|
||||
|
|
@ -42,9 +60,9 @@ async def make_call(
|
|||
response: Final = await client.post(api_base, headers=headers, data=data, stream=not fake_stream)
|
||||
|
||||
if streaming_decoder is not None:
|
||||
completion_stream: Any = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024))
|
||||
completion_stream = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024))
|
||||
elif fake_stream:
|
||||
model_response: Final = ModelResponse(**response.json())
|
||||
model_response: Final = _fake_streamed_model_response(response.json())
|
||||
completion_stream = MockResponseIterator(model_response=model_response)
|
||||
else:
|
||||
completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False)
|
||||
|
|
@ -82,7 +100,7 @@ def make_sync_call(
|
|||
if streaming_decoder is not None:
|
||||
completion_stream = streaming_decoder.iter_bytes(response.iter_bytes(chunk_size=1024))
|
||||
elif fake_stream:
|
||||
model_response: Final = ModelResponse(**response.json())
|
||||
model_response: Final = _fake_streamed_model_response(response.json())
|
||||
completion_stream = MockResponseIterator(model_response=model_response)
|
||||
else:
|
||||
completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True)
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import asyncio
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
|
|
@ -29,6 +31,16 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
class _RunwayMLTask(TypedDict, total=False):
|
||||
"""The RunwayML task payload returned by POST /v1/text_to_image and GET /v1/tasks/{id}."""
|
||||
|
||||
id: ReadOnly[str]
|
||||
status: ReadOnly[str]
|
||||
output: ReadOnly[Sequence[str | Mapping[str, str]]]
|
||||
failure: ReadOnly[str]
|
||||
failureCode: ReadOnly[str]
|
||||
|
||||
|
||||
class RunwayMLImageGenerationConfig(BaseImageGenerationConfig):
|
||||
"""
|
||||
Configuration for RunwayML image generation models.
|
||||
|
|
@ -80,7 +92,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig):
|
|||
|
||||
@staticmethod
|
||||
def _transform_runwayml_response_to_openai(
|
||||
response_data: dict[str, Any],
|
||||
response_data: _RunwayMLTask,
|
||||
model_response: ImageResponse,
|
||||
) -> ImageResponse:
|
||||
"""
|
||||
|
|
@ -155,7 +167,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig):
|
|||
raise TimeoutError(f"RunwayML task polling timed out after {timeout_secs} seconds")
|
||||
|
||||
@staticmethod
|
||||
def _check_task_status(response_data: dict[str, Any]) -> str:
|
||||
def _check_task_status(response_data: _RunwayMLTask) -> str:
|
||||
"""
|
||||
Check RunwayML task status from response.
|
||||
|
||||
|
|
@ -227,7 +239,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig):
|
|||
response = client.get(url=task_url, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
response_data = response.json()
|
||||
response_data: _RunwayMLTask = response.json()
|
||||
|
||||
# Check task status
|
||||
status = self._check_task_status(response_data=response_data)
|
||||
|
|
@ -276,7 +288,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig):
|
|||
response = await client.get(url=task_url, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
response_data = response.json()
|
||||
response_data: _RunwayMLTask = response.json()
|
||||
|
||||
# Check task status
|
||||
status = self._check_task_status(response_data=response_data)
|
||||
|
|
@ -322,7 +334,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig):
|
|||
}
|
||||
"""
|
||||
try:
|
||||
response_data = raw_response.json()
|
||||
response_data: _RunwayMLTask = raw_response.json()
|
||||
except Exception as e:
|
||||
raise self.get_error_class(
|
||||
error_message=f"Error transforming image generation response: {e}",
|
||||
|
|
@ -382,7 +394,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig):
|
|||
We need to poll the task until it completes (status SUCCEEDED) using async polling.
|
||||
"""
|
||||
try:
|
||||
response_data = raw_response.json()
|
||||
response_data: _RunwayMLTask = raw_response.json()
|
||||
except Exception as e:
|
||||
raise self.get_error_class(
|
||||
error_message=f"Error transforming image generation response: {e}",
|
||||
|
|
|
|||
|
|
@ -8,9 +8,10 @@ from dataclasses import dataclass
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, Protocol
|
||||
|
||||
import httpx
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -33,8 +34,8 @@ def _get_home() -> str:
|
|||
return os.getenv(HOME_PATH_ENV_VAR, DEFAULT_HOME_PATH)
|
||||
|
||||
|
||||
def _get_nested(d: dict[str, Any] | str, path: Sequence[str]) -> Any:
|
||||
cur: Any = d
|
||||
def _get_nested(d: object, path: Sequence[str]) -> object:
|
||||
cur: object = d
|
||||
if isinstance(cur, str):
|
||||
# This shouldn't happen if service keys are pre-parsed correctly
|
||||
try:
|
||||
|
|
@ -54,7 +55,7 @@ def _get_nested(d: dict[str, Any] | str, path: Sequence[str]) -> Any:
|
|||
return cur
|
||||
|
||||
|
||||
def _load_json_env(var_name: str) -> dict[str, Any] | None:
|
||||
def _load_json_env(var_name: str) -> dict[str, object] | None:
|
||||
raw: Final = os.environ.get(var_name)
|
||||
if not raw:
|
||||
return None
|
||||
|
|
@ -64,7 +65,7 @@ def _load_json_env(var_name: str) -> dict[str, Any] | None:
|
|||
return None
|
||||
|
||||
|
||||
def _str_or_none(value) -> str | None:
|
||||
def _str_or_none(value: object) -> str | None:
|
||||
try:
|
||||
return str(value) if value is not None else None
|
||||
except Exception:
|
||||
|
|
@ -124,7 +125,7 @@ CREDENTIAL_VALUES: Final[list[CredentialsValue]] = [
|
|||
]
|
||||
|
||||
|
||||
def init_conf(profile: str | None = None) -> dict[str, Any]:
|
||||
def init_conf(profile: str | None = None) -> dict[str, object]:
|
||||
"""
|
||||
Loads config JSON from:
|
||||
1) $AICORE_CONFIG if set, otherwise
|
||||
|
|
@ -191,7 +192,7 @@ def resolve_resource_group(sources: list[Source]) -> str | None:
|
|||
|
||||
def _parse_service_key_once(
|
||||
service_key: str | dict | None,
|
||||
) -> dict[str, Any] | None:
|
||||
) -> dict[str, object] | None:
|
||||
"""
|
||||
Pre-parse service_key if it's a string to avoid repeated JSON parsing.
|
||||
|
||||
|
|
@ -348,8 +349,33 @@ def validate_credentials(
|
|||
)
|
||||
|
||||
|
||||
class _TokenBody(TypedDict):
|
||||
"""Decoded body of the SAP AI Core OAuth2 token response."""
|
||||
|
||||
access_token: ReadOnly[str]
|
||||
expires_in: ReadOnly[NotRequired[int]]
|
||||
|
||||
|
||||
class _TokenResponse(Protocol):
|
||||
"""The token endpoint's HTTP response, read for the decoded token body it carries."""
|
||||
|
||||
def json(self) -> _TokenBody: ...
|
||||
|
||||
|
||||
def _bearer_token_and_expiry(response: _TokenResponse) -> tuple[str, datetime]:
|
||||
"""Read a token response into the Authorization header value and the token's absolute expiry."""
|
||||
payload: Final = response.json()
|
||||
expires_in: Final = int(payload.get("expires_in", 3600))
|
||||
access_token: Final = payload["access_token"]
|
||||
return f"Bearer {access_token}", datetime.now(timezone.utc) + timedelta(seconds=expires_in)
|
||||
|
||||
|
||||
def _request_token(
|
||||
client_id: str, auth_url: str, timeout: float, cert_pair=None, client_secret=None
|
||||
client_id: str,
|
||||
auth_url: str,
|
||||
timeout: float,
|
||||
cert_pair: tuple[str, str] | None = None,
|
||||
client_secret: str | None = None,
|
||||
) -> tuple[str, datetime]:
|
||||
data: Final = {"grant_type": "client_credentials", "client_id": client_id}
|
||||
if client_secret:
|
||||
|
|
@ -361,15 +387,10 @@ def _request_token(
|
|||
with httpx.Client(cert=cert_pair) as raw_client:
|
||||
handler = HTTPHandler(client=raw_client)
|
||||
resp = handler.post(auth_url, data=data, timeout=timeout)
|
||||
payload = resp.json()
|
||||
else:
|
||||
handler = _get_httpx_client()
|
||||
resp = handler.post(auth_url, data=data, timeout=timeout)
|
||||
payload = resp.json()
|
||||
access_token: Final = payload["access_token"]
|
||||
expires_in: Final = int(payload.get("expires_in", 3600))
|
||||
expiry_date: Final = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
|
||||
return f"Bearer {access_token}", expiry_date
|
||||
return _bearer_token_and_expiry(resp)
|
||||
handler = _get_httpx_client()
|
||||
resp = handler.post(auth_url, data=data, timeout=timeout)
|
||||
return _bearer_token_and_expiry(resp)
|
||||
except Exception as e:
|
||||
msg: Final = resp.text if resp is not None else getattr(e, "text", str(e))
|
||||
raise RuntimeError(f"Token request failed: {msg}") from e
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from urllib.parse import quote, unquote
|
|||
import httpx
|
||||
from httpx import Headers, Response
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
from typing_extensions import ReadOnly
|
||||
from typing_extensions import ReadOnly, Required
|
||||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -104,6 +104,27 @@ class _VertexBatchRow(TypedDict, total=False):
|
|||
processed_time: ReadOnly[str]
|
||||
|
||||
|
||||
class _VertexEmbeddingVector(TypedDict):
|
||||
values: ReadOnly[list[float]]
|
||||
|
||||
|
||||
class _VertexEmbeddingUsageMetadata(TypedDict, total=False):
|
||||
promptTokenCount: ReadOnly[int]
|
||||
|
||||
|
||||
class _VertexEmbeddingResponse(TypedDict, total=False):
|
||||
embedding: ReadOnly[Required[_VertexEmbeddingVector]]
|
||||
usageMetadata: ReadOnly[_VertexEmbeddingUsageMetadata]
|
||||
tokenCount: ReadOnly[int]
|
||||
|
||||
|
||||
class _VertexEmbeddingBatchRow(TypedDict, total=False):
|
||||
key: ReadOnly[str]
|
||||
request: ReadOnly[Mapping[str, object]]
|
||||
status: ReadOnly[Required[str]]
|
||||
response: ReadOnly[Required[_VertexEmbeddingResponse]]
|
||||
|
||||
|
||||
class _OpenAIBatchOutputError(TypedDict):
|
||||
code: ReadOnly[str]
|
||||
message: ReadOnly[str]
|
||||
|
|
@ -111,7 +132,7 @@ class _OpenAIBatchOutputError(TypedDict):
|
|||
|
||||
class _OpenAIBatchOutputResponse(TypedDict):
|
||||
status_code: ReadOnly[int]
|
||||
request_id: ReadOnly[str]
|
||||
request_id: ReadOnly[object]
|
||||
body: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
|
|
@ -218,7 +239,7 @@ def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object] | None
|
|||
return str(labels.get("litellm_custom_id", "unknown"))
|
||||
|
||||
|
||||
def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) -> bool:
|
||||
def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool:
|
||||
"""
|
||||
Whether a Vertex batch output row came from an `EmbedContentRequest`.
|
||||
|
||||
|
|
@ -237,7 +258,7 @@ def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any])
|
|||
|
||||
def _openai_batch_output_row(
|
||||
custom_id: str,
|
||||
body: Mapping[str, Any] | None = None,
|
||||
body: Mapping[str, object] | None = None,
|
||||
error_code: str | None = None,
|
||||
error_message: str = "",
|
||||
) -> _OpenAIBatchOutputRow:
|
||||
|
|
@ -259,7 +280,7 @@ def _openai_batch_output_row(
|
|||
}
|
||||
|
||||
|
||||
def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int, int]:
|
||||
def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[str, int, int]:
|
||||
"""
|
||||
Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch
|
||||
output row.
|
||||
|
|
@ -278,7 +299,7 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str,
|
|||
return unquote(match["custom_id"]), int(match["index"]), int(match["total"])
|
||||
|
||||
|
||||
def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int:
|
||||
def _embedding_prompt_token_count(vertex_response: _VertexEmbeddingResponse) -> int:
|
||||
"""
|
||||
Prompt tokens billed for one Vertex Gemini Embedding batch row.
|
||||
|
||||
|
|
@ -293,7 +314,7 @@ def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int:
|
|||
|
||||
def _vertex_embeddings_rows_to_openai_batch_output_row(
|
||||
custom_id: str,
|
||||
vertex_output_rows: tuple[Mapping[str, Any], ...],
|
||||
vertex_output_rows: tuple[_VertexEmbeddingBatchRow, ...],
|
||||
element_indices: tuple[int, ...],
|
||||
element_count: int,
|
||||
model: str | None,
|
||||
|
|
@ -348,7 +369,7 @@ def _vertex_embeddings_rows_to_openai_batch_output_row(
|
|||
|
||||
|
||||
def _transform_vertex_embeddings_batch_output_to_openai(
|
||||
vertex_output_rows: Iterable[Mapping[str, Any]],
|
||||
vertex_output_rows: Iterable[_VertexEmbeddingBatchRow],
|
||||
model: str | None,
|
||||
) -> tuple[_OpenAIBatchOutputRow, ...]:
|
||||
"""
|
||||
|
|
@ -388,7 +409,7 @@ def _model_from_managed_gcs_url(url: str) -> str | None:
|
|||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool:
|
||||
def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool:
|
||||
"""
|
||||
Whether an OpenAI batch JSONL line targets the embeddings endpoint.
|
||||
|
||||
|
|
@ -431,7 +452,7 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str:
|
|||
return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}"
|
||||
|
||||
|
||||
def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, Any]) -> Mapping[str, Any]:
|
||||
def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""
|
||||
One Vertex Gemini Embedding batch input row.
|
||||
|
||||
|
|
@ -453,8 +474,8 @@ def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str,
|
|||
|
||||
|
||||
def _openai_batch_jsonl_entry_to_vertex_embeddings_rows(
|
||||
openai_entry: Mapping[str, Any],
|
||||
) -> tuple[Mapping[str, Any], ...]:
|
||||
openai_entry: Mapping[str, object],
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""
|
||||
Transforms a single OpenAI `/v1/embeddings` batch entry into Vertex Gemini Embedding
|
||||
batch rows, one per requested embedding.
|
||||
|
|
@ -512,7 +533,7 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows(
|
|||
def _openai_batch_jsonl_entry_to_vertex_rows(
|
||||
openai_entry: dict[str, Any],
|
||||
map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]],
|
||||
) -> tuple[Mapping[str, Any], ...]:
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""
|
||||
Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to.
|
||||
|
||||
|
|
@ -533,7 +554,7 @@ def _openai_batch_jsonl_entry_to_vertex_rows(
|
|||
cached_content=None,
|
||||
)
|
||||
|
||||
custom_id: Final = openai_entry.get("custom_id")
|
||||
custom_id: Final[object] = openai_entry.get("custom_id")
|
||||
if custom_id is not None:
|
||||
if "labels" not in vertex_request_body:
|
||||
vertex_request_body["labels"] = {}
|
||||
|
|
|
|||
|
|
@ -250,7 +250,7 @@ def _gs_uri_requires_content_type_metadata(url: str) -> bool:
|
|||
|
||||
|
||||
def _image_url_payload_may_need_sync_gcs_metadata_fetch(
|
||||
raw_image_url: Any,
|
||||
raw_image_url: object,
|
||||
) -> bool:
|
||||
"""
|
||||
True when this image_url value (content-part image_url or assistant ``images[]``
|
||||
|
|
@ -326,7 +326,7 @@ def _openai_messages_may_need_sync_gcs_metadata_fetch(
|
|||
def _get_gcs_object_content_type(
|
||||
image_url: str,
|
||||
vertex_project: str | None = None,
|
||||
vertex_credentials: Any | None = None,
|
||||
vertex_credentials: object = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Resolve content type from GCS object metadata.
|
||||
|
|
@ -479,7 +479,7 @@ def _process_gemini_media(
|
|||
model: str | None = None,
|
||||
video_metadata: dict[str, Any] | None = None,
|
||||
vertex_project: str | None = None,
|
||||
vertex_credentials: Any | None = None,
|
||||
vertex_credentials: object = None,
|
||||
) -> PartType:
|
||||
"""
|
||||
Given a media URL (image, audio, or video), return the appropriate PartType for Gemini
|
||||
|
|
@ -1002,7 +1002,7 @@ def _gemini_convert_messages_with_history(
|
|||
if isinstance(_ss_invocations, list):
|
||||
for invocation in _ss_invocations:
|
||||
# Re-inject toolCall part
|
||||
tc_part: dict[str, Any] = {
|
||||
tc_part: dict[str, object] = {
|
||||
"toolCall": {
|
||||
"toolType": invocation.get("tool_type"),
|
||||
"id": invocation.get("id"),
|
||||
|
|
@ -1015,13 +1015,13 @@ def _gemini_convert_messages_with_history(
|
|||
|
||||
# Re-inject toolResponse part if response is present
|
||||
if "response" in invocation:
|
||||
tr_dict: dict[str, Any] = {
|
||||
tr_dict: dict[str, object] = {
|
||||
"id": invocation.get("id"),
|
||||
"response": invocation.get("response"),
|
||||
}
|
||||
if invocation.get("tool_type"):
|
||||
tr_dict["toolType"] = invocation["tool_type"]
|
||||
tr_part: dict[str, Any] = {"toolResponse": tr_dict}
|
||||
tr_part: dict[str, object] = {"toolResponse": tr_dict}
|
||||
if "response_thought_signature" in invocation:
|
||||
tr_part["thoughtSignature"] = invocation["response_thought_signature"]
|
||||
assistant_content.append(tr_part)
|
||||
|
|
@ -1090,7 +1090,7 @@ def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None:
|
|||
data_dict[k] = v
|
||||
|
||||
|
||||
def _has_google_maps_tool(tools: Any | None) -> bool:
|
||||
def _has_google_maps_tool(tools: object) -> bool:
|
||||
"""Return True if any tool object in the list has a 'googleMaps' key."""
|
||||
if not isinstance(tools, list):
|
||||
return False
|
||||
|
|
@ -1127,7 +1127,7 @@ def _rewrite_mime_type_to_response_format(generation_config: GenerationConfig) -
|
|||
schema = generation_config.pop("response_schema", None)
|
||||
generation_config.pop("response_mime_type", None)
|
||||
|
||||
response_format: Final[dict[str, Any]] = {"text": {"mimeType": "APPLICATION_JSON"}}
|
||||
response_format: Final[dict[str, dict[str, object]]] = {"text": {"mimeType": "APPLICATION_JSON"}}
|
||||
if schema is not None:
|
||||
response_format["text"]["schema"] = schema
|
||||
generation_config["responseFormat"] = response_format
|
||||
|
|
@ -1316,7 +1316,7 @@ async def async_transform_request_body(
|
|||
timeout: float | httpx.Timeout | None,
|
||||
extra_headers: dict | None,
|
||||
optional_params: dict,
|
||||
logging_obj: litellm.litellm_core_utils.litellm_logging.Logging,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
|
||||
litellm_params: dict,
|
||||
vertex_project: str | None,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import json
|
|||
import os
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import litellm
|
||||
|
|
@ -47,6 +47,21 @@ else:
|
|||
GoogleCredentialsObject = Any
|
||||
|
||||
|
||||
class _VertexCredentialsObject(Protocol):
|
||||
"""Structural view of the google-auth credentials handle that this class caches and refreshes."""
|
||||
|
||||
@property
|
||||
def token(self) -> object: ...
|
||||
|
||||
@property
|
||||
def quota_project_id(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def expired(self) -> object: ...
|
||||
|
||||
def refresh(self, request: object) -> None: ...
|
||||
|
||||
|
||||
class VertexBase:
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
|
@ -55,7 +70,7 @@ class VertexBase:
|
|||
self._credentials: GoogleCredentialsObject | None = None
|
||||
self._credentials_project_mapping: dict[
|
||||
tuple[VERTEX_CREDENTIALS_TYPES | None, str | None],
|
||||
tuple[GoogleCredentialsObject, str | None],
|
||||
tuple[_VertexCredentialsObject, str | None],
|
||||
] = {}
|
||||
self.project_id: str | None = None
|
||||
self.async_handler: AsyncHTTPHandler | None = None
|
||||
|
|
@ -109,7 +124,7 @@ class VertexBase:
|
|||
self,
|
||||
credentials: VERTEX_CREDENTIALS_TYPES | None,
|
||||
project_id: str | None,
|
||||
) -> tuple[Any, str]:
|
||||
) -> tuple[_VertexCredentialsObject | None, str]:
|
||||
if credentials is not None:
|
||||
if isinstance(credentials, str):
|
||||
_is_path: Final = os.path.exists(
|
||||
|
|
@ -209,7 +224,7 @@ class VertexBase:
|
|||
return creds, project_id
|
||||
|
||||
# Google Auth Helpers -- extracted for mocking purposes in tests
|
||||
def _credentials_from_identity_pool(self, json_obj, scopes):
|
||||
def _credentials_from_identity_pool(self, json_obj, scopes) -> _VertexCredentialsObject:
|
||||
try:
|
||||
from google.auth import identity_pool
|
||||
except ImportError:
|
||||
|
|
@ -220,7 +235,7 @@ class VertexBase:
|
|||
creds = creds.with_scopes(scopes)
|
||||
return creds
|
||||
|
||||
def _credentials_from_pluggable(self, json_obj, scopes):
|
||||
def _credentials_from_pluggable(self, json_obj, scopes) -> _VertexCredentialsObject:
|
||||
try:
|
||||
from google.auth import pluggable
|
||||
except ImportError:
|
||||
|
|
@ -231,7 +246,7 @@ class VertexBase:
|
|||
creds = creds.with_scopes(scopes)
|
||||
return creds
|
||||
|
||||
def _credentials_from_identity_pool_with_aws(self, json_obj, scopes):
|
||||
def _credentials_from_identity_pool_with_aws(self, json_obj, scopes) -> _VertexCredentialsObject:
|
||||
try:
|
||||
from google.auth import aws
|
||||
except ImportError:
|
||||
|
|
@ -242,7 +257,7 @@ class VertexBase:
|
|||
creds = creds.with_scopes(scopes)
|
||||
return creds
|
||||
|
||||
def _credentials_from_authorized_user(self, json_obj, scopes):
|
||||
def _credentials_from_authorized_user(self, json_obj, scopes) -> _VertexCredentialsObject:
|
||||
try:
|
||||
import google.oauth2.credentials
|
||||
except ImportError:
|
||||
|
|
@ -250,7 +265,7 @@ class VertexBase:
|
|||
|
||||
return google.oauth2.credentials.Credentials.from_authorized_user_info(json_obj, scopes=scopes)
|
||||
|
||||
def _credentials_from_service_account(self, json_obj, scopes):
|
||||
def _credentials_from_service_account(self, json_obj, scopes) -> _VertexCredentialsObject:
|
||||
try:
|
||||
import google.oauth2.service_account
|
||||
except ImportError:
|
||||
|
|
@ -258,7 +273,7 @@ class VertexBase:
|
|||
|
||||
return google.oauth2.service_account.Credentials.from_service_account_info(json_obj, scopes=scopes)
|
||||
|
||||
def _credentials_from_default_auth(self, scopes):
|
||||
def _credentials_from_default_auth(self, scopes) -> tuple[_VertexCredentialsObject, str | None]:
|
||||
try:
|
||||
import google.auth as google_auth
|
||||
except ImportError:
|
||||
|
|
@ -350,7 +365,7 @@ class VertexBase:
|
|||
)
|
||||
return api_base
|
||||
|
||||
def refresh_auth(self, credentials: Any) -> None:
|
||||
def refresh_auth(self, credentials: _VertexCredentialsObject) -> None:
|
||||
try:
|
||||
from google.auth.transport.requests import (
|
||||
Request,
|
||||
|
|
@ -426,7 +441,7 @@ class VertexBase:
|
|||
self,
|
||||
credential_cache_key: tuple,
|
||||
project_id: str | None,
|
||||
) -> tuple[str, str, "TokenState", Any, str | None] | None:
|
||||
) -> tuple[str, str, "TokenState", _VertexCredentialsObject, str | None] | None:
|
||||
"""
|
||||
Look up cached credentials and return usable token info for FRESH or
|
||||
STALE tokens (both are still valid for outbound requests). STALE
|
||||
|
|
@ -449,7 +464,9 @@ class VertexBase:
|
|||
return None
|
||||
return creds.token, resolved_project, token_state, creds, cached_project_id
|
||||
|
||||
def _unpack_cached_credentials(self, credential_cache_key: tuple) -> tuple[Any, str | None]:
|
||||
def _unpack_cached_credentials(
|
||||
self, credential_cache_key: tuple
|
||||
) -> tuple[_VertexCredentialsObject | None, str | None]:
|
||||
"""
|
||||
Return (credentials, project_id) from the cache, or (None, None) if
|
||||
not cached. Handles both tuple and legacy cache formats.
|
||||
|
|
@ -461,7 +478,7 @@ class VertexBase:
|
|||
return cached_entry
|
||||
return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None)
|
||||
|
||||
def _get_token_state(self, credentials: Any) -> "TokenState":
|
||||
def _get_token_state(self, credentials: _VertexCredentialsObject) -> "TokenState":
|
||||
"""
|
||||
Return the token state using google-auth's TokenState enum.
|
||||
|
||||
|
|
@ -485,7 +502,7 @@ class VertexBase:
|
|||
credentials: VERTEX_CREDENTIALS_TYPES | None,
|
||||
project_id: str | None,
|
||||
credential_cache_key: tuple,
|
||||
) -> tuple[Any, str | None]:
|
||||
) -> tuple[_VertexCredentialsObject, str | None]:
|
||||
"""Load credentials via load_auth (in thread) and cache the result."""
|
||||
try:
|
||||
_credentials, credential_project_id = await asyncify(self.load_auth)(
|
||||
|
|
@ -505,7 +522,7 @@ class VertexBase:
|
|||
|
||||
async def _background_refresh_credentials(
|
||||
self,
|
||||
credentials: Any,
|
||||
credentials: _VertexCredentialsObject,
|
||||
credential_cache_key: tuple,
|
||||
credential_project_id: str | None,
|
||||
) -> None:
|
||||
|
|
@ -557,7 +574,7 @@ class VertexBase:
|
|||
|
||||
def _schedule_background_refresh(
|
||||
self,
|
||||
credentials: Any,
|
||||
credentials: _VertexCredentialsObject,
|
||||
credential_cache_key: tuple,
|
||||
credential_project_id: str | None,
|
||||
) -> None:
|
||||
|
|
@ -575,7 +592,7 @@ class VertexBase:
|
|||
self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id)
|
||||
)
|
||||
|
||||
def _drop_background_refresh_task(_fut: asyncio.Future[Any]) -> None:
|
||||
def _drop_background_refresh_task(_fut: asyncio.Future[None]) -> None:
|
||||
if self._background_refresh_tasks.get(credential_cache_key) is _fut:
|
||||
self._background_refresh_tasks.pop(credential_cache_key, None)
|
||||
|
||||
|
|
@ -888,7 +905,7 @@ class VertexBase:
|
|||
# Convert dict credentials to string for caching
|
||||
cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials
|
||||
credential_cache_key: Final = (cache_credentials, project_id)
|
||||
_credentials: GoogleCredentialsObject | None = None
|
||||
_credentials: _VertexCredentialsObject | None = None
|
||||
|
||||
verbose_logger.debug("Checking cached credentials for project_id: %s", project_id)
|
||||
|
||||
|
|
|
|||
|
|
@ -6949,12 +6949,18 @@ def embedding(
|
|||
aembedding=aembedding,
|
||||
headers=headers,
|
||||
)
|
||||
elif custom_llm_provider == "dashscope":
|
||||
dashscope_key: Final = api_key or litellm.api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"):
|
||||
from litellm.llms.dashscope.common_utils import (
|
||||
missing_dashscope_family_key_message,
|
||||
resolve_dashscope_family_api_key,
|
||||
)
|
||||
|
||||
dashscope_key: Final = resolve_dashscope_family_api_key(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_key=api_key or litellm.api_key,
|
||||
)
|
||||
if dashscope_key is None:
|
||||
raise ValueError(
|
||||
"Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter."
|
||||
)
|
||||
raise ValueError(missing_dashscope_family_key_message(custom_llm_provider))
|
||||
if extra_headers is not None and isinstance(extra_headers, dict):
|
||||
headers = extra_headers
|
||||
else:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue