mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
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:
commit
22003cf27c
685 changed files with 33346 additions and 8733 deletions
6
.github/workflows/helm_unit_test.yml
vendored
6
.github/workflows/helm_unit_test.yml
vendored
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/end_user/",
|
||||
"/sso/",
|
||||
"/liteadmin/slack/connect/",
|
||||
"/moyai/connect/",
|
||||
"/login",
|
||||
"/v2/login",
|
||||
"/v3/login",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
model_list: []
|
||||
|
||||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
tracing:
|
||||
|
|
|
|||
73
deploy/lens/configure.py
Normal file
73
deploy/lens/configure.py
Normal 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()
|
||||
52
deploy/lens/test_configure.py
Normal file
52
deploy/lens/test_configure.py
Normal 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()
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -62,6 +62,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/rerank",
|
||||
"/v1/decisions",
|
||||
"/decisions",
|
||||
"/v1/systemone",
|
||||
"/systemone",
|
||||
"/v1/ocr",
|
||||
"/ocr",
|
||||
"/v1/rag/",
|
||||
|
|
|
|||
|
|
@ -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 -}}
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
98
helm/litellm-helm/templates/lens/clickhouse.yaml
Normal file
98
helm/litellm-helm/templates/lens/clickhouse.yaml
Normal 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 }}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
27
helm/litellm-helm/templates/lens/secrets.yaml
Normal file
27
helm/litellm-helm/templates/lens/secrets.yaml
Normal 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 }}
|
||||
84
helm/litellm-helm/tests/lens_endpoint_defaults_tests.yaml
Normal file
84
helm/litellm-helm/tests/lens_endpoint_defaults_tests.yaml
Normal 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
|
||||
45
helm/litellm-helm/tests/lens_saved_secrets_tests.yaml
Normal file
45
helm/litellm-helm/tests/lens_saved_secrets_tests.yaml
Normal 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=
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
160
helm/litellm-helm/tests/lens_setup_tests.yaml
Normal file
160
helm/litellm-helm/tests/lens_setup_tests.yaml
Normal 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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 -}}
|
||||
|
|
|
|||
98
helm/litellm/templates/lens/clickhouse.yaml
Normal file
98
helm/litellm/templates/lens/clickhouse.yaml
Normal 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 }}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
27
helm/litellm/templates/lens/secrets.yaml
Normal file
27
helm/litellm/templates/lens/secrets.yaml
Normal 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 }}
|
||||
86
helm/litellm/tests/lens_endpoint_defaults_tests.yaml
Normal file
86
helm/litellm/tests/lens_endpoint_defaults_tests.yaml
Normal 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
|
||||
47
helm/litellm/tests/lens_saved_secrets_tests.yaml
Normal file
47
helm/litellm/tests/lens_saved_secrets_tests.yaml
Normal 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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
162
helm/litellm/tests/lens_setup_tests.yaml
Normal file
162
helm/litellm/tests/lens_setup_tests.yaml
Normal 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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -42,5 +42,6 @@ prettyplease = "0.2"
|
|||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
wiremock.workspace = true
|
||||
uuid.workspace = true
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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>(
|
||||
|
|
|
|||
|
|
@ -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> {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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>(),
|
||||
|
|
|
|||
408
litellm-rust/crates/traces-clickhouse/tests/lens_feedback.rs
Normal file
408
litellm-rust/crates/traces-clickhouse/tests/lens_feedback.rs
Normal 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, ¶meters).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(())
|
||||
}
|
||||
|
|
@ -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") {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,9 @@ pub enum ReadQuery {
|
|||
Sample,
|
||||
Content,
|
||||
Evidence,
|
||||
FeedbackTarget,
|
||||
Feedback,
|
||||
FeedbackSummary,
|
||||
}
|
||||
|
||||
impl ReadQuery {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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"}}),
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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[
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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="",
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
from .transformation import DecisionsProviderConfig, JevCompatibleDecisionsEndpoint
|
||||
from .transformation import BaseDecisionsConfig
|
||||
|
||||
__all__ = ["DecisionsProviderConfig", "JevCompatibleDecisionsEndpoint"]
|
||||
__all__ = ["BaseDecisionsConfig"]
|
||||
|
|
|
|||
247
litellm/llms/base_llm/decisions/systemone.py
Normal file
247
litellm/llms/base_llm/decisions/systemone.py
Normal 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,
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
320
litellm/llms/openai/decisions/transformation.py
Normal file
320
litellm/llms/openai/decisions/transformation.py
Normal 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)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue