Merge remote-tracking branch 'origin/main' into litellm_remove_lit002_dict_ban_v2

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	litellm/types/decisions.py
This commit is contained in:
mateo 2026-10-08 05:00:01 +00:00
commit 22003cf27c
685 changed files with 33346 additions and 8733 deletions

View file

@ -22,6 +22,12 @@ jobs:
with:
persist-credentials: false
- name: Check Lens Compose configuration
run: |
python3 -m unittest discover -s deploy/lens -p 'test_*.py'
python3 deploy/lens/configure.py --version 1.2.3 --env-file "$RUNNER_TEMP/lens.env"
docker compose --env-file "$RUNNER_TEMP/lens.env" -f deploy/lens/stack.yaml config --quiet
- name: Set up Helm 3.11.1
uses: azure/setup-helm@1a275c3b69536ee54be43f2070a358922e12c8d4 # v4.3.1
with:

View file

@ -23,6 +23,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/end_user/",
"/sso/",
"/liteadmin/slack/connect/",
"/moyai/connect/",
"/login",
"/v2/login",
"/v3/login",

View file

@ -6,21 +6,20 @@ Agent exporters send traces directly to Lens. LiteLLM sends its optional request
## New local installation
Install Docker with Compose and Git, then build the gateway and Lens from one checkout:
Install Docker with Compose, Python 3.10 or later, and Git. Clone LiteLLM, select a published release that includes Lens, and start the existing Compose stack:
```bash
git clone https://github.com/BerriAI/litellm.git
cd litellm
export LITELLM_RELEASE_TAG="sha-$(git rev-parse HEAD)"
export LITELLM_MASTER_KEY="sk-$(openssl rand -hex 24)"
export LITELLM_LENS_SERVICE_TOKEN="$(openssl rand -hex 32)"
export OPENAI_API_KEY='<your-provider-key>'
docker compose -f docker/docker-compose.tracing.yml up -d --build
python3 deploy/lens/configure.py --version <release-version>
docker compose --env-file deploy/lens/.env -f deploy/lens/stack.yaml up -d --wait
```
Save the generated keys privately and reuse them when restarting or upgrading. This stack binds to localhost and uses development database passwords; use your normal secrets, TLS, backups, and ingress for a hosted deployment
The configuration command generates your keys and database passwords once, saves them in `deploy/lens/.env` with owner-only permissions, and preserves them on subsequent runs. Back up this file alongside your database volumes. Both images use the selected release; there is no local image build
Open `http://localhost:4002/ui/` and sign in as `admin` with `LITELLM_MASTER_KEY`. Under **Lens > Traces > Set up tracing**, generate a tracing key and copy the ingestion URL. Local exporters use `http://localhost:4318`. Model calls keep their existing LiteLLM URL and model key
Open `http://localhost:4000/ui/` and sign in as `admin` using `LITELLM_MASTER_KEY` from the saved file. Open **Lens**, select your framework, generate a tracing key, and copy the displayed configuration. The trace endpoint is already filled in. Keep your agent's existing model credentials; the tracing key only authorizes trace uploads
PostgreSQL and ClickHouse use persistent Docker volumes and have no host ports. The dashboard and trace listener bind to localhost. Use your normal TLS and ingress for a hosted deployment. Stop the stack with `docker compose --env-file deploy/lens/.env -f deploy/lens/stack.yaml down`; omit `-v` to retain data
Under **Lens > Investigations > Connect worker**, choose an analysis model and monthly budget. The deployed service connects automatically after you save these settings. There is no worker command or second token to copy
@ -68,14 +67,30 @@ Lens does not need provider credentials, PostgreSQL credentials, a GPU, or the L
### Kubernetes with Helm
Both `helm/litellm` and `helm/litellm-helm` support the Lens service. Keep your existing release, namespace, values, and database configuration. Create two Secrets through your normal secret manager: `litellm-lens-service` with key `service-token`, and `litellm-lens-clickhouse` with key `url`
Both `helm/litellm` and `helm/litellm-helm` support Lens. Keep your existing chart, release name, namespace, and values. Add:
```yaml
lensWorker:
enabled: true
```
Then run your usual Helm deployment command using the matching published chart. The chart supplies the matching Lens image, generates the shared service secret, starts a single ClickHouse instance with a persistent volume, and connects the services. Your cluster needs a default storage class, or set `lensWorker.clickhouse.storageClassName`. Bundled storage defaults to 20 GiB; set `lensWorker.clickhouse.storage` before installation to choose another size
When your chart manages an ingress with one hostname, the chart fills in the public tracing address and routes `/lens-ingest` directly to Lens. TLS is detected from `ingress.tls` or an ALB certificate annotation. With custom ingress, multiple hostnames, or TLS terminated elsewhere, set the address explicitly:
```yaml
lensWorker:
enabled: true
publicUrl: https://<your-litellm-host>/lens-ingest
```
For a dedicated trace hostname, configure `lensWorker.ingress.enabled`, `host`, `className`, and `tls`. Its hostname supplies the public address unless you override `publicUrl`. Internal Lens routes stay private
To use an existing ClickHouse database and secrets managed by your platform, keep these overrides:
```yaml
lensWorker:
enabled: true
image:
repository: <matching-worker-image-repository>
digest: sha256:<matching-worker-image-digest>
serviceTokenSecret:
name: litellm-lens-service
key: service-token
@ -84,21 +99,13 @@ lensWorker:
key: url
clickhouseDatabase: litellm
retentionDays: 14
publicUrl: https://<your-litellm-host>/lens-ingest
```
Set `clickhouseDatabase` and `retentionDays` to your existing database and retention before upgrading
Supplying `clickhouseSecret.name` uses that database and disables bundled storage. Keep your database name and retention policy. For GitOps tools that render Helm without cluster access, supply both existing secrets so rendering cannot regenerate credentials
When the chart's main ingress is enabled, it routes `/lens-ingest` directly to Lens. With a custom ingress, add that route yourself. For a dedicated hostname, use `lensWorker.ingress.enabled`, `host`, `className`, and `tls`, and set `publicUrl` to that hostname. The chart connects LiteLLM to Lens internally and gives both services the shared secret
Normal Helm upgrades reuse the generated credentials. Secrets are retained on uninstall, and the ClickHouse volume is retained by Kubernetes. Back them up together. Treat changing the database, storage class, or secret reference as an infrastructure change, not a routine version update
Update your existing component image overrides to matching builds, then use the chart from that checkout:
```bash
helm upgrade --install litellm ./helm/litellm \
--namespace litellm -f values.yaml --wait
```
Use `./helm/litellm-helm` if that is your existing chart. `lensWorker.replicaCount` scales ingestion and investigations. Each replica needs access to the same ClickHouse and gateway. Credentials refresh every 30 seconds; a newly created key may briefly receive a retryable 429. Revocations propagate on refresh, and a replica stops accepting traces when its credential snapshot reaches 90 seconds
After deployment, open **Lens**. If it was already open, click **Check setup**. The setup section moves to your framework and tracing key when Lens is reachable and storage is ready. Investigation setup asks for the analysis model and budget; the installed service connects automatically
## Upgrade
@ -106,13 +113,9 @@ Upgrade LiteLLM and Lens from the same source commit and release identity. For a
Keep the same databases, encryption keys, shared service secret, and public ingestion URL. Pause scheduled investigations and finish or cancel active runs, update both images through your usual deployment process, then check ingestion and run an investigation before resuming schedules. Do not run `docker compose down -v`
When upgrading from the Python worker, replace it with the Rust Lens service, move the existing ClickHouse connection to Lens, and configure the service URLs and secret on LiteLLM. Existing trace data remains in the same ClickHouse database; findings and settings remain in PostgreSQL. Stop the old worker. Generate dedicated tracing keys and change agent exporters to the ingestion URL. A virtual model key no longer authorizes uploads; the old gateway upload endpoints return 410 with setup guidance
If you retain an explicit `LENS_WORKER_TOKEN`, it remains an optional investigation credential. Normal setup uses the shared service connection and registers one managed worker identity. Configure the analysis model and billing key in the dashboard; provider keys stay on LiteLLM
## Development
`make lens-dev` starts LiteLLM, the Rust Lens service, and the hot-reload dashboard. Set `LENS_DEV_PROXY_PORT` and `LENS_DEV_UI_PORT` to change the local ports. For containers, pass the same release identity to both builds. Unversioned or incompatible workers are refused before claiming work
`make lens-dev` starts LiteLLM, Lens, and the hot-reload dashboard. Set `LENS_DEV_PROXY_PORT` and `LENS_DEV_UI_PORT` to change the local ports. For containers, pass the same release identity to both builds. Unversioned or incompatible workers are refused before claiming work
## Configure a lens

View file

@ -1,3 +1,5 @@
model_list: []
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
tracing:

73
deploy/lens/configure.py Normal file
View file

@ -0,0 +1,73 @@
from __future__ import annotations
import argparse
import os
import re
import secrets
import shlex
import sys
from pathlib import Path
from typing import Final
SECRET_NAMES: Final = (
"LITELLM_MASTER_KEY",
"LITELLM_SALT_KEY",
"LITELLM_LENS_SERVICE_TOKEN",
"POSTGRES_PASSWORD",
"CLICKHOUSE_PASSWORD",
)
def environment_content(path: Path) -> str:
if not path.exists():
return "".join(
f"{name}={'sk-' if name.endswith('KEY') else ''}{secrets.token_hex(32)}\n" for name in SECRET_NAMES
)
saved: Final = path.read_text().splitlines()
values: Final = dict(line.split("=", 1) for line in saved if "=" in line)
if any(not values.get(name) for name in SECRET_NAMES):
raise ValueError(f"{path} is incomplete. Restore your saved credentials before continuing")
return "\n".join(line for line in saved if not line.startswith("LITELLM_VERSION=")) + "\n"
def configure(path: Path, version: str) -> None:
release: Final = version.removeprefix("v")
if not re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+(?:[-.][a-zA-Z0-9.-]+)?", release):
raise ValueError("Use a published release version, such as 1.82.0 or 1.82.0-nightly")
if path.is_symlink():
raise ValueError(f"Refusing to replace a symlink: {path}")
existing: Final = path.exists()
content: Final = environment_content(path)
temporary: Final = path.with_name(f".{path.name}.{secrets.token_hex(8)}")
descriptor: Final = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
try:
with os.fdopen(descriptor, "w") as output:
output.write(content + f"LITELLM_VERSION={release}\n")
if existing:
os.replace(temporary, path)
else:
os.link(temporary, path)
finally:
temporary.unlink(missing_ok=True)
def main() -> None:
parser: Final = argparse.ArgumentParser(description="Create or update the configuration for the Lens Compose stack")
parser.add_argument("--version", required=True, help="Published LiteLLM release; Lens uses the matching version")
parser.add_argument("--env-file", type=Path, default=Path(__file__).with_name(".env"))
arguments: Final = parser.parse_args()
try:
configure(arguments.env_file, arguments.version)
except (OSError, ValueError) as error:
parser.exit(1, f"Could not configure Lens: {error}\n")
sys.stdout.write(
f"Saved {arguments.env_file}. Existing keys and database passwords are preserved\n"
f"Start with: docker compose --env-file {shlex.quote(str(arguments.env_file))} "
"-f deploy/lens/stack.yaml up -d --wait\n"
"Open http://localhost:4000/ui/ and sign in as admin with LITELLM_MASTER_KEY from the saved file\n"
"Back up this file with your database volumes. Do not commit it\n"
)
if __name__ == "__main__":
main()

View file

@ -0,0 +1,52 @@
import tempfile
import unittest
from pathlib import Path
from configure import SECRET_NAMES, configure
class ComposeConfigurationTests(unittest.TestCase):
def test_restart_and_upgrade_preserve_private_credentials_and_custom_settings(self) -> None:
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / ".env"
configure(path, "v1.2.3")
original = dict(line.split("=", 1) for line in path.read_text().splitlines())
self.assertEqual(path.stat().st_mode & 0o777, 0o600)
self.assertEqual(len({original[name] for name in SECRET_NAMES}), len(SECRET_NAMES))
self.assertTrue(all(len(original[name]) >= 64 for name in SECRET_NAMES))
with path.open("a") as output:
output.write("LITELLM_LENS_PUBLIC_URL=https://traces.example/prefix\n")
configure(path, "1.2.3")
configure(path, "v1.2.4-nightly")
updated = dict(line.split("=", 1) for line in path.read_text().splitlines())
self.assertEqual({name: updated[name] for name in SECRET_NAMES}, {name: original[name] for name in SECRET_NAMES})
self.assertEqual(updated["LITELLM_VERSION"], "1.2.4-nightly")
self.assertEqual(updated["LITELLM_LENS_PUBLIC_URL"], "https://traces.example/prefix")
self.assertEqual(path.stat().st_mode & 0o777, 0o600)
def test_incomplete_configuration_is_never_replaced_with_new_database_passwords(self) -> None:
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / ".env"
original = "POSTGRES_PASSWORD=existing\n"
path.write_text(original)
with self.assertRaisesRegex(ValueError, "incomplete"):
configure(path, "1.2.3")
self.assertEqual(path.read_text(), original)
def test_invalid_release_and_symlink_leave_existing_files_untouched(self) -> None:
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / ".env"
target = Path(directory) / "saved"
target.write_text("preserve")
path.symlink_to(target)
with self.assertRaisesRegex(ValueError, "symlink"):
configure(path, "1.2.3")
self.assertEqual(target.read_text(), "preserve")
path.unlink()
with self.assertRaisesRegex(ValueError, "published release"):
configure(path, "1.2.3\nPOSTGRES_PASSWORD=replaced")
self.assertFalse(path.exists())
if __name__ == "__main__":
unittest.main()

View file

@ -255,7 +255,7 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
_MANAGED_FILES_TARGET: Final = "managed_files"
class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
self.internal_usage_cache = internal_usage_cache
@ -2110,4 +2110,4 @@ class PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
verbose_logger.debug(
f"Converted file {file_id} from storage backend to base64 with format {content_type}"
)
_PROXY_LiteLLMManagedFiles = PROXY_LiteLLMManagedFiles
PROXY_LiteLLMManagedFiles = _PROXY_LiteLLMManagedFiles

View file

@ -38,7 +38,7 @@ else:
PrismaClient = Any
class PROXY_LiteLLMManagedVectorStores(
class _PROXY_LiteLLMManagedVectorStores(
CustomLogger, BaseManagedResource[VectorStoreCreateResponse]
):
"""
@ -462,4 +462,4 @@ class PROXY_LiteLLMManagedVectorStores(
parent_otel_span=parent_otel_span,
resource_id_key="vector_store_id",
)
_PROXY_LiteLLMManagedVectorStores = PROXY_LiteLLMManagedVectorStores
PROXY_LiteLLMManagedVectorStores = _PROXY_LiteLLMManagedVectorStores

View file

@ -62,6 +62,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/rerank",
"/v1/decisions",
"/decisions",
"/v1/systemone",
"/systemone",
"/v1/ocr",
"/ocr",
"/v1/rag/",

View file

@ -371,3 +371,30 @@ shutdown drain window.
{{- $_ := set $labels "app.kubernetes.io/name" (printf "%s-lens-worker" (include "litellm.name" . | trunc 51 | trimSuffix "-")) -}}
{{- toYaml $labels -}}
{{- end -}}
{{- define "litellm.lensWorker.serviceTokenSecretName" -}}
{{- .Values.lensWorker.serviceTokenSecret.name | default (printf "%s-lens-service" (include "litellm.fullname" .)) -}}
{{- end -}}
{{- define "litellm.lensWorker.bundledClickhouse" -}}
{{- if and .Values.lensWorker.enabled .Values.lensWorker.clickhouse.enabled (not .Values.lensWorker.clickhouseSecret.name) -}}true{{- end -}}
{{- end -}}
{{- define "litellm.lensWorker.publicUrl" -}}
{{- if .Values.lensWorker.publicUrl -}}
{{- .Values.lensWorker.publicUrl -}}
{{- else if .Values.lensWorker.ingress.enabled -}}
{{- $tls := or (not (empty .Values.lensWorker.ingress.tls)) (hasKey .Values.lensWorker.ingress.annotations "alb.ingress.kubernetes.io/certificate-arn") -}}
{{- printf "%s://%s" (ternary "https" "http" $tls) (required "lensWorker.ingress.host is required" .Values.lensWorker.ingress.host) -}}
{{- else if and .Values.ingress.enabled (eq (len .Values.ingress.hosts) 1) -}}
{{- $host := required "ingress.hosts[0].host is required" (first .Values.ingress.hosts).host -}}
{{- $tls := or (not (empty .Values.ingress.tls)) (hasKey .Values.ingress.annotations "alb.ingress.kubernetes.io/certificate-arn") -}}
{{- printf "%s://%s/lens-ingest" (ternary "https" "http" $tls) $host -}}
{{- else -}}
{{- fail "lensWorker.publicUrl is required when there is no single ingress hostname" -}}
{{- end -}}
{{- end -}}
{{- define "litellm.lensWorker.clickhouseName" -}}
{{- printf "%s-lens-clickhouse" (include "litellm.fullname" . | trunc 47 | trimSuffix "-") -}}
{{- end -}}

View file

@ -60,11 +60,11 @@ spec:
- name: LITELLM_LENS_URL
value: {{ printf "http://%s-lens-worker:%v" (include "litellm.fullname" .) .Values.lensWorker.service.port | quote }}
- name: LITELLM_LENS_PUBLIC_URL
value: {{ required "lensWorker.publicUrl is required" .Values.lensWorker.publicUrl | quote }}
value: {{ include "litellm.lensWorker.publicUrl" . | quote }}
- name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: {{ required "lensWorker.serviceTokenSecret.name is required" .Values.lensWorker.serviceTokenSecret.name | quote }}
name: {{ include "litellm.lensWorker.serviceTokenSecretName" . | quote }}
key: {{ .Values.lensWorker.serviceTokenSecret.key | quote }}
{{- end }}
{{- include "litellm.proxyEnv" . | nindent 12 }}

View file

@ -0,0 +1,98 @@
{{- if include "litellm.lensWorker.bundledClickhouse" . }}
{{- $name := include "litellm.lensWorker.clickhouseName" . }}
apiVersion: v1
kind: Service
metadata:
name: {{ $name }}
spec:
clusterIP: None
selector:
app.kubernetes.io/instance: {{ .Release.Name }}
app.kubernetes.io/component: lens-clickhouse
ports:
- name: http
port: 8123
targetPort: http
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
name: {{ $name }}
spec:
serviceName: {{ $name }}
replicas: 1
selector:
matchLabels:
app.kubernetes.io/instance: {{ .Release.Name }}
app.kubernetes.io/component: lens-clickhouse
template:
metadata:
labels:
app.kubernetes.io/instance: {{ .Release.Name }}
app.kubernetes.io/component: lens-clickhouse
spec:
automountServiceAccountToken: false
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
{{- end }}
securityContext:
runAsNonRoot: true
runAsUser: 101
runAsGroup: 101
fsGroup: 101
seccompProfile:
type: RuntimeDefault
containers:
- name: clickhouse
image: {{ .Values.lensWorker.clickhouse.image | quote }}
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop: [ALL]
env:
- name: CLICKHOUSE_USER
value: default
- name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ $name }}
key: password
- name: CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT
value: "1"
ports:
- name: http
containerPort: 8123
startupProbe:
httpGet:
path: /ping
port: http
periodSeconds: 5
timeoutSeconds: 3
failureThreshold: 60
readinessProbe:
httpGet:
path: /ping
port: http
livenessProbe:
httpGet:
path: /ping
port: http
timeoutSeconds: 3
resources:
{{- toYaml .Values.lensWorker.clickhouse.resources | nindent 12 }}
volumeMounts:
- name: data
mountPath: /var/lib/clickhouse
volumeClaimTemplates:
- metadata:
name: data
spec:
accessModes: [ReadWriteOnce]
{{- if ne .Values.lensWorker.clickhouse.storageClassName nil }}
storageClassName: {{ .Values.lensWorker.clickhouse.storageClassName | quote }}
{{- end }}
resources:
requests:
storage: {{ .Values.lensWorker.clickhouse.storage | quote }}
{{- end }}

View file

@ -45,13 +45,23 @@ spec:
- name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: {{ required "lensWorker.serviceTokenSecret.name is required" .Values.lensWorker.serviceTokenSecret.name | quote }}
name: {{ include "litellm.lensWorker.serviceTokenSecretName" . | quote }}
key: {{ .Values.lensWorker.serviceTokenSecret.key | quote }}
{{- if include "litellm.lensWorker.bundledClickhouse" . }}
- name: CLICKHOUSE_HOST
value: {{ include "litellm.lensWorker.clickhouseName" . | quote }}
- name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.lensWorker.clickhouseName" . | quote }}
key: password
{{- else }}
- name: CLICKHOUSE_URL
valueFrom:
secretKeyRef:
name: {{ required "lensWorker.clickhouseSecret.name is required" .Values.lensWorker.clickhouseSecret.name | quote }}
key: {{ .Values.lensWorker.clickhouseSecret.key | quote }}
{{- end }}
- name: CLICKHOUSE_DATABASE
value: {{ .Values.lensWorker.clickhouseDatabase | quote }}
- name: AGENT_TRACING_RETENTION_DAYS

View file

@ -0,0 +1,27 @@
{{- if and .Values.lensWorker.enabled (not .Values.lensWorker.serviceTokenSecret.name) }}
{{- $name := include "litellm.lensWorker.serviceTokenSecretName" . }}
{{- $existing := lookup "v1" "Secret" .Release.Namespace $name }}
apiVersion: v1
kind: Secret
metadata:
name: {{ $name }}
annotations:
helm.sh/resource-policy: keep
type: Opaque
data:
{{ .Values.lensWorker.serviceTokenSecret.key }}: {{ if $existing }}{{ required "Saved Lens service secret is missing its key" (index $existing.data .Values.lensWorker.serviceTokenSecret.key) | quote }}{{ else }}{{ randAlphaNum 64 | b64enc | quote }}{{ end }}
{{- end }}
{{- if include "litellm.lensWorker.bundledClickhouse" . }}
{{- $name := include "litellm.lensWorker.clickhouseName" . }}
{{- $existing := lookup "v1" "Secret" .Release.Namespace $name }}
---
apiVersion: v1
kind: Secret
metadata:
name: {{ $name }}
annotations:
helm.sh/resource-policy: keep
type: Opaque
data:
password: {{ if $existing }}{{ required "Saved Lens ClickHouse secret is missing its password" (index $existing.data "password") | quote }}{{ else }}{{ randAlphaNum 64 | b64enc | quote }}{{ end }}
{{- end }}

View file

@ -0,0 +1,84 @@
suite: Lens endpoint defaults
templates:
- deployment.yaml
- configmap-litellm.yaml
set:
lensWorker.enabled: true
ingress.enabled: true
ingress.hosts: [{host: gateway.example, paths: [{path: /, pathType: Prefix}]}]
ingress.tls:
- hosts: [gateway.example]
secretName: tls
tests:
- it: derives the trace endpoint from the deployment hostname and TLS
template: deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://gateway.example/lens-ingest
- it: preserves an explicitly configured public address
set:
lensWorker.publicUrl: https://custom.example/prefix
template: deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://custom.example/prefix
- it: uses the dedicated Lens ingress when configured
set:
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
lensWorker.ingress.tls:
- hosts: [traces.example]
secretName: traces-tls
template: deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://traces.example
- it: uses HTTPS for a dedicated ALB ingress with certificate annotations
set:
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
lensWorker.ingress.className: alb
lensWorker.ingress.tls: []
lensWorker.ingress.annotations:
alb.ingress.kubernetes.io/certificate-arn: arn:aws:acm:us-west-2:123456789012:certificate/test
alb.ingress.kubernetes.io/listen-ports: '[{"HTTP":80},{"HTTPS":443}]'
alb.ingress.kubernetes.io/ssl-redirect: "443"
template: deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://traces.example
- it: uses HTTP for a dedicated ingress without its own TLS
set:
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
lensWorker.ingress.tls: []
lensWorker.ingress.annotations: {}
template: deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: http://traces.example
- it: uses HTTP when ingress has no TLS
set:
ingress.tls: []
template: deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: http://gateway.example/lens-ingest

View file

@ -0,0 +1,45 @@
suite: Lens preserves credentials across Helm upgrades
templates:
- lens/secrets.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
release:
namespace: lens
name: lens-test
upgrade: true
kubernetesProvider:
scheme:
v1/Secret:
gvr:
version: v1
resource: secrets
namespaced: true
objects:
- apiVersion: v1
kind: Secret
metadata:
name: lens-test-lens-service
namespace: lens
data:
service-token: c2F2ZWQtc2VydmljZS10b2tlbg==
- apiVersion: v1
kind: Secret
metadata:
name: lens-test-lens-clickhouse
namespace: lens
data:
password: c2F2ZWQtZGF0YWJhc2UtcGFzc3dvcmQ=
tests:
- it: reuses the service credential instead of breaking running services
documentIndex: 0
asserts:
- equal:
path: data.service-token
value: c2F2ZWQtc2VydmljZS10b2tlbg==
- it: reuses the database password instead of locking out stored traces
documentIndex: 1
asserts:
- equal:
path: data.password
value: c2F2ZWQtZGF0YWJhc2UtcGFzc3dvcmQ=

View file

@ -117,7 +117,7 @@ tests:
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- failedTemplate:
errorMessage: lensWorker.publicUrl is required
errorMessage: lensWorker.publicUrl is required when there is no single ingress hostname
- it: omits Lens connection settings when disabled in deployment.yaml
template: deployment.yaml
asserts:

View file

@ -0,0 +1,160 @@
suite: Lens managed setup
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://traces.example
release:
namespace: lens
name: lens-test
templates:
- lens/deployment.yaml
- lens/clickhouse.yaml
- lens/secrets.yaml
tests:
- it: generates both private credentials for a new installation
template: lens/secrets.yaml
asserts:
- hasDocuments:
count: 2
- matchRegex:
path: data.service-token
pattern: '^[A-Za-z0-9+/]{86}==$'
documentIndex: 0
- matchRegex:
path: data.password
pattern: '^[A-Za-z0-9+/]{86}==$'
documentIndex: 1
- equal:
path: metadata.annotations["helm.sh/resource-policy"]
value: keep
- it: connects Lens to its private bundled storage
template: lens/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_HOST
value: lens-test-lens-clickhouse
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: lens-test-lens-clickhouse
key: password
- it: stores traces on a persistent volume with the chosen storage class
template: lens/clickhouse.yaml
documentIndex: 1
set:
lensWorker.clickhouse.storage: 40Gi
lensWorker.clickhouse.storageClassName: fast
asserts:
- equal:
path: spec.volumeClaimTemplates[0].spec
value:
accessModes: [ReadWriteOnce]
storageClassName: fast
resources:
requests:
storage: 40Gi
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: lens-test-lens-clickhouse
key: password
- it: keeps an existing external database instead of creating a new one
template: lens/clickhouse.yaml
set:
lensWorker.clickhouseSecret.name: external-clickhouse
asserts:
- hasDocuments:
count: 0
- it: leaves supplied secrets under their existing manager
template: lens/secrets.yaml
set:
lensWorker.clickhouseSecret.name: external-clickhouse
lensWorker.serviceTokenSecret.name: external-service
asserts:
- hasDocuments:
count: 0
- it: creates no credentials or database when Lens is disabled
templates:
- lens/secrets.yaml
- lens/clickhouse.yaml
set:
lensWorker.enabled: false
asserts:
- hasDocuments:
count: 0
- it: requires an external database when bundled storage is explicitly disabled
template: lens/deployment.yaml
set:
lensWorker.clickhouse.enabled: false
asserts:
- failedTemplate:
errorMessage: lensWorker.clickhouseSecret.name is required
- it: keeps storage names valid and uses the same name for the connection
set:
fullnameOverride: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
asserts:
- matchRegex:
path: metadata.name
pattern: '^[a-z]([-a-z0-9]{0,61}[a-z0-9])?$'
template: lens/clickhouse.yaml
documentIndex: 0
- matchRegex:
path: metadata.name
pattern: '^[a-z]([-a-z0-9]{0,61}[a-z0-9])?$'
template: lens/clickhouse.yaml
documentIndex: 1
- equal:
path: metadata.name
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/clickhouse.yaml
documentIndex: 0
- equal:
path: spec.serviceName
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/clickhouse.yaml
documentIndex: 1
- equal:
path: metadata.name
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/secrets.yaml
documentIndex: 1
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_HOST
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
key: password
template: lens/deployment.yaml
- it: allows cold storage startup and inherits registry credentials
template: lens/clickhouse.yaml
documentIndex: 1
set:
imagePullSecrets: [{name: registry-auth}]
asserts:
- equal:
path: spec.template.spec.imagePullSecrets
value: [{name: registry-auth}]
- equal:
path: spec.template.spec.containers[0].startupProbe
value:
httpGet: {path: /ping, port: http}
periodSeconds: 5
timeoutSeconds: 3
failureThreshold: 60

View file

@ -667,6 +667,17 @@ lensWorker:
serviceTokenSecret:
name: ""
key: service-token
clickhouse:
enabled: true
image: clickhouse/clickhouse-server:26.9.6.6
storage: 20Gi
storageClassName: null
resources:
requests:
cpu: 100m
memory: 512Mi
limits:
memory: 2Gi
clickhouseDatabase: litellm
retentionDays: 14
clickhouseSecret:

View file

@ -520,11 +520,11 @@ shutdown drain window.
- name: LITELLM_LENS_URL
value: {{ printf "http://%s-lens-worker:%v" (include "litellm.fullname" .) .Values.lensWorker.service.port | quote }}
- name: LITELLM_LENS_PUBLIC_URL
value: {{ required "lensWorker.publicUrl is required" .Values.lensWorker.publicUrl | quote }}
value: {{ include "litellm.lensWorker.publicUrl" . | quote }}
- name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: {{ required "lensWorker.serviceTokenSecret.name is required" .Values.lensWorker.serviceTokenSecret.name | quote }}
name: {{ include "litellm.lensWorker.serviceTokenSecretName" . | quote }}
key: {{ .Values.lensWorker.serviceTokenSecret.key | quote }}
{{- end }}
{{- end -}}
@ -534,3 +534,29 @@ shutdown drain window.
{{- $_ := set $labels "app.kubernetes.io/name" (printf "%s-lens-worker" (include "litellm.name" . | trunc 51 | trimSuffix "-")) -}}
{{- toYaml $labels -}}
{{- end -}}
{{- define "litellm.lensWorker.serviceTokenSecretName" -}}
{{- .Values.lensWorker.serviceTokenSecret.name | default (printf "%s-lens-service" (include "litellm.fullname" .)) -}}
{{- end -}}
{{- define "litellm.lensWorker.bundledClickhouse" -}}
{{- if and .Values.lensWorker.enabled .Values.lensWorker.clickhouse.enabled (not .Values.lensWorker.clickhouseSecret.name) -}}true{{- end -}}
{{- end -}}
{{- define "litellm.lensWorker.publicUrl" -}}
{{- if .Values.lensWorker.publicUrl -}}
{{- .Values.lensWorker.publicUrl -}}
{{- else if .Values.lensWorker.ingress.enabled -}}
{{- $tls := or (not (empty .Values.lensWorker.ingress.tls)) (hasKey .Values.lensWorker.ingress.annotations "alb.ingress.kubernetes.io/certificate-arn") -}}
{{- printf "%s://%s" (ternary "https" "http" $tls) (required "lensWorker.ingress.host is required" .Values.lensWorker.ingress.host) -}}
{{- else if and .Values.ingress.enabled .Values.ingress.host -}}
{{- $tls := or (not (empty .Values.ingress.tls)) (hasKey .Values.ingress.annotations "alb.ingress.kubernetes.io/certificate-arn") -}}
{{- printf "%s://%s/lens-ingest" (ternary "https" "http" $tls) .Values.ingress.host -}}
{{- else -}}
{{- fail "lensWorker.publicUrl is required when there is no single ingress hostname" -}}
{{- end -}}
{{- end -}}
{{- define "litellm.lensWorker.clickhouseName" -}}
{{- printf "%s-lens-clickhouse" (include "litellm.fullname" . | trunc 47 | trimSuffix "-") -}}
{{- end -}}

View file

@ -0,0 +1,98 @@
{{- if include "litellm.lensWorker.bundledClickhouse" . }}
{{- $name := include "litellm.lensWorker.clickhouseName" . }}
apiVersion: v1
kind: Service
metadata:
name: {{ $name }}
spec:
clusterIP: None
selector:
app.kubernetes.io/instance: {{ .Release.Name }}
app.kubernetes.io/component: lens-clickhouse
ports:
- name: http
port: 8123
targetPort: http
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
name: {{ $name }}
spec:
serviceName: {{ $name }}
replicas: 1
selector:
matchLabels:
app.kubernetes.io/instance: {{ .Release.Name }}
app.kubernetes.io/component: lens-clickhouse
template:
metadata:
labels:
app.kubernetes.io/instance: {{ .Release.Name }}
app.kubernetes.io/component: lens-clickhouse
spec:
automountServiceAccountToken: false
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
{{- end }}
securityContext:
runAsNonRoot: true
runAsUser: 101
runAsGroup: 101
fsGroup: 101
seccompProfile:
type: RuntimeDefault
containers:
- name: clickhouse
image: {{ .Values.lensWorker.clickhouse.image | quote }}
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop: [ALL]
env:
- name: CLICKHOUSE_USER
value: default
- name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ $name }}
key: password
- name: CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT
value: "1"
ports:
- name: http
containerPort: 8123
startupProbe:
httpGet:
path: /ping
port: http
periodSeconds: 5
timeoutSeconds: 3
failureThreshold: 60
readinessProbe:
httpGet:
path: /ping
port: http
livenessProbe:
httpGet:
path: /ping
port: http
timeoutSeconds: 3
resources:
{{- toYaml .Values.lensWorker.clickhouse.resources | nindent 12 }}
volumeMounts:
- name: data
mountPath: /var/lib/clickhouse
volumeClaimTemplates:
- metadata:
name: data
spec:
accessModes: [ReadWriteOnce]
{{- if ne .Values.lensWorker.clickhouse.storageClassName nil }}
storageClassName: {{ .Values.lensWorker.clickhouse.storageClassName | quote }}
{{- end }}
resources:
requests:
storage: {{ .Values.lensWorker.clickhouse.storage | quote }}
{{- end }}

View file

@ -45,13 +45,23 @@ spec:
- name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: {{ required "lensWorker.serviceTokenSecret.name is required" .Values.lensWorker.serviceTokenSecret.name | quote }}
name: {{ include "litellm.lensWorker.serviceTokenSecretName" . | quote }}
key: {{ .Values.lensWorker.serviceTokenSecret.key | quote }}
{{- if include "litellm.lensWorker.bundledClickhouse" . }}
- name: CLICKHOUSE_HOST
value: {{ include "litellm.lensWorker.clickhouseName" . | quote }}
- name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.lensWorker.clickhouseName" . | quote }}
key: password
{{- else }}
- name: CLICKHOUSE_URL
valueFrom:
secretKeyRef:
name: {{ required "lensWorker.clickhouseSecret.name is required" .Values.lensWorker.clickhouseSecret.name | quote }}
key: {{ .Values.lensWorker.clickhouseSecret.key | quote }}
{{- end }}
- name: CLICKHOUSE_DATABASE
value: {{ .Values.lensWorker.clickhouseDatabase | quote }}
- name: AGENT_TRACING_RETENTION_DAYS

View file

@ -0,0 +1,27 @@
{{- if and .Values.lensWorker.enabled (not .Values.lensWorker.serviceTokenSecret.name) }}
{{- $name := include "litellm.lensWorker.serviceTokenSecretName" . }}
{{- $existing := lookup "v1" "Secret" .Release.Namespace $name }}
apiVersion: v1
kind: Secret
metadata:
name: {{ $name }}
annotations:
helm.sh/resource-policy: keep
type: Opaque
data:
{{ .Values.lensWorker.serviceTokenSecret.key }}: {{ if $existing }}{{ required "Saved Lens service secret is missing its key" (index $existing.data .Values.lensWorker.serviceTokenSecret.key) | quote }}{{ else }}{{ randAlphaNum 64 | b64enc | quote }}{{ end }}
{{- end }}
{{- if include "litellm.lensWorker.bundledClickhouse" . }}
{{- $name := include "litellm.lensWorker.clickhouseName" . }}
{{- $existing := lookup "v1" "Secret" .Release.Namespace $name }}
---
apiVersion: v1
kind: Secret
metadata:
name: {{ $name }}
annotations:
helm.sh/resource-policy: keep
type: Opaque
data:
password: {{ if $existing }}{{ required "Saved Lens ClickHouse secret is missing its password" (index $existing.data "password") | quote }}{{ else }}{{ randAlphaNum 64 | b64enc | quote }}{{ end }}
{{- end }}

View file

@ -0,0 +1,86 @@
suite: Lens endpoint defaults
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
set:
lensWorker.enabled: true
ingress.enabled: true
ingress.host: gateway.example
ingress.tls:
- hosts: [gateway.example]
secretName: tls
tests:
- it: derives the trace endpoint from the deployment hostname and TLS
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://gateway.example/lens-ingest
- it: preserves an explicitly configured public address
set:
lensWorker.publicUrl: https://custom.example/prefix
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://custom.example/prefix
- it: uses the dedicated Lens ingress when configured
set:
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
lensWorker.ingress.tls:
- hosts: [traces.example]
secretName: traces-tls
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://traces.example
- it: uses HTTPS for a dedicated ALB ingress with certificate annotations
set:
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
lensWorker.ingress.className: alb
lensWorker.ingress.tls: []
lensWorker.ingress.annotations:
alb.ingress.kubernetes.io/certificate-arn: arn:aws:acm:us-west-2:123456789012:certificate/test
alb.ingress.kubernetes.io/listen-ports: '[{"HTTP":80},{"HTTPS":443}]'
alb.ingress.kubernetes.io/ssl-redirect: "443"
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: https://traces.example
- it: uses HTTP for a dedicated ingress without its own TLS
set:
lensWorker.ingress.enabled: true
lensWorker.ingress.host: traces.example
lensWorker.ingress.tls: []
lensWorker.ingress.annotations: {}
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: http://traces.example
- it: uses HTTP when ingress has no TLS
set:
ingress.tls: []
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_PUBLIC_URL
value: http://gateway.example/lens-ingest
values:
- ./values/required.yaml

View file

@ -0,0 +1,47 @@
suite: Lens preserves credentials across Helm upgrades
templates:
- lens/secrets.yaml
set:
fullnameOverride: lens-test
lensWorker.enabled: true
release:
namespace: lens
name: lens-test
upgrade: true
kubernetesProvider:
scheme:
v1/Secret:
gvr:
version: v1
resource: secrets
namespaced: true
objects:
- apiVersion: v1
kind: Secret
metadata:
name: lens-test-lens-service
namespace: lens
data:
service-token: c2F2ZWQtc2VydmljZS10b2tlbg==
- apiVersion: v1
kind: Secret
metadata:
name: lens-test-lens-clickhouse
namespace: lens
data:
password: c2F2ZWQtZGF0YWJhc2UtcGFzc3dvcmQ=
tests:
- it: reuses the service credential instead of breaking running services
documentIndex: 0
asserts:
- equal:
path: data.service-token
value: c2F2ZWQtc2VydmljZS10b2tlbg==
- it: reuses the database password instead of locking out stored traces
documentIndex: 1
asserts:
- equal:
path: data.password
value: c2F2ZWQtZGF0YWJhc2UtcGFzc3dvcmQ=
values:
- ./values/required.yaml

View file

@ -136,7 +136,7 @@ tests:
lensWorker.clickhouseSecret.name: lens-storage
asserts:
- failedTemplate:
errorMessage: lensWorker.publicUrl is required
errorMessage: lensWorker.publicUrl is required when there is no single ingress hostname
- it: omits Lens connection settings when disabled in gateway/deployment.yaml
template: gateway/deployment.yaml
asserts:

View file

@ -0,0 +1,162 @@
suite: Lens managed setup
set:
fullnameOverride: lens-test
lensWorker.enabled: true
lensWorker.publicUrl: https://traces.example
release:
namespace: lens
name: lens-test
templates:
- lens/deployment.yaml
- lens/clickhouse.yaml
- lens/secrets.yaml
tests:
- it: generates both private credentials for a new installation
template: lens/secrets.yaml
asserts:
- hasDocuments:
count: 2
- matchRegex:
path: data.service-token
pattern: '^[A-Za-z0-9+/]{86}==$'
documentIndex: 0
- matchRegex:
path: data.password
pattern: '^[A-Za-z0-9+/]{86}==$'
documentIndex: 1
- equal:
path: metadata.annotations["helm.sh/resource-policy"]
value: keep
- it: connects Lens to its private bundled storage
template: lens/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_HOST
value: lens-test-lens-clickhouse
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: lens-test-lens-clickhouse
key: password
- it: stores traces on a persistent volume with the chosen storage class
template: lens/clickhouse.yaml
documentIndex: 1
set:
lensWorker.clickhouse.storage: 40Gi
lensWorker.clickhouse.storageClassName: fast
asserts:
- equal:
path: spec.volumeClaimTemplates[0].spec
value:
accessModes: [ReadWriteOnce]
storageClassName: fast
resources:
requests:
storage: 40Gi
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: lens-test-lens-clickhouse
key: password
- it: keeps an existing external database instead of creating a new one
template: lens/clickhouse.yaml
set:
lensWorker.clickhouseSecret.name: external-clickhouse
asserts:
- hasDocuments:
count: 0
- it: leaves supplied secrets under their existing manager
template: lens/secrets.yaml
set:
lensWorker.clickhouseSecret.name: external-clickhouse
lensWorker.serviceTokenSecret.name: external-service
asserts:
- hasDocuments:
count: 0
- it: creates no credentials or database when Lens is disabled
templates:
- lens/secrets.yaml
- lens/clickhouse.yaml
set:
lensWorker.enabled: false
asserts:
- hasDocuments:
count: 0
- it: requires an external database when bundled storage is explicitly disabled
template: lens/deployment.yaml
set:
lensWorker.clickhouse.enabled: false
asserts:
- failedTemplate:
errorMessage: lensWorker.clickhouseSecret.name is required
- it: keeps storage names valid and uses the same name for the connection
set:
fullnameOverride: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
asserts:
- matchRegex:
path: metadata.name
pattern: '^[a-z]([-a-z0-9]{0,61}[a-z0-9])?$'
template: lens/clickhouse.yaml
documentIndex: 0
- matchRegex:
path: metadata.name
pattern: '^[a-z]([-a-z0-9]{0,61}[a-z0-9])?$'
template: lens/clickhouse.yaml
documentIndex: 1
- equal:
path: metadata.name
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/clickhouse.yaml
documentIndex: 0
- equal:
path: spec.serviceName
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/clickhouse.yaml
documentIndex: 1
- equal:
path: metadata.name
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/secrets.yaml
documentIndex: 1
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_HOST
value: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
template: lens/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: CLICKHOUSE_PASSWORD
valueFrom:
secretKeyRef:
name: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-lens-clickhouse
key: password
template: lens/deployment.yaml
- it: allows cold storage startup and inherits registry credentials
template: lens/clickhouse.yaml
documentIndex: 1
set:
imagePullSecrets: [{name: registry-auth}]
asserts:
- equal:
path: spec.template.spec.imagePullSecrets
value: [{name: registry-auth}]
- equal:
path: spec.template.spec.containers[0].startupProbe
value:
httpGet: {path: /ping, port: http}
periodSeconds: 5
timeoutSeconds: 3
failureThreshold: 60
values:
- ./values/required.yaml

View file

@ -77,13 +77,20 @@ tests:
asserts:
- hasDocuments:
count: 0
- it: requires a shared service secret when enabled
- it: uses the managed service secret when none is supplied
template: lens/deployment.yaml
set:
lensWorker.enabled: true
lensWorker.publicUrl: https://traces.example
asserts:
- failedTemplate:
errorMessage: lensWorker.serviceTokenSecret.name is required
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_LENS_SERVICE_TOKEN
valueFrom:
secretKeyRef:
name: RELEASE-NAME-litellm-lens-service
key: service-token
- it: uses the chart release and a secret without granting Kubernetes access
template: lens/deployment.yaml
chart:

View file

@ -644,6 +644,17 @@ lensWorker:
serviceTokenSecret:
name: ""
key: service-token
clickhouse:
enabled: true
image: clickhouse/clickhouse-server:26.9.6.6
storage: 20Gi
storageClassName: null
resources:
requests:
cpu: 100m
memory: 512Mi
limits:
memory: 2Gi
clickhouseDatabase: litellm
retentionDays: 14
clickhouseSecret:

View file

@ -42,5 +42,6 @@ prettyplease = "0.2"
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["test-util"] }
wiremock.workspace = true
uuid.workspace = true

View file

@ -46,6 +46,15 @@ pub fn response(content_type: Option<&str>, outcome: Result<(), Error>) -> Respo
.map(|_| StatusCode::OK)
.unwrap_or_else(|error| error.status());
let message = status.canonical_reason().unwrap_or("Trace request failed");
let rpc_code = match status {
StatusCode::OK => 0,
StatusCode::BAD_REQUEST => 3,
StatusCode::UNAUTHORIZED => 16,
StatusCode::PAYLOAD_TOO_LARGE | StatusCode::TOO_MANY_REQUESTS => 8,
StatusCode::CONFLICT => 10,
StatusCode::SERVICE_UNAVAILABLE => 14,
_ => 2,
};
let protobuf = content_type.is_some_and(|value| {
value
.split(';')
@ -58,7 +67,7 @@ pub fn response(content_type: Option<&str>, outcome: Result<(), Error>) -> Respo
Vec::new()
} else {
OtlpError {
code: 0,
code: rpc_code,
message: message.into(),
}
.encode_to_vec()
@ -70,7 +79,7 @@ pub fn response(content_type: Option<&str>, outcome: Result<(), Error>) -> Respo
if outcome.is_ok() {
b"{}".to_vec()
} else {
serde_json::json!({"code": 0, "message": message})
serde_json::json!({"code": rpc_code, "message": message})
.to_string()
.into_bytes()
},
@ -169,3 +178,67 @@ async fn store(
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::{OtlpError, response};
use crate::Error;
use axum::body::to_bytes;
use prost::Message;
use rstest::rstest;
#[rstest]
#[case::invalid_request(Error::InvalidRequest, 400, 3, false)]
#[case::unauthenticated(Error::Unauthorized, 401, 16, false)]
#[case::payload_too_large(Error::TooLarge, 413, 8, false)]
#[case::credentials_pending(Error::CredentialsPending, 429, 8, true)]
#[case::storage_unavailable(Error::Unavailable, 503, 14, true)]
#[case::conflict(Error::TraceChanged, 409, 10, false)]
#[tokio::test]
async fn rejected_batches_have_matching_http_and_rpc_errors(
#[case] error: Error,
#[case] http_status: u16,
#[case] rpc_code: i32,
#[case] retryable: bool,
#[values("application/json", "application/x-protobuf")] content_type: &str,
) {
let reply = response(Some(content_type), Err(error));
assert_eq!(reply.status().as_u16(), http_status);
assert_eq!(reply.headers()["content-type"], content_type);
assert_eq!(
reply
.headers()
.get("retry-after")
.map(|v| v.to_str().unwrap()),
retryable.then_some("5")
);
let message = reply.status().canonical_reason().unwrap();
let body = to_bytes(reply.into_body(), 1024).await.unwrap();
if content_type == "application/x-protobuf" {
let status = OtlpError::decode(body).unwrap();
assert_eq!(status.code, rpc_code);
assert_eq!(status.message, message);
} else {
let status: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(
status,
serde_json::json!({"code": rpc_code, "message": message})
);
}
}
#[rstest]
#[case::json("application/json", b"{}")]
#[case::protobuf("application/x-protobuf", b"")]
#[tokio::test]
async fn accepted_batches_keep_the_empty_export_response(
#[case] content_type: &str,
#[case] expected: &[u8],
) {
let reply = response(Some(content_type), Ok(()));
assert_eq!(reply.status(), 200);
assert_eq!(reply.headers()["content-type"], content_type);
assert!(!reply.headers().contains_key("retry-after"));
assert_eq!(to_bytes(reply.into_body(), 1024).await.unwrap(), expected);
}
}

View file

@ -26,6 +26,7 @@ use litellm_traces_clickhouse::InsertTable;
use serde_json::Value;
use std::{
collections::BTreeMap,
future::Future,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
@ -33,6 +34,7 @@ use std::{
time::Duration,
};
pub use storage::Storage;
use tokio::sync::Semaphore;
#[allow(
dead_code,
@ -46,7 +48,8 @@ pub use storage::Storage;
pub mod wire {
include!(concat!(env!("OUT_DIR"), "/wire.rs"));
}
use tokio::sync::Semaphore;
const READ_QUEUE_WAIT: Duration = Duration::from_secs(10);
pub struct State {
pub credentials: Arc<auth::Credentials>,
@ -80,6 +83,15 @@ impl State {
}
}
async fn wait_for_read_slot<P>(
acquire: impl Future<Output = Result<P, tokio::sync::AcquireError>>,
) -> Result<P, Error> {
tokio::time::timeout(READ_QUEUE_WAIT, acquire)
.await
.map_err(|_| Error::Unavailable)?
.map_err(|_| Error::Unavailable)
}
pub fn router(state: Arc<State>) -> Router {
let public = Router::new()
.route("/health/live", get(|| async { StatusCode::OK }))
@ -104,6 +116,7 @@ pub fn router(state: Arc<State>) -> Router {
Router::new()
.route("/internal/read", post(read))
.route("/internal/spend", post(spend))
.route("/internal/feedback", post(feedback))
.route("/internal/credentials", post(credentials))
.route("/internal/status", get(status)),
)
@ -125,10 +138,7 @@ async fn receipt(
) -> Result<Json<Value>, Error> {
let tenant = state.credentials.tenant(&headers)?;
state.require_storage()?;
let _permit = state
.read_slots
.try_acquire()
.map_err(|_| Error::Unavailable)?;
let _permit = wait_for_read_slot(state.read_slots.acquire()).await?;
let body = tokio::time::timeout(Duration::from_secs(5), to_bytes(body, 64 * 1024))
.await
.map_err(|_| Error::Unavailable)?
@ -206,11 +216,7 @@ async fn read(
) -> Result<Json<Value>, Error> {
auth::authorize_service(&headers, &state.service_token)?;
state.require_storage()?;
let permit = state
.read_slots
.clone()
.try_acquire_owned()
.map_err(|_| Error::Unavailable)?;
let permit = wait_for_read_slot(state.read_slots.clone().acquire_owned()).await?;
let body = tokio::time::timeout(Duration::from_secs(10), to_bytes(body, 1024 * 1024))
.await
.map_err(|_| Error::Unavailable)?
@ -228,6 +234,23 @@ async fn spend(
AppState(state): AppState<Arc<State>>,
headers: HeaderMap,
body: Body,
) -> Result<StatusCode, Error> {
insert(state, headers, body, InsertTable::SpendLogs).await
}
async fn feedback(
AppState(state): AppState<Arc<State>>,
headers: HeaderMap,
body: Body,
) -> Result<StatusCode, Error> {
insert(state, headers, body, InsertTable::LensFeedback).await
}
async fn insert(
state: Arc<State>,
headers: HeaderMap,
body: Body,
table: InsertTable,
) -> Result<StatusCode, Error> {
auth::authorize_service(&headers, &state.service_token)?;
state.require_storage()?;
@ -251,7 +274,7 @@ async fn spend(
&state.storage.client,
state.storage.config.storage().writer(),
state.storage.config.storage().database(),
InsertTable::SpendLogs,
table,
rows,
)
.await?;
@ -279,3 +302,39 @@ pub async fn provision(state: Arc<State>) {
tokio::time::sleep(Duration::from_secs(10)).await;
}
}
#[cfg(test)]
mod tests {
use super::{Error, READ_QUEUE_WAIT, wait_for_read_slot};
use std::sync::Arc;
use tokio::sync::Semaphore;
#[tokio::test]
async fn ninth_read_waits_for_a_permit_and_succeeds() {
let slots = Arc::new(Semaphore::new(8));
let permits = (0..8)
.map(|_| slots.clone().try_acquire_owned().expect("available permit"))
.collect::<Vec<_>>();
let waiting_slots = slots.clone();
let waiting =
tokio::spawn(async move { wait_for_read_slot(waiting_slots.acquire_owned()).await });
tokio::task::yield_now().await;
assert!(!waiting.is_finished());
drop(permits);
assert!(waiting.await.expect("joined read").is_ok());
}
#[tokio::test(start_paused = true)]
async fn read_queue_timeout_returns_unavailable() {
let slots = Arc::new(Semaphore::new(0));
let waiting = tokio::spawn(wait_for_read_slot(slots.acquire_owned()));
tokio::task::yield_now().await;
tokio::time::advance(READ_QUEUE_WAIT).await;
assert!(matches!(
waiting.await.expect("joined read"),
Err(Error::Unavailable)
));
}
}

View file

@ -116,6 +116,74 @@ async fn agent_picker_query_preserves_scope_through_the_internal_read_route() {
assert_eq!(response.json::<serde_json::Value>().await.unwrap(), result);
}
#[rstest]
#[tokio::test]
async fn feedback_summary_query_preserves_scope_through_the_internal_read_route() {
let store = MockServer::start().await;
let result = json!({"data": [{
"trace_id": "1234567890abcdef1234567890abcdef", "trace_ref": "REF",
"count": "2", "average": 5.5, "lowest": "2"
}]});
Mock::given(method("POST"))
.and(body_string_contains("FROM lens_feedback FINAL"))
.and(query_param("param_all_teams", "0"))
.and(query_param("param_team", "feedback-team"))
.and(query_param(
"param_trace_ids",
"['1234567890abcdef1234567890abcdef']",
))
.respond_with(ResponseTemplate::new(200).set_body_json(&result))
.expect(1)
.mount(&store)
.await;
let server = serve(&store.uri(), true).await;
let response = http_client()
.unwrap()
.post(format!("{}/internal/read", server.url))
.bearer_auth(SERVICE_TOKEN)
.json(&json!({
"operation": "query", "name": "feedback_summary", "parameters": {
"all_teams": 0, "team": "feedback-team", "key_hash": "",
"trace_ids": ["1234567890abcdef1234567890abcdef"]
}
}))
.send()
.await
.unwrap();
assert_eq!(response.status(), 200);
assert_eq!(response.json::<serde_json::Value>().await.unwrap(), result);
}
#[rstest]
#[tokio::test]
async fn feedback_rows_are_written_to_the_feedback_table() {
let store = MockServer::start().await;
Mock::given(method("POST"))
.and(query_param(
"query",
"INSERT INTO `litellm`.lens_feedback FORMAT JSONEachRow",
))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&store)
.await;
let server = serve(&store.uri(), true).await;
let response = http_client()
.unwrap()
.post(format!("{}/internal/feedback", server.url))
.bearer_auth(SERVICE_TOKEN)
.json(&json!([{
"TeamId": "team", "ApiKeyHash": "", "TraceId": "1234567890abcdef1234567890abcdef",
"Author": "customer-1042", "Score": 2, "Comment": "wrong command",
"CreatedAt": "2026-10-07T21:57:01.414Z", "UpdatedAt": "2026-10-07T21:57:01.414Z",
"IsDeleted": 0
}]))
.send()
.await
.unwrap();
assert_eq!(response.status(), 204);
}
#[rstest]
#[tokio::test]
async fn ingestion_confirms_storage_and_overwrites_exporter_tenant() {
@ -297,7 +365,7 @@ async fn ingestion_key_cannot_read_or_export_gateway_records() {
let store = MockServer::start().await;
let server = serve(&store.uri(), true).await;
let client = http_client().unwrap();
for path in ["/internal/read", "/internal/spend"] {
for path in ["/internal/read", "/internal/spend", "/internal/feedback"] {
let response = client
.post(format!("{}{path}", server.url))
.bearer_auth(KEY)

View file

@ -328,6 +328,7 @@ mod tests {
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
source_user: String::new(),
team_id: String::new(),
api_key_hash: String::new(),
user_id: String::new(),

View file

@ -289,6 +289,7 @@ fn span(index: usize) -> TraceSpansRow {
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
source_user: String::new(),
team_id: "team".into(),
api_key_hash: "key".into(),
user_id: "user".into(),

View file

@ -42,6 +42,7 @@ fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> Trac
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
source_user: String::new(),
team_id: String::new(),
api_key_hash: String::new(),
user_id: String::new(),

View file

@ -0,0 +1,18 @@
CREATE TABLE IF NOT EXISTS {database}.lens_feedback
(
TeamId LowCardinality(String),
ApiKeyHash String,
TraceId String CODEC(ZSTD(1)),
Author String,
Score UInt8,
Comment String CODEC(ZSTD(3)),
CreatedAt DateTime64(3),
UpdatedAt DateTime64(3),
IsDeleted UInt8,
EngineReceivedMs UInt64 DEFAULT 0,
INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1,
CONSTRAINT score_range CHECK Score <= 10
)
ENGINE = ReplacingMergeTree(UpdatedAt, IsDeleted)
ORDER BY (TeamId, ApiKeyHash, TraceId, Author)
SETTINGS materialize_ttl_recalculate_only = 1, non_replicated_deduplication_window = 1000

View file

@ -0,0 +1,12 @@
SELECT TraceId AS trace_id,
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref,
Author AS author, Score AS score, Comment AS comment,
formatDateTime(CreatedAt, '%FT%T.%fZ', 'UTC') AS created_at,
formatDateTime(UpdatedAt, '%FT%T.%fZ', 'UTC') AS updated_at
FROM lens_feedback FINAL
WHERE TraceId = {trace_id:String}
AND hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String}
AND ({all_teams:UInt8}=1 OR TeamId={team:String})
AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})
AND IsDeleted = 0
ORDER BY UpdatedAt DESC, Author

View file

@ -0,0 +1,9 @@
SELECT TraceId AS trace_id,
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref,
count() AS count, avg(Score) AS average, min(Score) AS lowest
FROM lens_feedback FINAL
WHERE TraceId IN {trace_ids:Array(String)}
AND ({all_teams:UInt8}=1 OR TeamId={team:String})
AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})
AND IsDeleted = 0
GROUP BY TeamId, ApiKeyHash, TraceId

View file

@ -0,0 +1,9 @@
SELECT TeamId AS team_id, ApiKeyHash AS key_hash,
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref
FROM agent_traces_by_key
WHERE TraceId = {trace_id:String}
AND ({all_teams:UInt8}=1 OR TeamId={team:String})
AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})
AND ({trace_ref:String}='' OR hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId)))={trace_ref:String})
GROUP BY TeamId, ApiKeyHash, TraceId
LIMIT 2

View file

@ -5,7 +5,6 @@ SELECT TraceId AS trace_id,
ifNull(any(RootName), '') AS name, any(ServiceName) AS service,
ifNull(any(RootInput), '') AS input_preview, ifNull(any(RootStatus), '') AS status,
toUnixTimestamp64Milli(min(StartTs)) AS start_ms,
min(StartTs) AS trace_start, max(EndTs) AS trace_end,
dateDiff('millisecond', min(StartTs), max(EndTs)) AS duration_ms,
sum(SpanCount) AS span_count,
sum(AgentCount) AS agent_invocations,
@ -27,7 +26,7 @@ HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64})
ORDER BY start_ms DESC, trace_ref DESC
LIMIT {limit:UInt32}
)
SELECT page.* EXCEPT (trace_start, trace_end),
SELECT page.*,
identities.agent_names AS agent_names, identities.agent_count AS agent_count,
identities.frameworks AS frameworks
FROM page
@ -37,10 +36,8 @@ LEFT JOIN (
arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks,
uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count
FROM otel_traces
WHERE Timestamp >= (SELECT min(trace_start) FROM page)
AND Timestamp <= (SELECT max(trace_end) FROM page)
WHERE Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64})
AND TraceId IN (SELECT trace_id FROM page)
AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page)
GROUP BY TeamId, ApiKeyHash, TraceId
) AS identities
ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash

View file

@ -14,7 +14,7 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.SpanAttributes['agent.source.title'] AS source_title, o.SpanAttributes['agent.source.user'] AS source_user,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64})

View file

@ -13,7 +13,7 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.SpanAttributes['agent.source.title'] AS source_title, o.SpanAttributes['agent.source.user'] AS source_user,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64})

View file

@ -14,7 +14,7 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.SpanAttributes['agent.source.title'] AS source_title, o.SpanAttributes['agent.source.user'] AS source_user,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.TraceId = {trace_id:String}

View file

@ -13,7 +13,7 @@ SELECT o.TraceId AS trace_id, o.SpanAttributes['lens.original_trace_id'] AS orig
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
AS tool_call_id,
o.SpanAttributes['agent.source.type'] AS source_type, o.SpanAttributes['agent.source.url'] AS source_url,
o.SpanAttributes['agent.source.title'] AS source_title,
o.SpanAttributes['agent.source.title'] AS source_title, o.SpanAttributes['agent.source.user'] AS source_user,
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.TraceId = {trace_id:String}

View file

@ -33,6 +33,7 @@ pub type InsertRow = BTreeMap<String, Shared<Value>>;
pub enum InsertTable {
OtelTraces,
SpendLogs,
LensFeedback,
}
impl InsertTable {
@ -40,6 +41,7 @@ impl InsertTable {
match value {
"otel_traces" => Ok(Self::OtelTraces),
"spend_logs" => Ok(Self::SpendLogs),
"lens_feedback" => Ok(Self::LensFeedback),
_ => Err(Error::InvalidTable),
}
}
@ -48,6 +50,7 @@ impl InsertTable {
match self {
Self::OtelTraces => "otel_traces",
Self::SpendLogs => "spend_logs",
Self::LensFeedback => "lens_feedback",
}
}
}

View file

@ -6,13 +6,16 @@ const SAMPLE_READ_LIMITS: ReadLimits = ReadLimits {
..litellm_storage_clickhouse::READ_LIMITS
};
pub const LENS_QUERIES: [litellm_traces::ReadQuery; 6] = [
pub const LENS_QUERIES: [litellm_traces::ReadQuery; 9] = [
litellm_traces::ReadQuery::TraceAgents,
litellm_traces::ReadQuery::Availability,
litellm_traces::ReadQuery::Agents,
litellm_traces::ReadQuery::Sample,
litellm_traces::ReadQuery::Content,
litellm_traces::ReadQuery::Evidence,
litellm_traces::ReadQuery::FeedbackTarget,
litellm_traces::ReadQuery::Feedback,
litellm_traces::ReadQuery::FeedbackSummary,
];
#[macro_rules_attribute::apply(wire_type)]
@ -339,3 +342,108 @@ impl Query for LensEvidence {
const SQL: &'static str = include_str!("../../query/lens_evidence.sql");
}
pub struct LensFeedbackTarget;
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[serde(deny_unknown_fields)]
pub struct LensFeedbackTargetParams {
#[serde(flatten)]
pub access: LensAccessParams,
pub trace_id: String,
pub trace_ref: String,
}
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[cfg_attr(feature = "schema", schemars(rename = "FeedbackTargetRow"))]
pub struct LensFeedbackTargetRow {
pub team_id: String,
pub key_hash: String,
pub trace_ref: String,
}
impl Query for LensFeedbackTarget {
type Params = LensFeedbackTargetParams;
type Row = LensFeedbackTargetRow;
const SQL: &'static str = include_str!("../../query/lens_feedback_target.sql");
}
pub struct LensFeedback;
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[serde(deny_unknown_fields)]
pub struct LensFeedbackParams {
#[serde(flatten)]
pub access: LensAccessParams,
pub trace_id: String,
pub trace_ref: String,
}
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[cfg_attr(feature = "schema", schemars(rename = "FeedbackRow"))]
pub struct LensFeedbackRow {
pub trace_id: String,
pub trace_ref: String,
pub author: String,
#[serde(deserialize_with = "super::number::deserialize")]
#[cfg_attr(
feature = "schema",
schemars(schema_with = "crate::wire_schema::u64_number")
)]
pub score: u64,
pub comment: String,
pub created_at: String,
pub updated_at: String,
}
impl Query for LensFeedback {
type Params = LensFeedbackParams;
type Row = LensFeedbackRow;
const SQL: &'static str = include_str!("../../query/lens_feedback.sql");
}
pub struct LensFeedbackSummary;
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[serde(deny_unknown_fields)]
pub struct LensFeedbackSummaryParams {
#[serde(flatten)]
pub access: LensAccessParams,
pub trace_ids: Vec<String>,
}
#[macro_rules_attribute::apply(wire_type)]
#[derive(Debug)]
#[cfg_attr(feature = "schema", schemars(rename = "FeedbackSummaryRow"))]
pub struct LensFeedbackSummaryRow {
pub trace_id: String,
pub trace_ref: String,
#[serde(deserialize_with = "super::number::deserialize")]
#[cfg_attr(
feature = "schema",
schemars(schema_with = "crate::wire_schema::u64_number")
)]
pub count: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub average: f64,
#[serde(deserialize_with = "super::number::deserialize")]
#[cfg_attr(
feature = "schema",
schemars(schema_with = "crate::wire_schema::u64_number")
)]
pub lowest: u64,
}
impl Query for LensFeedbackSummary {
type Params = LensFeedbackSummaryParams;
type Row = LensFeedbackSummaryRow;
const SQL: &'static str = include_str!("../../query/lens_feedback_summary.sql");
}

View file

@ -134,6 +134,8 @@ struct TraceSpansRowEncoding {
pub source_url: String,
#[serde(default)]
pub source_title: String,
#[serde(default)]
pub source_user: String,
pub team_id: String,
pub api_key_hash: String,
pub user_id: String,
@ -355,7 +357,7 @@ mod tests {
quoted,
);
round_trip::<TraceSpansRow>(
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "source_type": "slack", "source_url": "https://acme.slack.com/archives/C1/p1", "source_title": "thread", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "source_type": "slack", "source_url": "https://acme.slack.com/archives/C1/p1", "source_title": "thread", "source_user": "tin@berri.ai", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
quoted,
);
round_trip::<SpanDetailRow>(

View file

@ -14,10 +14,11 @@ static MIGRATOR: Migrator = Migrator {
..sqlx::migrate!("./migrations")
};
const RETENTION: [(&str, &str); 3] = [
const RETENTION: [(&str, &str); 4] = [
("otel_traces", "toDateTime(Timestamp)"),
("agent_traces_by_key", "toDateTime(StartTs)"),
("spend_logs", "toDateTime(start_time)"),
("lens_feedback", "toDateTime(CreatedAt)"),
];
fn validate_schema(database: &str, retention_days: u32) -> Result<(), Error> {

View file

@ -4,7 +4,7 @@
use std::{future::Future, marker::PhantomData};
use litellm_http::Client;
use litellm_storage_clickhouse::{Query, fetch};
use litellm_storage_clickhouse::{Query, ReadLimits, fetch};
use litellm_traces::query::named as contracts;
use litellm_traces_cache::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, StoreError};
use serde::{Serialize, de::DeserializeOwned};
@ -14,7 +14,12 @@ use crate::{
query::named::{SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow},
};
const PAGE_SIZE: u32 = 256;
const PAGE_SIZE: u32 = 8192;
const SPAN_BATCH_READ_LIMITS: ReadLimits = ReadLimits {
result_rows: PAGE_SIZE as u64,
response_bytes: 16 * 1024 * 1024,
..litellm_storage_clickhouse::READ_LIMITS
};
#[derive(Default)]
struct ReadBudget {
@ -67,6 +72,7 @@ struct Paged<K>(PhantomData<K>);
impl<K: Keyset> Query for Paged<K> {
type Params = Batch<K>;
type Row = K::Row;
const READ_LIMITS: ReadLimits = SPAN_BATCH_READ_LIMITS;
const SQL: &'static str = K::SQL;
}
@ -316,8 +322,8 @@ mod tests {
}
#[rstest]
#[case::fits(1000, PAGE_SIZE, &[256, 256, 256, 256])]
#[case::uniform_large_rows(1000, 100, &[256, 128, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64])]
#[case::fits(1000, PAGE_SIZE, &[8192])]
#[case::uniform_large_rows(200, 100, &[8192, 4096, 2048, 1024, 512, 256, 128, 64, 64, 64, 64])]
#[tokio::test]
async fn a_rejected_page_size_is_not_retried(
#[case] total: u32,
@ -334,6 +340,13 @@ mod tests {
assert_eq!(table.requests.lock().unwrap().as_slice(), requests);
}
#[test]
fn span_batches_use_larger_read_limits() {
let limits = <Paged<Numbers> as Query>::READ_LIMITS;
assert_eq!(limits.result_rows, 8192);
assert_eq!(limits.response_bytes, 16 * 1024 * 1024);
}
#[rstest]
#[tokio::test]
async fn a_single_oversized_row_fails_the_read() {
@ -346,7 +359,9 @@ mod tests {
assert!(matches!(result, Err(StoreError::TooLarge)), "{result:?}");
assert_eq!(
table.requests.lock().unwrap().as_slice(),
&[256, 128, 64, 32, 16, 8, 4, 2, 1]
&[
8192, 4096, 2048, 1024, 512, 256, 128, 64, 32, 16, 8, 4, 2, 1
]
);
}
}

View file

@ -37,6 +37,13 @@ pub async fn execute_named_read(
ReadQuery::Sample => named_json::<LensSample>(client, connection, parameters).await,
ReadQuery::Content => named_json::<LensContent>(client, connection, parameters).await,
ReadQuery::Evidence => named_json::<LensEvidence>(client, connection, parameters).await,
ReadQuery::FeedbackTarget => {
named_json::<LensFeedbackTarget>(client, connection, parameters).await
}
ReadQuery::Feedback => named_json::<LensFeedback>(client, connection, parameters).await,
ReadQuery::FeedbackSummary => {
named_json::<LensFeedbackSummary>(client, connection, parameters).await
}
}
}

View file

@ -81,6 +81,18 @@ pub fn schemas() -> BTreeMap<&'static str, Schema> {
("LensSampleParams", received::<lens::LensSampleParams>()),
("LensContentParams", received::<lens::LensContentParams>()),
("LensEvidenceParams", received::<lens::LensEvidenceParams>()),
(
"LensFeedbackTargetParams",
received::<lens::LensFeedbackTargetParams>(),
),
("LensFeedbackParams", received::<lens::LensFeedbackParams>()),
(
"LensFeedbackSummaryParams",
received::<lens::LensFeedbackSummaryParams>(),
),
("FeedbackTargetRow", received::<lens::LensFeedbackTargetRow>()),
("FeedbackRow", received::<lens::LensFeedbackRow>()),
("FeedbackSummaryRow", received::<lens::LensFeedbackSummaryRow>()),
(
"ActivityAvailability",
received::<lens::LensAvailabilityRow>(),

View file

@ -0,0 +1,408 @@
use std::collections::BTreeMap;
use litellm_traces_clickhouse::{
Connection, InsertTable, Parameter, ReadQuery, ensure_schema, execute_named_read, execute_read,
insert_rows,
};
use rstest::rstest;
use serde_json::{Value, json};
use time::{Duration, OffsetDateTime, format_description::well_known::Rfc3339};
mod support;
use support::{ClickHouseDatabase, TestResult, database};
struct Feedback<'a> {
team: &'a str,
key: &'a str,
trace: &'a str,
author: &'a str,
score: u8,
comment: &'a str,
edited_after_seconds: i64,
deleted: bool,
}
const SAVED: Feedback<'static> = Feedback {
team: "team-a",
key: "key-a",
trace: "trace-1",
author: "alice",
score: 3,
comment: "missed the file",
edited_after_seconds: 0,
deleted: false,
};
fn now() -> TestResult<OffsetDateTime> {
Ok(OffsetDateTime::now_utc().replace_millisecond(0)?)
}
fn iso(at: OffsetDateTime) -> TestResult<String> {
Ok(at.format(&Rfc3339)?)
}
fn row(feedback: &Feedback, created: OffsetDateTime) -> TestResult<BTreeMap<String, Value>> {
let updated = created + Duration::seconds(feedback.edited_after_seconds);
Ok(serde_json::from_value(json!({
"TeamId": feedback.team, "ApiKeyHash": feedback.key, "TraceId": feedback.trace,
"Author": feedback.author, "Score": feedback.score, "Comment": feedback.comment,
"CreatedAt": iso(created)?, "UpdatedAt": iso(updated)?,
"IsDeleted": u8::from(feedback.deleted)
}))?)
}
async fn ready(database: &ClickHouseDatabase) -> TestResult<Connection> {
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
Ok(writer)
}
async fn save(
database: &ClickHouseDatabase,
writer: &Connection,
created: OffsetDateTime,
feedback: &[Feedback<'_>],
) -> TestResult {
let rows = feedback
.iter()
.map(|entry| row(entry, created))
.collect::<TestResult<Vec<_>>>()?;
insert_rows(
&database.client,
writer,
"trace_test",
InsertTable::LensFeedback,
rows,
)
.await?;
Ok(())
}
fn access(team: &str) -> BTreeMap<String, Parameter> {
BTreeMap::from([
(
"all_teams".into(),
Parameter::Integer(i64::from(team.is_empty())),
),
("team".into(), Parameter::Text(team.into())),
("key_hash".into(), Parameter::Text(String::new())),
])
}
async fn read(
database: &ClickHouseDatabase,
query: ReadQuery,
parameters: BTreeMap<String, Parameter>,
) -> TestResult<Vec<Value>> {
let reader = Connection::configured(&database.url, "trace_test", "default", "")?;
let body: Value = serde_json::from_str(
&execute_named_read(&database.client, &reader, query, &parameters).await?,
)?;
Ok(body["data"]
.as_array()
.cloned()
.ok_or_else(|| format!("no data in {body}"))?)
}
async fn trace_ref(
database: &ClickHouseDatabase,
team: &str,
key: &str,
trace: &str,
) -> TestResult<String> {
let reader = Connection::configured(&database.url, "trace_test", "default", "")?;
let sql = format!(
"SELECT hex(SHA256(concat('{team}', char(0), '{key}', char(0), '{trace}'))) AS trace_ref"
);
let response: Value = serde_json::from_str(
&execute_read(&database.client, &reader, &sql, &BTreeMap::new()).await?,
)?;
Ok(response["data"][0]["trace_ref"]
.as_str()
.unwrap_or_default()
.to_owned())
}
async fn feedback(
database: &ClickHouseDatabase,
team: &str,
trace: &str,
trace_ref: &str,
) -> TestResult<Vec<Value>> {
let mut parameters = access(team);
parameters.insert("trace_id".into(), Parameter::Text(trace.into()));
parameters.insert("trace_ref".into(), Parameter::Text(trace_ref.into()));
read(database, ReadQuery::Feedback, parameters).await
}
fn timestamp(row: &Value, field: &str) -> TestResult<OffsetDateTime> {
Ok(OffsetDateTime::parse(
row[field].as_str().unwrap_or_default(),
&Rfc3339,
)?)
}
#[rstest]
#[tokio::test]
async fn a_later_save_replaces_the_authors_feedback_and_keeps_other_authors(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = ready(&database).await?;
let created = now()?;
save(&database, &writer, created, &[SAVED]).await?;
save(
&database,
&writer,
created,
&[
Feedback {
score: 8,
comment: "fine after retry",
edited_after_seconds: 300,
..SAVED
},
Feedback {
author: "bob",
score: 10,
comment: "",
..SAVED
},
],
)
.await?;
let reference = trace_ref(&database, "team-a", "key-a", "trace-1").await?;
let rows = feedback(&database, "", "trace-1", &reference).await?;
let shown: Vec<(&str, u64, &str)> = rows
.iter()
.map(|row| {
(
row["author"].as_str().unwrap_or_default(),
row["score"].as_u64().unwrap_or(99),
row["comment"].as_str().unwrap_or_default(),
)
})
.collect();
assert_eq!(shown, [("alice", 8, "fine after retry"), ("bob", 10, "")]);
assert_eq!(timestamp(&rows[0], "created_at")?, created);
assert_eq!(
timestamp(&rows[0], "updated_at")?,
created + Duration::seconds(300)
);
Ok(())
}
#[rstest]
#[tokio::test]
async fn a_deleted_version_hides_the_authors_feedback(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = ready(&database).await?;
let created = now()?;
save(
&database,
&writer,
created,
&[
SAVED,
Feedback {
author: "bob",
score: 6,
..SAVED
},
],
)
.await?;
save(
&database,
&writer,
created,
&[Feedback {
deleted: true,
edited_after_seconds: 540,
..SAVED
}],
)
.await?;
let reference = trace_ref(&database, "team-a", "key-a", "trace-1").await?;
let authors: Vec<Value> = feedback(&database, "", "trace-1", &reference)
.await?
.into_iter()
.map(|row| row["author"].clone())
.collect();
assert_eq!(authors, [json!("bob")]);
Ok(())
}
#[rstest]
#[tokio::test]
async fn summaries_aggregate_live_feedback_per_trace_and_respect_team_scope(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = ready(&database).await?;
save(
&database,
&writer,
now()?,
&[
Feedback { score: 2, ..SAVED },
Feedback {
author: "bob",
score: 7,
..SAVED
},
Feedback {
author: "carol",
score: 0,
deleted: true,
..SAVED
},
Feedback {
team: "team-b",
key: "key-b",
trace: "trace-2",
score: 9,
..SAVED
},
],
)
.await?;
let summary = |team: &str| {
let mut parameters = access(team);
parameters.insert(
"trace_ids".into(),
Parameter::Strings(vec![
"trace-1".into(),
"trace-2".into(),
"trace-unrated".into(),
]),
);
parameters
};
let everyone = read(&database, ReadQuery::FeedbackSummary, summary("")).await?;
let team_a = read(&database, ReadQuery::FeedbackSummary, summary("team-a")).await?;
let by_trace: BTreeMap<&str, (u64, f64, u64)> = everyone
.iter()
.map(|row| {
(
row["trace_id"].as_str().unwrap_or_default(),
(
row["count"].as_u64().unwrap_or(0),
row["average"].as_f64().unwrap_or(-1.0),
row["lowest"].as_u64().unwrap_or(99),
),
)
})
.collect();
assert_eq!(
by_trace,
BTreeMap::from([("trace-1", (2, 4.5, 2)), ("trace-2", (1, 9.0, 9))])
);
assert_eq!(
team_a
.iter()
.map(|row| row["trace_id"].clone())
.collect::<Vec<_>>(),
[json!("trace-1")]
);
Ok(())
}
#[rstest]
#[tokio::test]
async fn feedback_for_another_teams_trace_is_not_readable(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = ready(&database).await?;
save(&database, &writer, now()?, &[SAVED]).await?;
let reference = trace_ref(&database, "team-a", "key-a", "trace-1").await?;
assert!(
feedback(&database, "team-b", "trace-1", &reference)
.await?
.is_empty()
);
assert_eq!(
feedback(&database, "team-a", "trace-1", &reference)
.await?
.len(),
1
);
Ok(())
}
#[rstest]
#[tokio::test]
async fn the_table_rejects_scores_above_ten(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = ready(&database).await?;
let rejected = save(
&database,
&writer,
now()?,
&[Feedback { score: 11, ..SAVED }],
)
.await;
assert!(rejected.is_err());
Ok(())
}
#[rstest]
#[tokio::test]
async fn target_resolves_team_and_key_for_a_visible_trace_only(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
let writer = ready(&database).await?;
let span: BTreeMap<String, Value> = serde_json::from_value(json!({
"Timestamp": now()?.unix_timestamp_nanos() as i64, "TraceId": "trace-1", "SpanId": "root",
"ParentSpanId": "", "ServiceName": "agent", "SpanName": "run",
"ResourceAttributes": {"litellm.team_id": "team-a", "litellm.api_key_hash": "key-a"}
}))?;
insert_rows(
&database.client,
&writer,
"trace_test",
InsertTable::OtelTraces,
vec![span],
)
.await?;
let target = |team: &str| {
let mut parameters = access(team);
parameters.insert("trace_id".into(), Parameter::Text("trace-1".into()));
parameters.insert("trace_ref".into(), Parameter::Text(String::new()));
parameters
};
let visible = read(&database, ReadQuery::FeedbackTarget, target("team-a")).await?;
let hidden = read(&database, ReadQuery::FeedbackTarget, target("team-b")).await?;
assert_eq!(visible.len(), 1);
assert_eq!(
(
visible[0]["team_id"].as_str(),
visible[0]["key_hash"].as_str()
),
(Some("team-a"), Some("key-a"))
);
assert_eq!(
visible[0]["trace_ref"].as_str().map(str::to_owned),
Some(trace_ref(&database, "team-a", "key-a", "trace-1").await?)
);
assert!(hidden.is_empty());
Ok(())
}

View file

@ -335,10 +335,10 @@ async fn concurrent_schema_setup_succeeds(
&database,
"SELECT count() AS tables FROM system.tables \
WHERE database = 'trace_test' AND name IN \
('otel_traces', 'agent_traces_by_key', 'spend_logs')",
('otel_traces', 'agent_traces_by_key', 'spend_logs', 'lens_feedback')",
)
.await?;
assert_eq!(tables["data"][0]["tables"].as_u64(), Some(3));
assert_eq!(tables["data"][0]["tables"].as_u64(), Some(4));
assert_eq!(
migration_ledger_versions(&database).await?,
migration_versions()
@ -1040,6 +1040,7 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
tables["data"],
serde_json::json!([
{"name": "agent_traces_by_key"},
{"name": "lens_feedback"},
{"name": "otel_traces"},
{"name": "spend_logs"}
])
@ -1057,7 +1058,13 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
"start_time": old_timestamp_ms, "end_time": old_timestamp_ms + 1000
}))?;
insert_rows(&database, "otel_traces", vec![span]).await?;
let old_iso = old_time.format(&time::format_description::well_known::Rfc3339)?;
let feedback = serde_json::from_value(serde_json::json!({
"TeamId": "team-1", "ApiKeyHash": "", "TraceId": "expired", "Author": "admin",
"Score": 4, "Comment": "", "CreatedAt": old_iso, "UpdatedAt": old_iso, "IsDeleted": 0
}))?;
insert_rows(&database, "spend_logs", vec![spend]).await?;
insert_rows(&database, "lens_feedback", vec![feedback]).await?;
assert_eq!(table_rows(&database, "agent_traces_by_key").await?, 1);
ensure_schema(&database.client, &writer, "trace_test", 14).await?;
let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
@ -1087,6 +1094,8 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
)
.await?;
execute_write(&database, "OPTIMIZE TABLE trace_test.spend_logs FINAL").await?;
execute_write(&database, "OPTIMIZE TABLE trace_test.lens_feedback FINAL").await?;
assert_eq!(table_rows(&database, "lens_feedback").await?, 0);
assert_eq!(table_rows(&database, "otel_traces").await?, 0);
assert_eq!(table_rows(&database, "agent_traces_by_key").await?, 0);
assert_eq!(table_rows(&database, "spend_logs").await?, 0);
@ -1109,7 +1118,7 @@ async fn retention_reconciliation_updates_each_table_ttl(
&database,
"SELECT name, create_table_query FROM system.tables \
WHERE database = 'trace_test' AND name IN \
('otel_traces', 'agent_traces_by_key', 'spend_logs') ORDER BY name",
('otel_traces', 'agent_traces_by_key', 'spend_logs', 'lens_feedback') ORDER BY name",
)
.await?;
let ttl_queries = ttl_queries["data"].as_array().expect("retention tables");
@ -1118,7 +1127,12 @@ async fn retention_reconciliation_updates_each_table_ttl(
.iter()
.map(|row| row["name"].as_str().expect("table name"))
.collect::<Vec<_>>(),
["agent_traces_by_key", "otel_traces", "spend_logs"]
[
"agent_traces_by_key",
"lens_feedback",
"otel_traces",
"spend_logs"
]
);
for row in ttl_queries {
let query = row["create_table_query"]
@ -1135,7 +1149,7 @@ async fn retention_reconciliation_updates_each_table_ttl(
&database,
"SELECT name, create_table_query FROM system.tables \
WHERE database = 'trace_test' AND name IN \
('otel_traces', 'agent_traces_by_key', 'spend_logs') ORDER BY name",
('otel_traces', 'agent_traces_by_key', 'spend_logs', 'lens_feedback') ORDER BY name",
)
.await?;
for row in ttl_queries["data"].as_array().expect("retention tables") {

View file

@ -13,6 +13,7 @@ pub const DATABASE: &str = "trace_test";
pub struct SeededDatabase {
pub database: ClickHouseDatabase,
#[allow(dead_code)] // dead_code: also shared with reads.rs, which uses a direct storage reader
pub readers: QueryReaders,
}

View file

@ -3,9 +3,7 @@ use std::collections::BTreeMap;
use litellm_http::Client;
use litellm_traces::query::named::ReadAccessParams;
use litellm_traces_cache::{ReadError, TraceReader};
use litellm_traces_clickhouse::{
ClickHouseTraces, Connection, InsertTable, QueryScope, insert_rows,
};
use litellm_traces_clickhouse::{ClickHouseTraces, Connection, InsertTable, insert_rows};
use rstest::rstest;
use serde_json::json;
@ -83,10 +81,7 @@ async fn list_costs_match_each_run_when_response_ids_are_reused(
.collect(),
)
.await?;
let connection = fixture
.readers
.connection(client, &QueryScope::All, "fixture-secret")
.await?;
let connection = Connection::reader(&fixture.database.url, DATABASE)?;
let (reader, store) = make_reader(client, connection);
let access = ReadAccessParams {
all_teams: false,
@ -212,10 +207,7 @@ async fn large_runs_remain_complete_under_default_reader_limits(
.collect::<Vec<_>>();
insert_rows(client, &writer, DATABASE, InsertTable::SpendLogs, costs).await?;
}
let connection = fixture
.readers
.connection(client, &QueryScope::All, "fixture-secret")
.await?;
let connection = Connection::reader(&fixture.database.url, DATABASE)?;
let (reader, store) = make_reader(client, connection);
let access = ReadAccessParams {
all_teams: false,
@ -365,10 +357,7 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive(
) -> TestResult {
let fixture = seeded_database?;
let client = &fixture.database.client;
let connection = fixture
.readers
.connection(client, &QueryScope::All, "fixture-secret")
.await?;
let connection = Connection::reader(&fixture.database.url, DATABASE)?;
let (reader, store) = make_reader(client, connection.clone());
let access = ReadAccessParams {
all_teams: true,
@ -528,10 +517,7 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals(
) -> TestResult {
let fixture = seeded_database?;
let client = &fixture.database.client;
let connection = fixture
.readers
.connection(client, &QueryScope::All, "fixture-secret")
.await?;
let connection = Connection::reader(&fixture.database.url, DATABASE)?;
let (reader, store) = make_reader(client, connection.clone());
let access = ReadAccessParams {
all_teams: true,
@ -557,10 +543,7 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals(
("TraceId".into(), json!(run.trace_id)),
("SpanId".into(), json!("oversized-child")),
("ParentSpanId".into(), json!("0101010101010101")),
(
"SpanName".into(),
json!("x".repeat(litellm_storage_clickhouse::READ_LIMITS.response_bytes + 1)),
),
("SpanName".into(), json!("x".repeat(16 * 1024 * 1024 + 1))),
("ObservationType".into(), json!("tool")),
("TeamId".into(), json!("team-a")),
("ApiKeyHash".into(), json!("key-a")),
@ -699,10 +682,7 @@ async fn assigned_call_ids_require_shared_ownership_through_detail_and_batch_rea
.collect(),
)
.await?;
let connection = fixture
.readers
.connection(client, &QueryScope::All, "fixture-secret")
.await?;
let connection = Connection::reader(&fixture.database.url, DATABASE)?;
let (reader, store) = make_reader(client, connection);
let access = ReadAccessParams {
all_teams: false,
@ -815,10 +795,7 @@ async fn native_cost_correlation_survives_session_grouping_and_excludes_other_ow
.collect(),
)
.await?;
let connection = fixture
.readers
.connection(client, &QueryScope::All, "fixture-secret")
.await?;
let connection = Connection::reader(&fixture.database.url, DATABASE)?;
let (reader, store) = make_reader(client, connection);
let access = ReadAccessParams {
all_teams: true,

View file

@ -17,6 +17,9 @@ pub enum ReadQuery {
Sample,
Content,
Evidence,
FeedbackTarget,
Feedback,
FeedbackSummary,
}
impl ReadQuery {

View file

@ -116,6 +116,8 @@ pub struct TraceSpansRow {
pub source_url: String,
#[serde(default)]
pub source_title: String,
#[serde(default)]
pub source_user: String,
pub team_id: String,
pub api_key_hash: String,
pub user_id: String,

View file

@ -146,6 +146,7 @@ fn source(row: &TraceSpansRow) -> Option<RunSource> {
.unwrap_or(RunSourceType::Custom),
url: row.source_url.clone(),
title: row.source_title.clone(),
user: row.source_user.clone(),
})
}

View file

@ -88,6 +88,10 @@ pub struct RunSource {
pub kind: RunSourceType,
pub url: String,
pub title: String,
/// Who started the conversation, e.g. the Slack user's email.
#[serde(default, skip_serializing_if = "String::is_empty")]
#[cfg_attr(feature = "schema", schemars(extend("x-python-optional" = true)))]
pub user: String,
}
#[macro_rules_attribute::apply(response_type)]

View file

@ -210,6 +210,7 @@ fn trace_span(span: DecodedSpan) -> TraceSpansRow {
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
source_user: String::new(),
team_id: "fixture-team".into(),
api_key_hash: "fixture-key".into(),
user_id: "fixture-user".into(),
@ -297,6 +298,7 @@ fn unrelated_transport(call: &TraceSpansRow) -> TraceSpansRow {
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
source_user: String::new(),
team_id: call.team_id.clone(),
api_key_hash: call.api_key_hash.clone(),
user_id: call.user_id.clone(),

View file

@ -1650,6 +1650,9 @@ fn session_capture_joins_native_logs_and_traces_across_turns_without_changing_sp
let first = litellm_traces::decode_otlp_logs(&logs.encode_to_vec(), None).unwrap();
let second = decode_otlp(&request.encode_to_vec(), None).unwrap();
assert_eq!(first[0].trace_id, second[0].trace_id);
// Lens feedback resolves session ids the same way; keep in sync with
// litellm/proxy/lens/feedback_repository.py::session_trace_id.
assert_eq!(second[0].trace_id, "5fddf060372c8501dca4f331b9da882b");
assert_eq!(
first[0].attributes["lens.original_trace_id"],
"01".repeat(16)

View file

@ -14,6 +14,9 @@ use rstest::rstest;
#[case::sample("sample", ReadQuery::Sample)]
#[case::content("content", ReadQuery::Content)]
#[case::evidence("evidence", ReadQuery::Evidence)]
#[case::feedback_target("feedback_target", ReadQuery::FeedbackTarget)]
#[case::feedback("feedback", ReadQuery::Feedback)]
#[case::feedback_summary("feedback_summary", ReadQuery::FeedbackSummary)]
fn names_select_the_public_query(#[case] name: &str, #[case] query: ReadQuery) {
assert_eq!(ReadQuery::parse(name).unwrap(), query);
assert_eq!(query.as_ref(), name);

View file

@ -53,7 +53,7 @@ fn result_contracts_preserve_public_field_names() {
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["framework"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}),
);
round_trip::<TraceSpansRow>(
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "framework", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "source_type": "slack", "source_url": "https://acme.slack.com/archives/C1/p1", "source_title": "thread", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
json!({"trace_id": "trace", "original_trace_id": "original", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "framework", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "source_type": "slack", "source_url": "https://acme.slack.com/archives/C1/p1", "source_title": "thread", "source_user": "tin@berri.ai", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
);
round_trip::<SpanDetailRow>(
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),

View file

@ -36,6 +36,7 @@ fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> Trac
source_type: String::new(),
source_url: String::new(),
source_title: String::new(),
source_user: String::new(),
team_id: "team".into(),
api_key_hash: "key".into(),
user_id: String::new(),
@ -227,6 +228,16 @@ fn summary_source_type_picks_the_app(#[case] source_type: &str, #[case] expected
assert_eq!(source.map(|source| source.kind), Some(expected));
}
#[rstest]
#[case::set("tin@berri.ai")]
#[case::missing("")]
fn summary_source_carries_who_started_it(#[case] user: &str) {
let mut root = sourced(row("root", "", "agent", "agent", "agent"), THREAD, "t");
root.source_user = user.into();
let source = resolve_trace("t", "", &[root], &[]).unwrap().summary.source;
assert_eq!(source.map(|source| source.user), Some(user.to_owned()));
}
#[rstest]
fn spans_are_offset_from_the_trace_start() {
let trace = resolve_trace("t1", "", &deep_agent(1), &[]).unwrap();

View file

@ -1667,6 +1667,24 @@ if TYPE_CHECKING:
from .llms.hosted_vllm.rerank.transformation import (
HostedVLLMRerankConfig as HostedVLLMRerankConfig,
)
from .llms.perplexity.decisions.transformation import (
PerplexityDecisionsConfig as PerplexityDecisionsConfig,
)
from .llms.typesafe.decisions.transformation import (
TypeSafeDecisionsConfig as TypeSafeDecisionsConfig,
)
from .llms.openrouter.decisions.transformation import (
OpenRouterDecisionsConfig as OpenRouterDecisionsConfig,
)
from .llms.cloudflare.decisions.transformation import (
CloudflareDecisionsConfig as CloudflareDecisionsConfig,
)
from .llms.strands_decider.decisions.transformation import (
StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig,
)
from .llms.openai.decisions.transformation import (
OpenAIDecisionsConfig as OpenAIDecisionsConfig,
)
from .llms.nvidia_nim.rerank.transformation import (
NvidiaNimRerankConfig as NvidiaNimRerankConfig,
)

View file

@ -156,6 +156,12 @@ LLM_CONFIG_NAMES: Final = (
"ScalewayRerankConfig",
"DeepinfraRerankConfig",
"HostedVLLMRerankConfig",
"PerplexityDecisionsConfig",
"TypeSafeDecisionsConfig",
"OpenRouterDecisionsConfig",
"CloudflareDecisionsConfig",
"StrandsDeciderDecisionsConfig",
"OpenAIDecisionsConfig",
"NvidiaNimRerankConfig",
"NvidiaNimRankingConfig",
"VertexAIRerankConfig",
@ -707,6 +713,15 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.hosted_vllm.rerank.transformation",
"HostedVLLMRerankConfig",
),
"PerplexityDecisionsConfig": (".llms.perplexity.decisions.transformation", "PerplexityDecisionsConfig"),
"TypeSafeDecisionsConfig": (".llms.typesafe.decisions.transformation", "TypeSafeDecisionsConfig"),
"OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"),
"CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"),
"StrandsDeciderDecisionsConfig": (
".llms.strands_decider.decisions.transformation",
"StrandsDeciderDecisionsConfig",
),
"OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"),
"NvidiaNimRerankConfig": (
".llms.nvidia_nim.rerank.transformation",
"NvidiaNimRerankConfig",

View file

@ -416,11 +416,11 @@ def _provider_output_file_id(output_file_id: str) -> str:
llm_output_file_id, model-encoded ids decode to the raw provider id, raw ids pass through.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
get_original_file_id,
is_base64_encoded_unified_file_id,
)
unified_file_id: Final = _is_base64_encoded_unified_file_id(output_file_id)
unified_file_id: Final = is_base64_encoded_unified_file_id(output_file_id)
if not unified_file_id:
return get_original_file_id(output_file_id)
try:

View file

@ -19,7 +19,7 @@ from openai.types.responses.response_input_param import (
from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam
from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm import ModelResponse
@ -46,6 +46,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolParamFunctionChunk,
PromptCacheBreakpoint,
Reasoning,
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
@ -238,6 +239,19 @@ def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Li
return "length"
_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(PromptCacheBreakpoint)
def _prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object:
if marker is None or not drop_params:
return marker
try:
return _PROMPT_CACHE_BREAKPOINT.validate_python(marker)
except ValidationError:
verbose_logger.debug("Chat provider: dropping malformed prompt_cache_breakpoint %r under drop_params", marker)
return None
def _input_file_from_file_value(file_value: object) -> dict[str, object]:
if not isinstance(file_value, dict):
return {"type": "input_file"}
@ -400,9 +414,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
self,
messages: list["AllMessageValues"],
*,
drop_params: bool = False,
keep_prompt_cache_breakpoints: bool = False,
) -> tuple[list[object], str | None]:
converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(messages)
converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(
messages, drop_params=drop_params
)
return (
converted_input_items
if keep_prompt_cache_breakpoints
@ -411,7 +428,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
def _convert_chat_completion_messages_to_responses_input(
self, messages: list["AllMessageValues"]
self, messages: list["AllMessageValues"], *, drop_params: bool = False
) -> tuple[list[object], str | None]:
input_items: Final[list[object]] = []
instructions: str | None = None
@ -452,6 +469,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
"content": self._convert_content_to_responses_format(
content,
role,
drop_params=drop_params,
),
}
)
@ -470,6 +488,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
tool_output = self._convert_content_to_responses_format(
content,
"user", # Use "user" role to get input_* types
drop_params=drop_params,
)
else:
# Fallback: convert unexpected types to input_text
@ -497,7 +516,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
{
"type": "message",
"role": "assistant",
"content": self._convert_content_to_responses_format(content, "assistant"),
"content": self._convert_content_to_responses_format(
content, "assistant", drop_params=drop_params
),
}
)
for tool_call in tool_calls:
@ -531,7 +552,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
{
"type": "message",
"role": role,
"content": self._convert_content_to_responses_format(content, cast(str, role)),
"content": self._convert_content_to_responses_format(
content, cast(str, role), drop_params=drop_params
),
}
)
elif role == "assistant":
@ -647,6 +670,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
converted_input_items, converted_instructions = self.convert_chat_completion_messages_to_responses_api(
messages,
drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint,
)
# OpenAI's Responses API rejects an empty input. For a system-only
@ -1126,6 +1150,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
]
| None,
role: str,
drop_params: bool = False,
) -> list[dict[str, object]]:
"""Convert chat completion content to responses API format"""
from litellm.types.llms.openai import ChatCompletionImageObject
@ -1152,7 +1177,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if original_type == "text":
converted = with_prompt_cache_breakpoint(
self._convert_content_str_to_input_text(item.get("text", ""), role),
item.get("prompt_cache_breakpoint"),
_prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
)
result.append(converted)
verbose_logger.debug("Chat provider: text -> %s", converted)
@ -1165,7 +1190,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
cast(ChatCompletionImageObject, item), role
),
),
item.get("prompt_cache_breakpoint"),
_prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
)
result.append(converted)
verbose_logger.debug("Chat provider: image_url -> %s", converted)
@ -1181,7 +1206,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
_input_file_from_file_value(
cast("ChatCompletionFileObject", item).get("file"), # cast-ok: type tag checked
),
item.get("prompt_cache_breakpoint"),
_prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
)
result.append(converted)
verbose_logger.debug("Chat provider: file -> %s", converted)
@ -1203,7 +1228,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
verbose_logger.debug("Chat provider: passthrough -> %s", item)
else:
# Default to input_text for unknown types
converted = self._convert_content_str_to_input_text(str(item.get("text", item)), role)
converted = with_prompt_cache_breakpoint(
self._convert_content_str_to_input_text(str(item.get("text", item)), role),
_prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params),
)
result.append(converted)
verbose_logger.debug("Chat provider: unknown(%s) -> %s", original_type, converted)
verbose_logger.debug("Chat provider: Final converted content: %s", result)

View file

@ -64,6 +64,8 @@ AGENT_TRACING_AGENT_LIST_LIMIT: Final = get_env_int("AGENT_TRACING_AGENT_LIST_LI
LENS_DATASET_MAX_CASES: Final = get_env_int("LENS_DATASET_MAX_CASES", 200)
LENS_DATASET_MAX_CASE_CHARS: Final = get_env_int("LENS_DATASET_MAX_CASE_CHARS", 20_000)
LENS_DATASET_TRACE_PAGE_SIZE: Final = 500
LENS_FEEDBACK_MAX_COMMENT_CHARS: Final = 10_000
LENS_FEEDBACK_MAX_SCORE: Final = 10
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))
DEFAULT_S3_MAX_CONCURRENT_UPLOADS: Final = int(os.getenv("DEFAULT_S3_MAX_CONCURRENT_UPLOADS", "16"))

View file

@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast
from httpx import Response
from pydantic import BaseModel
from typing_extensions import ReadOnly, TypedDict
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
import litellm._logging
@ -100,7 +100,7 @@ from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_ro
from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.decisions import DecisionsResponse, DecisionsUsage
from litellm.types.decisions import DecisionsResponse, DecisionsUsage, OpenAIDecisionResponse, OpenAIDecisionUsage
from litellm.types.llms.base import CachedTokensDetails, LiteLLMBaseModel
from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
@ -1034,6 +1034,8 @@ def get_usage_object(
),
)
if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
return None if completion_response.usage is None else _decisions_usage(completion_response.usage)
if usage_obj is None:
return None
if isinstance(usage_obj, Usage):
@ -1064,12 +1066,37 @@ def get_usage_object(
return None
def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper:
match usage:
case DecisionsUsage():
return PromptTokensDetailsWrapper(
cached_tokens=usage.cached_tokens, cache_write_tokens=usage.cache_write_tokens
)
case OpenAIDecisionUsage():
return PromptTokensDetailsWrapper(
cached_tokens=usage.input_tokens_details.cached_tokens,
cache_write_tokens=usage.input_tokens_details.cache_write_tokens,
)
case _:
assert_never(usage)
def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage:
return Usage(
prompt_tokens=usage.input_tokens,
completion_tokens=usage.output_tokens,
total_tokens=usage.input_tokens + usage.output_tokens,
prompt_tokens_details=_decisions_prompt_tokens_details(usage),
)
def _is_known_usage_objects(usage_obj):
"""Returns True if the usage obj is a known Usage type"""
return (
isinstance(usage_obj, litellm.Usage)
or isinstance(usage_obj, ResponseAPIUsage)
or isinstance(usage_obj, DecisionsUsage)
or isinstance(usage_obj, OpenAIDecisionUsage)
or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj)
)
@ -1481,11 +1508,8 @@ def completion_cost(
"usage",
litellm.Usage(**_usage_for_dump.model_dump()),
)
if isinstance(usage_obj, DecisionsUsage):
_usage = {
"prompt_tokens": usage_obj.input_tokens,
"completion_tokens": usage_obj.output_tokens,
}
if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)):
_usage = _decisions_usage(usage_obj).model_dump()
elif usage_obj is None:
_usage = {}
elif isinstance(usage_obj, BaseModel):
@ -1977,7 +2001,8 @@ def response_cost_calculator(
| OpenAIModerationResponse
| Response
| SearchResponse
| DecisionsResponse,
| DecisionsResponse
| OpenAIDecisionResponse,
model: str,
custom_llm_provider: str | None,
call_type: Literal[

View file

@ -1,114 +1,147 @@
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import Final
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.decisions.transformation import DecisionsProviderConfig
from litellm.llms.cloudflare.decisions.transformation import CLOUDFLARE_DECISIONS_ENDPOINT
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client, get_httpx_client
from litellm.llms.openrouter.decisions.transformation import OPENROUTER_DECISIONS_ENDPOINT
from litellm.llms.perplexity.decisions.transformation import PERPLEXITY_DECISIONS_ENDPOINT
from litellm.llms.strands_decider.decisions.transformation import STRANDS_DECIDER_DECISIONS_ENDPOINT
from litellm.llms.typesafe.decisions.transformation import TYPESAFE_DECISIONS_ENDPOINT
from litellm.secret_managers.main import get_secret_str
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.decisions.transformation import (
BaseDecisionsConfig,
ir_to_systemone_response,
systemone_request_to_ir,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.openai.decisions.transformation import ir_to_openai_response, openai_request_to_ir
from litellm.types.decisions import (
DecisionQuestion,
DecisionsIRRequest,
DecisionsIRResponse,
DecisionsJSON,
DecisionsRequest,
DecisionsRequestBody,
DecisionsResponse,
OpenAIDecisionInput,
OpenAIDecisionQuestion,
OpenAIDecisionRequestBody,
OpenAIDecisionResponse,
UnsupportedDecisionsRequest,
)
from litellm.utils import client
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager, client
DECISIONS_ENDPOINTS: Final[Mapping[str, DecisionsProviderConfig]] = MappingProxyType(
{
"perplexity": PERPLEXITY_DECISIONS_ENDPOINT,
"typesafe": TYPESAFE_DECISIONS_ENDPOINT,
"openrouter": OPENROUTER_DECISIONS_ENDPOINT,
"cloudflare": CLOUDFLARE_DECISIONS_ENDPOINT,
"strands_decider": STRANDS_DECIDER_DECISIONS_ENDPOINT,
}
DecisionsQuestions: TypeAlias = (
Mapping[str, DecisionQuestion | Mapping[str, object]] | Sequence[OpenAIDecisionQuestion | Mapping[str, object]]
)
DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequestBody
_DECISIONS_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequest]] = TypeAdapter(DecisionsRequest)
_DECISIONS_PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object)
_DECISIONS_RESPONSE_ADAPTER: Final[TypeAdapter[DecisionsResponse]] = TypeAdapter(DecisionsResponse)
_SYSTEMONE_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody)
_OPENAI_REQUEST_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody)
_HANDLER: Final = BaseLLMHTTPHandler()
@dataclass(frozen=True, slots=True, repr=False)
class _PreparedDecisionsRequest:
config: DecisionsProviderConfig
provider: str
upstream_model: str
url: str
api_key: str | None = field(repr=False)
headers: Mapping[str, str] = field(repr=False)
class _DecisionsCall:
model: str
requested_model: str
custom_llm_provider: str
provider_config: BaseDecisionsConfig
request: DecisionsRequestFormat
ir_request: DecisionsIRRequest
body: Mapping[str, object] = field(repr=False)
api_base: str
api_key: str | None = field(repr=False)
logging_obj: LiteLLMLoggingObj | None
headers: Mapping[str, str]
timeout: float | httpx.Timeout | None
def _resolve_provider_model(model: str, custom_llm_provider: str | None) -> tuple[str, str]:
provider: Final = model.partition("/")[0] if custom_llm_provider is None else custom_llm_provider
if provider not in DECISIONS_ENDPOINTS:
supported: Final = ", ".join(DECISIONS_ENDPOINTS)
def _supported_providers() -> tuple[str, ...]:
return tuple(
provider.value
for provider in LlmProviders
if ProviderConfigManager.get_provider_decisions_config(model="", provider=provider) is not None
)
def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig:
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
provider_config: Final = (
None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider)
)
if provider_config is None:
supported: Final = ", ".join(_supported_providers())
raise litellm.BadRequestError(
message=f"Unknown Decisions provider '{provider}'. Supported providers: {supported}",
message=f"Unknown Decisions provider '{custom_llm_provider}'. Supported providers: {supported}",
model=model,
llm_provider=provider,
llm_provider=custom_llm_provider,
)
upstream_model: Final = model.removeprefix(f"{provider}/")
return provider_config
def _validate_request(
*,
state: DecisionsJSON | None,
questions: DecisionsQuestions | None,
decision_input: OpenAIDecisionInput | None,
safety_identifier: str | None,
) -> DecisionsRequestFormat:
if decision_input is None:
return _SYSTEMONE_REQUEST_ADAPTER.validate_python({"state": state, "questions": questions})
return _OPENAI_REQUEST_ADAPTER.validate_python(
{"input": decision_input, "questions": questions, "safety_identifier": safety_identifier}
)
def _ir_request(request: DecisionsRequestFormat) -> DecisionsIRRequest:
match request:
case DecisionsRequestBody():
return systemone_request_to_ir(request)
case OpenAIDecisionRequestBody():
return openai_request_to_ir(request)
case _:
assert_never(request)
def _prepare_call(
*,
model: str,
state: DecisionsJSON | None,
questions: DecisionsQuestions | None,
decision_input: OpenAIDecisionInput | None,
safety_identifier: str | None,
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> _DecisionsCall:
upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
provider_config: Final = _provider_config(upstream_model, provider)
canonical_model: Final = provider_config.canonical_model(upstream_model)
if not upstream_model:
raise litellm.BadRequestError(
message="A model name is required for the Decisions API",
model=model,
llm_provider=provider,
)
return provider, upstream_model
def _resolve_api_key(
*,
provider: str,
model: str,
endpoint: DecisionsProviderConfig,
api_key: str | None,
) -> str | None:
if api_key is not None:
return api_key
server_api_key: Final = next(
(key for key in (get_secret_str(name) for name in endpoint.api_key_env) if key),
None,
)
if server_api_key is None:
if not endpoint.api_key_required:
return None
raise litellm.AuthenticationError(
message=f"Missing API key for Decisions provider '{provider}'",
if state is not None and decision_input is not None:
raise litellm.BadRequestError(
message="Pass either state (System One format) or input (OpenAI format) to the Decisions API, not both",
model=model,
llm_provider=provider,
)
return server_api_key
def _prepare_request(
*,
model: str,
state: DecisionsJSON,
questions: Mapping[str, DecisionQuestion | Mapping[str, object]],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, str] | None,
) -> _PreparedDecisionsRequest:
provider, upstream_model = _resolve_provider_model(model, custom_llm_provider)
try:
validated_request: Final = _DECISIONS_REQUEST_ADAPTER.validate_python(
{"model": model, "state": state, "questions": questions}
request: Final = _validate_request(
state=state, questions=questions, decision_input=decision_input, safety_identifier=safety_identifier
)
except ValidationError as error:
raise litellm.BadRequestError(
@ -117,109 +150,94 @@ def _prepare_request(
llm_provider=provider,
) from error
endpoint: Final = DECISIONS_ENDPOINTS[provider]
env_api_base: Final = get_secret_str(endpoint.api_base_env)
default_api_base: Final = endpoint.default_api_base()
resolved_api_base: Final = api_base or env_api_base or default_api_base
resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
if resolved_api_base is None:
raise litellm.BadRequestError(
message=endpoint.missing_api_base_message(provider),
message=provider_config.missing_api_base_message(provider),
model=model,
llm_provider=provider,
)
resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
if resolved_api_key is None and provider_config.api_key_required:
raise litellm.AuthenticationError(
message=f"Missing API key for Decisions provider '{provider}'",
model=model,
llm_provider=provider,
)
resolved_api_key: Final = _resolve_api_key(
provider=provider,
model=model,
endpoint=endpoint,
api_key=api_key,
ir_request: Final = _ir_request(request)
body: Final = provider_config.transform_decisions_request(
model=canonical_model, request=ir_request, custom_llm_provider=provider
)
if isinstance(body, UnsupportedDecisionsRequest):
raise litellm.BadRequestError(
message=f"Decisions provider '{provider}' cannot serve this request: {body.reason}",
model=model,
llm_provider=provider,
)
canonical_model: Final = endpoint.canonical_model(upstream_model)
outbound_headers: Final = MappingProxyType(
{
**{
name: value
for name, value in (extra_headers or {}).items()
if name.lower() not in {"authorization", "content-type"}
},
**({"Authorization": f"Bearer {resolved_api_key}"} if resolved_api_key is not None else {}),
"Content-Type": "application/json",
}
)
body: Final = MappingProxyType(
{
"model": endpoint.request_model(canonical_model),
"state": validated_request.state,
"questions": {
name: question.model_dump(mode="json", exclude_none=True)
for name, question in validated_request.questions.items()
},
}
)
return _PreparedDecisionsRequest(
config=endpoint,
provider=provider,
upstream_model=canonical_model,
url=endpoint.endpoint_url(resolved_api_base, canonical_model),
api_key=resolved_api_key,
headers=outbound_headers,
body=body,
)
def _log_request(
prepared: _PreparedDecisionsRequest,
kwargs: Mapping[str, object],
) -> LiteLLMLoggingObj | None:
logging_obj: Final = kwargs.get("litellm_logging_obj")
if not isinstance(logging_obj, LiteLLMLoggingObj):
return None
logging_obj.update_from_kwargs(
kwargs=dict(kwargs),
model=prepared.upstream_model,
litellm_params={
"litellm_call_id": kwargs.get("litellm_call_id"),
"api_base": prepared.url,
},
custom_llm_provider=prepared.provider,
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_from_kwargs(
kwargs=dict(kwargs),
model=canonical_model,
litellm_params={
"litellm_call_id": kwargs.get("litellm_call_id"),
"api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
},
custom_llm_provider=provider,
)
return _DecisionsCall(
model=canonical_model,
requested_model=model,
custom_llm_provider=provider,
provider_config=provider_config,
request=request,
ir_request=ir_request,
body=body,
api_base=resolved_api_base,
api_key=resolved_api_key,
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
headers=extra_headers or {},
timeout=timeout,
)
request_body: Final = dict(prepared.body)
request_headers: Final = dict(prepared.headers)
logging_obj.pre_call(
input=request_body,
api_key=prepared.api_key,
model=prepared.upstream_model,
additional_args={
"api_base": prepared.url,
"complete_input_dict": request_body,
"headers": request_headers,
},
)
return logging_obj
def _parse_response(
response: httpx.Response,
prepared: _PreparedDecisionsRequest,
) -> DecisionsResponse:
response.raise_for_status()
payload: Final[object] = _DECISIONS_PAYLOAD_ADAPTER.validate_json(response.content)
result: Final = _DECISIONS_RESPONSE_ADAPTER.validate_python(prepared.config.unwrap_response(payload))
result.hidden_params.update(
def _format_response(response: DecisionsIRResponse, call: _DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
formatted: Final = _formatted_response(response, call)
formatted.set_hidden_params(
{
"model": f"{prepared.provider}/{prepared.upstream_model}",
"custom_llm_provider": prepared.provider,
"provider_response_model": f"{prepared.provider}/{prepared.upstream_model}",
"model": f"{call.custom_llm_provider}/{call.model}",
"custom_llm_provider": call.custom_llm_provider,
"provider_response_model": f"{call.custom_llm_provider}/{call.model}",
}
)
return result
return formatted
def _map_upstream_exception(error: Exception, prepared: _PreparedDecisionsRequest) -> Exception:
def _formatted_response(
response: DecisionsIRResponse, call: _DecisionsCall
) -> DecisionsResponse | OpenAIDecisionResponse:
match call.request:
case DecisionsRequestBody():
return ir_to_systemone_response(response, call.ir_request)
case OpenAIDecisionRequestBody():
return ir_to_openai_response(response, call.ir_request, call.requested_model)
case _:
assert_never(call.request)
def _map_upstream_exception(error: Exception, call: _DecisionsCall) -> Exception:
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
provider_label: Final = f"{call.custom_llm_provider[0].upper()}{call.custom_llm_provider[1:]}Exception"
return litellm.APIConnectionError(
message=f"{provider_label} - {error.message}",
llm_provider=call.custom_llm_provider,
model=f"{call.custom_llm_provider}/{call.model}",
)
return litellm.exception_type(
model=f"{prepared.provider}/{prepared.upstream_model}",
custom_llm_provider=prepared.provider,
model=f"{call.custom_llm_provider}/{call.model}",
custom_llm_provider=call.custom_llm_provider,
original_exception=error,
)
@ -227,73 +245,91 @@ def _map_upstream_exception(error: Exception, prepared: _PreparedDecisionsReques
@client
async def adecisions(
model: str,
state: DecisionsJSON,
questions: Mapping[str, DecisionQuestion | Mapping[str, object]],
state: DecisionsJSON | None = None,
questions: DecisionsQuestions | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
input: OpenAIDecisionInput | None = None,
safety_identifier: str | None = None,
**kwargs: object,
) -> DecisionsResponse:
prepared: Final = _prepare_request(
) -> DecisionsResponse | OpenAIDecisionResponse:
call: Final = _prepare_call(
model=model,
state=state,
questions=questions,
decision_input=input,
safety_identifier=safety_identifier,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
logging_obj: Final = _log_request(prepared, kwargs)
try:
handler: Final = get_async_httpx_client(llm_provider=prepared.provider)
response: Final = await handler.post(
prepared.url,
json=dict(prepared.body),
headers=dict(prepared.headers),
timeout=timeout,
logging_obj=logging_obj,
response: Final = await _HANDLER.adecisions(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
request=call.ir_request,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
)
return _parse_response(response=response, prepared=prepared)
except Exception as error:
raise _map_upstream_exception(error, prepared) from error
raise _map_upstream_exception(error, call) from error
return _format_response(response, call)
@client
def decisions(
model: str,
state: DecisionsJSON,
questions: Mapping[str, DecisionQuestion | Mapping[str, object]],
state: DecisionsJSON | None = None,
questions: DecisionsQuestions | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
input: OpenAIDecisionInput | None = None,
safety_identifier: str | None = None,
**kwargs: object,
) -> DecisionsResponse:
prepared: Final = _prepare_request(
) -> DecisionsResponse | OpenAIDecisionResponse:
call: Final = _prepare_call(
model=model,
state=state,
questions=questions,
decision_input=input,
safety_identifier=safety_identifier,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
logging_obj: Final = _log_request(prepared, kwargs)
try:
handler: Final = get_httpx_client()
response: Final = handler.post(
prepared.url,
json=dict(prepared.body),
headers=dict(prepared.headers),
timeout=timeout,
logging_obj=logging_obj,
response: Final = _HANDLER.decisions(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
request=call.ir_request,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
)
return _parse_response(response=response, prepared=prepared)
except Exception as error:
raise _map_upstream_exception(error, prepared) from error
raise _map_upstream_exception(error, call) from error
return _format_response(response, call)
__all__ = ["DECISIONS_ENDPOINTS", "adecisions", "decisions"]
__all__ = ["adecisions", "decisions"]

View file

@ -90,7 +90,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
end_time: Final = datetime.now()
asyncio.create_task(
PassThroughStreamingHandler._route_streaming_logging_to_handler(
PassThroughStreamingHandler.route_streaming_logging_to_handler(
litellm_logging_obj=self.litellm_logging_obj,
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
url_route="/v1/generateContent",

View file

@ -1845,7 +1845,7 @@ Model Info:
try:
from litellm.proxy.spend_tracking.spend_management_endpoints import (
_get_spend_report_for_time_range,
get_spend_report_for_time_range,
)
# Parse the time range
@ -1862,7 +1862,7 @@ Model Info:
if await self.internal_usage_cache.async_get_cache(key=_event_cache_key):
return
_resp: Final = await _get_spend_report_for_time_range(
_resp: Final = await get_spend_report_for_time_range(
start_date=start_date.strftime("%Y-%m-%d"),
end_date=todays_date.strftime("%Y-%m-%d"),
)
@ -1909,7 +1909,7 @@ Model Info:
from calendar import monthrange
from litellm.proxy.spend_tracking.spend_management_endpoints import (
_get_spend_report_for_time_range,
get_spend_report_for_time_range,
)
todays_date: Final = datetime.datetime.now().date()
@ -1921,7 +1921,7 @@ Model Info:
if await self.internal_usage_cache.async_get_cache(key=_event_cache_key):
return
_resp: Final = await _get_spend_report_for_time_range(
_resp: Final = await get_spend_report_for_time_range(
start_date=first_day_of_month.strftime("%Y-%m-%d"),
end_date=last_day_of_month.strftime("%Y-%m-%d"),
)

View file

@ -44,9 +44,9 @@ class GcsPubSubLogger(CustomBatchLogger):
topic_id (str): Pub/Sub topic ID
credentials_path (str, optional): Path to Google Cloud credentials JSON file
"""
from litellm.proxy.utils import _premium_user_check
from litellm.proxy.utils import premium_user_check
_premium_user_check()
premium_user_check()
self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
@ -107,9 +107,9 @@ class GcsPubSubLogger(CustomBatchLogger):
from litellm.proxy.spend_tracking.spend_tracking_utils import (
get_logging_payload,
)
from litellm.proxy.utils import _premium_user_check
from litellm.proxy.utils import premium_user_check
_premium_user_check()
premium_user_check()
try:
verbose_logger.debug("PubSub: Logging - Enters logging function for model %s", kwargs)

View file

@ -3833,7 +3833,7 @@ class PrometheusLogger(CustomLogger):
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy.management_endpoints.key_management_endpoints import (
_list_key_helper,
list_key_helper,
)
from litellm.proxy.proxy_server import prisma_client
@ -3847,7 +3847,7 @@ class PrometheusLogger(CustomLogger):
list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken],
int | None,
]:
key_list_response: Final = await _list_key_helper(
key_list_response: Final = await list_key_helper(
prisma_client=prisma_client,
page=page,
size=page_size,

View file

@ -633,9 +633,9 @@ async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
from litellm.exceptions import BudgetExceededError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
_team_max_budget_check,
_virtual_key_max_budget_check,
get_team_object,
team_max_budget_check,
virtual_key_max_budget_check,
)
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
except ImportError:
@ -645,7 +645,7 @@ async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
if not isinstance(auth, UserAPIKeyAuth):
return False
try:
await _virtual_key_max_budget_check(valid_token=auth, proxy_logging_obj=proxy_logging_obj)
await virtual_key_max_budget_check(valid_token=auth, proxy_logging_obj=proxy_logging_obj)
if auth.team_id:
team: Final = await get_team_object(
team_id=auth.team_id,
@ -653,7 +653,7 @@ async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
user_api_key_cache=user_api_key_cache,
check_cache_only=True,
)
await _team_max_budget_check(team_object=team, valid_token=auth, proxy_logging_obj=proxy_logging_obj)
await team_max_budget_check(team_object=team, valid_token=auth, proxy_logging_obj=proxy_logging_obj)
except BudgetExceededError:
return True
except Exception as e: # noqa: BLE001 # advisory gate: a failed read must not block sampling

View file

@ -53,8 +53,14 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i
VectorStorePreCallHook,
)
from litellm.integrations.zerobus import ZerobusLogger
from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import _PROXY_DynamicRateLimitHandlerV3
from litellm.proxy.hooks.dynamic_rate_limiter import ( # noqa: F401 # legacy module exports
PROXY_DynamicRateLimitHandler,
_PROXY_DynamicRateLimitHandler, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
)
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( # noqa: F401 # legacy module exports
PROXY_DynamicRateLimitHandlerV3,
_PROXY_DynamicRateLimitHandlerV3, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
)
class CustomLoggerRegistry:
@ -101,8 +107,8 @@ class CustomLoggerRegistry:
"pointfive": PointFiveLogger,
"zerobus": ZerobusLogger,
"aws_sqs": SQSLogger,
"dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler,
"dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3,
"dynamic_rate_limiter": PROXY_DynamicRateLimitHandler,
"dynamic_rate_limiter_v3": PROXY_DynamicRateLimitHandlerV3,
"vector_store_pre_call_hook": VectorStorePreCallHook,
"dotprompt": DotpromptManager,
"bitbucket": BitBucketPromptManager,

View file

@ -2147,7 +2147,10 @@ def _map_openrouter_exception(
exception_provider: str,
extra_information: str,
) -> None:
if hasattr(original_exception, "status_code"):
received_status: Final = hasattr(original_exception, "status_code") and not getattr(
original_exception, "status_code_is_synthesized", False
)
if received_status:
if original_exception.status_code == 400:
raise BadRequestError(
message=f"{exception_provider} - {error_str}",

View file

@ -120,7 +120,7 @@ from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.containers.main import ContainerObject
from litellm.types.decisions import DecisionsResponse
from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
from litellm.types.interactions import (
InteractionsAPIResponse,
@ -2650,6 +2650,7 @@ class Logging(LiteLLMLoggingBaseClass):
or isinstance(logging_result, OCRResponse) # OCR
or isinstance(logging_result, SearchResponse) # Search API
or isinstance(logging_result, DecisionsResponse)
or isinstance(logging_result, OpenAIDecisionResponse)
or (
isinstance(logging_result, InteractionsAPIResponse)
and logging_result.usage is not None
@ -4969,17 +4970,17 @@ def _init_custom_logger_compatible_class(
return _otel_logger
elif logging_integration == "dynamic_rate_limiter":
from litellm.proxy.hooks.dynamic_rate_limiter import (
_PROXY_DynamicRateLimitHandler,
PROXY_DynamicRateLimitHandler,
)
for callback in _in_memory_loggers:
if isinstance(callback, _PROXY_DynamicRateLimitHandler):
if isinstance(callback, PROXY_DynamicRateLimitHandler):
return callback
if internal_usage_cache is None:
raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}")
dynamic_rate_limiter_obj: Final = _PROXY_DynamicRateLimitHandler(internal_usage_cache=internal_usage_cache)
dynamic_rate_limiter_obj: Final = PROXY_DynamicRateLimitHandler(internal_usage_cache=internal_usage_cache)
if llm_router is not None and isinstance(llm_router, litellm.Router):
dynamic_rate_limiter_obj.update_variables(llm_router=llm_router)
@ -4987,17 +4988,19 @@ def _init_custom_logger_compatible_class(
return dynamic_rate_limiter_obj
elif logging_integration == "dynamic_rate_limiter_v3":
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3,
PROXY_DynamicRateLimitHandlerV3,
)
for callback in _in_memory_loggers:
if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3):
if isinstance(callback, PROXY_DynamicRateLimitHandlerV3):
return callback
if internal_usage_cache is None:
raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}")
dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=internal_usage_cache)
dynamic_rate_limiter_obj_v3: Final = PROXY_DynamicRateLimitHandlerV3(
internal_usage_cache=internal_usage_cache
)
if llm_router is not None and isinstance(llm_router, litellm.Router):
dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router)
@ -5546,19 +5549,19 @@ def get_custom_logger_compatible_class(
elif logging_integration == "dynamic_rate_limiter":
from litellm.proxy.hooks.dynamic_rate_limiter import (
_PROXY_DynamicRateLimitHandler,
PROXY_DynamicRateLimitHandler,
)
for callback in _in_memory_loggers:
if isinstance(callback, _PROXY_DynamicRateLimitHandler):
if isinstance(callback, PROXY_DynamicRateLimitHandler):
return callback
elif logging_integration == "dynamic_rate_limiter_v3":
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3,
PROXY_DynamicRateLimitHandlerV3,
)
for callback in _in_memory_loggers:
if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3):
if isinstance(callback, PROXY_DynamicRateLimitHandlerV3):
return callback
elif logging_integration == "langtrace":

View file

@ -565,9 +565,9 @@ def update_messages_with_model_file_ids(
}
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
convert_b64_uid_to_unified_uid,
get_original_file_id,
is_base64_encoded_unified_file_id,
is_model_embedded_id,
)
@ -603,7 +603,7 @@ def update_messages_with_model_file_ids(
if model_file_id_mapping and model_id is not None
else None
)
if not provider_file_id and _is_base64_encoded_unified_file_id(file_id):
if not provider_file_id and is_base64_encoded_unified_file_id(file_id):
unified_file_id = convert_b64_uid_to_unified_uid(file_id)
if "llm_output_file_id," in unified_file_id:
provider_file_id = unified_file_id.split("llm_output_file_id,")[1].split(";")[0]
@ -634,9 +634,9 @@ def update_responses_input_with_model_file_ids(
Format: {"litellm_file_id": {"model_id": "provider_file_id"}}
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
convert_b64_uid_to_unified_uid,
get_original_file_id,
is_base64_encoded_unified_file_id,
is_model_embedded_id,
)
@ -671,7 +671,7 @@ def update_responses_input_with_model_file_ids(
updated_content.append(updated_content_item)
else:
# Check if this is a base64-encoded unified file ID without mapping
is_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
is_unified_file_id = is_base64_encoded_unified_file_id(file_id)
if is_unified_file_id:
# Fallback: decode unified file ID
unified_file_id = convert_b64_uid_to_unified_uid(file_id)

View file

@ -206,36 +206,108 @@ def _handle_ollama_system_message(messages: list, prompt: str, msg_i: int) -> tu
return system_content_str, msg_i
_OLLAMA_USER_ROLES: Final = frozenset({"user", "tool", "function"})
def _ollama_bad_message(model: str, message: AllMessageValues, msg_i: int, detail: str) -> "litellm.BadRequestError":
return litellm.BadRequestError(
message=BAD_MESSAGE_ERROR_STR + f"the {message['role']} message at index {msg_i} {detail}",
model=model,
llm_provider="ollama",
)
def _ollama_content_part(model: str, message: AllMessageValues, part: object, msg_i: int) -> tuple[str, str]:
match part:
case {"type": "text", "text": str() as text}:
return text, ""
case {"type": "text", "text": bad_text}:
raise _ollama_bad_message(
model, message, msg_i, f"has a {type(bad_text).__name__} text part; text must be a string"
)
case {"type": "text"}:
raise _ollama_bad_message(model, message, msg_i, "has a text part with no text; text must be a string")
case {"type": "image_url", "image_url": str() as image_url}:
return "", image_url
case {"type": "image_url", "image_url": {"url": str() as image_url}}:
return "", image_url
case {"type": "image_url", "image_url": dict()}:
raise _ollama_bad_message(
model,
message,
msg_i,
"has an image_url object without a url string; image_url must be a URL string or an object with a url",
)
case {"type": "image_url", "image_url": bad_image_url}:
raise _ollama_bad_message(
model,
message,
msg_i,
f"has a {type(bad_image_url).__name__} image_url; image_url must be a URL string or an object with a url",
)
case {"type": "image_url"}:
raise _ollama_bad_message(
model,
message,
msg_i,
"has an image_url part with no image_url; image_url must be a URL string or an object with a url",
)
case Mapping():
return "", ""
case _:
raise _ollama_bad_message(
model, message, msg_i, f"has a {type(part).__name__} content part; content parts must be objects"
)
def _ollama_user_message_parts(
model: str, message: AllMessageValues, msg_i: int
) -> tuple[tuple[str, ...], tuple[str, ...]]:
msg_content: Final = message.get("content")
if msg_content is None:
return (), ()
if isinstance(msg_content, str):
return ((msg_content,) if msg_content else ()), ()
if not isinstance(msg_content, list):
raise _ollama_bad_message(
model,
message,
msg_i,
f"has {type(msg_content).__name__} content; content must be a string or a list of content parts",
)
parts: Final = tuple(_ollama_content_part(model, message, part, msg_i) for part in msg_content)
texts: Final = tuple(text for text, _ in parts if text)
image_urls: Final = tuple(image_url for _, image_url in parts if image_url)
return texts, image_urls
def _ollama_user_turn(model: str, messages: Sequence[AllMessageValues], msg_i: int) -> tuple[str, tuple[str, ...], int]:
user_run: Final = tuple(
itertools.takewhile(
lambda message: message["role"] in _OLLAMA_USER_ROLES, (messages[i] for i in range(msg_i, len(messages)))
)
)
user_parts: Final = tuple(
_ollama_user_message_parts(model, message, msg_i + offset) for offset, message in enumerate(user_run)
)
user_content_str: Final = "\n".join("\n".join(texts) for texts, _ in user_parts if texts)
image_urls: Final = tuple(itertools.chain.from_iterable(urls for _, urls in user_parts))
return user_content_str, image_urls, len(user_run)
def ollama_pt(
model: str, messages: list
) -> (
str | OllamaVisionModelObject
): # https://github.com/ollama/ollama/blob/af4cf55884ac54b9e637cd71dadfe9b7a5685877/docs/modelfile.md#template
user_message_types: Final = {"user", "tool", "function"}
msg_i = 0
images: Final = []
prompt = ""
while msg_i < len(messages):
init_msg_i = msg_i
user_content_str = ""
## MERGE CONSECUTIVE USER CONTENT ##
while msg_i < len(messages) and messages[msg_i]["role"] in user_message_types:
msg_content = messages[msg_i].get("content")
if msg_content:
if isinstance(msg_content, list):
for m in msg_content:
if m.get("type", "") == "image_url":
if isinstance(m["image_url"], str):
images.append(m["image_url"])
elif isinstance(m["image_url"], dict):
images.append(m["image_url"]["url"])
elif m.get("type", "") == "text":
user_content_str += m["text"]
else:
# Tool message content will always be a string
user_content_str += msg_content
msg_i += 1
user_content_str, image_urls, user_run_len = _ollama_user_turn(model, messages, msg_i)
images.extend(image_urls)
msg_i += user_run_len
if user_content_str:
prompt += f"### User:\n{user_content_str}\n\n"

View file

@ -322,7 +322,7 @@ class AnthropicMessagesHandler(BaseTranslation):
if not chunks:
return None
try:
return AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks(
return AnthropicPassthroughLoggingHandler.build_usage_only_response_from_chunks(
all_chunks=chunks,
model=str((request_data or {}).get("model") or ""),
)
@ -1226,7 +1226,7 @@ class AnthropicMessagesHandler(BaseTranslation):
has_ended: Final = self._check_streaming_has_ended(responses_so_far)
if has_ended:
# build the model response from the responses_so_far
built_response: Final = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
built_response: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response(
all_chunks=responses_so_far,
litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj),
model="",

View file

@ -69,6 +69,12 @@ def _error_status_and_message(exc: Exception) -> tuple[int, str]:
return 500, str(exc) or "Upstream stream ended before completion"
def _provider_error(exc: Exception) -> Exception:
if isinstance(exc, MidStreamFallbackError) and exc.original_exception is not None:
return exc.original_exception
return exc
def _mid_stream_error_sse_event(exc: Exception) -> bytes:
from litellm.anthropic_interface.exceptions.exception_mapping_utils import (
anthropic_error_sse_frame,
@ -336,6 +342,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
self._message_id: str = f"msg_{uuid.uuid4()}"
if litellm_logging_obj is not None:
litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
self.litellm_logging_obj = litellm_logging_obj
# Mapping of truncated tool names to original names (for OpenAI's 64-char limit)
self.tool_name_mapping = tool_name_mapping or {}
# Polyfill applied_edits on final message_delta.
@ -1041,7 +1048,25 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
else:
yield chunk
except Exception as e: # noqa: BLE001 # boundary before the socket: any upstream failure becomes an Anthropic error event
verbose_logger.exception("Anthropic Adapter - mid-stream error, emitting Anthropic error event: %s", e)
verbose_logger.exception("Anthropic Adapter - mid-stream error: %s", e)
logging_obj: Final = self.litellm_logging_obj
provider_error: Final = _provider_error(e)
if logging_obj is not None and logging_obj.on_detached_stream_failure is not None:
if provider_error is e:
raise
raise provider_error from e
if logging_obj is not None:
try:
await logging_obj.dispatch_failure_handlers(
exception=provider_error,
traceback_exception=traceback.format_exc(),
prefer_async_handlers=True,
)
except Exception as failure_handler_error: # noqa: BLE001 # a failing failure handler must not also drop the error frame
verbose_logger.exception(
"Anthropic Adapter - failure handler raised while reporting a mid-stream error: %s",
failure_handler_error,
)
yield _mid_stream_error_sse_event(e)
def _increment_content_block_index(self):

View file

@ -228,7 +228,7 @@ async def _check_summary_model_access(
try:
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.auth_checks import (
_can_object_call_model,
can_object_call_model,
can_project_access_model,
can_user_call_model,
get_project_object,
@ -258,7 +258,7 @@ async def _check_summary_model_access(
if not models:
continue
try:
_can_object_call_model(
can_object_call_model(
model=summary_model,
llm_router=llm_router,
models=models,
@ -370,7 +370,7 @@ async def _check_summary_model_access(
)
if member_allowed_models:
try:
_can_object_call_model(
can_object_call_model(
model=summary_model,
llm_router=llm_router,
models=list(member_allowed_models),
@ -558,7 +558,7 @@ async def _check_summary_model_rate_limit(
limiter: Final[object] = get_proxy_hook("parallel_request_limiter") if get_proxy_hook is not None else None
should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None)
create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr(
limiter, "_create_rate_limit_descriptors", None
limiter, "create_rate_limit_descriptors", None
)
add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None

View file

@ -421,7 +421,7 @@ class BaseAnthropicMessagesStreamingIterator:
if self.completion_start_time is not None:
self.litellm_logging_obj.completion_start_time = self.completion_start_time
self.litellm_logging_obj.model_call_details["completion_start_time"] = self.completion_start_time
logging_coroutine: Final = PassThroughStreamingHandler._route_streaming_logging_to_handler(
logging_coroutine: Final = PassThroughStreamingHandler.route_streaming_logging_to_handler(
litellm_logging_obj=self.litellm_logging_obj,
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
url_route="/v1/messages",

View file

@ -65,6 +65,7 @@ _IDENTITY_SOURCE_PARAM: Final = "anthropic_identity_source"
_IDENTITY_SOURCE_ENV: Final = "ANTHROPIC_IDENTITY_SOURCE"
_IDENTITY_TOKEN_FILE_PARAM: Final = "anthropic_identity_token_file"
_IDENTITY_TOKEN_PARAM: Final = "anthropic_identity_token"
_LEGACY_REF_PARAMS: Final = (_IDENTITY_TOKEN_FILE_PARAM, _IDENTITY_TOKEN_PARAM)
# litellm_params key -> InternalIssuerSource/KeycloakSource field name. Every key here must
# also be listed in ANTHROPIC_WIF_KWARGS_KEYS (types/workload_identity.py), which is what makes it
@ -139,7 +140,7 @@ def resolve_anthropic_wif_params(litellm_params: Mapping[str, object] | None) ->
)
organization_id: Final = _config_value(litellm_params, "anthropic_organization_id", "ANTHROPIC_ORGANIZATION_ID")
if federation_rule_id is None or organization_id is None:
_raise_if_identity_source_configured(litellm_params, federation_rule_id, organization_id)
_raise_if_federation_requested(litellm_params, federation_rule_id, organization_id)
return None
identity_source: Final = _resolve_identity_source(litellm_params)
if identity_source is None:
@ -204,17 +205,15 @@ def _raise_unknown_source_kind(source_kind: str) -> NoReturn:
)
def _raise_if_identity_source_configured(
def _raise_if_federation_requested(
litellm_params: Mapping[str, object] | None, federation_rule_id: str | None, organization_id: str | None
) -> None:
"""A configured identity source is an explicit request to federate, so a missing rule or
organization id fails closed with the ids named, rather than silently skipping federation
and surfacing later as a missing API key."""
source_kind: Final = _resolve_source_kind(litellm_params)
if source_kind is None:
"""A configured identity source, or a token file or inline token set on the deployment, is an
explicit request to federate, so a missing rule or organization id fails closed with the ids
named, rather than silently skipping federation and surfacing later as a missing API key."""
request: Final = _explicit_federation_request(litellm_params)
if request is None:
return
if source_kind not in {kind.value for kind in AnthropicIdentitySourceKind}:
_raise_unknown_source_kind(source_kind)
missing: Final = tuple(
param
for param, value in (
@ -225,7 +224,7 @@ def _raise_if_identity_source_configured(
)
raise litellm.AuthenticationError(
message=(
f"{_IDENTITY_SOURCE_PARAM} is {source_kind!r}, but {' and '.join(missing)} "
f"{request}, but {' and '.join(missing)} "
f"{'is' if len(missing) == 1 else 'are'} not set. {_MISSING_IDS_HINT}"
),
llm_provider="anthropic",
@ -233,14 +232,25 @@ def _raise_if_identity_source_configured(
)
def _explicit_federation_request(litellm_params: Mapping[str, object] | None) -> str | None:
source_kind: Final = _resolve_source_kind(litellm_params)
if source_kind is None:
legacy_ref_param: Final = _legacy_ref_param(litellm_params)
return None if legacy_ref_param is None else f"{legacy_ref_param} is set"
if source_kind not in {kind.value for kind in AnthropicIdentitySourceKind}:
_raise_unknown_source_kind(source_kind)
return f"{_IDENTITY_SOURCE_PARAM} is {source_kind!r}"
def _resolve_source_kind(litellm_params: Mapping[str, object] | None) -> str | None:
param_kind: Final = _param_str(litellm_params, _IDENTITY_SOURCE_PARAM)
if param_kind is not None:
return param_kind
has_param_legacy_ref: Final = any(
_param_str(litellm_params, key) is not None for key in (_IDENTITY_TOKEN_FILE_PARAM, _IDENTITY_TOKEN_PARAM)
)
return None if has_param_legacy_ref else _env_str(_IDENTITY_SOURCE_ENV)
return None if _legacy_ref_param(litellm_params) is not None else _env_str(_IDENTITY_SOURCE_ENV)
def _legacy_ref_param(litellm_params: Mapping[str, object] | None) -> str | None:
return next((key for key in _LEGACY_REF_PARAMS if _param_str(litellm_params, key) is not None), None)
def _reject_foreign_variant_fields(

View file

@ -1,3 +1,3 @@
from .transformation import DecisionsProviderConfig, JevCompatibleDecisionsEndpoint
from .transformation import BaseDecisionsConfig
__all__ = ["DecisionsProviderConfig", "JevCompatibleDecisionsEndpoint"]
__all__ = ["BaseDecisionsConfig"]

View file

@ -0,0 +1,247 @@
"""The Jev / System One wire shape and its translation to and from the OpenAI Decisions shape.
System One (TypeSafe, Perplexity, OpenRouter, Cloudflare Clef, Strands Decider) takes
{"model", "state", "questions": {name: question}} and answers with {"model", "answers": {name: answer}, "usage"}.
Predicates are `noul` questions, choice options are a `criteria` map, score levels are a `criteria` list.
"""
import itertools
from collections.abc import Mapping, Sequence
from typing import Final, Literal, TypeAlias
from pydantic import ConfigDict, TypeAdapter
from typing_extensions import assert_never
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.base import LiteLLMPydanticObjectBase
from litellm.types.openai_decisions import (
ChoiceAnswer,
ChoiceProbability,
ChoiceQuestion,
DecisionAnswer,
DecisionChoice,
DecisionInput,
DecisionInputMessage,
DecisionInputPart,
DecisionInputTokensDetails,
DecisionOutputTokensDetails,
DecisionQuestion,
DecisionsRequest,
DecisionsRequestBody,
DecisionsResponse,
DecisionUsage,
PredicateAnswer,
PredicateQuestion,
ScoreAnswer,
ScoreProbability,
ScoreQuestion,
)
class SystemOneObjectBase(LiteLLMPydanticObjectBase):
model_config = ConfigDict(extra="allow", frozen=True)
class SystemOneNoulAnswer(SystemOneObjectBase):
type: Literal["noul"]
noul: float
class SystemOneChoiceAnswer(SystemOneObjectBase):
type: Literal["choice"]
choice: str
confidence: float
probabilities: Mapping[str, float]
class SystemOneScoreAnswer(SystemOneObjectBase):
type: Literal["score"]
score: float
confidence: float
probabilities: Mapping[str, float]
SystemOneAnswer: TypeAlias = SystemOneNoulAnswer | SystemOneChoiceAnswer | SystemOneScoreAnswer
class SystemOneUsage(SystemOneObjectBase):
input_tokens: int = 0
output_tokens: int = 0
class SystemOneResponse(SystemOneObjectBase):
model: str | None = None
answers: Mapping[str, SystemOneAnswer]
usage: SystemOneUsage | None = None
SYSTEM_ONE_RESPONSE_ADAPTER: Final[TypeAdapter[SystemOneResponse]] = TypeAdapter(SystemOneResponse)
def _unsupported(what: str, custom_llm_provider: str) -> BaseLLMException:
return BaseLLMException(
status_code=400,
message=f"Decisions provider '{custom_llm_provider}' does not support {what}",
)
def to_system_one_request(model: str, body: DecisionsRequestBody, custom_llm_provider: str) -> dict[str, object]:
keys: Final = question_keys(body.questions, custom_llm_provider)
return {
"model": model,
"state": _state(body.input, custom_llm_provider),
"questions": {
key: _question(question, custom_llm_provider) for key, question in zip(keys, body.questions, strict=True)
},
}
def question_keys(questions: Sequence[DecisionQuestion], custom_llm_provider: str) -> tuple[str, ...]:
"""System One keys questions and answers by name, so unnamed questions get a positional key."""
names: Final = tuple(question.name for question in questions if question.name is not None)
if len(set(names)) != len(names):
raise BaseLLMException(
status_code=400,
message=f"Decisions provider '{custom_llm_provider}' requires a unique name per question",
)
taken: Final = frozenset(names)
return tuple(
question.name if question.name is not None else _positional_key(index, taken)
for index, question in enumerate(questions)
)
def _positional_key(index: int, taken: frozenset[str]) -> str:
candidates: Final = (f"{'_' * depth}q{index}" for depth in itertools.count())
return next(key for key in candidates if key not in taken)
def _state(input_value: DecisionInput, custom_llm_provider: str) -> str:
if isinstance(input_value, str):
return input_value
return "\n".join(_message_text(message, custom_llm_provider) for message in input_value)
def _message_text(message: DecisionInputMessage, custom_llm_provider: str) -> str:
if isinstance(message.content, str):
return message.content
return "\n".join(_part_text(part, custom_llm_provider) for part in message.content)
def _part_text(part: DecisionInputPart, custom_llm_provider: str) -> str:
if part.type != "input_text":
raise _unsupported("input_image parts", custom_llm_provider)
return part.text
def _question(question: DecisionQuestion, custom_llm_provider: str) -> dict[str, object]:
match question:
case PredicateQuestion():
return {"type": "noul", "instructions": question.instructions}
case ChoiceQuestion():
return {
"type": "choice",
"instructions": question.instructions,
"criteria": _choice_criteria(question, custom_llm_provider),
}
case ScoreQuestion():
return {
"type": "score",
"instructions": question.instructions,
"criteria": [_level_text(level.label, level.description) for level in question.levels],
}
case _:
assert_never(question)
def _choice_criteria(question: ChoiceQuestion, custom_llm_provider: str) -> dict[str, str | None]:
values: Final = tuple(_choice_key(choice, custom_llm_provider) for choice in question.choices)
if len(set(values)) != len(values):
raise _unsupported("repeated choice values", custom_llm_provider)
return {value: choice.description for value, choice in zip(values, question.choices, strict=True)}
def _choice_key(choice: DecisionChoice, custom_llm_provider: str) -> str:
if not isinstance(choice.value, str):
raise _unsupported("boolean choice values", custom_llm_provider)
return choice.value
def _level_text(label: str, description: str | None) -> str:
return description if description is not None else label
def to_decisions_response(
system_one: SystemOneResponse,
request: DecisionsRequest,
custom_llm_provider: str,
) -> DecisionsResponse:
keys: Final = question_keys(request.body.questions, custom_llm_provider)
return DecisionsResponse(
model=system_one.model if system_one.model is not None else request.model,
answers=[
_answer(key, question, system_one.answers, custom_llm_provider)
for key, question in zip(keys, request.body.questions, strict=True)
],
usage=_usage(system_one.usage),
)
def _answer(
key: str,
question: DecisionQuestion,
answers: Mapping[str, SystemOneAnswer],
custom_llm_provider: str,
) -> DecisionAnswer:
match question, answers.get(key):
case PredicateQuestion(), SystemOneNoulAnswer() as answer:
return PredicateAnswer(type="predicate", name=question.name, probability=answer.noul)
case ChoiceQuestion(), SystemOneChoiceAnswer() as answer:
return ChoiceAnswer(
type="choice",
name=question.name,
choice=answer.choice,
probabilities=_choice_probabilities(question, answer),
confidence=answer.confidence,
)
case ScoreQuestion(), SystemOneScoreAnswer() as answer:
return ScoreAnswer(
type="score",
name=question.name,
score=answer.score,
probabilities=_score_probabilities(question, answer),
confidence=answer.confidence,
)
case _:
raise BaseLLMException(
status_code=500,
message=(
f"Decisions provider '{custom_llm_provider}' returned no {question.type} answer "
f"for question '{key}'"
),
)
def _choice_probabilities(question: ChoiceQuestion, answer: SystemOneChoiceAnswer) -> list[ChoiceProbability]:
return [
ChoiceProbability(value=choice.value, probability=answer.probabilities.get(str(choice.value), 0.0))
for choice in question.choices
]
def _score_probabilities(question: ScoreQuestion, answer: SystemOneScoreAnswer) -> list[ScoreProbability]:
return [
ScoreProbability(value=index, label=level.label, probability=answer.probabilities.get(str(index), 0.0))
for index, level in enumerate(question.levels)
]
def _usage(usage: SystemOneUsage | None) -> DecisionUsage:
counted: Final = usage if usage is not None else SystemOneUsage()
return DecisionUsage(
input_tokens=counted.input_tokens,
input_tokens_details=DecisionInputTokensDetails(cached_tokens=0, cache_write_tokens=0),
output_tokens=counted.output_tokens,
output_tokens_details=DecisionOutputTokensDetails(reasoning_tokens=0),
total_tokens=counted.input_tokens + counted.output_tokens,
)

View file

@ -1,20 +1,322 @@
from dataclasses import dataclass
from typing import Protocol
import json
from abc import ABC
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
import httpx
from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.decisions import (
ChoiceAnswer,
ChoiceQuestion,
DecisionAnswer,
DecisionQuestion,
DecisionsIRAnswer,
DecisionsIRChoiceAnswer,
DecisionsIRChoiceOption,
DecisionsIRChoiceProbability,
DecisionsIRChoiceQuestion,
DecisionsIRMessages,
DecisionsIRPredicateAnswer,
DecisionsIRPredicateQuestion,
DecisionsIRQuestion,
DecisionsIRRefusal,
DecisionsIRRequest,
DecisionsIRResponse,
DecisionsIRScoreAnswer,
DecisionsIRScoreLevel,
DecisionsIRScoreProbability,
DecisionsIRScoreQuestion,
DecisionsIRState,
DecisionsIRUsage,
DecisionsJSON,
DecisionsRequestBody,
DecisionsResponse,
DecisionsUsage,
NoulAnswer,
NoulQuestion,
OpenAIDecisionInputImage,
OpenAIDecisionInputMessage,
OpenAIDecisionInputText,
ScoreAnswer,
ScoreQuestion,
UnsupportedDecisionsRequest,
systemone_choice_key,
)
_PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object)
_SYSTEMONE_RESPONSE_ADAPTER: Final[TypeAdapter[DecisionsResponse]] = TypeAdapter(DecisionsResponse)
_RESERVED_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "content-type"})
_TEXT_ONLY: Final = UnsupportedDecisionsRequest(
reason="input_image content parts are not supported because System One providers accept text input only"
)
@dataclass(frozen=True, slots=True)
class JevCompatibleDecisionsEndpoint:
default_api_base_value: str | None
path: str
api_key_env: tuple[str, ...]
api_base_env: str
def decisions_text(value: DecisionsJSON) -> str:
return value if isinstance(value, str) else json.dumps(value)
def systemone_keys(questions: Sequence[DecisionsIRQuestion]) -> tuple[str, ...]:
names: Final = tuple(question.name for question in questions if question.name is not None)
if len(frozenset(names)) == len(questions):
return names
return tuple(str(index) for index in range(len(questions)))
def _ir_question(name: str, question: DecisionQuestion) -> DecisionsIRQuestion:
extra: Final = MappingProxyType(question.model_extra or {})
match question:
case NoulQuestion():
return DecisionsIRPredicateQuestion(
name=name, instructions=question.instructions, criteria=question.criteria, extra=extra
)
case ChoiceQuestion():
return DecisionsIRChoiceQuestion(
name=name,
instructions=question.instructions,
choices=tuple(
DecisionsIRChoiceOption(value=value, description=description)
for value, description in question.criteria.items()
),
extra=extra,
)
case ScoreQuestion():
return DecisionsIRScoreQuestion(
name=name,
instructions=question.instructions,
levels=tuple(
DecisionsIRScoreLevel(label=criterion, description=None) for criterion in question.criteria
),
extra=extra,
)
case _:
assert_never(question)
def systemone_request_to_ir(request: DecisionsRequestBody) -> DecisionsIRRequest:
return DecisionsIRRequest(
input=DecisionsIRState(state=request.state),
questions=tuple(_ir_question(name, question) for name, question in request.questions.items()),
)
def _has_image(message: OpenAIDecisionInputMessage) -> bool:
return not isinstance(message.content, str) and any(
isinstance(part, OpenAIDecisionInputImage) for part in message.content
)
def _message_text(message: OpenAIDecisionInputMessage) -> str:
if isinstance(message.content, str):
return message.content
return "\n\n".join(part.text for part in message.content if isinstance(part, OpenAIDecisionInputText))
def _systemone_state(
decision_input: DecisionsIRState | DecisionsIRMessages,
) -> DecisionsJSON | UnsupportedDecisionsRequest:
match decision_input:
case DecisionsIRState():
return decision_input.state
case DecisionsIRMessages():
if any(_has_image(message) for message in decision_input.messages):
return _TEXT_ONLY
return "\n\n".join(_message_text(message) for message in decision_input.messages)
case _:
assert_never(decision_input)
def _optional(key: str, value: object) -> Mapping[str, object]:
return {} if value is None else {key: value}
def _level_criterion(level: DecisionsIRScoreLevel) -> DecisionsJSON:
return level.label if level.description is None else f"{decisions_text(level.label)}: {level.description}"
def _systemone_question(question: DecisionsIRQuestion) -> Mapping[str, object]:
match question:
case DecisionsIRPredicateQuestion():
return {
"type": "noul",
**_optional("instructions", question.instructions),
**_optional("criteria", question.criteria),
**question.extra,
}
case DecisionsIRChoiceQuestion():
return {
"type": "choice",
**_optional("instructions", question.instructions),
"criteria": {systemone_choice_key(option.value): option.description for option in question.choices},
**question.extra,
}
case DecisionsIRScoreQuestion():
return {
"type": "score",
**_optional("instructions", question.instructions),
"criteria": [_level_criterion(level) for level in question.levels],
**question.extra,
}
case _:
assert_never(question)
def ir_to_systemone_request(
model: str, request: DecisionsIRRequest
) -> Mapping[str, object] | UnsupportedDecisionsRequest:
state: Final = _systemone_state(request.input)
if isinstance(state, UnsupportedDecisionsRequest):
return state
keyed_questions: Final = zip(systemone_keys(request.questions), request.questions, strict=True)
return {
"model": model,
"state": state,
"questions": {key: _systemone_question(question) for key, question in keyed_questions},
}
def _ir_choice_answer(question: DecisionsIRChoiceQuestion, answer: ChoiceAnswer) -> DecisionsIRChoiceAnswer:
typed_values: Final = {systemone_choice_key(option.value): option.value for option in question.choices}
return DecisionsIRChoiceAnswer(
choice=typed_values.get(answer.choice, answer.choice),
confidence=answer.confidence,
probabilities=tuple(
DecisionsIRChoiceProbability(value=typed_values.get(key, key), probability=probability)
for key, probability in answer.probabilities.items()
),
extra=MappingProxyType(answer.model_extra or {}),
)
def _ir_score_answer(question: DecisionsIRScoreQuestion, answer: ScoreAnswer) -> DecisionsIRScoreAnswer:
return DecisionsIRScoreAnswer(
score=answer.score,
confidence=answer.confidence,
probabilities=tuple(
DecisionsIRScoreProbability(
value=index,
label=answer.legend.get(str(index), level.label),
probability=answer.probabilities.get(str(index), 0.0),
)
for index, level in enumerate(question.levels)
),
extra=MappingProxyType(answer.model_extra or {}),
)
def _ir_answer(question: DecisionsIRQuestion, answer: DecisionAnswer | None) -> DecisionsIRAnswer:
match question, answer:
case DecisionsIRPredicateQuestion(), NoulAnswer():
return DecisionsIRPredicateAnswer(probability=answer.noul, extra=MappingProxyType(answer.model_extra or {}))
case DecisionsIRChoiceQuestion(), ChoiceAnswer():
return _ir_choice_answer(question, answer)
case DecisionsIRScoreQuestion(), ScoreAnswer():
return _ir_score_answer(question, answer)
case _:
return DecisionsIRRefusal()
def systemone_response_to_ir(response: DecisionsResponse, request: DecisionsIRRequest) -> DecisionsIRResponse:
usage: Final = response.usage or DecisionsUsage()
keyed_questions: Final = zip(systemone_keys(request.questions), request.questions, strict=True)
return DecisionsIRResponse(
model=response.model,
answers=tuple(_ir_answer(question, response.answers.get(key)) for key, question in keyed_questions),
usage=DecisionsIRUsage(
input_tokens=usage.input_tokens,
output_tokens=usage.output_tokens,
cached_tokens=usage.cached_tokens,
cache_write_tokens=usage.cache_write_tokens,
extra=MappingProxyType(usage.model_extra or {}),
),
extra=MappingProxyType(response.model_extra or {}),
)
def parse_systemone_response(payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse:
return systemone_response_to_ir(_SYSTEMONE_RESPONSE_ADAPTER.validate_python(payload), request)
def _systemone_answer(answer: DecisionsIRAnswer) -> DecisionAnswer | None:
match answer:
case DecisionsIRPredicateAnswer():
return NoulAnswer.model_validate({**answer.extra, "type": "noul", "noul": answer.probability})
case DecisionsIRChoiceAnswer():
return ChoiceAnswer.model_validate(
{
**answer.extra,
"type": "choice",
"choice": systemone_choice_key(answer.choice),
"confidence": answer.confidence,
"probabilities": {
systemone_choice_key(item.value): item.probability for item in answer.probabilities
},
}
)
case DecisionsIRScoreAnswer():
return ScoreAnswer.model_validate(
{
**answer.extra,
"type": "score",
"score": answer.score,
"confidence": answer.confidence,
"legend": {str(item.value): item.label for item in answer.probabilities},
"probabilities": {str(item.value): item.probability for item in answer.probabilities},
}
)
case DecisionsIRRefusal():
return None
case _:
assert_never(answer)
def ir_to_systemone_response(response: DecisionsIRResponse, request: DecisionsIRRequest) -> DecisionsResponse:
keyed_answers: Final = zip(
systemone_keys(request.questions), (_systemone_answer(answer) for answer in response.answers), strict=True
)
return DecisionsResponse.model_validate(
{
**response.extra,
"model": response.model,
"answers": {key: answer for key, answer in keyed_answers if answer is not None},
"usage": DecisionsUsage.model_validate(
{
**response.usage.extra,
"input_tokens": response.usage.input_tokens,
"output_tokens": response.usage.output_tokens,
"cached_tokens": response.usage.cached_tokens,
"cache_write_tokens": response.usage.cache_write_tokens,
}
),
}
)
class BaseDecisionsConfig(ABC):
path: str = "/v1/systemone"
api_key_env: tuple[str, ...] = ()
api_base_env: tuple[str, ...] = ()
api_key_required: bool = True
def default_api_base(self) -> str | None:
return self.default_api_base_value
def get_default_api_base(self) -> str | None:
return None
def missing_api_base_message(self, provider: str) -> str:
return f"api_base is required for Decisions provider '{provider}'"
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return f"api_base is required for Decisions provider '{custom_llm_provider}'"
def resolve_api_base(self, api_base: str | None) -> str | None:
return api_base or self._first_secret(self.api_base_env) or self.get_default_api_base()
def resolve_api_key(self, api_key: str | None) -> str | None:
return api_key or self._first_secret(self.api_key_env)
@staticmethod
def _first_secret(names: tuple[str, ...]) -> str | None:
return next((value for value in (get_secret_str(name) for name in names) if value), None)
def canonical_model(self, model: str) -> str:
return model
@ -22,31 +324,50 @@ class JevCompatibleDecisionsEndpoint:
def request_model(self, model: str) -> str:
return model
def endpoint_url(self, api_base: str, model: str) -> str:
def validate_environment(self, headers: Mapping[str, str], model: str, api_key: str | None) -> dict[str, str]:
return {
**{name: value for name, value in headers.items() if name.lower() not in _RESERVED_HEADERS},
**({"Authorization": f"Bearer {api_key}"} if api_key is not None else {}),
"Content-Type": "application/json",
}
def get_complete_url(self, api_base: str, model: str) -> str:
return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}"
def transform_decisions_request(
self,
model: str,
request: DecisionsIRRequest,
custom_llm_provider: str,
) -> Mapping[str, object] | UnsupportedDecisionsRequest:
return ir_to_systemone_request(self.request_model(model), request)
def unwrap_response(self, payload: object) -> object:
return payload
def parse_response(self, payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse:
return parse_systemone_response(self.unwrap_response(payload), request)
class DecisionsProviderConfig(Protocol):
@property
def api_key_env(self) -> tuple[str, ...]: ...
def transform_decisions_response(
self,
model: str,
custom_llm_provider: str,
raw_response: httpx.Response,
request: DecisionsIRRequest,
) -> DecisionsIRResponse:
payload: Final[object] = _PAYLOAD_ADAPTER.validate_json(raw_response.content)
try:
return self.parse_response(payload, request)
except ValidationError as error:
raise BaseLLMException(
status_code=500,
message=f"Decisions provider '{custom_llm_provider}' returned an unexpected response: {error}",
) from error
@property
def api_base_env(self) -> str: ...
@property
def api_key_required(self) -> bool: ...
def default_api_base(self) -> str | None: ...
def missing_api_base_message(self, provider: str) -> str: ...
def canonical_model(self, model: str) -> str: ...
def request_model(self, model: str) -> str: ...
def endpoint_url(self, api_base: str, model: str) -> str: ...
def unwrap_response(self, payload: object) -> object: ...
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, str] | httpx.Headers,
) -> BaseLLMException:
return BaseLLMException(status_code=status_code, message=error_message, headers=headers)

View file

@ -25,7 +25,7 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI: Final[
] = {
"Submitted": "validating",
"Validating": "validating",
"Scheduled": "validating",
"Scheduled": "in_progress",
"InProgress": "in_progress",
"Stopping": "cancelling",
"Stopped": "cancelled",

View file

@ -1,9 +1,9 @@
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.secret_managers.main import (
get_secret_str,
normalize_nonempty_secret_str,
@ -12,19 +12,17 @@ from litellm.secret_managers.main import (
_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
@dataclass(frozen=True, slots=True)
class CloudflareDecisionsEndpoint:
api_key_env: tuple[str, ...] = ("CLOUDFLARE_API_KEY",)
api_base_env: str = "CLOUDFLARE_API_BASE"
api_key_required: bool = True
class CloudflareDecisionsConfig(BaseDecisionsConfig):
api_key_env = ("CLOUDFLARE_API_KEY",)
api_base_env = ("CLOUDFLARE_API_BASE",)
def default_api_base(self) -> str | None:
def get_default_api_base(self) -> str | None:
account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
if account_id is None:
return None
return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
def missing_api_base_message(self, provider: str) -> str:
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly"
def canonical_model(self, model: str) -> str:
@ -35,7 +33,7 @@ class CloudflareDecisionsEndpoint:
def request_model(self, model: str) -> str:
return model.rsplit("/", maxsplit=1)[-1]
def endpoint_url(self, api_base: str, model: str) -> str:
def get_complete_url(self, api_base: str, model: str) -> str:
normalized_api_base: Final = api_base.rstrip("/")
if normalized_api_base.endswith("/ai/v1"):
return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}"
@ -44,15 +42,13 @@ class CloudflareDecisionsEndpoint:
return f"{normalized_api_base}/ai/run/{model}"
def unwrap_response(self, payload: object) -> object:
if not isinstance(payload, Mapping):
try:
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
except ValidationError:
return payload
response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload)
if "answers" in response_mapping:
return payload
result: Final = response_mapping.get("result")
if isinstance(result, Mapping):
return result
return payload
CLOUDFLARE_DECISIONS_ENDPOINT: Final[CloudflareDecisionsEndpoint] = CloudflareDecisionsEndpoint()
return response_mapping
try:
return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result"))
except ValidationError:
return response_mapping

View file

@ -450,6 +450,18 @@ def get_shared_realtime_ssl_context() -> bool | str | ssl.SSLContext:
return _shared_realtime_ssl_context
def realtime_ssl_for_url(url: str) -> bool | str | ssl.SSLContext | None:
if url.startswith("ws://"):
return None
shared: Final = get_shared_realtime_ssl_context()
if shared is not False:
return shared
unverified: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
unverified.check_hostname = False
unverified.verify_mode = ssl.CERT_NONE
return unverified
def mask_sensitive_info(error_message):
# Find the start of the key parameter
if isinstance(error_message, str):

View file

@ -70,6 +70,7 @@ from litellm.llms.base_llm.base_model_iterator import (
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig
from litellm.llms.base_llm.files.transformation import (
@ -122,6 +123,7 @@ from litellm.types.containers.main import (
ContainerObject,
DeleteContainerResult,
)
from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse
from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
from litellm.types.integrations.custom_logger import (
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
@ -197,7 +199,7 @@ def _rust_responses_websocket_enabled(
return decision(context) is not Decision.PYTHON
from .http_handler import get_shared_realtime_ssl_context
from .http_handler import get_shared_realtime_ssl_context, realtime_ssl_for_url
if TYPE_CHECKING:
from aiohttp import ClientSession
@ -258,7 +260,7 @@ class _WebsocketsModule(Protocol):
*,
additional_headers: Mapping[str, str],
max_size: int | None,
ssl: bool | str | ssl.SSLContext,
ssl: bool | str | ssl.SSLContext | None,
open_timeout: float,
) -> Awaitable["ClientConnection"]: ...
@ -1481,6 +1483,100 @@ class BaseLLMHTTPHandler:
request_data=request_data,
)
def _prepare_decisions_request(
self,
model: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseDecisionsConfig,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
) -> tuple[str, dict[str, str], dict[str, object]]:
outbound_headers: Final = provider_config.validate_environment(headers=headers, model=model, api_key=api_key)
url: Final = provider_config.get_complete_url(api_base=api_base, model=model)
data: Final = dict(body)
if logging_obj is not None:
logging_obj.pre_call(
input=data,
api_key=api_key,
model=model,
additional_args={"api_base": url, "complete_input_dict": data, "headers": outbound_headers},
)
return url, outbound_headers, data
def decisions(
self,
model: str,
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseDecisionsConfig,
request: DecisionsIRRequest,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
client: HTTPHandler | None = None,
) -> DecisionsIRResponse:
url, outbound_headers, data = self._prepare_decisions_request(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
api_base=api_base,
api_key=api_key,
headers=headers,
)
sync_httpx_client: Final = client if client is not None else get_httpx_client()
try:
response: Final = sync_httpx_client.post(
url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
return provider_config.transform_decisions_response(
model=model, custom_llm_provider=custom_llm_provider, raw_response=response, request=request
)
async def adecisions(
self,
model: str,
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj | None,
provider_config: BaseDecisionsConfig,
request: DecisionsIRRequest,
body: Mapping[str, object],
api_base: str,
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
client: AsyncHTTPHandler | None = None,
) -> DecisionsIRResponse:
url, outbound_headers, data = self._prepare_decisions_request(
model=model,
logging_obj=logging_obj,
provider_config=provider_config,
body=body,
api_base=api_base,
api_key=api_key,
headers=headers,
)
async_httpx_client: Final = (
client
if client is not None
else get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider))
)
try:
response: Final = await async_httpx_client.post(
url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
return provider_config.transform_decisions_response(
model=model, custom_llm_provider=custom_llm_provider, raw_response=response, request=request
)
def _prepare_audio_transcription_request(
self,
model: str,
@ -6156,6 +6252,7 @@ class BaseLLMHTTPHandler:
e: Exception,
provider_config: Union[
BaseConfig,
BaseDecisionsConfig,
BaseRerankConfig,
BaseResponsesAPIConfig,
BaseImageEditConfig,
@ -6257,7 +6354,7 @@ class BaseLLMHTTPHandler:
websockets_module: _WebsocketsModule,
url: str,
headers: dict,
ssl_context: bool | str | ssl.SSLContext,
ssl_context: bool | str | ssl.SSLContext | None,
*,
open_timeout: float = 8.0,
max_attempts: int = 3,
@ -6329,12 +6426,7 @@ class BaseLLMHTTPHandler:
)
try:
ssl_context = get_shared_realtime_ssl_context()
if url.startswith("wss://") and ssl_context is False:
# Keep TLS for wss:// while honoring SSL_VERIFY=False semantics.
ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE
ssl_context: Final = realtime_ssl_for_url(url)
provider_backend: Final = await provider_config.open_backend(url, headers)
backend_ws: Final = (
provider_backend

View file

@ -39,6 +39,7 @@ from litellm.types.utils import (
ModelResponseStream,
ProviderField,
StreamingChoices,
generate_id,
)
from ..common_utils import OllamaError, OllamaModelInfo, convert_image
@ -524,6 +525,7 @@ class OllamaConfig(BaseConfig):
class OllamaTextCompletionResponseIterator(BaseModelResponseIterator):
def __init__(self, streaming_response, sync_stream: bool, json_mode: bool | None = False):
super().__init__(streaming_response, sync_stream, json_mode)
self.response_id: Final[str] = generate_id()
self.started_reasoning_content: bool = False
self.finished_reasoning_content: bool = False
self.streamed_content: bool = False
@ -622,6 +624,7 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator):
content = self._hold_json_object_start(text)
return ModelResponseStream(
id=self.response_id,
choices=[
StreamingChoices(
index=0,
@ -641,24 +644,26 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator):
# Return reasoning content as ModelResponseStream so UIs can render it
thinking_content: Final = chunk.get("thinking") or ""
return ModelResponseStream(
id=self.response_id,
choices=[
StreamingChoices(
index=0,
delta=Delta(reasoning_content=thinking_content),
)
]
],
)
else:
# In this case, 'thinking' is not present in the chunk, chunk["done"] is false,
# and chunk["response"] is falsy (None or empty string),
# but Ollama is just starting to stream, so it should be processed as a normal dict
return ModelResponseStream(
id=self.response_id,
choices=[
StreamingChoices(
index=0,
delta=Delta(reasoning_content=""),
)
]
],
)
# raise Exception(f"Unable to parse ollama chunk - {chunk}")
except Exception as e:

View file

@ -0,0 +1,320 @@
from collections.abc import Mapping, Sequence
from typing import Final
from pydantic import TypeAdapter
from typing_extensions import assert_never
import litellm
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, decisions_text
from litellm.secret_managers.main import get_secret_str
from litellm.types.decisions import (
DecisionsIRAnswer,
DecisionsIRChoiceAnswer,
DecisionsIRChoiceOption,
DecisionsIRChoiceProbability,
DecisionsIRChoiceQuestion,
DecisionsIRMessages,
DecisionsIRPredicateAnswer,
DecisionsIRPredicateQuestion,
DecisionsIRQuestion,
DecisionsIRRefusal,
DecisionsIRRequest,
DecisionsIRResponse,
DecisionsIRScoreAnswer,
DecisionsIRScoreLevel,
DecisionsIRScoreProbability,
DecisionsIRScoreQuestion,
DecisionsIRState,
DecisionsIRUsage,
DecisionsJSON,
OpenAIChoiceAnswer,
OpenAIChoiceProbability,
OpenAIChoiceQuestion,
OpenAIDecisionAnswer,
OpenAIDecisionInput,
OpenAIDecisionInputTokensDetails,
OpenAIDecisionOutputTokensDetails,
OpenAIDecisionQuestion,
OpenAIDecisionRequestBody,
OpenAIDecisionResponse,
OpenAIDecisionUsage,
OpenAIPredicateAnswer,
OpenAIPredicateQuestion,
OpenAIRefusalAnswer,
OpenAIScoreAnswer,
OpenAIScoreProbability,
OpenAIScoreQuestion,
UnsupportedDecisionsRequest,
systemone_choice_key,
)
_OPENAI_RESPONSE_ADAPTER: Final[TypeAdapter[OpenAIDecisionResponse]] = TypeAdapter(OpenAIDecisionResponse)
_SINGLE_OPTION: Final = UnsupportedDecisionsRequest(
reason="OpenAI needs at least 2 choices or levels on every choice or score question"
)
def _ir_input(decision_input: OpenAIDecisionInput) -> DecisionsIRState | DecisionsIRMessages:
if isinstance(decision_input, str):
return DecisionsIRState(state=decision_input)
return DecisionsIRMessages(messages=tuple(decision_input))
def _ir_question(question: OpenAIDecisionQuestion) -> DecisionsIRQuestion:
match question:
case OpenAIPredicateQuestion():
return DecisionsIRPredicateQuestion(name=question.name, instructions=question.instructions)
case OpenAIChoiceQuestion():
return DecisionsIRChoiceQuestion(
name=question.name,
instructions=question.instructions,
choices=tuple(
DecisionsIRChoiceOption(value=option.value, description=option.description)
for option in question.choices
),
)
case OpenAIScoreQuestion():
return DecisionsIRScoreQuestion(
name=question.name,
instructions=question.instructions,
levels=tuple(
DecisionsIRScoreLevel(label=level.label, description=level.description) for level in question.levels
),
)
case _:
assert_never(question)
def openai_request_to_ir(request: OpenAIDecisionRequestBody) -> DecisionsIRRequest:
return DecisionsIRRequest(
input=_ir_input(request.input),
questions=tuple(_ir_question(question) for question in request.questions),
safety_identifier=request.safety_identifier,
)
def _text_field(key: str, value: DecisionsJSON | None) -> Mapping[str, str]:
return {} if value is None else {key: decisions_text(value)}
def _instructions(value: DecisionsJSON | None, default: str) -> str:
return default if value is None else decisions_text(value)
def _predicate_instructions(question: DecisionsIRPredicateQuestion) -> str:
instructions: Final = () if question.instructions is None else (decisions_text(question.instructions),)
criteria: Final = tuple(
f"Answer {answer} when: {decisions_text(rule)}"
for answer, rule in (question.criteria or {}).items()
if rule is not None
)
return "\n\n".join((*instructions, *criteria))
def _openai_question(question: DecisionsIRQuestion) -> Mapping[str, object]:
match question:
case DecisionsIRPredicateQuestion():
return {
"type": "predicate",
**_text_field("name", question.name),
"instructions": _predicate_instructions(question) or "Is this true of the input?",
}
case DecisionsIRChoiceQuestion():
return {
"type": "choice",
**_text_field("name", question.name),
"instructions": _instructions(question.instructions, "Which choice best fits the input?"),
"choices": [
{"value": option.value, **_text_field("description", option.description)}
for option in question.choices
],
}
case DecisionsIRScoreQuestion():
return {
"type": "score",
**_text_field("name", question.name),
"instructions": _instructions(question.instructions, "Which level best fits the input?"),
"levels": [
{"label": decisions_text(level.label), **_text_field("description", level.description)}
for level in question.levels
],
}
case _:
assert_never(question)
def _openai_input(decision_input: DecisionsIRState | DecisionsIRMessages) -> str | Sequence[Mapping[str, object]]:
match decision_input:
case DecisionsIRState():
return decisions_text(decision_input.state)
case DecisionsIRMessages():
return [message.model_dump(mode="json", exclude_none=True) for message in decision_input.messages]
case _:
assert_never(decision_input)
def _has_one_option(question: DecisionsIRQuestion) -> bool:
match question:
case DecisionsIRChoiceQuestion():
return len(question.choices) < 2
case DecisionsIRScoreQuestion():
return len(question.levels) < 2
case _:
return False
def ir_to_openai_request(model: str, request: DecisionsIRRequest) -> Mapping[str, object] | UnsupportedDecisionsRequest:
if any(_has_one_option(question) for question in request.questions):
return _SINGLE_OPTION
return {
"model": model,
"input": _openai_input(request.input),
"questions": [_openai_question(question) for question in request.questions],
**_text_field("safety_identifier", request.safety_identifier),
}
def _level_label(question: DecisionsIRScoreQuestion, probability: OpenAIScoreProbability) -> DecisionsJSON:
if 0 <= probability.value < len(question.levels):
return question.levels[probability.value].label
return probability.label
def _ir_answer(question: DecisionsIRQuestion, answer: OpenAIDecisionAnswer | None) -> DecisionsIRAnswer:
match question, answer:
case DecisionsIRPredicateQuestion(), OpenAIPredicateAnswer():
return DecisionsIRPredicateAnswer(probability=answer.probability)
case DecisionsIRChoiceQuestion(), OpenAIChoiceAnswer():
return DecisionsIRChoiceAnswer(
choice=answer.choice,
confidence=answer.confidence,
probabilities=tuple(
DecisionsIRChoiceProbability(value=item.value, probability=item.probability)
for item in answer.probabilities
),
)
case DecisionsIRScoreQuestion(), OpenAIScoreAnswer():
return DecisionsIRScoreAnswer(
score=answer.score,
confidence=answer.confidence,
probabilities=tuple(
DecisionsIRScoreProbability(
value=item.value, label=_level_label(question, item), probability=item.probability
)
for item in answer.probabilities
),
)
case _:
return DecisionsIRRefusal()
def openai_response_to_ir(response: OpenAIDecisionResponse, request: DecisionsIRRequest) -> DecisionsIRResponse:
answers: Final = response.answers
return DecisionsIRResponse(
model=response.model,
answers=tuple(
_ir_answer(question, answers[index] if index < len(answers) else None)
for index, question in enumerate(request.questions)
),
usage=DecisionsIRUsage(
input_tokens=response.usage.input_tokens,
output_tokens=response.usage.output_tokens,
cached_tokens=response.usage.input_tokens_details.cached_tokens,
cache_write_tokens=response.usage.input_tokens_details.cache_write_tokens,
reasoning_tokens=response.usage.output_tokens_details.reasoning_tokens,
),
)
def _openai_choice_answer(question: DecisionsIRChoiceQuestion, answer: DecisionsIRChoiceAnswer) -> OpenAIChoiceAnswer:
probabilities: Final = {systemone_choice_key(item.value): item.probability for item in answer.probabilities}
return OpenAIChoiceAnswer(
name=question.name,
choice=answer.choice,
probabilities=tuple(
OpenAIChoiceProbability(
value=option.value, probability=probabilities.get(systemone_choice_key(option.value), 0.0)
)
for option in question.choices
),
confidence=answer.confidence,
)
def _openai_score_answer(question: DecisionsIRScoreQuestion, answer: DecisionsIRScoreAnswer) -> OpenAIScoreAnswer:
probabilities: Final = {item.value: item.probability for item in answer.probabilities}
return OpenAIScoreAnswer(
name=question.name,
score=answer.score,
probabilities=tuple(
OpenAIScoreProbability(
value=index, label=decisions_text(level.label), probability=probabilities.get(index, 0.0)
)
for index, level in enumerate(question.levels)
),
confidence=answer.confidence,
)
def _openai_answer(question: DecisionsIRQuestion, answer: DecisionsIRAnswer) -> OpenAIDecisionAnswer:
match question, answer:
case DecisionsIRPredicateQuestion(), DecisionsIRPredicateAnswer():
return OpenAIPredicateAnswer(name=question.name, probability=answer.probability)
case DecisionsIRChoiceQuestion(), DecisionsIRChoiceAnswer():
return _openai_choice_answer(question, answer)
case DecisionsIRScoreQuestion(), DecisionsIRScoreAnswer():
return _openai_score_answer(question, answer)
case _:
return OpenAIRefusalAnswer(name=question.name)
def ir_to_openai_response(
response: DecisionsIRResponse, request: DecisionsIRRequest, requested_model: str
) -> OpenAIDecisionResponse:
return OpenAIDecisionResponse(
model=response.model or requested_model,
answers=tuple(
_openai_answer(question, answer)
for question, answer in zip(request.questions, response.answers, strict=True)
),
usage=OpenAIDecisionUsage(
input_tokens=response.usage.input_tokens,
input_tokens_details=OpenAIDecisionInputTokensDetails(
cached_tokens=response.usage.cached_tokens,
cache_write_tokens=response.usage.cache_write_tokens,
),
output_tokens=response.usage.output_tokens,
output_tokens_details=OpenAIDecisionOutputTokensDetails(reasoning_tokens=response.usage.reasoning_tokens),
total_tokens=response.usage.input_tokens + response.usage.output_tokens,
),
)
class OpenAIDecisionsConfig(BaseDecisionsConfig):
path = "/v1/decisions"
def get_default_api_base(self) -> str | None:
return "https://api.openai.com"
def resolve_api_base(self, api_base: str | None) -> str | None:
return (
api_base
or litellm.api_base
or get_secret_str("OPENAI_BASE_URL")
or get_secret_str("OPENAI_API_BASE")
or self.get_default_api_base()
)
def resolve_api_key(self, api_key: str | None) -> str | None:
return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
def transform_decisions_request(
self,
model: str,
request: DecisionsIRRequest,
custom_llm_provider: str,
) -> Mapping[str, object] | UnsupportedDecisionsRequest:
return ir_to_openai_request(self.request_model(model), request)
def parse_response(self, payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse:
return openai_response_to_ir(_OPENAI_RESPONSE_ADAPTER.validate_python(payload), request)

View file

@ -1,10 +1,10 @@
from typing import Final
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.decisions.transformation import JevCompatibleDecisionsEndpoint
OPENROUTER_DECISIONS_ENDPOINT: Final[JevCompatibleDecisionsEndpoint] = JevCompatibleDecisionsEndpoint(
default_api_base_value="https://openrouter.ai/api",
path="/alpha/decisions",
api_key_env=("OPENROUTER_API_KEY",),
api_base_env="OPENROUTER_API_BASE",
)
class OpenRouterDecisionsConfig(BaseDecisionsConfig):
path = "/alpha/decisions"
api_key_env = ("OPENROUTER_API_KEY",)
api_base_env = ("OPENROUTER_API_BASE",)
def get_default_api_base(self) -> str | None:
return "https://openrouter.ai/api"

View file

@ -1,10 +1,10 @@
from typing import Final
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.decisions.transformation import JevCompatibleDecisionsEndpoint
PERPLEXITY_DECISIONS_ENDPOINT: Final[JevCompatibleDecisionsEndpoint] = JevCompatibleDecisionsEndpoint(
default_api_base_value="https://api.perplexity.ai",
path="/v1/decisions",
api_key_env=("PERPLEXITYAI_API_KEY", "PERPLEXITY_API_KEY"),
api_base_env="PERPLEXITY_API_BASE",
)
class PerplexityDecisionsConfig(BaseDecisionsConfig):
path = "/v1/decisions"
api_key_env = ("PERPLEXITYAI_API_KEY", "PERPLEXITY_API_KEY")
api_base_env = ("PERPLEXITY_API_BASE",)
def get_default_api_base(self) -> str | None:
return "https://api.perplexity.ai"

View file

@ -1,11 +1,7 @@
from typing import Final
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.decisions.transformation import JevCompatibleDecisionsEndpoint
STRANDS_DECIDER_DECISIONS_ENDPOINT: Final[JevCompatibleDecisionsEndpoint] = JevCompatibleDecisionsEndpoint(
default_api_base_value=None,
path="/v1/systemone",
api_key_env=("STRANDS_DECIDER_API_KEY",),
api_base_env="STRANDS_DECIDER_API_BASE",
api_key_required=False,
)
class StrandsDeciderDecisionsConfig(BaseDecisionsConfig):
api_key_env = ("STRANDS_DECIDER_API_KEY",)
api_base_env = ("STRANDS_DECIDER_API_BASE",)
api_key_required = False

View file

@ -1,10 +1,9 @@
from typing import Final
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.decisions.transformation import JevCompatibleDecisionsEndpoint
TYPESAFE_DECISIONS_ENDPOINT: Final[JevCompatibleDecisionsEndpoint] = JevCompatibleDecisionsEndpoint(
default_api_base_value="https://api.typesafe.ai",
path="/v1/systemone",
api_key_env=("TYPESAFE_API_KEY",),
api_base_env="TYPESAFE_API_BASE",
)
class TypeSafeDecisionsConfig(BaseDecisionsConfig):
api_key_env = ("TYPESAFE_API_KEY",)
api_base_env = ("TYPESAFE_API_BASE",)
def get_default_api_base(self) -> str | None:
return "https://api.typesafe.ai"

Some files were not shown because too many files have changed in this diff Show more