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

This commit is contained in:
yucheng 2026-10-01 18:12:45 +00:00
commit 6454a525a0
280 changed files with 6927 additions and 2536 deletions

View file

@ -107,7 +107,7 @@ legacy_paths() {
echo tests/unit/proxy/test_update_spend.py
echo tests/unit/skills/test_skills_db.py ;;
proxy-db-endpoints-and-responses)
echo tests/unit/proxy/engine
echo tests/unit/proxy/lens
echo tests/unit/proxy/auth/test_models_fallback_endpoint.py
echo tests/unit/proxy/common_utils/test_check_batch_cost.py
echo tests/unit/proxy/common_utils/test_check_responses_cost.py

View file

@ -11,6 +11,7 @@ UNSUPPORTED: Final = re.compile(
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
r"|^tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e\.py$"
r"|^tests/e2e/logging/test_langsmith_batch_serialization_e2e\.py$"
r"|^tests/e2e/logging/test_s3_log_e2e\.py$"
r"|^tests/e2e/secret_manager/"
)
HARNESS: Final = re.compile(

View file

@ -5,13 +5,13 @@ on:
branches: [main, litellm_oss_branch, "litellm_**"]
paths:
- deploy/lens/**
- litellm/proxy/engine/**
- litellm/proxy/lens/**
- .github/workflows/lens-worker.yml
push:
branches: [main, litellm_agent_engine]
branches: [main]
paths:
- deploy/lens/**
- litellm/proxy/engine/**
- litellm/proxy/lens/**
- .github/workflows/lens-worker.yml
workflow_dispatch:
@ -41,8 +41,8 @@ jobs:
--security-opt no-new-privileges --entrypoint python \
lens-worker:${{ github.sha }} -c '
import os
import engine.worker
from engine.trace_store import trace_store
import lens.worker
from lens.trace_store import trace_store
assert os.getuid() == 65532
with trace_store() as store:
assert store.count() == 0

View file

@ -140,42 +140,46 @@ jobs:
- shard: proxy-endpoints
artifact-name: proxy-endpoints
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/unit/proxy/analytics_endpoints
tests/unit/proxy/management_endpoints
tests/test_litellm/proxy/list_api
tests/test_litellm/proxy/memory
tests/unit/proxy/list_api
tests/unit/proxy/memory
tests/unit/proxy/guardrails
tests/unit/proxy/management_helpers
--ignore=tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py
--ignore=tests/unit/proxy/management_endpoints/test_key_generate_prisma.py
--ignore=tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py
--ignore=tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
tests/test_litellm/proxy/anthropic_endpoints
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/container_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/ocr_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/credential_endpoints
tests/test_litellm/proxy/discovery_endpoints
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/shutdown
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/rerank_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
tests/test_litellm/proxy/utils
--ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py
--ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py
--ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py
--ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py
tests/unit/proxy/anthropic_endpoints
tests/unit/proxy/google_endpoints
tests/unit/proxy/openai_files_endpoint
tests/unit/proxy/batches_endpoints
tests/unit/proxy/container_endpoints
tests/unit/proxy/fine_tuning_endpoints
tests/unit/proxy/vector_store_files_endpoints
tests/unit/proxy/video_endpoints
tests/unit/proxy/response_api_endpoints
tests/unit/proxy/image_endpoints
tests/unit/proxy/ocr_endpoints
tests/unit/proxy/vector_store_endpoints
tests/unit/proxy/agent_endpoints
tests/unit/proxy/a2a
tests/unit/proxy/credential_endpoints
tests/unit/proxy/discovery_endpoints
tests/unit/proxy/health_endpoints
tests/unit/proxy/shutdown
tests/unit/proxy/public_endpoints
tests/unit/proxy/prompts
tests/unit/proxy/rag_endpoints
tests/unit/proxy/rerank_endpoints
tests/unit/proxy/realtime_endpoints
tests/unit/proxy/ui_crud_endpoints
tests/unit/proxy/config_resolvers
tests/unit/proxy/utils
workers: 4
reruns: 2
timeout-minutes: 20

View file

@ -327,7 +327,7 @@ test-unit-proxy-core: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/auth tests/unit/proxy/client tests/test_litellm/proxy/db tests/unit/proxy/hooks tests/unit/proxy/policy_engine --tb=short -vv -n 4 --durations=20
test-unit-proxy-misc: install-test-deps
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py tests/unit/proxy/test_proxy_server_endpoints_and_startup.py tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/unit/proxy/agent_endpoints tests/unit/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/unit/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/unit/proxy/google_endpoints tests/unit/proxy/health_endpoints tests/unit/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/unit/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/unit/proxy/prompts tests/unit/proxy/public_endpoints tests/unit/proxy/response_api_endpoints tests/unit/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/unit/proxy/ui_crud_endpoints tests/unit/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py tests/unit/proxy/test_proxy_server_endpoints_and_startup.py tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py --ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py --ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py --ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py --ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py --tb=short -vv -n 4 --durations=20
test-unit-integrations: install-test-deps
$(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20

View file

@ -81,7 +81,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
# Spend / analytics
"/spend/",
"/analytics/",
"/engine/",
"/lens/",
"/v1/traces",
"/global/",
"/user_agent",
@ -146,7 +146,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
{
"/",
"/routes",
"/engine",
"/lens",
"/openapi.json",
"/docs",
"/docs/oauth2-redirect",

View file

@ -1,6 +1,6 @@
FROM python:3.12-slim
WORKDIR /app
RUN pip install --no-cache-dir httpx==0.28.1 pydantic==2.11.7
COPY litellm/proxy/engine/__init__.py litellm/proxy/engine/models.py litellm/proxy/engine/trace_store.py litellm/proxy/engine/analysis.py litellm/proxy/engine/worker.py /app/engine/
COPY litellm/proxy/lens/__init__.py litellm/proxy/lens/models.py litellm/proxy/lens/trace_store.py litellm/proxy/lens/analysis.py litellm/proxy/lens/worker.py /app/lens/
USER 65532:65532
CMD ["python", "-m", "engine.worker"]
CMD ["python", "-m", "lens.worker"]

View file

@ -1,8 +1,8 @@
**
!litellm/
!litellm/proxy/
!litellm/proxy/engine/
!litellm/proxy/engine/__init__.py
!litellm/proxy/engine/models.py
!litellm/proxy/engine/analysis.py
!litellm/proxy/engine/worker.py
!litellm/proxy/lens/
!litellm/proxy/lens/__init__.py
!litellm/proxy/lens/models.py
!litellm/proxy/lens/analysis.py
!litellm/proxy/lens/worker.py

View file

@ -74,7 +74,7 @@ V1 requires ClickHouse for both sources. It does not reconstruct sessions from u
The UI and API use the same scan lifecycle. Authenticate with a proxy administrator credential for writes, or a proxy-admin viewer credential for reads. Worker credentials are only for worker operations
```bash
curl "$LITELLM_URL/engine" -H "Authorization: Bearer $LITELLM_API_KEY" \
curl "$LITELLM_URL/lens" -H "Authorization: Bearer $LITELLM_API_KEY" \
-H 'Content-Type: application/json' -d '{
"name": "Research quality", "model": "your-model-alias",
"context": "Answer the requested question using cited, retrieved evidence.",
@ -83,14 +83,14 @@ curl "$LITELLM_URL/engine" -H "Authorization: Bearer $LITELLM_API_KEY" \
"enabled": true, "interval_minutes": 1440, "monthly_budget": 50
}'
curl "$LITELLM_URL/engine/$LENS_ID/runs" -X POST \
curl "$LITELLM_URL/lens/$LENS_ID/runs" -X POST \
-H "Authorization: Bearer $LITELLM_API_KEY" -H 'Content-Type: application/json' -d '{}'
curl "$LITELLM_URL/engine/$LENS_ID/runs?offset=0" -H "Authorization: Bearer $LITELLM_API_KEY"
curl "$LITELLM_URL/engine/$LENS_ID/runs/$BATCH_ID" -H "Authorization: Bearer $LITELLM_API_KEY"
curl "$LITELLM_URL/lens/$LENS_ID/runs?offset=0" -H "Authorization: Bearer $LITELLM_API_KEY"
curl "$LITELLM_URL/lens/$LENS_ID/runs/$BATCH_ID" -H "Authorization: Bearer $LITELLM_API_KEY"
```
Creation queues the first batch. Posting to `/engine/{id}/runs` queues another, or returns the existing active batch. The run response contains its ID under `jobs[0].id`. Poll the batch URL for status, findings and assessments. List responses omit large result payloads; request a batch to retrieve them. Supply an optional complete `settings` object on the runs POST for a one-off override; the saved lens stays unchanged. Selection accepts `team_id`, exact `filters`, and opaque `execution_ids` returned by `/engine/preview/sample`. Preview accepts `offset` and `as_of` to keep the time window fixed while paging. Feedback uses `PATCH /engine/{id}/findings/{finding_id}` with `status` and `reason`
Creation queues the first batch. Posting to `/lens/{id}/runs` queues another, or returns the existing active batch. The run response contains its ID under `jobs[0].id`. Poll the batch URL for status, findings and assessments. List responses omit large result payloads; request a batch to retrieve them. Supply an optional complete `settings` object on the runs POST for a one-off override; the saved lens stays unchanged. Selection accepts `team_id`, exact `filters`, and opaque `execution_ids` returned by `/lens/preview/sample`. Preview accepts `offset` and `as_of` to keep the time window fixed while paging. Feedback uses `PATCH /lens/{id}/findings/{finding_id}` with `status` and `reason`
## Quality evaluation
@ -107,3 +107,11 @@ Set `LITELLM_API_KEY` privately. This makes paid model calls. Inspect missed and
The worker uses temporary disk space for trace content while reviewing it, and removes those files after each review. The Docker command supplies a writable temporary mount while keeping the application filesystem read-only
To check that accepted behavior stays accepted without hiding new problems, run the evaluator with `--dataset tests/proxy_behavior/lens/feedback_cases.json`. Reports include elapsed time, model call count, reported cost when the proxy provides it, missed checks, unexpected checks, and inconclusive candidates
## Upgrading from the original Lens API
The Lens API now uses `/lens` instead of `/engine`, list responses use `lenses`, and worker claims use `lens_id`. Upgrade the proxy and recreate every worker with the image shown by the upgraded dashboard before starting new scans. Update API clients to the new paths and response fields. Old worker images cannot poll the renamed API
Stop workers and let active scans finish before upgrading. Deploy proxy instances together: older proxies cannot use the renamed database tables. The schema migration renames the three Lens tables and the run-history identifier column in place, preserving saved investigations, findings, history, worker credentials, and billing assignments. Existing migration files retain their original names and checksums
Upgrades using `--use_prisma_db_push` stop before schema changes if any legacy Lens table exists, preventing Prisma from dropping saved data. Apply `litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql` to the configured database schema before retrying. Deployments already using migration history can instead start without `--use_prisma_db_push` to apply the shipped migration normally. Fresh databases and databases already using the renamed tables can continue using database push

View file

@ -1,6 +1,6 @@
services:
lens-worker:
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:c41e932eaf3e4efbcaf8cc5027c7e93021e5b2823f21cb8785cd107e37b91c9a}
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:a8e8731d954916594eea462969946b9292fb771681ff515a9fd296b53f856c77}
environment:
LITELLM_URL: ${LITELLM_URL:?Set the URL reachable from this container}
LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:?Create a worker credential in the Lens UI}

Binary file not shown.

Before

Width:  |  Height:  |  Size: 95 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 6.9 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 89 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 80 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 70 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 132 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 59 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 54 KiB

View file

@ -0,0 +1,18 @@
DO $$
BEGIN
ALTER TABLE IF EXISTS "LiteLLM_Engine" RENAME TO "LiteLLM_Lens";
ALTER TABLE IF EXISTS "LiteLLM_EngineRun" RENAME TO "LiteLLM_LensRun";
ALTER TABLE IF EXISTS "LiteLLM_EngineWorker" RENAME TO "LiteLLM_LensWorker";
IF EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_LensRun"')
AND attname = 'engine_id' AND NOT attisdropped
) THEN
ALTER TABLE "LiteLLM_LensRun" RENAME COLUMN "engine_id" TO "lens_id";
END IF;
ALTER INDEX IF EXISTS "LiteLLM_Engine_pkey" RENAME TO "LiteLLM_Lens_pkey";
ALTER INDEX IF EXISTS "LiteLLM_EngineRun_pkey" RENAME TO "LiteLLM_LensRun_pkey";
ALTER INDEX IF EXISTS "LiteLLM_EngineWorker_pkey" RENAME TO "LiteLLM_LensWorker_pkey";
ALTER INDEX IF EXISTS "LiteLLM_EngineWorker_token_hash_key" RENAME TO "LiteLLM_LensWorker_token_hash_key";
ALTER INDEX IF EXISTS "LiteLLM_EngineRun_engine_id_created_at_idx" RENAME TO "LiteLLM_LensRun_lens_id_created_at_idx";
END $$;

View file

@ -1895,22 +1895,22 @@ model LiteLLM_WorkflowMessage {
@@index([run_id])
}
model LiteLLM_Engine {
model LiteLLM_Lens {
id String @id
version Int @default(0)
data Json
}
model LiteLLM_EngineRun {
model LiteLLM_LensRun {
id String @id
engine_id String
lens_id String
created_at DateTime
data Json
@@index([engine_id, created_at])
@@index([lens_id, created_at])
}
model LiteLLM_EngineWorker {
model LiteLLM_LensWorker {
id String @id
token_hash String @unique
data Json

View file

@ -590,6 +590,36 @@ class ProxyExtrasDBManager:
f"Failed to resolve migration {migration_name}: {e.stderr}"
)
@staticmethod
def raise_if_lens_rename_pending() -> None:
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return
try:
import psycopg
except ImportError as exc:
raise RuntimeError("Install psycopg to verify Lens data safety before prisma db push.") from exc
try:
with psycopg.connect(
ProxyExtrasDBManager._strip_prisma_query_params(database_url), connect_timeout=10, autocommit=True
) as connection:
legacy: Final = connection.execute(
"SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid=c.relnamespace "
"WHERE n.nspname=%s AND c.relname IN ('LiteLLM_Engine', 'LiteLLM_EngineRun', 'LiteLLM_EngineWorker') "
"LIMIT 1",
(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",),
).fetchone()
except psycopg.Error as exc:
raise RuntimeError(
"Cannot verify Lens data safety; refusing prisma db push. Check database connectivity and psycopg installation."
) from exc
if legacy is not None:
raise RuntimeError(
"Legacy Lens tables exist. prisma db push would drop saved Lens data. "
"Apply the shipped 20261001100000_rename_lens migration to this database schema before retrying. "
"Deployments using migration history can upgrade without --use_prisma_db_push instead."
)
@staticmethod
def spend_logs_is_partitioned() -> bool:
"""True when the connected database's LiteLLM_SpendLogs is a
@ -895,6 +925,7 @@ class ProxyExtrasDBManager:
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
if not use_migrate:
ProxyExtrasDBManager.raise_if_lens_rename_pending()
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
original_dir = os.getcwd()
@ -1398,6 +1429,7 @@ class ProxyExtrasDBManager:
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
# Use prisma db push with increased timeout
ProxyExtrasDBManager.raise_if_lens_rename_pending()
prisma_toolchain.run_prisma(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),

View file

@ -4,6 +4,10 @@ version = "0.4.103"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
dependencies = [
"psycopg>=3.2,<4.0",
"psycopg-binary>=3.2,<4.0",
]
license = "MIT"
license-files = ["LICENSE"]
authors = [

View file

@ -69,6 +69,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY"
S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY"
MAX_FILE_LIST_LIMIT: Final = 10000
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))

View file

@ -498,6 +498,13 @@
"ui_name": "Log Prompts Only",
"description": "Log request messages to S3 but drop the model response from each logged object",
"required": false
},
"s3_partition_granularity": {
"type": "select",
"ui_name": "Folder Partitioning",
"description": "day writes one folder per date, hour adds an hour folder below each date (s3_v2 only)",
"options": ["day", "hour"],
"required": false
}
},
"description": "S3 Bucket (AWS) Logging Integration"

View file

@ -16,8 +16,10 @@ from litellm.constants import (
MAX_S3_OBJECT_KEY_BYTES,
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES,
S3_LOG_PROMPTS_ONLY_ENV_VAR,
S3_PARTITION_GRANULARITY_ENV_VAR,
S3_PREFIX_DIGEST_CHARS,
)
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
from litellm.types.utils import StandardLoggingPayload
_S3_BOOL: Final = TypeAdapter(bool)
@ -36,6 +38,18 @@ def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] |
return True
def resolve_s3_partition_granularity(
configured: object, environ: Mapping[str, str] | None = None
) -> S3PartitionGranularity:
env: Final = os.environ if environ is None else environ
raw: Final = env.get(S3_PARTITION_GRANULARITY_ENV_VAR) if configured is None else configured
if raw == "hour":
return "hour"
if raw is not None and raw not in ("", "day"):
verbose_logger.warning("s3 logging: s3_partition_granularity=%r is not one of day, hour, using day", raw)
return "day"
def _resolve_positive_int(setting: str, configured: object, fallback: int, *, reject_bool: bool) -> int:
if configured is None or configured == "":
return fallback
@ -371,10 +385,11 @@ def get_s3_object_key(
prefix: str,
start_time: datetime,
s3_file_name: str,
partition_granularity: S3PartitionGranularity = "day",
) -> str:
sanitized_s3_file_name: Final = s3_file_name.replace("/", "_").replace(":", "_")
configured_prefix: Final = (s3_path.rstrip("/") + "/" if s3_path else "") + prefix
date_segment: Final = start_time.strftime("%Y-%m-%d") + "/"
date_segment: Final = start_time.strftime("%Y-%m-%d/%H/" if partition_granularity == "hour" else "%Y-%m-%d/")
# we need the s3 key to include the time, so we log cache hits too
s3_object_key: Final = configured_prefix + date_segment + sanitized_s3_file_name + ".json"
if len(s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES:

View file

@ -9,6 +9,7 @@ NOTE 1: S3 does not provide a BATCH PUT API endpoint; by default each element is
import asyncio
import contextvars
import logging
import os
import re
import time
from collections.abc import Awaitable, Callable, Mapping
@ -28,6 +29,7 @@ from litellm.constants import (
DEFAULT_S3_FLUSH_INTERVAL_SECONDS,
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY,
DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
S3_PARTITION_GRANULARITY_ENV_VAR,
)
from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample
from litellm.integrations.s3 import (
@ -42,6 +44,7 @@ from litellm.integrations.s3 import (
resolve_s3_max_concurrent_uploads,
resolve_s3_max_queue_size,
resolve_s3_max_retry_age_seconds,
resolve_s3_partition_granularity,
resolve_sse_params,
)
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
@ -53,7 +56,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
from litellm.types.integrations.s3_v2 import S3PartitionGranularity, s3BatchLoggingElement
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
from .custom_batch_logger import CustomBatchLogger
@ -119,6 +122,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
_upload_limiter: asyncio.Semaphore | AdaptiveConcurrencyLimiter | None = None
s3_drop_on_terminal_error: bool = True
s3_max_retry_age_seconds: int | None = 3600
s3_partition_granularity: object = None
_partition_granularity_cache: tuple[object, S3PartitionGranularity] | None = None
def __init__(
self,
@ -147,6 +152,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_server_side_encryption: str | None = None,
s3_sse_kms_key_id: str | None = None,
s3_log_prompts_only: bool | None = None,
s3_partition_granularity: str | None = None,
s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
s3_max_queue_size: int | None = None,
s3_max_retry_age_seconds: int | None = 3600,
@ -195,6 +201,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_server_side_encryption=s3_server_side_encryption,
s3_sse_kms_key_id=s3_sse_kms_key_id,
s3_log_prompts_only=s3_log_prompts_only,
s3_partition_granularity=s3_partition_granularity,
s3_max_concurrent_uploads=s3_max_concurrent_uploads,
s3_max_queue_size=s3_max_queue_size,
s3_max_retry_age_seconds=s3_max_retry_age_seconds,
@ -271,6 +278,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_server_side_encryption: str | None = None,
s3_sse_kms_key_id: str | None = None,
s3_log_prompts_only: bool | None = None,
s3_partition_granularity: str | None = None,
s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
s3_max_queue_size: int | None = None,
s3_max_retry_age_seconds: int | None = 3600,
@ -331,6 +339,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only
)
self.s3_partition_granularity = (
params.get("s3_partition_granularity") if s3_partition_granularity is None else s3_partition_granularity
)
self._partition_granularity_cache = None
self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
params.get("s3_server_side_encryption") or s3_server_side_encryption,
params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id,
@ -482,6 +495,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
"audit_logs/",
now,
f"{now.strftime('%H-%M-%S')}_{audit_log_id}",
partition_granularity=self.resolve_partition_granularity(),
)
element: Final = s3BatchLoggingElement(
@ -758,6 +772,19 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
),
)
def resolve_partition_granularity(self) -> S3PartitionGranularity:
raw: Final = (
os.environ.get(S3_PARTITION_GRANULARITY_ENV_VAR)
if self.s3_partition_granularity is None
else self.s3_partition_granularity
)
cached: Final = self._partition_granularity_cache
if cached is not None and cached[0] == raw:
return cached[1]
resolved: Final = resolve_s3_partition_granularity(raw)
self._partition_granularity_cache = (raw, resolved)
return resolved
def create_s3_batch_logging_element(
self,
start_time: datetime,
@ -803,11 +830,27 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
prefix_path,
s3_file_name,
)
s3_object_key: Final = get_s3_object_key(
s3_path=cast(str | None, self.s3_path) or "",
prefix=prefix_path,
start_time=start_time,
s3_file_name=s3_file_name,
def object_key(partition_granularity: S3PartitionGranularity) -> str:
return get_s3_object_key(
s3_path=cast(str | None, self.s3_path) or "",
prefix=prefix_path,
start_time=start_time,
s3_file_name=s3_file_name,
partition_granularity=partition_granularity,
)
metadata: Final = standard_logging_payload.get("metadata")
cold_storage_object_key: Final = (
metadata.get("cold_storage_object_key")
if metadata is not None and litellm.cold_storage_custom_logger == "s3_v2"
else None
)
s3_object_key: Final = (
cold_storage_object_key
if cold_storage_object_key is not None
and cold_storage_object_key in (object_key("day"), object_key("hour"))
else object_key(self.resolve_partition_granularity())
)
verbose_logger.debug("s3_object_key=%s", s3_object_key)

View file

@ -117,6 +117,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.integrations.s3_v2 import S3PartitionGranularity
from litellm.types.interactions import (
InteractionsAPIResponse,
InteractionsAPIStreamingResponse,
@ -6059,6 +6060,7 @@ class StandardLoggingPayloadSetup:
# Get the actual s3_path from the configured cold storage logger instance
s3_path = "" # default value
partition_granularity: S3PartitionGranularity = "day"
# Try to get the actual logger instance from the logger name
try:
@ -6067,6 +6069,8 @@ class StandardLoggingPayloadSetup:
)
if custom_logger and hasattr(custom_logger, "s3_path") and getattr(custom_logger, "s3_path"):
s3_path = getattr(custom_logger, "s3_path")
if isinstance(custom_logger, S3V2Logger):
partition_granularity = custom_logger.resolve_partition_granularity()
except Exception:
# If any error occurs in getting the logger instance, use default empty s3_path
pass
@ -6076,6 +6080,7 @@ class StandardLoggingPayloadSetup:
prefix="", # Don't split by team alias for cold storage
start_time=start_time,
s3_file_name=s3_file_name,
partition_granularity=partition_granularity,
)
return s3_object_key

View file

@ -34666,52 +34666,6 @@
]
}
},
"/engine/workers/register": {
"post": {
"operationId": "register_worker_engine_workers_register_post",
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/WorkerName"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/WorkerCreated"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Register Worker",
"tags": [
"mcp_discoverable"
]
}
},
"/guardrails/register": {
"post": {
"description": "Register a guardrail for onboarding (team submission).\n\nAccepts a guardrail config in the\n[Generic Guardrail API](https://docs.litellm.ai/docs/adding_provider/generic_guardrail_api) format.\nThe submission is stored with status `pending_review` until an admin approves it.",
@ -34804,6 +34758,52 @@
]
}
},
"/lens/workers/register": {
"post": {
"operationId": "register_worker_lens_workers_register_post",
"requestBody": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/WorkerName"
}
}
},
"required": true
},
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/WorkerCreated"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Register Worker",
"tags": [
"mcp_discoverable"
]
}
},
"/register": {
"post": {
"operationId": "register_client_register_post",

View file

@ -520,16 +520,16 @@ class LiteLLMRoutes(enum.Enum):
"/v1/rag/ingest",
"/rag/query",
"/v1/rag/query",
"/engine",
"/engine/{engine_id}",
"/engine/{engine_id}/runs",
"/engine/{engine_id}/runs/{job_id}",
"/engine/{engine_id}/executions/{execution_id}",
"/engine/{engine_id}/cancel",
"/engine/{engine_id}/findings/{finding_id}",
"/engine/preview/sample",
"/engine/workers/register",
"/engine/workers/{worker_id}",
"/lens",
"/lens/{lens_id}",
"/lens/{lens_id}/runs",
"/lens/{lens_id}/runs/{job_id}",
"/lens/{lens_id}/executions/{execution_id}",
"/lens/{lens_id}/cancel",
"/lens/{lens_id}/findings/{finding_id}",
"/lens/preview/sample",
"/lens/workers/register",
"/lens/workers/{worker_id}",
"/v1/traces",
"/v1/traces/{trace_id}",
"/v1/traces/{trace_id}/spans/{span_id}",
@ -3957,6 +3957,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
"AWS_SECRET_ACCESS_KEY",
"AWS_REGION_NAME",
"S3_LOG_PROMPTS_ONLY",
"S3_PARTITION_GRANULARITY",
],
)

View file

@ -953,6 +953,9 @@ class PrismaManager:
verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
return False
from litellm_proxy_extras.utils import ProxyExtrasDBManager
ProxyExtrasDBManager.raise_if_lens_rename_pending()
PrismaManager._raise_if_partitioned_spend_logs()
run_prisma(
[

View file

@ -13,17 +13,17 @@ from pydantic import AwareDatetime, BaseModel, Field, TypeAdapter
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
from litellm.proxy.engine.billing import validate_key
from litellm.proxy.engine.models import (
from litellm.proxy.lens.billing import validate_key
from litellm.proxy.lens.models import (
Claim,
Engine,
EngineList,
EngineSettings,
Execution,
ExecutionContent,
FindingDraft,
FindingUpdate,
Job,
Lens,
LensList,
LensSettings,
ModelRequest,
ModelResult,
Progress,
@ -34,9 +34,9 @@ from litellm.proxy.engine.models import (
Worker,
WorkerCreated,
)
from litellm.proxy.engine.repository import EngineRepository, WriterDatabase
from litellm.proxy.engine.sources import SourceReader, parse_execution
from litellm.proxy.engine.state import (
from litellm.proxy.lens.repository import LensRepository, WriterDatabase
from litellm.proxy.lens.sources import SourceReader, parse_execution
from litellm.proxy.lens.state import (
can_access,
claim_job,
current_job,
@ -46,17 +46,17 @@ from litellm.proxy.engine.state import (
snapshot_finding,
)
router: Final = APIRouter(prefix="/engine", tags=["Lens"]) # mutable-ok: FastAPI requires list
router: Final = APIRouter(prefix="/lens", tags=["Lens"]) # mutable-ok: FastAPI requires list
_bearer: Final = HTTPBearer()
Auth: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
def repository() -> EngineRepository:
def repository() -> LensRepository:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(503, "Lens needs a connected Postgres database")
return EngineRepository(WriterDatabase(writer_wrapper(prisma_client.db)))
return LensRepository(WriterDatabase(writer_wrapper(prisma_client.db)))
def source_reader() -> SourceReader:
@ -73,11 +73,11 @@ def user_scope(auth: UserAPIKeyAuth, write: bool = False) -> Scope:
raise HTTPException(403, "Lens requires proxy administrator access")
async def get_engine(engine_id: str, scope: Scope) -> Engine:
engine: Final = await repository().get(engine_id)
if engine is None or not can_access(scope, engine.scope):
async def get_lens(lens_id: str, scope: Scope) -> Lens:
lens: Final = await repository().get(lens_id)
if lens is None or not can_access(scope, lens.scope):
raise HTTPException(404, "Lens not found")
return engine
return lens
async def worker_auth(credentials: Annotated[HTTPAuthorizationCredentials, Depends(_bearer)]) -> Worker:
@ -90,9 +90,9 @@ async def worker_auth(credentials: Annotated[HTTPAuthorizationCredentials, Depen
WorkerAuth: TypeAlias = Annotated[Worker, Depends(worker_auth)]
async def assigned(engine_id: str, job_id: str, worker: Worker) -> tuple[Engine, Job]:
engine: Final = await get_engine(engine_id, worker.scope)
job: Final = current_job(engine)
async def assigned(lens_id: str, job_id: str, worker: Worker) -> tuple[Lens, Job]:
lens: Final = await get_lens(lens_id, worker.scope)
job: Final = current_job(lens)
if (
job is None
or job.id != job_id
@ -102,16 +102,16 @@ async def assigned(engine_id: str, job_id: str, worker: Worker) -> tuple[Engine,
or job.lease_until <= datetime.now(timezone.utc)
):
raise HTTPException(409, "This worker no longer owns the job")
return engine, job
return lens, job
def required(engine: Engine | None) -> Engine:
if engine is None:
def required(lens: Lens | None) -> Lens:
if lens is None:
raise HTTPException(409, "Lens changed concurrently; retry the operation")
return engine
return lens
def validate_selection(settings: EngineSettings) -> None:
def validate_selection(settings: LensSettings) -> None:
for identity in settings.execution_ids:
try:
source, _, _, _ = parse_execution(identity)
@ -121,7 +121,7 @@ def validate_selection(settings: EngineSettings) -> None:
raise HTTPException(422, "Choose execution IDs returned by the activity preview")
def validate_model(settings: EngineSettings, auth: UserAPIKeyAuth) -> None:
def validate_model(settings: LensSettings, auth: UserAPIKeyAuth) -> None:
from litellm.proxy.proxy_server import llm_router
validate_selection(settings)
@ -137,24 +137,24 @@ def validate_model(settings: EngineSettings, auth: UserAPIKeyAuth) -> None:
raise HTTPException(403, "This key does not have access to the analysis model")
@router.get("", response_model=EngineList)
async def list_engines(auth: Auth) -> EngineList:
@router.get("", response_model=LensList)
async def list_lenses(auth: Auth) -> LensList:
from litellm.proxy import tracing_endpoints
scope: Final = user_scope(auth)
return EngineList(
engines=tuple(e for e in await repository().engines() if can_access(scope, e.scope)),
return LensList(
lenses=tuple(e for e in await repository().lenses() if can_access(scope, e.scope)),
workers=tuple(w for w in await repository().workers() if can_access(scope, w.scope)),
tracing_enabled=tracing_endpoints.receiver is not None,
)
@router.post("", response_model=Engine)
async def create_engine(settings: EngineSettings, auth: Auth) -> Engine:
@router.post("", response_model=Lens)
async def create_lens(settings: LensSettings, auth: Auth) -> Lens:
scope: Final = user_scope(auth, write=True)
validate_model(settings, auth)
now: Final = datetime.now(timezone.utc)
engine: Final = Engine(
lens: Final = Lens(
id=str(uuid4()),
scope=scope,
settings=settings,
@ -162,16 +162,16 @@ async def create_engine(settings: EngineSettings, auth: Auth) -> Engine:
next_run_at=now,
budget_month=now.strftime("%Y-%m"),
)
return await repository().create(queue_job(engine, now, str(uuid4())))
return await repository().create(queue_job(lens, now, str(uuid4())))
@router.put("/{engine_id}", response_model=Engine)
async def update_engine(engine_id: str, settings: EngineSettings, auth: Auth) -> Engine:
await get_engine(engine_id, user_scope(auth, write=True))
@router.put("/{lens_id}", response_model=Lens)
async def update_lens(lens_id: str, settings: LensSettings, auth: Auth) -> Lens:
await get_lens(lens_id, user_scope(auth, write=True))
validate_model(settings, auth)
return required(
await repository().update(
engine_id,
lens_id,
lambda e: e.model_copy(
update=MappingProxyType(
{
@ -184,47 +184,47 @@ async def update_engine(engine_id: str, settings: EngineSettings, auth: Auth) ->
)
@router.post("/{engine_id}/runs", response_model=Engine)
async def run_engine(engine_id: str, body: RunRequest, auth: Auth) -> Engine:
await get_engine(engine_id, user_scope(auth, write=True))
@router.post("/{lens_id}/runs", response_model=Lens)
async def run_lens(lens_id: str, body: RunRequest, auth: Auth) -> Lens:
await get_lens(lens_id, user_scope(auth, write=True))
if body.settings is not None:
validate_model(body.settings, auth)
now: Final = datetime.now(timezone.utc)
job_id: Final = str(uuid4())
return required(
await repository().update(engine_id, lambda e: queue_job(e, now, job_id, body.lookback_hours, body.settings))
await repository().update(lens_id, lambda e: queue_job(e, now, job_id, body.lookback_hours, body.settings))
)
@router.get("/{engine_id}", response_model=Engine)
async def read_engine(engine_id: str, auth: Auth) -> Engine:
return await get_engine(engine_id, user_scope(auth))
@router.get("/{lens_id}", response_model=Lens)
async def read_lens(lens_id: str, auth: Auth) -> Lens:
return await get_lens(lens_id, user_scope(auth))
@router.get("/{engine_id}/runs", response_model=tuple[Job, ...])
async def list_runs(engine_id: str, auth: Auth, offset: int = Query(default=0, ge=0)) -> tuple[Job, ...]:
await get_engine(engine_id, user_scope(auth))
@router.get("/{lens_id}/runs", response_model=tuple[Job, ...])
async def list_runs(lens_id: str, auth: Auth, offset: int = Query(default=0, ge=0)) -> tuple[Job, ...]:
await get_lens(lens_id, user_scope(auth))
return tuple(
j.model_copy(update=MappingProxyType({"sample": None, "findings": None, "assessments": ()}))
for j in await repository().jobs(engine_id, offset)
for j in await repository().jobs(lens_id, offset)
)
@router.get("/{engine_id}/runs/{job_id}", response_model=Job)
async def read_run(engine_id: str, job_id: str, auth: Auth) -> Job:
await get_engine(engine_id, user_scope(auth))
job: Final = await repository().job(engine_id, job_id)
@router.get("/{lens_id}/runs/{job_id}", response_model=Job)
async def read_run(lens_id: str, job_id: str, auth: Auth) -> Job:
await get_lens(lens_id, user_scope(auth))
job: Final = await repository().job(lens_id, job_id)
if job is None:
raise HTTPException(404, "Investigation not found")
return job
@router.post("/{engine_id}/cancel", response_model=Engine)
async def cancel_engine(engine_id: str, auth: Auth) -> Engine:
await get_engine(engine_id, user_scope(auth, write=True))
@router.post("/{lens_id}/cancel", response_model=Lens)
async def cancel_lens(lens_id: str, auth: Auth) -> Lens:
await get_lens(lens_id, user_scope(auth, write=True))
now: Final = datetime.now(timezone.utc)
def cancel(e: Engine) -> Engine:
def cancel(e: Lens) -> Lens:
job: Final = current_job(e)
if job is None:
return e
@ -235,15 +235,15 @@ async def cancel_engine(engine_id: str, auth: Auth) -> Engine:
update=MappingProxyType({"next_run_at": now + timedelta(minutes=e.settings.interval_minutes)})
)
return required(await repository().update(engine_id, cancel))
return required(await repository().update(lens_id, cancel))
@router.patch("/{engine_id}/findings/{finding_id}", response_model=Engine)
async def update_finding(engine_id: str, finding_id: str, body: FindingUpdate, auth: Auth) -> Engine:
await get_engine(engine_id, user_scope(auth, write=True))
@router.patch("/{lens_id}/findings/{finding_id}", response_model=Lens)
async def update_finding(lens_id: str, finding_id: str, body: FindingUpdate, auth: Auth) -> Lens:
await get_lens(lens_id, user_scope(auth, write=True))
return required(
await repository().update(
engine_id,
lens_id,
lambda e: e.model_copy(
update=MappingProxyType(
{
@ -260,7 +260,7 @@ async def update_finding(engine_id: str, finding_id: str, body: FindingUpdate, a
class Preview(BaseModel):
as_of: AwareDatetime | None = None
offset: int = Field(default=0, ge=0)
settings: EngineSettings
settings: LensSettings
lookback_hours: int = Field(default=24, ge=1, le=720)
@ -335,7 +335,7 @@ async def claim(worker: WorkerAuth, protocol_version: int = 1) -> Claim | None:
raise HTTPException(409, "Assign an analysis key to this worker in Lens setup")
now: Final = datetime.now(timezone.utc)
await repository().heartbeat(worker.id, now.isoformat())
for candidate in await repository().engines():
for candidate in await repository().lenses():
if not can_access(worker.scope, candidate.scope):
continue
if claimed := await claim_candidate(candidate, worker, now):
@ -343,12 +343,12 @@ async def claim(worker: WorkerAuth, protocol_version: int = 1) -> Claim | None:
return None
@router.post("/worker/{engine_id}/{job_id}/progress", response_model=bool)
async def progress(engine_id: str, job_id: str, body: Progress, worker: WorkerAuth) -> bool:
await assigned(engine_id, job_id, worker)
@router.post("/worker/{lens_id}/{job_id}/progress", response_model=bool)
async def progress(lens_id: str, job_id: str, body: Progress, worker: WorkerAuth) -> bool:
await assigned(lens_id, job_id, worker)
now: Final = datetime.now(timezone.utc)
def renew(e: Engine) -> Engine:
def renew(e: Lens) -> Lens:
job: Final = current_job(e)
if job is None or job.id != job_id or job.worker_id != worker.id:
return e
@ -361,21 +361,21 @@ async def progress(engine_id: str, job_id: str, body: Progress, worker: WorkerAu
),
)
required(await repository().update(engine_id, renew))
required(await repository().update(lens_id, renew))
await repository().heartbeat(worker.id, now.isoformat())
return True
@router.get("/worker/{engine_id}/{job_id}/sample", response_model=Sample)
async def sample(engine_id: str, job_id: str, worker: WorkerAuth) -> Sample:
engine, job = await assigned(engine_id, job_id, worker)
@router.get("/worker/{lens_id}/{job_id}/sample", response_model=Sample)
async def sample(lens_id: str, job_id: str, worker: WorkerAuth) -> Sample:
lens, job = await assigned(lens_id, job_id, worker)
if job.sample is not None:
return job.sample
pages: list[Sample] = [] # mutable-ok: freeze selection after stable cursor traversal
cursor = "" # rebind-ok: advance by immutable identity, never by shifting row positions
while True:
page = await source_reader().sample(
engine.scope,
lens.scope,
job.settings,
int(job.start.timestamp() * 1000),
int(job.end.timestamp() * 1000),
@ -390,7 +390,7 @@ async def sample(engine_id: str, job_id: str, worker: WorkerAuth) -> Sample:
) # comprehension-ok: flatten query pages
selected: Final = Sample(executions=executions, eligible=pages[0].eligible, selected=len(executions))
def freeze(e: Engine) -> Engine:
def freeze(e: Lens) -> Lens:
active: Final = current_job(e)
if active is None or active.id != job_id or active.worker_id != worker.id:
raise HTTPException(409, "Job was cancelled or reassigned")
@ -400,45 +400,45 @@ async def sample(engine_id: str, job_id: str, worker: WorkerAuth) -> Sample:
else e
)
updated: Final = required(await repository().update(engine_id, freeze))
updated: Final = required(await repository().update(lens_id, freeze))
frozen: Final = next(j for j in updated.jobs if j.id == job_id).sample
if frozen is None:
raise HTTPException(409, "Could not freeze the sample")
return frozen
@router.get("/worker/{engine_id}/{job_id}/content", response_model=ExecutionContent)
@router.get("/worker/{lens_id}/{job_id}/content", response_model=ExecutionContent)
async def content(
engine_id: str,
lens_id: str,
job_id: str,
execution_id: str,
worker: WorkerAuth,
cursor: str = "",
offset: int = Query(default=0, ge=0),
) -> ExecutionContent:
engine, job = await assigned(engine_id, job_id, worker)
lens, job = await assigned(lens_id, job_id, worker)
selected: Final = job.sample or Sample(executions=(), eligible=0)
execution: Final = next((e for e in selected.executions if e.id == execution_id), None)
if execution is None:
raise HTTPException(404, "Execution is outside this job's sample")
return await source_reader().content(engine.scope, execution, cursor, offset)
return await source_reader().content(lens.scope, execution, cursor, offset)
@router.post("/worker/{engine_id}/{job_id}/model", response_model=ModelResult)
async def model(engine_id: str, job_id: str, body: ModelRequest, worker: WorkerAuth, request: Request) -> ModelResult:
from litellm.proxy.engine.inference import analyze
@router.post("/worker/{lens_id}/{job_id}/model", response_model=ModelResult)
async def model(lens_id: str, job_id: str, body: ModelRequest, worker: WorkerAuth, request: Request) -> ModelResult:
from litellm.proxy.lens.inference import analyze
engine, job = await assigned(engine_id, job_id, worker)
return await analyze(repository(), engine, job, worker, body, request)
lens, job = await assigned(lens_id, job_id, worker)
return await analyze(repository(), lens, job, worker, body, request)
@router.post("/worker/{engine_id}/{job_id}/result", response_model=Engine)
async def result(engine_id: str, job_id: str, body: Result, worker: WorkerAuth) -> Engine:
engine: Final = await get_engine(engine_id, worker.scope)
old: Final = next((j for j in engine.jobs if j.id == job_id), None)
@router.post("/worker/{lens_id}/{job_id}/result", response_model=Lens)
async def result(lens_id: str, job_id: str, body: Result, worker: WorkerAuth) -> Lens:
lens: Final = await get_lens(lens_id, worker.scope)
old: Final = next((j for j in lens.jobs if j.id == job_id), None)
if old and old.status in ("completed", "failed") and old.worker_id == worker.id:
return engine
_, job = await assigned(engine_id, job_id, worker)
return lens
_, job = await assigned(lens_id, job_id, worker)
now: Final = datetime.now(timezone.utc)
selected: Final = job.sample or Sample(executions=(), eligible=0)
allowed: Final = frozenset(e.id for e in selected.executions)
@ -455,9 +455,9 @@ async def result(engine_id: str, job_id: str, body: Result, worker: WorkerAuth)
raise HTTPException(422, "Finding references evidence outside the job")
for finding in body.findings:
await validate_finding(engine, selected, finding)
await validate_finding(lens, selected, finding)
def finish(e: Engine) -> Engine:
def finish(e: Lens) -> Lens:
active: Final = current_job(e)
if active is None or active.id != job_id or active.worker_id != worker.id:
return e
@ -488,29 +488,29 @@ async def result(engine_id: str, job_id: str, body: Result, worker: WorkerAuth)
)
)
return required(await repository().update(engine_id, finish))
return required(await repository().update(lens_id, finish))
def merge_results(engine: Engine, result: Result, revision: int, now: datetime) -> Engine:
def merge_one(current: Engine, draft: FindingDraft) -> Engine:
def merge_results(lens: Lens, result: Result, revision: int, now: datetime) -> Lens:
def merge_one(current: Lens, draft: FindingDraft) -> Lens:
finding: Final = merge_finding(current, draft, revision, now)
return current.model_copy(
update=MappingProxyType({"findings": (finding, *(f for f in current.findings if f.id != finding.id))})
)
return reduce(merge_one, result.findings, engine)
return reduce(merge_one, result.findings, lens)
@router.post("/worker/{engine_id}/{job_id}/heartbeat", response_model=bool)
async def heartbeat(engine_id: str, job_id: str, worker: WorkerAuth) -> bool:
_, job = await assigned(engine_id, job_id, worker)
return await progress(engine_id, job_id, Progress(stage=job.stage, coverage=job.coverage), worker)
@router.post("/worker/{lens_id}/{job_id}/heartbeat", response_model=bool)
async def heartbeat(lens_id: str, job_id: str, worker: WorkerAuth) -> bool:
_, job = await assigned(lens_id, job_id, worker)
return await progress(lens_id, job_id, Progress(stage=job.stage, coverage=job.coverage), worker)
async def claim_candidate(candidate: Engine, worker: Worker, now: datetime) -> Claim | None:
async def claim_candidate(candidate: Lens, worker: Worker, now: datetime) -> Claim | None:
job_id: Final = str(uuid4())
def schedule(e: Engine) -> Engine:
def schedule(e: Lens) -> Lens:
scheduled: Final = queue_job(e, now, job_id) if e.settings.enabled and e.next_run_at <= now else e
return claim_job(scheduled, worker, now)
@ -519,31 +519,31 @@ async def claim_candidate(candidate: Engine, worker: Worker, now: datetime) -> C
return None
job: Final = current_job(updated)
if job and job.worker_id == worker.id and job.status == "running" and job != current_job(candidate):
return Claim(engine_id=updated.id, job=job, findings=updated.findings)
return Claim(lens_id=updated.id, job=job, findings=updated.findings)
return None
async def validate_finding(engine: Engine, selected: Sample, finding: FindingDraft) -> None:
previous: Final = next((f for f in engine.findings if f.id == finding.existing_finding_id), None)
async def validate_finding(lens: Lens, selected: Sample, finding: FindingDraft) -> None:
previous: Final = next((f for f in lens.findings if f.id == finding.existing_finding_id), None)
if finding.existing_finding_id and (previous is None or previous.check_id != finding.check_id):
raise HTTPException(422, "Existing finding must belong to the same check")
for evidence in finding.evidence:
if not await source_reader().verify_evidence(
engine.scope, next(e for e in selected.executions if e.id == evidence.execution_id), evidence
lens.scope, next(e for e in selected.executions if e.id == evidence.execution_id), evidence
):
raise HTTPException(422, "Evidence quote does not match stored content")
@router.get("/{engine_id}/executions/{execution_id}", response_model=ExecutionContent)
@router.get("/{lens_id}/executions/{execution_id}", response_model=ExecutionContent)
async def evidence_content(
engine_id: str, execution_id: str, auth: Auth, cursor: str = "", offset: int = Query(default=0, ge=0)
lens_id: str, execution_id: str, auth: Auth, cursor: str = "", offset: int = Query(default=0, ge=0)
) -> ExecutionContent:
engine: Final = await get_engine(engine_id, user_scope(auth))
lens: Final = await get_lens(lens_id, user_scope(auth))
try:
source, team, trace_id, trace_ref = parse_execution(execution_id)
except ValueError:
raise HTTPException(404, "Execution not found")
if source not in ("traces", "requests") or (not engine.scope.all_teams and team != engine.scope.team_id):
if source not in ("traces", "requests") or (not lens.scope.all_teams and team != lens.scope.team_id):
raise HTTPException(404, "Execution not found")
execution: Final = Execution(
id=execution_id,
@ -556,4 +556,4 @@ async def evidence_content(
span_count=1,
root_seen=source == "requests",
)
return await source_reader().content(engine.scope, execution, cursor, offset)
return await source_reader().content(lens.scope, execution, cursor, offset)

View file

@ -8,10 +8,10 @@ from pydantic import BaseModel, ConfigDict, Field
import litellm
from litellm.integrations.clickhouse.context import lens_analysis
from litellm.litellm_core_utils.initialize_dynamic_callback_params import inherit_message_logging_privacy
from litellm.proxy.engine.billing import complete, validate_key
from litellm.proxy.engine.models import Engine, Job, ModelRequest, ModelResult, Worker
from litellm.proxy.engine.repository import EngineRepository
from litellm.proxy.engine.state import current_job, renew_budget, replace_job
from litellm.proxy.lens.billing import complete, validate_key
from litellm.proxy.lens.models import Job, Lens, ModelRequest, ModelResult, Worker
from litellm.proxy.lens.repository import LensRepository
from litellm.proxy.lens.state import current_job, renew_budget, replace_job
from litellm.types.utils import CostPerToken, ModelResponse
@ -87,7 +87,7 @@ def quote(deployments: tuple[Deployment, ...], prompt: str) -> float:
async def analyze(
repo: EngineRepository, engine: Engine, job: Job, worker: Worker, body: ModelRequest, request: Request
repo: LensRepository, lens: Lens, job: Job, worker: Worker, body: ModelRequest, request: Request
) -> ModelResult:
from litellm.proxy.proxy_server import llm_router
@ -106,7 +106,7 @@ async def analyze(
estimate: Final = quote(deployments, body.prompt)
now: Final = datetime.now(timezone.utc)
def reserve(e: Engine) -> Engine:
def reserve(e: Lens) -> Lens:
current: Final = renew_budget(e, now)
active: Final = current_job(current)
if (
@ -124,7 +124,7 @@ async def analyze(
).model_copy(update=MappingProxyType({"spent": current.spent + estimate}))
async def reserve_budget() -> None:
if await repo.update(engine.id, reserve) is None:
if await repo.update(lens.id, reserve) is None:
raise HTTPException(409, "Could not reserve analysis budget")
data: Final[dict[str, object]] = { # mutable-ok: proxy processing enriches request data
@ -140,8 +140,8 @@ async def analyze(
"disable_fallbacks": True,
"response_format": {"type": "json_object"}, # mutable-ok: provider response-format JSON
"metadata": { # mutable-ok: request processing enriches metadata
"tags": ["litellm-engine"], # mutable-ok: logging callbacks require a list
"lens_id": engine.id,
"tags": ["litellm-lens"], # mutable-ok: logging callbacks require a list
"lens_id": lens.id,
"lens_run_id": job.id,
"lens_worker_id": worker.id,
"user_api_key_team_id": team_id,
@ -153,7 +153,7 @@ async def analyze(
parsed: Final = Completion.model_validate_json(response.model_dump_json())
cost: Final = billed_cost if billed_cost is not None else completion_charge(deployments, response, estimate)
def settle(e: Engine) -> Engine:
def settle(e: Lens) -> Lens:
charged: Final = next((j for j in e.jobs if j.id == job.id), None)
adjusted: Final = (
e.model_copy(update=MappingProxyType({"spent": max(0, e.spent - estimate + cost)}))
@ -168,7 +168,7 @@ async def analyze(
else adjusted
)
await repo.update(engine.id, settle)
await repo.update(lens.id, settle)
return ModelResult(content=parsed.choices[0].message.content or "{}", cost=cost)

View file

@ -25,7 +25,7 @@ class Check(Record):
enabled: bool = True
class EngineSettings(Record):
class LensSettings(Record):
name: str = Field(min_length=1, max_length=100)
context: str = Field(default="", max_length=6000)
source: Literal["traces", "requests", "both"] = "traces"
@ -44,7 +44,7 @@ class EngineSettings(Record):
monthly_budget: float = Field(default=20, gt=0, le=100000, allow_inf_nan=False)
@model_validator(mode="after")
def unique_checks(self) -> "EngineSettings":
def unique_checks(self) -> "LensSettings":
if len(frozenset(c.id for c in self.checks)) != len(self.checks):
raise ValueError("Each check must have a unique ID")
if not self.context.strip() and not any(c.enabled for c in self.checks):
@ -163,7 +163,7 @@ class Job(Record):
created_at: datetime
start: datetime
end: datetime
settings: EngineSettings
settings: LensSettings
revision: int
worker_id: str | None = None
lease_until: datetime | None = None
@ -177,10 +177,10 @@ class Job(Record):
assessments: tuple[RunAssessment, ...] = ()
class Engine(Record):
class Lens(Record):
id: str
scope: Scope
settings: EngineSettings
settings: LensSettings
revision: int = 1
version: int = 0
created_at: datetime
@ -206,14 +206,14 @@ class WorkerCreated(Record):
token: str
class EngineList(Record):
engines: tuple[Engine, ...]
class LensList(Record):
lenses: tuple[Lens, ...]
workers: tuple[Worker, ...]
tracing_enabled: bool
class RunRequest(Record):
settings: EngineSettings | None = None
settings: LensSettings | None = None
lookback_hours: int | None = Field(default=None, ge=1, le=720)
@ -223,7 +223,7 @@ class FindingUpdate(Record):
class Claim(Record):
engine_id: str
lens_id: str
job: Job
findings: tuple[Finding, ...]

View file

@ -5,7 +5,7 @@ from typing import Final, Protocol
from pydantic import BaseModel, JsonValue, TypeAdapter
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.engine.models import Engine, Job, Worker
from litellm.proxy.lens.models import Job, Lens, Worker
class Database(Protocol):
@ -20,44 +20,44 @@ class Row(BaseModel):
_ROWS: Final = TypeAdapter(tuple[Row, ...])
class EngineRepository:
class LensRepository:
def __init__(self, db: Database) -> None:
self.db: Final = db
async def engines(self) -> tuple[Engine, ...]:
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_Engine" ORDER BY id'))
return tuple(Engine.model_validate(row.data) for row in rows)
async def lenses(self) -> tuple[Lens, ...]:
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_Lens" ORDER BY id'))
return tuple(Lens.model_validate(row.data) for row in rows)
async def get(self, engine_id: str) -> Engine | None:
async def get(self, lens_id: str) -> Lens | None:
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
'SELECT data FROM "LiteLLM_Engine" WHERE id=$1',
engine_id,
'SELECT data FROM "LiteLLM_Lens" WHERE id=$1',
lens_id,
)
)
return Engine.model_validate(rows[0].data) if rows else None
return Lens.model_validate(rows[0].data) if rows else None
async def create(self, engine: Engine) -> Engine:
async def create(self, lens: Lens) -> Lens:
await self.db.execute_raw(
'INSERT INTO "LiteLLM_Engine" (id, version, data) VALUES ($1,0,$2::jsonb)',
engine.id,
engine.model_dump_json(),
'INSERT INTO "LiteLLM_Lens" (id, version, data) VALUES ($1,0,$2::jsonb)',
lens.id,
lens.model_dump_json(),
)
return engine
return lens
async def update(
self, engine_id: str, transform: Callable[[Engine], Engine], attempts: int = 8, *, changed_only: bool = False
) -> Engine | None:
self, lens_id: str, transform: Callable[[Lens], Lens], attempts: int = 8, *, changed_only: bool = False
) -> Lens | None:
for _ in range(attempts):
completed, updated = await self._try_update(engine_id, transform, changed_only)
completed, updated = await self._try_update(lens_id, transform, changed_only)
if completed:
return updated
return None
async def _try_update(
self, engine_id: str, transform: Callable[[Engine], Engine], changed_only: bool
) -> tuple[bool, Engine | None]:
previous: Final = await self.get(engine_id)
self, lens_id: str, transform: Callable[[Lens], Lens], changed_only: bool
) -> tuple[bool, Lens | None]:
previous: Final = await self.get(lens_id)
if previous is None:
return True, None
candidate: Final = transform(previous)
@ -67,12 +67,12 @@ class EngineRepository:
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
"""WITH previous AS MATERIALIZED (
SELECT data FROM "LiteLLM_Engine" WHERE id=$2 AND version=$3 FOR UPDATE
SELECT data FROM "LiteLLM_Lens" WHERE id=$2 AND version=$3 FOR UPDATE
), updated AS (
UPDATE "LiteLLM_Engine" SET data=$1::jsonb, version=version+1
UPDATE "LiteLLM_Lens" SET data=$1::jsonb, version=version+1
WHERE id=$2 AND version=$3 AND EXISTS (SELECT 1 FROM previous) RETURNING id
)
, archived AS (INSERT INTO "LiteLLM_EngineRun" (id, engine_id, created_at, data)
, archived AS (INSERT INTO "LiteLLM_LensRun" (id, lens_id, created_at, data)
SELECT job->>'id', $2, (job->>'created_at')::timestamp, job
FROM previous, jsonb_array_elements(previous.data->'jobs') AS job
WHERE EXISTS (SELECT 1 FROM updated)
@ -81,46 +81,46 @@ class EngineRepository:
ON CONFLICT (id) DO NOTHING)
SELECT to_jsonb(count(*)) AS data FROM updated""",
updated.model_dump_json(),
engine_id,
lens_id,
previous.version,
)
)
return bool(rows and rows[0].data == 1), updated
async def jobs(self, engine_id: str, offset: int = 0) -> tuple[Job, ...]:
async def jobs(self, lens_id: str, offset: int = 0) -> tuple[Job, ...]:
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
"""SELECT data FROM (
SELECT data FROM "LiteLLM_EngineRun" WHERE engine_id=$1
SELECT data FROM "LiteLLM_LensRun" WHERE lens_id=$1
UNION ALL
SELECT jsonb_array_elements(data->'jobs') AS data FROM "LiteLLM_Engine" WHERE id=$1
SELECT jsonb_array_elements(data->'jobs') AS data FROM "LiteLLM_Lens" WHERE id=$1
) AS jobs ORDER BY data->>'created_at' DESC, data->>'id' DESC LIMIT 50 OFFSET $2""",
engine_id,
lens_id,
offset,
)
)
return tuple(Job.model_validate(row.data) for row in rows)
async def job(self, engine_id: str, job_id: str) -> Job | None:
async def job(self, lens_id: str, job_id: str) -> Job | None:
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
"""SELECT data FROM "LiteLLM_EngineRun" WHERE engine_id=$1 AND id=$2
UNION ALL SELECT job AS data FROM "LiteLLM_Engine", jsonb_array_elements(data->'jobs') AS job
"""SELECT data FROM "LiteLLM_LensRun" WHERE lens_id=$1 AND id=$2
UNION ALL SELECT job AS data FROM "LiteLLM_Lens", jsonb_array_elements(data->'jobs') AS job
WHERE id=$1 AND job->>'id'=$2 LIMIT 1""",
engine_id,
lens_id,
job_id,
)
)
return Job.model_validate(rows[0].data) if rows else None
async def workers(self) -> tuple[Worker, ...]:
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_EngineWorker"'))
rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_LensWorker"'))
return tuple(Worker.model_validate(row.data) for row in rows)
async def worker(self, token_hash: str) -> Worker | None:
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
'SELECT data FROM "LiteLLM_EngineWorker" WHERE token_hash=$1',
'SELECT data FROM "LiteLLM_LensWorker" WHERE token_hash=$1',
token_hash,
)
)
@ -129,20 +129,20 @@ class EngineRepository:
async def save_worker(self, worker: Worker, token_hash: str | None = None) -> None:
if token_hash is not None:
await self.db.execute_raw(
'INSERT INTO "LiteLLM_EngineWorker" (id,token_hash,data) VALUES ($1,$2,$3::jsonb)',
'INSERT INTO "LiteLLM_LensWorker" (id,token_hash,data) VALUES ($1,$2,$3::jsonb)',
worker.id,
token_hash,
worker.model_dump_json(),
)
return
await self.db.execute_raw(
'UPDATE "LiteLLM_EngineWorker" SET data=$1::jsonb WHERE id=$2', worker.model_dump_json(), worker.id
'UPDATE "LiteLLM_LensWorker" SET data=$1::jsonb WHERE id=$2', worker.model_dump_json(), worker.id
)
async def set_worker_billing(self, worker_id: str, key_id: str) -> Worker | None:
rows: Final = _ROWS.validate_python(
await self.db.query_raw(
"""UPDATE "LiteLLM_EngineWorker"
"""UPDATE "LiteLLM_LensWorker"
SET data=jsonb_set(data, '{analysis_key_id}', to_jsonb($1::text))
WHERE id=$2 AND COALESCE((data->>'revoked')::boolean, false)=false RETURNING data""",
key_id,
@ -153,13 +153,13 @@ class EngineRepository:
async def revoke_worker(self, worker_id: str) -> None:
await self.db.execute_raw(
"""UPDATE "LiteLLM_EngineWorker" SET data=jsonb_set(data, '{revoked}', 'true') WHERE id=$1""",
"""UPDATE "LiteLLM_LensWorker" SET data=jsonb_set(data, '{revoked}', 'true') WHERE id=$1""",
worker_id,
)
async def heartbeat(self, worker_id: str, now: str) -> None:
await self.db.execute_raw(
"""UPDATE "LiteLLM_EngineWorker" SET data=jsonb_set(data, '{last_seen}', to_jsonb($1::text)) WHERE id=$2""",
"""UPDATE "LiteLLM_LensWorker" SET data=jsonb_set(data, '{last_seen}', to_jsonb($1::text)) WHERE id=$2""",
now,
worker_id,
)

View file

@ -6,11 +6,11 @@ from typing import Final, Literal, Protocol
from pydantic import BaseModel, TypeAdapter
from litellm.proxy.engine.models import (
EngineSettings,
from litellm.proxy.lens.models import (
Evidence,
Execution,
ExecutionContent,
LensSettings,
MetadataFilter,
Sample,
Scope,
@ -93,7 +93,7 @@ class SourceReader:
async def sample(
self,
scope: Scope,
settings: EngineSettings,
settings: LensSettings,
start: int,
end: int,
offset: int = 0,

View file

@ -3,7 +3,7 @@ from datetime import datetime, timedelta
from types import MappingProxyType
from typing import Final
from litellm.proxy.engine.models import Engine, EngineSettings, Finding, FindingDraft, Job, Scope, Worker
from litellm.proxy.lens.models import Finding, FindingDraft, Job, Lens, LensSettings, Scope, Worker
def can_access(viewer: Scope, target: Scope) -> bool:
@ -14,46 +14,46 @@ def can_access(viewer: Scope, target: Scope) -> bool:
)
def current_job(engine: Engine) -> Job | None:
return next((job for job in engine.jobs if job.status in ("queued", "running")), None)
def current_job(lens: Lens) -> Job | None:
return next((job for job in lens.jobs if job.status in ("queued", "running")), None)
def replace_job(engine: Engine, job: Job) -> Engine:
return engine.model_copy(
update=MappingProxyType({"jobs": tuple(job if old.id == job.id else old for old in engine.jobs)})
def replace_job(lens: Lens, job: Job) -> Lens:
return lens.model_copy(
update=MappingProxyType({"jobs": tuple(job if old.id == job.id else old for old in lens.jobs)})
)
def queue_job(
engine: Engine,
lens: Lens,
now: datetime,
job_id: str,
lookback_hours: int | None = None,
settings: EngineSettings | None = None,
) -> Engine:
if current_job(engine):
return engine
selected: Final = settings or engine.settings
settings: LensSettings | None = None,
) -> Lens:
if current_job(lens):
return lens
selected: Final = settings or lens.settings
job: Final = Job(
id=job_id,
created_at=now,
start=now - timedelta(hours=lookback_hours if lookback_hours is not None else selected.lookback_hours),
end=now - timedelta(minutes=2),
settings=selected,
revision=engine.revision,
revision=lens.revision,
)
return engine.model_copy(update=MappingProxyType({"jobs": (job,)}))
return lens.model_copy(update=MappingProxyType({"jobs": (job,)}))
def claim_job(engine: Engine, worker: Worker, now: datetime) -> Engine:
job: Final = current_job(engine)
if job is None or not can_access(worker.scope, engine.scope):
return engine
def claim_job(lens: Lens, worker: Worker, now: datetime) -> Lens:
job: Final = current_job(lens)
if job is None or not can_access(worker.scope, lens.scope):
return lens
if job.status == "running" and job.lease_until is not None and job.lease_until > now:
return engine
return lens
if job.attempts >= 3:
return replace_job(
engine,
lens,
job.model_copy(
update=MappingProxyType(
{
@ -64,11 +64,9 @@ def claim_job(engine: Engine, worker: Worker, now: datetime) -> Engine:
}
)
),
).model_copy(
update=MappingProxyType({"next_run_at": now + timedelta(minutes=engine.settings.interval_minutes)})
)
).model_copy(update=MappingProxyType({"next_run_at": now + timedelta(minutes=lens.settings.interval_minutes)}))
return replace_job(
engine,
lens,
job.model_copy(
update=MappingProxyType(
{
@ -83,23 +81,23 @@ def claim_job(engine: Engine, worker: Worker, now: datetime) -> Engine:
)
def renew_budget(engine: Engine, now: datetime) -> Engine:
def renew_budget(lens: Lens, now: datetime) -> Lens:
month: Final = now.strftime("%Y-%m")
if engine.budget_month == month:
return engine
return engine.model_copy(update=MappingProxyType({"budget_month": month, "spent": 0}))
if lens.budget_month == month:
return lens
return lens.model_copy(update=MappingProxyType({"budget_month": month, "spent": 0}))
def merge_finding(engine: Engine, draft: FindingDraft, revision: int, now: datetime) -> Finding:
legacy_identity: Final = hashlib.sha256(f"{engine.id}:{draft.check_id}:{draft.title.lower()}".encode()).hexdigest()[
def merge_finding(lens: Lens, draft: FindingDraft, revision: int, now: datetime) -> Finding:
legacy_identity: Final = hashlib.sha256(f"{lens.id}:{draft.check_id}:{draft.title.lower()}".encode()).hexdigest()[
:24
]
identity: Final = hashlib.sha256(
f"{engine.id}:{draft.check_id}:{draft.kind}:{draft.title.lower()}".encode()
f"{lens.id}:{draft.check_id}:{draft.kind}:{draft.title.lower()}".encode()
).hexdigest()[:24]
identities: Final = (draft.existing_finding_id, identity, legacy_identity)
previous: Final = next(
(f for f in engine.findings if f.id in identities and f.kind == draft.kind and f.check_id == draft.check_id),
(f for f in lens.findings if f.id in identities and f.kind == draft.kind and f.check_id == draft.check_id),
None,
)
occurrences: Final = tuple(sorted(frozenset(e.execution_id for e in draft.evidence if e.role == "support")))
@ -137,8 +135,8 @@ def merge_finding(engine: Engine, draft: FindingDraft, revision: int, now: datet
)
def snapshot_finding(engine: Engine, draft: FindingDraft, revision: int, now: datetime) -> Finding:
merged: Final = merge_finding(engine, draft, revision, now)
def snapshot_finding(lens: Lens, draft: FindingDraft, revision: int, now: datetime) -> Finding:
merged: Final = merge_finding(lens, draft, revision, now)
return Finding.model_validate(
MappingProxyType(
{

View file

@ -12,10 +12,10 @@ import httpx
from .analysis import analyze_sample
from .models import Claim, Coverage, ExecutionContent, ModelRequest, ModelResult, Progress, Result, Sample
logger: Final = logging.getLogger("litellm.engine.worker")
logger: Final = logging.getLogger("litellm.lens.worker")
class EngineWorker:
class LensWorker:
def __init__(self, client: httpx.AsyncClient, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep) -> None:
self.client: Final = client
self.sleep: Final = sleep
@ -38,14 +38,12 @@ class EngineWorker:
return await self.model_request(path, body, attempt + 1)
async def run_once(self) -> bool:
response: Final = await self.client.post(
"/engine/worker/claim", params=MappingProxyType({"protocol_version": 2})
)
response: Final = await self.client.post("/lens/worker/claim", params=MappingProxyType({"protocol_version": 2}))
response.raise_for_status()
if response.json() is None:
return False
claim: Final = Claim.model_validate(response.json())
prefix: Final = f"/engine/worker/{claim.engine_id}/{claim.job.id}"
prefix: Final = f"/lens/worker/{claim.lens_id}/{claim.job.id}"
async def model(body: ModelRequest) -> ModelResult:
return await self.model_request(prefix + "/model", body)
@ -111,7 +109,7 @@ async def main() -> None:
async with httpx.AsyncClient(
base_url=url, headers=MappingProxyType({"Authorization": f"Bearer {token}"}), timeout=180
) as client:
worker: Final = EngineWorker(client)
worker: Final = LensWorker(client)
while True:
try:
await worker.run_once()

View file

@ -541,7 +541,6 @@ from litellm.proxy.discovery_endpoints import (
agent_skills_discovery_router,
ui_discovery_endpoints_router,
)
from litellm.proxy.engine.endpoints import router as engine_router
from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router
from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config
from litellm.proxy.google_endpoints.endpoints import router as google_router
@ -565,6 +564,7 @@ from litellm.proxy.hooks.prompt_injection_detection import (
)
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_spend_event
from litellm.proxy.image_endpoints.endpoints import router as image_router
from litellm.proxy.lens.endpoints import router as lens_router
from litellm.proxy.list_api.common import (
ManagementProblem,
problem_response,
@ -19988,7 +19988,7 @@ app.include_router(auto_router_management_router)
app.include_router(tag_management_router)
app.include_router(workflow_management_router)
app.include_router(memory_router)
app.include_router(engine_router)
app.include_router(lens_router)
app.include_router(plugin_router)
app.include_router(cost_tracking_settings_router)
app.include_router(prompt_caching_requests_router)

View file

@ -1895,22 +1895,22 @@ model LiteLLM_WorkflowMessage {
@@index([run_id])
}
model LiteLLM_Engine {
model LiteLLM_Lens {
id String @id
version Int @default(0)
data Json
}
model LiteLLM_EngineRun {
model LiteLLM_LensRun {
id String @id
engine_id String
lens_id String
created_at DateTime
data Json
@@index([engine_id, created_at])
@@index([lens_id, created_at])
}
model LiteLLM_EngineWorker {
model LiteLLM_LensWorker {
id String @id
token_hash String @unique
data Json

View file

@ -9,14 +9,23 @@ Pure functions, no I/O. Two steps:
"""
import json
from collections.abc import Callable, Mapping
from collections.abc import Mapping
from itertools import accumulate
from types import MappingProxyType
from typing import Any, Final
from typing import Final
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm.constants import OTLP_MAX_ATTRIBUTE_VALUE_BYTES, OTLP_MAX_BODY_BYTES
from litellm.rust_bridge.traces import DecodedSpan
from litellm.rust_bridge.traces import decode_otlp as native_decode_otlp
from litellm.tracing.types import SpanRow, SpanType
from litellm.tracing.normalizers import select_normalizer
from litellm.tracing.normalizers.base import to_int
from litellm.tracing.types import SpanRow
_MESSAGE_LIST: Final = TypeAdapter(tuple[dict[str, JsonValue], ...])
_MAX_JSON_ESCAPE_BYTES: Final = 6
# attributes whose content we lift into Input/Output and drop from SpanAttributes
_HEAVY_ATTRIBUTES: Final = frozenset(
@ -30,18 +39,6 @@ _HEAVY_ATTRIBUTES: Final = frozenset(
"output.value",
}
)
# LangChain / Deep Agents middleware wrappers: real spans, but noise in the UI
_FRAMEWORK_SUFFIXES: Final = (
".wrap_model_call",
".wrap_tool_call",
".before_agent",
".after_agent",
".before_model",
".after_model",
)
_LLM_OPERATIONS: Final = frozenset({"chat", "text_completion", "generate_content"})
_LC_ROLES: Final = MappingProxyType({"human": "user", "ai": "assistant", "system": "system", "tool": "tool"})
_OPENINFERENCE_TYPES: Final[Mapping[str, SpanType]] = MappingProxyType({"AGENT": "agent", "LLM": "llm", "TOOL": "tool"})
class InvalidOTLPPayloadError(ValueError):
@ -60,6 +57,83 @@ def _truncate(value: str) -> str:
return f"{kept}…[truncated {size - OTLP_MAX_ATTRIBUTE_VALUE_BYTES} bytes]"
def _size(value: str) -> int:
return len(value.encode("utf-8"))
class _ElisionMarker(TypedDict):
role: ReadOnly[str]
content: ReadOnly[str]
def _elided(count: int) -> str:
marker: Final[_ElisionMarker] = {"role": "system", "content": f"…[{count} earlier messages truncated]"}
return json.dumps(marker)
def _with_content(message: Mapping[str, JsonValue], content: str) -> str:
return json.dumps(MappingProxyType({**message, "content": content}), default=lambda proxy: proxy.copy())
def _shrunk_message(message: Mapping[str, JsonValue], budget: int) -> str:
"""One message cut to `budget` bytes, as valid JSON.
Shortens `content` first; if other fields (e.g. huge tool_calls) still don't fit, keeps only role + content.
"""
content: Final = message.get("content")
text: Final = content if isinstance(content, str) else json.dumps(content)
role_only: Final = MappingProxyType({"role": message.get("role", "user")})
attempts: Final = (
_cut_content(message, text, budget, 1),
_cut_content(role_only, text, budget, 1),
_cut_content(role_only, text, budget, _MAX_JSON_ESCAPE_BYTES),
)
return next((attempt for attempt in attempts if _size(attempt) <= budget), attempts[-1])
def _cut_content(message: Mapping[str, JsonValue], text: str, budget: int, escape_factor: int) -> str:
overhead: Final = _size(_with_content(message, ""))
room: Final = max(0, budget - overhead - 48) // escape_factor
kept: Final = text.encode("utf-8")[:room].decode("utf-8", "ignore")
return _with_content(message, f"{kept}…[truncated {_size(text) - _size(kept)} bytes]")
def _newest_that_fit(encoded: tuple[str, ...], budget: int) -> int:
"""How many trailing messages fit in `budget` bytes (comma separators included), scanning newest first."""
sizes: Final = tuple(_size(m) + 1 for m in reversed(encoded))
totals: Final = tuple(accumulate(sizes))
return next((count for count, total in enumerate(totals) if total > budget), len(totals))
def _truncate_payload(value: str) -> str:
"""Message arrays keep the first message, an elision marker and the newest messages that fit.
The result is always valid JSON: if even those don't fit, the first and last messages are shortened.
Anything that isn't a message array is byte-truncated as before.
"""
if _size(value) <= OTLP_MAX_ATTRIBUTE_VALUE_BYTES or not value.startswith("["):
return _truncate(value)
try:
messages: Final = _MESSAGE_LIST.validate_json(value)
except ValidationError:
return _truncate(value)
if len(messages) < 2:
return _truncate(value)
encoded: Final = tuple(json.dumps(m) for m in messages)
marker_budget: Final = _size(_elided(len(messages))) + 1
budget: Final = OTLP_MAX_ATTRIBUTE_VALUE_BYTES - 2 - _size(encoded[0]) - 1 - marker_budget
kept: Final = min(_newest_that_fit(encoded[1:], budget), len(messages) - 2)
if kept > 0:
tail: Final = encoded[len(encoded) - kept :]
return "[" + ", ".join((encoded[0], _elided(len(messages) - 1 - kept), *tail)) + "]"
half: Final = (OTLP_MAX_ATTRIBUTE_VALUE_BYTES - marker_budget - 4) // 2
middle: Final = (_elided(len(messages) - 2),) if len(messages) > 2 else ()
shrunk: Final = (
"[" + ", ".join((_shrunk_message(messages[0], half), *middle, _shrunk_message(messages[-1], half))) + "]"
)
return shrunk if _size(shrunk) <= OTLP_MAX_ATTRIBUTE_VALUE_BYTES else "[" + _elided(len(messages)) + "]"
def decode_otlp(
body: bytes, content_type: str | None = None, content_encoding: str | None = None
) -> tuple[SpanRow, ...]:
@ -116,157 +190,17 @@ def _span_row(span: DecodedSpan) -> SpanRow:
row["SpanAttributes"] = { # mutable-ok: the Rust JSON bridge requires a plain dict for span attributes
k: _truncate(v) for k, v in attributes.items() if k not in _HEAVY_ATTRIBUTES
}
row["Input"], row["Output"] = _truncate(row["Input"]), _truncate(row["Output"])
row["Input"], row["Output"] = _truncate_payload(row["Input"]), _truncate(row["Output"])
return row
def _loads(value: str) -> object:
try:
return json.loads(value)
except (ValueError, TypeError):
return None
def _lc_message(message: Mapping[str, Any]) -> dict[str, Any]:
"""LangChain serialized message (or plain {role, content}) -> {role, content, tool_calls?}."""
kwargs = message.get("kwargs", message)
role = _LC_ROLES.get(kwargs.get("type") or kwargs.get("role"), kwargs.get("role") or kwargs.get("type") or "")
content = kwargs.get("content", "")
out: dict[str, Any] = { # mutable-ok: the framework message is built for JSON serialization
"role": role,
"content": content if isinstance(content, str) else json.dumps(content),
}
if kwargs.get("tool_calls"):
out["tool_calls"] = tuple(
{"name": t.get("name"), "args": t.get("args")} # mutable-ok: JSON tool calls need object payloads
for t in kwargs["tool_calls"]
)
if role == "tool" and kwargs.get("name"):
out["name"] = kwargs["name"]
return out
def _langsmith_type(row: SpanRow, attributes: Mapping[str, str]) -> SpanType:
kind = attributes.get("langsmith.span.kind", "chain")
name = row["SpanName"]
if not row["ParentSpanId"] or name == attributes.get("langsmith.metadata.lc_agent_name"):
return "agent"
if kind in ("llm", "tool"):
return kind
if name.endswith(_FRAMEWORK_SUFFIXES):
return "framework"
return "chain"
def _langsmith_io(row: SpanRow, attributes: Mapping[str, str]) -> None:
prompt = _loads(attributes.get("gen_ai.prompt", ""))
completion = _loads(attributes.get("gen_ai.completion", ""))
prompt_payload = prompt if isinstance(prompt, dict) else MappingProxyType({})
if row["ObservationType"] == "llm" and isinstance(completion, dict):
messages = prompt_payload.get("messages") or ((),)
batch = messages[0] if messages and isinstance(messages[0], list) else messages
row["Input"] = (
json.dumps(tuple(_lc_message(m) for m in batch if isinstance(m, dict)))
if isinstance(batch, (list, tuple))
else ""
)
generations: Final = completion.get("generations")
first: Final = generations[0] if isinstance(generations, list) and generations else None
item: Final = first[0] if isinstance(first, list) and first else None
message: Final = item.get("message") if isinstance(item, dict) else None
generation: Final = message.get("kwargs") if isinstance(message, dict) else None
if isinstance(generation, dict):
row["Output"] = json.dumps(_lc_message(generation))
metadata: Final = generation.get("response_metadata")
row["LiteLLMRequestId"] = metadata.get("id", "") if isinstance(metadata, dict) else ""
else:
row["Output"] = attributes.get("gen_ai.completion", "")
return
if row["ObservationType"] == "tool":
output = completion.get("output", completion) if isinstance(completion, dict) else completion
if isinstance(output, dict) and "update" in output: # LangGraph Command, e.g. Deep Agents `task`
update: Final = output.get("update")
update_messages = update.get("messages") or () if isinstance(update, dict) else ()
output = update_messages[-1] if update_messages else output
if isinstance(output, dict):
output = output.get("content", output)
row["Input"] = attributes.get("gen_ai.prompt", "")
row["Output"] = output if isinstance(output, str) else json.dumps(output)
return
if row["ObservationType"] == "agent":
input_messages = prompt.get("messages") if isinstance(prompt, dict) else None
output_messages = completion.get("messages") if isinstance(completion, dict) else None
# agents built with @traceable take arbitrary args, not a message list: keep the raw payload then
row["Input"] = (
json.dumps(tuple(_lc_message(m) for m in input_messages if isinstance(m, dict)))
if input_messages
else attributes.get("gen_ai.prompt", "")
)
row["Output"] = (
json.dumps(_lc_message(output_messages[-1]))
if output_messages and isinstance(output_messages[-1], dict)
else attributes.get("gen_ai.completion", "")
)
return
row["Input"] = attributes.get("gen_ai.prompt", "")
row["Output"] = attributes.get("gen_ai.completion", "")
def normalize_langsmith(row: SpanRow, attributes: Mapping[str, str]) -> None:
row["ObservationType"] = _langsmith_type(row, attributes)
row["AgentName"] = attributes.get("langsmith.metadata.lc_agent_name", "")
row["Model"] = attributes.get("gen_ai.request.model", "")
_langsmith_io(row, attributes)
def normalize_genai(row: SpanRow, attributes: Mapping[str, str]) -> None:
operation = attributes.get("gen_ai.operation.name", "")
if operation == "invoke_agent" or not row["ParentSpanId"]:
row["ObservationType"] = "agent"
elif operation in _LLM_OPERATIONS:
row["ObservationType"] = "llm"
elif operation == "execute_tool":
row["ObservationType"] = "tool"
row["AgentName"] = attributes.get("gen_ai.agent.name", "")
row["Model"] = attributes.get("gen_ai.request.model") or attributes.get("gen_ai.response.model", "")
row["LiteLLMRequestId"] = attributes.get("gen_ai.response.id", "")
row["Input"] = attributes.get("gen_ai.input.messages") or attributes.get("gen_ai.tool.call.arguments", "")
row["Output"] = attributes.get("gen_ai.output.messages") or attributes.get("gen_ai.tool.call.result", "")
def normalize_openinference(row: SpanRow, attributes: Mapping[str, str]) -> None:
kind = attributes.get("openinference.span.kind", "").upper()
row["ObservationType"] = _OPENINFERENCE_TYPES.get(kind, "agent" if not row["ParentSpanId"] else "chain")
row["AgentName"] = attributes.get("agent.name", "")
row["Model"] = attributes.get("llm.model_name", "")
row["Input"] = attributes.get("input.value", "")
row["Output"] = attributes.get("output.value", "")
row["InputTokens"] = _to_int(attributes.get("llm.token_count.prompt"))
row["OutputTokens"] = _to_int(attributes.get("llm.token_count.completion"))
def _set_tokens(row: SpanRow, attributes: Mapping[str, str]) -> None:
row["InputTokens"] = _to_int(attributes.get("gen_ai.usage.input_tokens"))
row["OutputTokens"] = _to_int(attributes.get("gen_ai.usage.output_tokens"))
def _to_int(value: str | None) -> int:
try:
return int(value) if value else 0
except ValueError:
return 0
def select_normalizer(scope_name: str, attributes: Mapping[str, str]) -> Callable[[SpanRow, Mapping[str, str]], None]:
if scope_name == "langsmith" or "langsmith.span.kind" in attributes:
return normalize_langsmith
if "openinference.span.kind" in attributes:
return normalize_openinference
return normalize_genai
row["InputTokens"] = to_int(attributes.get("gen_ai.usage.input_tokens"))
row["OutputTokens"] = to_int(attributes.get("gen_ai.usage.output_tokens"))
def normalize(row: SpanRow, attributes: Mapping[str, str]) -> None:
select_normalizer(row["ScopeName"], attributes)(row, attributes)
select_normalizer(row["ScopeName"], attributes).normalize(row, attributes)
if not row["InputTokens"] and not row["OutputTokens"]:
_set_tokens(row, attributes)

View file

@ -0,0 +1,32 @@
"""Per-convention span normalizers, tried in order: the first whose `matches()` is true wins."""
from collections.abc import Mapping, Sequence
from typing import Final
from litellm.tracing.normalizers.base import SpanNormalizer
from litellm.tracing.normalizers.genai import GenAISemconvNormalizer
from litellm.tracing.normalizers.langsmith import LangSmithNormalizer
from litellm.tracing.normalizers.openinference import OpenInferenceNormalizer
NORMALIZERS: Final[tuple[SpanNormalizer, ...]] = (
LangSmithNormalizer(),
OpenInferenceNormalizer(),
GenAISemconvNormalizer(),
)
_FALLBACK: Final[SpanNormalizer] = GenAISemconvNormalizer()
def select_normalizer(
scope_name: str, attributes: Mapping[str, str], registry: Sequence[SpanNormalizer] = NORMALIZERS
) -> SpanNormalizer:
return next((n for n in registry if n.matches(scope_name, attributes)), _FALLBACK)
__all__ = (
"NORMALIZERS",
"GenAISemconvNormalizer",
"LangSmithNormalizer",
"OpenInferenceNormalizer",
"SpanNormalizer",
"select_normalizer",
)

View file

@ -0,0 +1,22 @@
from collections.abc import Mapping
from typing import Protocol
from litellm.tracing.types import SpanRow
class SpanNormalizer(Protocol):
"""Maps one tracing convention's span attributes onto the LiteLLM `SpanRow` columns."""
@property
def name(self) -> str: ...
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool: ...
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None: ...
def to_int(value: str | None) -> int:
try:
return int(value) if value else 0
except ValueError:
return 0

View file

@ -0,0 +1,31 @@
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final
from litellm.tracing.types import SpanRow
_LLM_OPERATIONS: Final = frozenset({"chat", "text_completion", "generate_content"})
@dataclass(frozen=True, slots=True)
class GenAISemconvNormalizer:
"""OTEL `gen_ai.*` semantic conventions. Matches every span, so it belongs last as the fallback."""
name: str = "genai"
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool:
return True
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None:
operation: Final = attributes.get("gen_ai.operation.name", "")
if operation == "invoke_agent" or not row["ParentSpanId"]:
row["ObservationType"] = "agent"
elif operation in _LLM_OPERATIONS:
row["ObservationType"] = "llm"
elif operation == "execute_tool":
row["ObservationType"] = "tool"
row["AgentName"] = attributes.get("gen_ai.agent.name", "")
row["Model"] = attributes.get("gen_ai.request.model") or attributes.get("gen_ai.response.model", "")
row["LiteLLMRequestId"] = attributes.get("gen_ai.response.id", "")
row["Input"] = attributes.get("gen_ai.input.messages") or attributes.get("gen_ai.tool.call.arguments", "")
row["Output"] = attributes.get("gen_ai.output.messages") or attributes.get("gen_ai.tool.call.result", "")

View file

@ -0,0 +1,115 @@
"""LangSmith OTEL mode, which LangChain, LangGraph and Deep Agents export through."""
import json
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm.tracing.normalizers.messages import lc_message
from litellm.tracing.types import SpanRow, SpanType
# LangChain / Deep Agents middleware wrappers: real spans, but noise in the UI
_FRAMEWORK_SUFFIXES: Final = (
".wrap_model_call",
".wrap_tool_call",
".before_agent",
".after_agent",
".before_model",
".after_model",
)
def _loads(value: str) -> object:
try:
return json.loads(value)
except (ValueError, TypeError):
return None
def _span_type(row: SpanRow, attributes: Mapping[str, str]) -> SpanType:
kind: Final = attributes.get("langsmith.span.kind", "chain")
name: Final = row["SpanName"]
if not row["ParentSpanId"] or name == attributes.get("langsmith.metadata.lc_agent_name"):
return "agent"
if kind in ("llm", "tool"):
return kind
if name.endswith(_FRAMEWORK_SUFFIXES):
return "framework"
return "chain"
def _tool_output(completion: object) -> object:
raw: Final = completion.get("output", completion) if isinstance(completion, dict) else completion
update: Final = raw.get("update") if isinstance(raw, dict) else None
update_messages: Final = update.get("messages") or () if isinstance(update, dict) else ()
is_command: Final = isinstance(raw, dict) and "update" in raw
# LangGraph Command (e.g. the Deep Agents `task` tool): the result is the last update message
output: Final = update_messages[-1] if is_command and update_messages else raw
return output.get("content", output) if isinstance(output, dict) else output
def _set_agent_io(row: SpanRow, attributes: Mapping[str, str], prompt: object, completion: object) -> None:
input_messages: Final = prompt.get("messages") if isinstance(prompt, dict) else None
output_messages: Final = completion.get("messages") if isinstance(completion, dict) else None
# agents built with @traceable take arbitrary args, not a message list: keep the raw payload then
row["Input"] = (
json.dumps(tuple(lc_message(m) for m in input_messages if isinstance(m, dict)))
if input_messages
else attributes.get("gen_ai.prompt", "")
)
row["Output"] = (
json.dumps(lc_message(output_messages[-1]))
if output_messages and isinstance(output_messages[-1], dict)
else attributes.get("gen_ai.completion", "")
)
def _set_io(row: SpanRow, attributes: Mapping[str, str]) -> None:
prompt: Final = _loads(attributes.get("gen_ai.prompt", ""))
completion: Final = _loads(attributes.get("gen_ai.completion", ""))
if row["ObservationType"] == "llm" and isinstance(completion, dict):
prompt_payload: Final = prompt if isinstance(prompt, dict) else MappingProxyType({})
messages: Final = prompt_payload.get("messages") or ((),)
batch: Final = messages[0] if messages and isinstance(messages[0], list) else messages
row["Input"] = (
json.dumps(tuple(lc_message(m) for m in batch if isinstance(m, dict)))
if isinstance(batch, (list, tuple))
else ""
)
generations: Final = completion.get("generations")
first: Final = generations[0] if isinstance(generations, list) and generations else None
item: Final = first[0] if isinstance(first, list) and first else None
message: Final = item.get("message") if isinstance(item, dict) else None
generation: Final = message.get("kwargs") if isinstance(message, dict) else None
if isinstance(generation, dict):
row["Output"] = json.dumps(lc_message(generation))
metadata: Final = generation.get("response_metadata")
row["LiteLLMRequestId"] = metadata.get("id", "") if isinstance(metadata, dict) else ""
return
row["Output"] = attributes.get("gen_ai.completion", "")
return
if row["ObservationType"] == "tool":
output: Final = _tool_output(completion)
row["Input"] = attributes.get("gen_ai.prompt", "")
row["Output"] = output if isinstance(output, str) else json.dumps(output)
return
if row["ObservationType"] == "agent":
_set_agent_io(row, attributes, prompt, completion)
return
row["Input"] = attributes.get("gen_ai.prompt", "")
row["Output"] = attributes.get("gen_ai.completion", "")
@dataclass(frozen=True, slots=True)
class LangSmithNormalizer:
name: str = "langsmith"
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool:
return scope_name == "langsmith" or "langsmith.span.kind" in attributes
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None:
row["ObservationType"] = _span_type(row, attributes)
row["AgentName"] = attributes.get("langsmith.metadata.lc_agent_name", "")
row["Model"] = attributes.get("gen_ai.request.model", "")
_set_io(row, attributes)

View file

@ -0,0 +1,59 @@
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
ChatRole: TypeAlias = Literal["system", "user", "assistant", "tool"]
MESSAGE_ROLES: Final[Mapping[str, ChatRole]] = MappingProxyType(
{"human": "user", "user": "user", "ai": "assistant", "assistant": "assistant", "system": "system", "tool": "tool"}
)
class _ContentBlock(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
type: str = ""
text: str | None = None
_CONTENT_BLOCKS: Final = TypeAdapter(tuple[_ContentBlock, ...])
_NON_TEXT_BLOCKS: Final = frozenset(
{"reasoning", "thinking", "redacted_thinking", "function_call", "tool_use", "tool_call"}
)
def content_text(content: object) -> str:
"""Message content as display text: Responses-style block lists keep only their text blocks."""
if content is None:
return ""
if isinstance(content, str):
return content
try:
blocks: Final = _CONTENT_BLOCKS.validate_python(content)
except ValidationError:
return json.dumps(content)
if not all(block.text is not None or block.type in _NON_TEXT_BLOCKS for block in blocks):
return json.dumps(content)
return "\n\n".join(block.text for block in blocks if block.text is not None)
def lc_message(message: Mapping[str, Any]) -> dict[str, Any]:
"""LangChain serialized message (or plain {role, content}) -> {role, content, tool_calls?}."""
kwargs: Final = message.get("kwargs", message)
role: Final = MESSAGE_ROLES.get(
kwargs.get("type") or kwargs.get("role"), kwargs.get("role") or kwargs.get("type") or ""
)
out: Final[dict[str, Any]] = { # mutable-ok: the framework message is built for JSON serialization
"role": role,
"content": content_text(kwargs.get("content", "")),
}
if kwargs.get("tool_calls"):
out["tool_calls"] = tuple(
{"name": t.get("name"), "args": t.get("args")} # mutable-ok: JSON tool calls need object payloads
for t in kwargs["tool_calls"]
)
if role == "tool" and kwargs.get("name"):
out["name"] = kwargs["name"]
return out

View file

@ -0,0 +1,27 @@
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm.tracing.normalizers.base import to_int
from litellm.tracing.types import SpanRow, SpanType
_OPENINFERENCE_TYPES: Final[Mapping[str, SpanType]] = MappingProxyType({"AGENT": "agent", "LLM": "llm", "TOOL": "tool"})
@dataclass(frozen=True, slots=True)
class OpenInferenceNormalizer:
name: str = "openinference"
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool:
return "openinference.span.kind" in attributes
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None:
kind: Final = attributes.get("openinference.span.kind", "").upper()
row["ObservationType"] = _OPENINFERENCE_TYPES.get(kind, "agent" if not row["ParentSpanId"] else "chain")
row["AgentName"] = attributes.get("agent.name", "")
row["Model"] = attributes.get("llm.model_name", "")
row["Input"] = attributes.get("input.value", "")
row["Output"] = attributes.get("output.value", "")
row["InputTokens"] = to_int(attributes.get("llm.token_count.prompt"))
row["OutputTokens"] = to_int(attributes.get("llm.token_count.completion"))

View file

@ -28,6 +28,7 @@ from litellm.tracing.types import (
TraceScope,
TraceSummary,
)
from litellm.tracing.ui_format import to_ui_content
NANOS_PER_MS: Final = 1_000_000
SPEND_WINDOW_MS: Final = 30 * 60 * 1000
@ -344,5 +345,7 @@ class ClickHouseTraceStore:
span_id=rows[0]["span_id"],
input=rows[0]["input"],
output=rows[0]["output"],
input_ui=to_ui_content(rows[0]["input"]),
output_ui=to_ui_content(rows[0]["output"]),
attributes=rows[0]["attributes"],
)

View file

@ -14,6 +14,8 @@ from typing import Literal
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.tracing.ui_format import UIContent
SpanType = Literal["agent", "llm", "tool", "chain", "framework"]
SpanStatus = Literal["ok", "error", "unset"]
@ -84,6 +86,8 @@ class SpanDetail(TypedDict):
span_id: ReadOnly[str]
input: ReadOnly[str]
output: ReadOnly[str]
input_ui: ReadOnly[UIContent]
output_ui: ReadOnly[UIContent]
attributes: ReadOnly[dict[str, str]]

View file

@ -0,0 +1,158 @@
"""The LiteLLM UI content format: span input / output reduced to messages, key/value fields or plain text."""
import json
from collections.abc import Mapping, Sequence
from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.tracing.normalizers.messages import MESSAGE_ROLES, ChatRole, content_text
class UIToolCall(TypedDict):
name: ReadOnly[str]
arguments: ReadOnly[str]
class UIMessage(TypedDict):
role: ReadOnly[ChatRole]
content: ReadOnly[str]
name: ReadOnly[NotRequired[str]]
tool_calls: ReadOnly[NotRequired[tuple[UIToolCall, ...]]]
class UIField(TypedDict):
key: ReadOnly[str]
value: ReadOnly[str]
class UIMessages(TypedDict):
kind: ReadOnly[Literal["messages"]]
messages: ReadOnly[tuple[UIMessage, ...]]
class UIFields(TypedDict):
kind: ReadOnly[Literal["fields"]]
fields: ReadOnly[tuple[UIField, ...]]
class UIText(TypedDict):
kind: ReadOnly[Literal["text"]]
text: ReadOnly[str]
UIContent: TypeAlias = UIMessages | UIFields | UIText
class _ToolFunction(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
name: str = ""
arguments: JsonValue = None
class _RawToolCall(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
name: str = ""
args: JsonValue = None
arguments: JsonValue = None
function: _ToolFunction | None = None
class _RawMessage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
role: str | None = None
type: str | None = None
content: JsonValue = None
name: str | None = None
tool_calls: tuple[_RawToolCall, ...] | None = None
kwargs: "_RawMessage | None" = None
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_MESSAGE: Final = TypeAdapter(_RawMessage)
_MESSAGES: Final = TypeAdapter(tuple[_RawMessage, ...])
def _unwrapped(message: _RawMessage) -> _RawMessage:
return message.kwargs if message.kwargs is not None else message
def _is_message(message: _RawMessage) -> bool:
has_role: Final = message.role is not None or message.type in MESSAGE_ROLES
return has_role and ("content" in message.model_fields_set or bool(message.tool_calls))
def _arguments_text(arguments: JsonValue) -> str:
match arguments:
case str():
return arguments
case None:
return "{}"
case _:
return json.dumps(arguments)
def _tool_call(call: _RawToolCall) -> UIToolCall:
if call.function is not None:
return UIToolCall(name=call.function.name or call.name, arguments=_arguments_text(call.function.arguments))
return UIToolCall(name=call.name, arguments=_arguments_text(call.arguments if call.args is None else call.args))
def _role(message: _RawMessage, has_tool_calls: bool) -> ChatRole:
"""Known roles and LangChain types map directly; any other role is the assistant when it calls tools, else the user."""
known: Final = MESSAGE_ROLES.get(message.role or message.type or "")
if known is not None:
return known
return "assistant" if has_tool_calls else "user"
def _ui_message(message: _RawMessage) -> UIMessage:
calls: Final = tuple(_tool_call(call) for call in message.tool_calls or ())
role: Final = _role(message, bool(calls))
content: Final = content_text(message.content)
match (message.name or None, calls):
case (None, ()):
return UIMessage(role=role, content=content)
case (None, _):
return UIMessage(role=role, content=content, tool_calls=calls)
case (str() as name, ()):
return UIMessage(role=role, content=content, name=name)
case (str() as name, _):
return UIMessage(role=role, content=content, name=name, tool_calls=calls)
def _messages(parsed: Sequence[JsonValue] | Mapping[str, JsonValue]) -> tuple[_RawMessage, ...] | None:
try:
raw: Final = (
(_MESSAGE.validate_python(parsed),) if isinstance(parsed, Mapping) else _MESSAGES.validate_python(parsed)
)
except ValidationError:
return None
unwrapped: Final = tuple(_unwrapped(message) for message in raw)
return unwrapped if unwrapped and all(_is_message(message) for message in unwrapped) else None
def _field_value(value: JsonValue) -> str:
return value if isinstance(value, str) else json.dumps(value)
def _parsed(raw: str) -> JsonValue:
try:
return _JSON.validate_json(raw)
except ValidationError:
return raw
def to_ui_content(raw: str) -> UIContent:
if not raw:
return UIText(kind="text", text="")
parsed: Final = _parsed(raw)
if not isinstance(parsed, list | dict):
return UIText(kind="text", text=parsed if isinstance(parsed, str) else raw)
messages: Final = _messages(parsed)
if messages is not None:
return UIMessages(kind="messages", messages=tuple(_ui_message(message) for message in messages))
if isinstance(parsed, dict):
return UIFields(kind="fields", fields=tuple(UIField(key=k, value=_field_value(v)) for k, v in parsed.items()))
return UIText(kind="text", text=raw)

View file

@ -1,5 +1,9 @@
from typing import Literal
from pydantic import BaseModel
S3PartitionGranularity = Literal["day", "hour"]
class s3BatchLoggingElement(BaseModel):
"""

View file

@ -1895,22 +1895,22 @@ model LiteLLM_WorkflowMessage {
@@index([run_id])
}
model LiteLLM_Engine {
model LiteLLM_Lens {
id String @id
version Int @default(0)
data Json
}
model LiteLLM_EngineRun {
model LiteLLM_LensRun {
id String @id
engine_id String
lens_id String
created_at DateTime
data Json
@@index([engine_id, created_at])
@@index([lens_id, created_at])
}
model LiteLLM_EngineWorker {
model LiteLLM_LensWorker {
id String @id
token_hash String @unique
data Json

View file

@ -3,8 +3,8 @@ import os
ALLOWED_FILES = [
# The standalone Lens process reuses one client for its entire lifetime, without importing the proxy SDK.
"../../litellm/proxy/engine/worker.py",
"./litellm/proxy/engine/worker.py",
"../../litellm/proxy/lens/worker.py",
"./litellm/proxy/lens/worker.py",
# local files
"../../litellm/__init__.py",
"../../litellm/llms/custom_httpx/http_handler.py",

View file

@ -1,6 +1,7 @@
# Logging integration delivery (behavior features). Grounded in litellm/integrations/.
- {id: logging.s3.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/s3_v2.py", rationale: "Primary audit trail; batch flush no-drop"}
- {id: logging.s3.failure.writes_object, module: logging, tier: P0, event: failure, assertions: [writes_object], exercised_on: [chat_completions, messages], source: "integrations/s3_v2.py", rationale: "Failed calls persisted for compliance"}
- {id: logging.s3.success.partition_layout, module: logging, tier: P1, event: success, assertions: [object_key_layout], exercised_on: [chat_completions], source: "integrations/s3_v2.py / LIT-8985", rationale: "s3_partition_granularity picks the date or date/hour folder every downstream query and lifecycle rule reads"}
- {id: logging.gcs_bucket.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/gcs_bucket/gcs_bucket.py", rationale: "GCS parallel to S3"}
- {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"}
- {id: logging.datadog.stream.exports_metric, module: logging, tier: P0, event: stream, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses], source: "integrations/datadog/datadog.py", rationale: "Streaming aggregates usage after the last chunk; delivery and cost must survive that path"}

View file

@ -106,6 +106,7 @@ UI_BASE_URL = os.environ.get("E2E_UI_BASE_URL", PROXY_BASE_URL).rstrip("/")
CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5")
CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
S3_PARTITION_GRANULARITY = os.environ.get("E2E_S3_PARTITION_GRANULARITY", "day")
LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp")
LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "")

View file

@ -22,11 +22,12 @@ alias per test turns the poll into a cheap prefix listing.
from __future__ import annotations
import math
import re
import time
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_config import CHEAP_ANTHROPIC_MODEL, S3_PARTITION_GRANULARITY, unique_marker
from lifecycle import ResourceManager
from logging_client import (
INVALID_UPSTREAM_API_KEY,
@ -106,6 +107,46 @@ class TestS3LogDelivery:
record.response_cost, outcome.response_cost, rel_tol=1e-9
), f"payload response_cost {record.response_cost!r} must equal the header cost {outcome.response_cost}"
@pytest.mark.covers("logging.s3.success.partition_layout", exercised_on=["chat_completions"])
def test_chat_completions_object_key_follows_the_partition_granularity(
self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager
) -> None:
"""The one object a call writes must sit in the folder layout the proxy's
s3_partition_granularity names: {alias}/{date}/ for day and
{alias}/{date}/{HH}/ for hour, where HH is the hour the object's own
time- file name records. E2E_S3_PARTITION_GRANULARITY tells the test
which one the proxy under test runs."""
_assert_s3_configured(client)
alias = f"s3-layout-{unique_marker()}"
key = client.key_with_alias(alias, models=[CHEAP_ANTHROPIC_MODEL])
resources.defer(lambda: client.delete_key(key))
outcome = first_ok(
client,
lambda: client.chat_raw(
key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {unique_marker()}", max_tokens=16
),
)
body_id = completion_response_id(outcome.body)
assert body_id is not None, "the completion body must carry an id (it names the s3 object)"
records = s3_logs.poll_records(prefix=f"{alias}/", predicate=lambda r: r.id == body_id)
assert len(records) == 1, f"expected exactly ONE s3 object for response {body_id}, got {len(records)}"
file_id = body_id.replace("/", "_").replace(":", "_")
hour_folder = r"(?P<folder_hour>\d{2})/" if S3_PARTITION_GRANULARITY == "hour" else ""
layout = re.compile(
rf"{re.escape(alias)}/\d{{4}}-\d{{2}}-\d{{2}}/{hour_folder}"
rf"time-(?P<file_hour>\d{{2}})-\d{{2}}-\d{{2}}-\d{{6}}_{re.escape(file_id)}\.json"
)
keys = [object_key for object_key in s3_logs.list_keys(f"{alias}/") if file_id in object_key]
assert len(keys) == 1, f"expected one object key for response {body_id}, got {keys}"
match = layout.fullmatch(keys[0])
assert match is not None, f"{keys[0]!r} is outside the {S3_PARTITION_GRANULARITY} layout {layout.pattern!r}"
assert S3_PARTITION_GRANULARITY != "hour" or match.group("folder_hour") == match.group("file_hour"), (
f"the hour folder must be the hour the object's file name records: {keys[0]!r}"
)
@pytest.mark.covers("logging.s3.failure.writes_object", exercised_on=["chat_completions"])
def test_chat_completions_failure_writes_one_object(
self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager

View file

@ -29,6 +29,7 @@ class DatabaseRelay:
self._armed: Final = threading.Event()
self.tripped: Final = threading.Event()
self.refused = 0
self.reconnected: Final = threading.Event()
self._tripped_at = 0.0
self._writers: tuple[asyncio.StreamWriter, ...] = ()
self._ready: Final = threading.Event()
@ -61,6 +62,8 @@ class DatabaseRelay:
self.refused += 1
client_writer.close()
return
if self.tripped.is_set():
self.reconnected.set()
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
self._writers = (*self._writers, client_writer, server_writer)

View file

@ -1,65 +0,0 @@
import asyncio
import os
from collections.abc import AsyncIterator
from datetime import datetime, timezone
from typing import Final
from uuid import uuid4
import pytest
import pytest_asyncio
from prisma import Prisma
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.engine.models import Check, Engine, EngineSettings, Scope, Worker
from litellm.proxy.engine.repository import EngineRepository, WriterDatabase
from litellm.proxy.engine.state import claim_job, queue_job
@pytest_asyncio.fixture(loop_scope="function")
async def engine_db() -> AsyncIterator[Prisma]:
async with Prisma(datasource={"url": os.environ["DATABASE_URL"]}) as db:
yield db
@pytest.mark.asyncio
async def test_concurrent_workers_cannot_both_acquire_the_same_job(engine_db: Prisma) -> None:
now: Final = datetime.now(timezone.utc)
scope: Final = Scope(team_id=uuid4().hex)
repo: Final = EngineRepository(WriterDatabase(PrismaWrapper(engine_db)))
engine: Final = Engine(
id=uuid4().hex,
scope=scope,
settings=EngineSettings(name="Lease test", model="test", checks=(Check(id="c", instruction="Find retries"),)),
created_at=now,
next_run_at=now,
budget_month=now.strftime("%Y-%m"),
)
await repo.create(queue_job(engine, now, uuid4().hex))
try:
workers: Final = tuple(Worker(id=uuid4().hex, name="worker", scope=scope, last_seen=now) for _ in range(2))
results: Final = await asyncio.gather(
*(repo.update(engine.id, lambda e, w=w: claim_job(e, w, now)) for w in workers)
)
stored: Final = await repo.get(engine.id)
assert stored is not None
assert stored.jobs[0].attempts == 1
assert stored.jobs[0].worker_id in tuple(w.id for w in workers)
assert tuple(r.jobs[0].worker_id for r in results if r) == (stored.jobs[0].worker_id, stored.jobs[0].worker_id)
finally:
await engine_db.execute_raw('DELETE FROM "LiteLLM_Engine" WHERE id=$1', engine.id)
@pytest.mark.asyncio
async def test_heartbeat_never_restores_revoked_access(engine_db: Prisma) -> None:
now: Final = datetime.now(timezone.utc)
repo: Final = EngineRepository(WriterDatabase(PrismaWrapper(engine_db)))
worker: Final = Worker(id=uuid4().hex, name="worker", scope=Scope(team_id=uuid4().hex), last_seen=now)
token_hash: Final = uuid4().hex
await repo.save_worker(worker, token_hash)
try:
await repo.save_worker(worker.model_copy(update={"revoked": True}))
await repo.heartbeat(worker.id, now.isoformat())
stored: Final = await repo.worker(token_hash)
assert stored is not None and stored.revoked is True
finally:
await engine_db.execute_raw('DELETE FROM "LiteLLM_EngineWorker" WHERE id=$1', worker.id)

View file

@ -0,0 +1,173 @@
import asyncio
import os
from collections.abc import AsyncIterator
from datetime import datetime, timezone
from pathlib import Path
from typing import Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from uuid import uuid4
import psycopg
import pytest
import pytest_asyncio
from prisma import Prisma
from psycopg import sql
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.lens.models import Check, Lens, LensSettings, Scope, Worker
from litellm.proxy.lens.repository import LensRepository, WriterDatabase
from litellm.proxy.lens.state import claim_job, queue_job
@pytest_asyncio.fixture(loop_scope="function")
async def lens_db() -> AsyncIterator[Prisma]:
async with Prisma(datasource={"url": os.environ["DATABASE_URL"]}) as db:
yield db
@pytest.mark.asyncio
async def test_concurrent_workers_cannot_both_acquire_the_same_job(lens_db: Prisma) -> None:
now: Final = datetime.now(timezone.utc)
scope: Final = Scope(team_id=uuid4().hex)
repo: Final = LensRepository(WriterDatabase(PrismaWrapper(lens_db)))
lens: Final = Lens(
id=uuid4().hex,
scope=scope,
settings=LensSettings(name="Lease test", model="test", checks=(Check(id="c", instruction="Find retries"),)),
created_at=now,
next_run_at=now,
budget_month=now.strftime("%Y-%m"),
)
await repo.create(queue_job(lens, now, uuid4().hex))
try:
workers: Final = tuple(Worker(id=uuid4().hex, name="worker", scope=scope, last_seen=now) for _ in range(2))
results: Final = await asyncio.gather(
*(repo.update(lens.id, lambda e, w=w: claim_job(e, w, now)) for w in workers)
)
stored: Final = await repo.get(lens.id)
assert stored is not None
assert stored.jobs[0].attempts == 1
assert stored.jobs[0].worker_id in tuple(w.id for w in workers)
assert tuple(r.jobs[0].worker_id for r in results if r) == (stored.jobs[0].worker_id, stored.jobs[0].worker_id)
finally:
await lens_db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', lens.id)
@pytest.mark.asyncio
async def test_heartbeat_never_restores_revoked_access(lens_db: Prisma) -> None:
now: Final = datetime.now(timezone.utc)
repo: Final = LensRepository(WriterDatabase(PrismaWrapper(lens_db)))
worker: Final = Worker(id=uuid4().hex, name="worker", scope=Scope(team_id=uuid4().hex), last_seen=now)
token_hash: Final = uuid4().hex
await repo.save_worker(worker, token_hash)
try:
await repo.save_worker(worker.model_copy(update={"revoked": True}))
await repo.heartbeat(worker.id, now.isoformat())
stored: Final = await repo.worker(token_hash)
assert stored is not None and stored.revoked is True
finally:
await lens_db.execute_raw('DELETE FROM "LiteLLM_LensWorker" WHERE id=$1', worker.id)
@pytest.mark.parametrize("populated", (False, True))
@pytest.mark.parametrize("preceding_schema", (False, True))
def test_lens_rename_preserves_saved_data_and_worker_credentials(populated: bool, preceding_schema: bool) -> None:
migrations: Final = (
Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
)
schema: Final = f"lens_migration_{uuid4().hex}"
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
try:
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
for name in ("20260930000000_agent_engine", "20261001000000_lens_run_history"):
connection.execute(sql.SQL((migrations / name / "migration.sql").read_text()))
if populated:
connection.execute(
"""INSERT INTO "LiteLLM_Engine" VALUES ('lens', 7, '{"findings":[{"id":"finding"}]}');
INSERT INTO "LiteLLM_EngineWorker" VALUES ('worker', 'token-hash', '{"analysis_key_id":"key"}');
INSERT INTO "LiteLLM_EngineRun" VALUES ('batch', 'lens', '2026-01-01', '{"cost":1.25}')"""
)
if preceding_schema:
first_schema: Final = f"lens_first_{uuid4().hex}"
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(first_schema)))
connection.execute(
sql.SQL("SET LOCAL search_path TO {}, {}").format(
sql.Identifier(first_schema), sql.Identifier(schema)
)
)
connection.execute(sql.SQL((migrations / "20261001100000_rename_lens" / "migration.sql").read_text()))
connection.execute(sql.SQL((migrations / "20261001100000_rename_lens" / "migration.sql").read_text()))
assert connection.execute('SELECT id, version, data FROM "LiteLLM_Lens"').fetchall() == (
[("lens", 7, {"findings": [{"id": "finding"}]})] if populated else []
)
assert connection.execute('SELECT id, token_hash, data FROM "LiteLLM_LensWorker"').fetchall() == (
[("worker", "token-hash", {"analysis_key_id": "key"})] if populated else []
)
assert connection.execute('SELECT id, lens_id, data FROM "LiteLLM_LensRun"').fetchall() == (
[("batch", "lens", {"cost": 1.25})] if populated else []
)
finally:
connection.rollback()
@pytest.mark.parametrize("entrypoint", ("proxy", "extras-v1", "extras-v2"))
@pytest.mark.parametrize("legacy_table", ("LiteLLM_Engine", "LiteLLM_EngineRun", "LiteLLM_EngineWorker"))
def test_db_push_refuses_legacy_lens_data(monkeypatch: pytest.MonkeyPatch, entrypoint: str, legacy_table: str) -> None:
from litellm_proxy_extras.utils import ProxyExtrasDBManager
from litellm.proxy.db.prisma_client import PrismaManager
database_url: Final = os.environ["DATABASE_URL"]
schema: Final = f"lens_push_{uuid4().hex}"
parsed: Final = urlsplit(database_url)
scoped: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})))
with psycopg.connect(database_url, autocommit=True) as connection:
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
try:
connection.execute(
sql.SQL("CREATE TABLE {} (id TEXT PRIMARY KEY, data JSONB)").format(
sql.Identifier(schema, legacy_table)
)
)
connection.execute(
sql.SQL("INSERT INTO {} VALUES ('saved', '{{\"keep\":true}}')").format(
sql.Identifier(schema, legacy_table)
)
)
monkeypatch.setenv("DATABASE_URL", scoped)
setup: Final = (
PrismaManager.setup_database if entrypoint == "proxy" else ProxyExtrasDBManager.setup_database
)
with pytest.raises(RuntimeError, match="Legacy Lens tables exist"):
setup(use_migrate=False, use_v2_resolver=entrypoint == "extras-v2")
assert connection.execute(
sql.SQL("SELECT id, data FROM {}").format(sql.Identifier(schema, legacy_table))
).fetchall() == [("saved", {"keep": True})]
finally:
connection.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
def test_db_push_creates_fresh_lens_tables_and_preserves_them_on_restart(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy.db.prisma_client import PrismaManager
database_url: Final = os.environ["DATABASE_URL"]
schema: Final = f"lens_fresh_push_{uuid4().hex}"
parsed: Final = urlsplit(database_url)
scoped: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})))
with psycopg.connect(database_url, autocommit=True) as connection:
connection.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
try:
monkeypatch.setenv("DATABASE_URL", scoped)
assert PrismaManager.setup_database(use_migrate=False)
connection.execute(
sql.SQL("INSERT INTO {} (id, data) VALUES ('saved', '{{\"keep\":true}}')").format(
sql.Identifier(schema, "LiteLLM_Lens")
)
)
assert PrismaManager.setup_database(use_migrate=False)
assert connection.execute(
sql.SQL("SELECT id, data FROM {}").format(sql.Identifier(schema, "LiteLLM_Lens"))
).fetchall() == [("saved", {"keep": True})]
finally:
connection.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))

File diff suppressed because it is too large Load diff

View file

@ -12,10 +12,10 @@ from tests.integration._support.process import owned_proxy
from tests.integration.pricing.test_off_peak_pricing import off_peak_window
def delete_lens(engine_id: str) -> None:
write_rows('DELETE FROM "LiteLLM_EngineRun" WHERE engine_id=%s', (engine_id,))
write_rows('DELETE FROM "LiteLLM_Engine" WHERE id=%s', (engine_id,))
assert read_rows('SELECT id FROM "LiteLLM_Engine" WHERE id=%s', (engine_id,)) == []
def delete_lens(lens_id: str) -> None:
write_rows('DELETE FROM "LiteLLM_LensRun" WHERE lens_id=%s', (lens_id,))
write_rows('DELETE FROM "LiteLLM_Lens" WHERE id=%s', (lens_id,))
assert read_rows('SELECT id FROM "LiteLLM_Lens" WHERE id=%s', (lens_id,)) == []
@pytest.mark.parametrize("off_peak", (False, True))
@ -37,12 +37,12 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
key: Final = scenario.key(models=[model], max_budget=1)
key_id: Final = sha256(key.encode()).hexdigest()
worker: Final = gateway.post(
"/engine/workers/register", {"name": "Billing regression", "analysis_key_id": key_id}
"/lens/workers/register", {"name": "Billing regression", "analysis_key_id": key_id}
)
worker_id: Final = string_value(object_value(worker["worker"])["id"])
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_EngineWorker" WHERE id=%s', (worker_id,))
engine: Final = gateway.post(
"/engine",
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_LensWorker" WHERE id=%s', (worker_id,))
lens: Final = gateway.post(
"/lens",
{
"name": "Billing regression",
"model": model,
@ -51,17 +51,17 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
"source": "requests",
},
)
engine_id: Final = string_value(engine["id"])
scenario.cleanups.callback(delete_lens, engine_id)
lens_id: Final = string_value(lens["id"])
scenario.cleanups.callback(delete_lens, lens_id)
worker_key: Final = string_value(worker["token"])
unauthorized: Final = gateway.request(
"POST", "/engine/workers/register", {"name": "Denied", "analysis_key_id": key_id}, key=key
"POST", "/lens/workers/register", {"name": "Denied", "analysis_key_id": key_id}, key=key
)
assert unauthorized.status_code == 403, unauthorized.text
with ThreadPoolExecutor(max_workers=8) as pool:
claims: Final = tuple(
pool.map(
lambda _: gateway.request("POST", "/engine/worker/claim?protocol_version=2", {}, key=worker_key),
lambda _: gateway.request("POST", "/lens/worker/claim?protocol_version=2", {}, key=worker_key),
range(8),
)
)
@ -69,9 +69,9 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
winners: Final = tuple(response.json() for response in claims if response.json() is not None)
assert len(winners) == 1
claim: Final = object_value(winners[0])
assert claim["engine_id"] == engine_id
assert claim["lens_id"] == lens_id
job_id: Final = string_value(object_value(claim["job"])["id"])
path: Final = f"/engine/worker/{engine_id}/{job_id}/model"
path: Final = f"/lens/worker/{lens_id}/{job_id}/model"
result: Final = gateway.post(path, {"prompt": "Inspect this run", "purpose": "extract"}, key=worker_key)
expected: Final = (20 * 0.000001 + 20 * 0.000002) * (0.5 if off_peak else 1)
assert result["cost"] == pytest.approx(expected)
@ -81,7 +81,7 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
seconds=70,
)
assert rows[0]["spend"] == pytest.approx(expected)
assert gateway.get(f"/engine/{engine_id}")["spent"] == pytest.approx(expected)
assert gateway.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected)
raw_hash: Final = gateway.request(
"POST",
"/v1/chat/completions",
@ -105,11 +105,11 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
gateway.post("/key/block", {"key": key})
blocked: Final = gateway.request("POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key)
assert blocked.status_code == 400, blocked.text
assert gateway.get(f"/engine/{engine_id}")["spent"] == pytest.approx(expected)
assert gateway.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected)
replacement: Final = scenario.key(models=[model], rpm_limit=1)
replacement_id: Final = sha256(replacement.encode()).hexdigest()
changed: Final = gateway.request(
"PUT", f"/engine/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
"PUT", f"/lens/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
)
assert changed.status_code == 200, changed.text
billed_replacement: Final = gateway.post(
@ -124,17 +124,17 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway,
seconds=70,
)
assert second_rows[0]["spend"] == pytest.approx(expected)
revoked: Final = gateway.request("DELETE", f"/engine/workers/{worker_id}")
revoked: Final = gateway.request("DELETE", f"/lens/workers/{worker_id}")
assert revoked.status_code == 200, revoked.text
denied_worker: Final = gateway.request(
"POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key
)
assert denied_worker.status_code == 401, denied_worker.text
forbidden_change: Final = gateway.request(
"PUT", f"/engine/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
"PUT", f"/lens/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id}
)
assert forbidden_change.status_code == 409, forbidden_change.text
gateway.post(f"/engine/{engine_id}/cancel", {})
gateway.post(f"/lens/{lens_id}/cancel", {})
@pytest.mark.parametrize("cancel_on_disconnect", (False, True))
@ -174,19 +174,19 @@ def test_worker_spend_logs_do_not_expose_investigation_content(
seconds=70,
)
assert marker in str(retained[0]), "Control must prove this proxy retains ordinary prompts"
worker: Final = isolated.post("/engine/workers/register", {"analysis_key_id": key_id})
worker: Final = isolated.post("/lens/workers/register", {"analysis_key_id": key_id})
worker_id: Final = string_value(object_value(worker["worker"])["id"])
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_EngineWorker" WHERE id=%s', (worker_id,))
engine: Final = isolated.post(
"/engine", {"name": "Log privacy", "model": model, "enabled": False, "context": "Find problems"}
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_LensWorker" WHERE id=%s', (worker_id,))
lens: Final = isolated.post(
"/lens", {"name": "Log privacy", "model": model, "enabled": False, "context": "Find problems"}
)
engine_id: Final = string_value(engine["id"])
scenario.cleanups.callback(delete_lens, engine_id)
lens_id: Final = string_value(lens["id"])
scenario.cleanups.callback(delete_lens, lens_id)
worker_token: Final = string_value(worker["token"])
claim: Final = isolated.post("/engine/worker/claim?protocol_version=2", {}, key=worker_token)
claim: Final = isolated.post("/lens/worker/claim?protocol_version=2", {}, key=worker_token)
job_id: Final = string_value(object_value(claim["job"])["id"])
result: Final = isolated.post(
f"/engine/worker/{engine_id}/{job_id}/model", {"prompt": marker, "purpose": "extract"}, key=worker_token
f"/lens/worker/{lens_id}/{job_id}/model", {"prompt": marker, "purpose": "extract"}, key=worker_token
)
assert result["content"], "The worker must still receive model output"
rows: Final = eventually(
@ -200,4 +200,4 @@ def test_worker_spend_logs_do_not_expose_investigation_content(
assert float(rows[0]["spend"]) == pytest.approx(result["cost"])
assert marker not in str(rows[0])
assert result["content"] not in str(rows[0]["response"])
isolated.post(f"/engine/{engine_id}/cancel", {})
isolated.post(f"/lens/{lens_id}/cancel", {})

View file

@ -13,13 +13,13 @@ from typing import Final
import httpx
from pydantic import BaseModel
from litellm.proxy.engine.analysis import analyze_sample
from litellm.proxy.engine.inference import _SYSTEM
from litellm.proxy.engine.models import (
from litellm.proxy.lens.analysis import analyze_sample
from litellm.proxy.lens.inference import _SYSTEM
from litellm.proxy.lens.models import (
Check,
Claim,
Coverage,
EngineSettings,
LensSettings,
Execution,
ExecutionContent,
Finding,
@ -90,7 +90,7 @@ async def evaluate(
feedback: tuple[Finding, ...] = (),
) -> dict[str, object]:
records: Final = MappingProxyType({case.name: fixtures(case) for case in cases})
settings: Final = EngineSettings(
settings: Final = LensSettings(
name="Quality evaluation",
model=model_name,
checks=checks,
@ -100,7 +100,7 @@ async def evaluate(
)
now: Final = datetime.now(timezone.utc)
claim: Final = Claim(
engine_id="evaluation",
lens_id="evaluation",
findings=feedback,
job=Job(id="evaluation", created_at=now, start=now, end=now, settings=settings, revision=1),
)

View file

@ -15,8 +15,8 @@ from litellm import Router
from litellm.proxy import proxy_server
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.engine import endpoints
from litellm.proxy.engine.models import Check, Coverage, EngineSettings, ModelRequest, Progress, Result, RunRequest
from litellm.proxy.lens import endpoints
from litellm.proxy.lens.models import Check, Coverage, LensSettings, ModelRequest, Progress, Result, RunRequest
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -61,13 +61,13 @@ async def lens_database() -> AsyncIterator[PrismaClient]:
@pytest.mark.asyncio
async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database: PrismaClient) -> None:
admin: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
settings: Final = EngineSettings(
settings: Final = LensSettings(
name="Lifecycle regression",
model="lens-test-analysis",
enabled=False,
checks=(Check(id="retries", instruction="Find unrecovered retries"),),
)
engine: Final = await endpoints.create_engine(settings, admin)
lens: Final = await endpoints.create_lens(settings, admin)
key_id: Final = hashlib.sha256(uuid4().bytes).hexdigest()
await lens_database.db.litellm_verificationtoken.create(data={"token": key_id, "models": ["lens-test-analysis"]})
registration: Final = await endpoints.register_worker(
@ -76,17 +76,17 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
credentials: Final = HTTPAuthorizationCredentials(scheme="Bearer", credentials=registration.token)
worker: Final = await endpoints.worker_auth(credentials)
try:
assert engine.jobs[0].status == "queued"
assert lens.jobs[0].status == "queued"
stored_worker: Final = await endpoints.repository().worker(
hashlib.sha256(registration.token.encode()).hexdigest()
)
assert stored_worker is not None and stored_worker.id == worker.id
assert worker.id == registration.worker.id
listing: Final = await endpoints.list_engines(admin)
assert engine.id in tuple(e.id for e in listing.engines)
listing: Final = await endpoints.list_lenses(admin)
assert lens.id in tuple(e.id for e in listing.lenses)
assert worker.id in tuple(w.id for w in listing.workers)
claims: Final = await asyncio.gather(
*(endpoints.claim_candidate(engine, worker, datetime.now(timezone.utc)) for _ in range(8))
*(endpoints.claim_candidate(lens, worker, datetime.now(timezone.utc)) for _ in range(8))
)
winners: Final = tuple(claim for claim in claims if claim is not None)
assert len(winners) == 1
@ -94,16 +94,16 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
assert claimed.job.worker_id == worker.id
assert (
await endpoints.claim_candidate(
await endpoints.get_engine(engine.id, worker.scope), worker, datetime.now(timezone.utc)
await endpoints.get_lens(lens.id, worker.scope), worker, datetime.now(timezone.utc)
)
is None
)
assert await endpoints.progress(
engine.id, claimed.job.id, Progress(stage="Reviewing", coverage=Coverage(screened=2)), worker
lens.id, claimed.job.id, Progress(stage="Reviewing", coverage=Coverage(screened=2)), worker
)
assert await endpoints.heartbeat(engine.id, claimed.job.id, worker)
assert await endpoints.heartbeat(lens.id, claimed.job.id, worker)
response: Final = await endpoints.model(
engine.id,
lens.id,
claimed.job.id,
ModelRequest(prompt="Return an empty observations list", purpose="extract"),
worker,
@ -111,7 +111,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
{
"type": "http",
"scheme": "http",
"path": "/engine/worker/model",
"path": "/lens/worker/model",
"headers": [],
"client": ("127.0.0.1", 1234),
}
@ -120,7 +120,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
assert '"observations"' in response.content
with pytest.raises(HTTPException) as denied_ip:
await endpoints.model(
engine.id,
lens.id,
claimed.job.id,
ModelRequest(prompt="Must not run", purpose="extract"),
worker,
@ -128,7 +128,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
{
"type": "http",
"scheme": "http",
"path": "/engine/worker/model",
"path": "/lens/worker/model",
"headers": [(b"x-forwarded-for", b"127.0.0.1")],
"client": ("192.0.2.1", 1234),
}
@ -136,7 +136,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
)
assert denied_ip.value.status_code == 403
forwarded: Final = await endpoints.model(
engine.id,
lens.id,
claimed.job.id,
ModelRequest(prompt="Return an empty observations list", purpose="extract"),
worker,
@ -144,7 +144,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
{
"type": "http",
"scheme": "http",
"path": "/engine/worker/model",
"path": "/lens/worker/model",
"headers": [(b"x-forwarded-for", b"127.0.0.1")],
"client": ("192.0.2.100", 1234),
}
@ -153,7 +153,7 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
assert '"observations"' in forwarded.content
with pytest.raises(HTTPException) as spoofed_chain:
await endpoints.model(
engine.id,
lens.id,
claimed.job.id,
ModelRequest(prompt="Must not run", purpose="extract"),
worker,
@ -161,14 +161,14 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
{
"type": "http",
"scheme": "http",
"path": "/engine/worker/model",
"path": "/lens/worker/model",
"headers": [(b"x-forwarded-for", b"127.0.0.1, 192.0.2.1")],
"client": ("192.0.2.100", 1234),
}
),
)
assert spoofed_chain.value.status_code == 403
charged: Final = await endpoints.get_engine(engine.id, worker.scope)
charged: Final = await endpoints.get_lens(lens.id, worker.scope)
assert charged.spent == pytest.approx(response.cost + forwarded.cost)
assert charged.jobs[0].cost == pytest.approx(response.cost + forwarded.cost)
legacy: Final = worker.model_copy(update={"analysis_key_id": None})
@ -178,37 +178,37 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
with pytest.raises(HTTPException) as needs_billing:
await endpoints.claim(authenticated_legacy, protocol_version=2)
assert needs_billing.value.status_code == 409
assert await endpoints.heartbeat(engine.id, claimed.job.id, authenticated_legacy)
assert await endpoints.heartbeat(lens.id, claimed.job.id, authenticated_legacy)
finished: Final = await endpoints.result(
engine.id, claimed.job.id, Result(coverage=Coverage(screened=2)), authenticated_legacy
lens.id, claimed.job.id, Result(coverage=Coverage(screened=2)), authenticated_legacy
)
assert finished.jobs[0].status == "completed"
assert finished.jobs[0].coverage.screened == 2
assert finished.last_scan_at == claimed.job.end
assert finished.next_run_at > finished.jobs[0].finished_at
assert await endpoints.result(engine.id, claimed.job.id, Result(coverage=Coverage()), worker) == finished
assert await endpoints.result(lens.id, claimed.job.id, Result(coverage=Coverage()), worker) == finished
with pytest.raises(HTTPException) as stale:
await endpoints.heartbeat(engine.id, claimed.job.id, worker)
await endpoints.heartbeat(lens.id, claimed.job.id, worker)
assert stale.value.status_code == 409
edited: Final = await endpoints.update_engine(
engine.id, settings.model_copy(update={"interval_minutes": 7}), admin
edited: Final = await endpoints.update_lens(
lens.id, settings.model_copy(update={"interval_minutes": 7}), admin
)
assert edited.revision == engine.revision + 1
rerun: Final = await endpoints.run_engine(engine.id, RunRequest(lookback_hours=3), admin)
assert edited.revision == lens.revision + 1
rerun: Final = await endpoints.run_lens(lens.id, RunRequest(lookback_hours=3), admin)
assert rerun.jobs[0].settings.interval_minutes == 7
assert rerun.jobs[0].created_at - rerun.jobs[0].start == timedelta(hours=3)
history: Final = await endpoints.list_runs(engine.id, admin, offset=0)
history: Final = await endpoints.list_runs(lens.id, admin, offset=0)
assert {job.id for job in history} == {claimed.job.id, rerun.jobs[0].id}
archived: Final = await endpoints.read_run(engine.id, claimed.job.id, admin)
archived: Final = await endpoints.read_run(lens.id, claimed.job.id, admin)
assert archived == finished.jobs[0]
assert archived.settings.interval_minutes == 15
assert archived.findings == ()
with pytest.raises(HTTPException) as foreign_history:
await endpoints.read_run(engine.id, claimed.job.id, UserAPIKeyAuth(team_id="other"))
await endpoints.read_run(lens.id, claimed.job.id, UserAPIKeyAuth(team_id="other"))
assert foreign_history.value.status_code == 403
cancelled: Final = await endpoints.cancel_engine(engine.id, admin)
cancelled: Final = await endpoints.cancel_lens(lens.id, admin)
assert cancelled.jobs[0].status == "cancelled"
assert await endpoints.cancel_engine(engine.id, admin) == cancelled
assert await endpoints.cancel_lens(lens.id, admin) == cancelled
assert await endpoints.revoke_worker(worker.id, admin)
assert await endpoints.repository().set_worker_billing(worker.id, key_id) is None
with pytest.raises(HTTPException) as revoked_billing:
@ -218,10 +218,10 @@ async def test_scan_lifecycle_persists_results_and_revokes_worker(lens_database:
await endpoints.worker_auth(credentials)
assert revoked.value.status_code == 401
with pytest.raises(HTTPException) as foreign:
await endpoints.get_engine(engine.id, endpoints.Scope(team_id="other"))
await endpoints.get_lens(lens.id, endpoints.Scope(team_id="other"))
assert foreign.value.status_code == 404
finally:
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_EngineRun" WHERE engine_id=$1', engine.id)
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_Engine" WHERE id=$1', engine.id)
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_EngineWorker" WHERE id=$1', worker.id)
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_LensRun" WHERE lens_id=$1', lens.id)
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', lens.id)
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_LensWorker" WHERE id=$1', worker.id)
await lens_database.db.execute_raw('DELETE FROM "LiteLLM_VerificationToken" WHERE token=$1', key_id)

View file

@ -6,9 +6,9 @@ from queue import SimpleQueue
from typing import Final
import httpx
from engine.models import (
from lens.models import (
Claim,
EngineSettings,
LensSettings,
Execution,
ExecutionContent,
Job,
@ -17,7 +17,7 @@ from engine.models import (
Sample,
TracePart,
)
from engine.worker import EngineWorker
from lens.worker import LensWorker
async def main() -> None:
@ -25,7 +25,7 @@ async def main() -> None:
claims: Final = iter(("full", "healthy"))
saved: Final = SimpleQueue[Result]()
pages: Final = SimpleQueue[str]()
settings: Final = EngineSettings(name="Storage recovery", model="unused", context="Finish the task", concurrency=1)
settings: Final = LensSettings(name="Storage recovery", model="unused", context="Finish the task", concurrency=1)
execution: Final = Execution(
id="run", source="traces", trace_id="trace", team_id="", name="Task", start_time="", span_count=10000
)
@ -34,7 +34,7 @@ async def main() -> None:
path: Final = request.url.path
if path.endswith("/claim"):
claim: Final = Claim(
engine_id="lens",
lens_id="lens",
job=Job(id=next(claims), created_at=now, start=now, end=now, settings=settings, revision=1),
findings=(),
)
@ -71,7 +71,7 @@ async def main() -> None:
return httpx.Response(200, json=True)
async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client:
worker: Final = EngineWorker(client)
worker: Final = LensWorker(client)
assert await worker.run_once()
failed: Final = saved.get_nowait()
assert failed.error.startswith("Worker temporary storage failed.")

View file

@ -2,7 +2,7 @@
Behavior tests for the cache analytics queries against a real Postgres. The info-route
exclusion and the Unknown grouping live in SQL, so these tests are the ones that exercise
them; the endpoint wiring is unit-tested in
tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py.
tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py.
"""
import json

View file

@ -11,7 +11,8 @@ from fastapi.testclient import TestClient
from litellm.proxy import tracing_endpoints
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.tracing import TracingPayloadTooLargeError
from litellm.tracing import TraceReceiver, TracingPayloadTooLargeError
from litellm.tracing.store import ClickHouseTraceStore
TEAM_KEY = UserAPIKeyAuth(
token="hashed-key", team_id="team-research", org_id="org-1", user_role=LitellmUserRoles.INTERNAL_USER
@ -148,6 +149,24 @@ def test_get_span_404_and_200(client, receiver):
receiver.get_span.assert_awaited_with("t1", "s1", {"team_ids": ("team-research",), "api_key_hash": ""}, "")
def test_get_span_serves_ui_content_from_stored_payloads(client, monkeypatch):
storage = MagicMock()
stored_output = '{"role": "ai", "content": "", "tool_calls": [{"name": "lookup", "args": {"id": 7}}]}'
storage.query = AsyncMock(
return_value=[{"span_id": "s1", "input": '{"city": "Paris"}', "output": stored_output, "attributes": {}}]
)
monkeypatch.setattr(tracing_endpoints, "receiver", TraceReceiver(ClickHouseTraceStore(storage)))
body = client.get("/v1/traces/t1/spans/s1").json()
assert body["output"] == stored_output
assert body["input_ui"] == {"kind": "fields", "fields": [{"key": "city", "value": "Paris"}]}
assert body["output_ui"] == {
"kind": "messages",
"messages": [
{"role": "assistant", "content": "", "tool_calls": [{"name": "lookup", "arguments": '{"id": 7}'}]}
],
}
def test_trace_detail_passes_scoped_reference(client, receiver):
receiver.get_trace.return_value = {"summary": {"trace_id": "t1"}, "agents": [], "spans": []}
assert client.get("/v1/traces/t1?trace_ref=run-one").status_code == 200

View file

@ -0,0 +1,68 @@
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm.tracing.normalizers import (
NORMALIZERS,
GenAISemconvNormalizer,
LangSmithNormalizer,
OpenInferenceNormalizer,
select_normalizer,
)
from litellm.tracing.types import SpanRow
_NO_ATTRIBUTES: Final[Mapping[str, str]] = MappingProxyType({})
def test_langsmith_scope_selects_langsmith_without_any_attributes():
assert isinstance(select_normalizer("langsmith", _NO_ATTRIBUTES), LangSmithNormalizer)
def test_langsmith_kind_attribute_selects_langsmith_under_any_scope():
assert isinstance(select_normalizer("other", MappingProxyType({"langsmith.span.kind": "llm"})), LangSmithNormalizer)
def test_langsmith_wins_over_openinference_when_both_markers_present():
attributes: Final = MappingProxyType({"langsmith.span.kind": "llm", "openinference.span.kind": "LLM"})
assert isinstance(select_normalizer("other", attributes), LangSmithNormalizer)
def test_openinference_kind_attribute_selects_openinference():
assert isinstance(
select_normalizer("other", MappingProxyType({"openinference.span.kind": "LLM"})), OpenInferenceNormalizer
)
def test_unmarked_span_falls_back_to_genai():
assert isinstance(
select_normalizer("other", MappingProxyType({"gen_ai.operation.name": "chat"})), GenAISemconvNormalizer
)
def test_empty_registry_falls_back_to_genai():
assert isinstance(select_normalizer("langsmith", _NO_ATTRIBUTES, registry=()), GenAISemconvNormalizer)
def test_registry_names_are_unique():
names: Final = tuple(n.name for n in NORMALIZERS)
assert len(names) == len(frozenset(names))
@dataclass(frozen=True, slots=True)
class _CustomNormalizer:
name: str = "custom"
def matches(self, scope_name: str, attributes: Mapping[str, str]) -> bool:
return scope_name == "custom-sdk"
def normalize(self, row: SpanRow, attributes: Mapping[str, str]) -> None:
return None
def test_normalizer_inserted_ahead_in_custom_registry_wins_only_where_it_matches():
registry: Final = (_CustomNormalizer(), *NORMALIZERS)
assert isinstance(
select_normalizer("custom-sdk", MappingProxyType({"langsmith.span.kind": "llm"}), registry), _CustomNormalizer
)
assert isinstance(select_normalizer("langsmith", _NO_ATTRIBUTES, registry), LangSmithNormalizer)

View file

@ -125,6 +125,49 @@ def test_incomplete_langsmith_completion_preserves_the_export(completion):
assert rows[0]["Output"] == completion
def test_llm_block_list_content_keeps_only_text():
reasoning = {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAB-opaque"}
history = [reasoning, {"type": "text", "text": "Earlier answer", "annotations": []}]
answer = [reasoning, {"type": "text", "text": "Part one"}, {"type": "text", "text": "Part two"}]
prompt = {
"messages": [
[
{"kwargs": {"type": "human", "content": "refund please"}},
{"kwargs": {"type": "ai", "content": history}},
{"kwargs": {"type": "ai", "content": [reasoning]}},
]
]
}
completion = {"generations": [[{"message": {"kwargs": {"type": "ai", "content": answer}}}]]}
span = _span(
"ChatOpenAI",
b"\x03" * 8,
b"\x02" * 8,
langsmith__span__kind="llm",
gen_ai__prompt=json.dumps(prompt),
gen_ai__completion=json.dumps(completion),
)
rows = decode_otlp(_export(span, scope="langsmith"), "application/x-protobuf")
assert [m["content"] for m in json.loads(rows[0]["Input"])] == ["refund please", "Earlier answer", ""]
assert json.loads(rows[0]["Output"])["content"] == "Part one\n\nPart two"
assert "encrypted_content" not in rows[0]["Input"] + rows[0]["Output"]
def test_llm_unrecognized_list_content_is_kept_as_json():
content = [{"type": "image_url", "image_url": {"url": "https://x.test/a.png"}}]
completion = {"generations": [[{"message": {"kwargs": {"type": "ai", "content": content}}}]]}
span = _span(
"ChatOpenAI",
b"\x03" * 8,
b"\x02" * 8,
langsmith__span__kind="llm",
gen_ai__prompt='{"messages": [[{"kwargs": {"type": "human", "content": "hi"}}]]}',
gen_ai__completion=json.dumps(completion),
)
rows = decode_otlp(_export(span, scope="langsmith"), "application/x-protobuf")
assert json.loads(json.loads(rows[0]["Output"])["content"]) == content
def test_task_tool_output_is_subagent_final_message_text(rows_by_name):
task = rows_by_name["task"]
assert json.loads(task["Input"])["subagent_type"] == "researcher"
@ -189,6 +232,64 @@ def test_long_values_are_truncated_with_marker():
assert len(task["Input"].split("…")[0].encode()) <= 100
def test_long_message_history_drops_middle_messages_and_stays_valid_json():
history = [{"kwargs": {"type": "human", "content": f"turn {i} " + "x" * 60}} for i in range(12)]
prompt = json.dumps({"messages": [[{"kwargs": {"type": "system", "content": "be brief"}}, *history]]})
completion = json.dumps({"generations": [[{"message": {"kwargs": {"type": "ai", "content": "ok"}}}]]})
span = _span(
"ChatOpenAI",
b"\x03" * 8,
b"\x02" * 8,
langsmith__span__kind="llm",
gen_ai__prompt=prompt,
gen_ai__completion=completion,
)
with patch.object(decode, "OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 400):
rows = decode_otlp(_export(span, scope="langsmith"), "application/x-protobuf")
messages = json.loads(rows[0]["Input"])
assert len(rows[0]["Input"].encode()) <= 400
assert messages[0]["content"] == "be brief"
assert "earlier messages truncated" in messages[1]["content"]
assert messages[-1]["content"].startswith("turn 11 ")
kept = int(messages[1]["content"].split("[")[1].split()[0])
assert kept + len(messages) - 2 == 12
@pytest.mark.parametrize(
"messages",
[
[{"role": "system", "content": "s" * 2000}, {"role": "user", "content": "short question"}],
[{"role": "user", "content": "a" * 900}, {"role": "assistant", "content": "b" * 900}],
[
{"role": "system", "content": "s" * 900},
{"role": "user", "content": "middle"},
{"role": "user", "content": "q" * 900},
],
],
ids=["huge-first-message", "two-messages", "huge-first-and-last"],
)
def test_oversized_message_arrays_are_shortened_not_cut(messages):
with patch.object(decode, "OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 400):
out = decode._truncate_payload(json.dumps(messages))
assert len(out.encode()) <= 400
kept = json.loads(out)
assert kept[0]["role"] == messages[0]["role"]
assert kept[-1]["role"] == messages[-1]["role"]
assert all(isinstance(m["content"], str) for m in kept)
def test_oversized_non_content_fields_still_fit_the_limit():
heavy = {"role": "assistant", "content": "x", "tool_calls": [{"name": "t", "args": {"blob": "z" * 3000}}]}
messages = [heavy, {"role": "user", "content": "—" * 900}]
with patch.object(decode, "OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 400):
out = decode._truncate_payload(json.dumps(messages))
kept = json.loads(out)
assert len(out.encode()) <= 400
assert [m["role"] for m in kept] == ["assistant", "user"]
assert kept[0]["content"].startswith("x")
assert kept[1]["content"].startswith("\u2014")
# ---------------------------------------------------------------- status / exceptions

View file

@ -306,11 +306,16 @@ async def test_get_span_not_found_and_found():
store = ClickHouseTraceStore(client)
scope: TraceScope = {"team_ids": (), "api_key_hash": ""}
assert await store.get_span("t", "s", scope) is None
client.query = AsyncMock(return_value=[{"span_id": "s", "input": "i", "output": "o", "attributes": {"k": "v"}}])
stored_input = '[{"role": "user", "content": "hi"}]'
client.query = AsyncMock(
return_value=[{"span_id": "s", "input": stored_input, "output": '{"ok": true}', "attributes": {"k": "v"}}]
)
assert await store.get_span("t", "s", scope) == {
"span_id": "s",
"input": "i",
"output": "o",
"input": stored_input,
"output": '{"ok": true}',
"input_ui": {"kind": "messages", "messages": ({"role": "user", "content": "hi"},)},
"output_ui": {"kind": "fields", "fields": ({"key": "ok", "value": "true"},)},
"attributes": {"k": "v"},
}

View file

@ -0,0 +1,119 @@
import json
import pytest
from litellm.tracing.ui_format import to_ui_content
def test_message_array_maps_roles_and_keeps_order():
raw = json.dumps(
[
{"role": "system", "content": "be brief"},
{"role": "human", "content": "hi"},
{"role": "tool", "name": "lookup", "content": "42"},
{"role": "narrator", "content": "aside"},
]
)
assert to_ui_content(raw) == {
"kind": "messages",
"messages": (
{"role": "system", "content": "be brief"},
{"role": "user", "content": "hi"},
{"role": "tool", "content": "42", "name": "lookup"},
{"role": "user", "content": "aside"},
),
}
@pytest.mark.parametrize(
"call",
[
{"name": "get_plan", "args": {"customer_id": "c-1"}},
{"name": "get_plan", "arguments": '{"customer_id": "c-1"}'},
{"id": "call_1", "type": "function", "function": {"name": "get_plan", "arguments": '{"customer_id": "c-1"}'}},
],
)
def test_single_assistant_message_with_tool_call(call: dict[str, object]):
content = to_ui_content(json.dumps({"role": "assistant", "content": None, "tool_calls": [call]}))
assert content["kind"] == "messages"
(message,) = content["messages"]
assert message["role"] == "assistant"
assert message["content"] == ""
calls = message.get("tool_calls")
assert calls is not None and len(calls) == 1
assert calls[0]["name"] == "get_plan"
assert json.loads(calls[0]["arguments"]) == {"customer_id": "c-1"}
def test_unknown_role_with_tool_calls_is_assistant():
content = to_ui_content(json.dumps({"role": "model", "content": "", "tool_calls": [{"name": "f", "args": None}]}))
assert content == {
"kind": "messages",
"messages": ({"role": "assistant", "content": "", "tool_calls": ({"name": "f", "arguments": "{}"},)},),
}
def test_block_list_content_keeps_text_and_drops_reasoning():
raw = json.dumps(
{
"role": "assistant",
"content": [
{"type": "reasoning", "encrypted_content": "opaque"},
{"type": "thinking", "thinking": "hidden chain"},
{"type": "text", "text": "first"},
{"type": "text", "text": "second"},
],
}
)
assert to_ui_content(raw) == {
"kind": "messages",
"messages": ({"role": "assistant", "content": "first\n\nsecond"},),
}
def test_langchain_kwargs_shape():
raw = json.dumps(
[
{"lc": 1, "type": "constructor", "kwargs": {"type": "human", "content": "question"}},
{"kwargs": {"type": "ai", "content": "", "tool_calls": [{"name": "search", "args": {"q": "x"}}]}},
]
)
content = to_ui_content(raw)
assert content["kind"] == "messages"
human, ai = content["messages"]
assert human == {"role": "user", "content": "question"}
assert ai["role"] == "assistant"
assert ai.get("tool_calls") == ({"name": "search", "arguments": '{"q": "x"}'},)
def test_plain_object_becomes_fields_in_key_order():
raw = json.dumps({"zeta": "plain", "alpha": {"nested": [1, 2]}, "count": 3, "missing": None})
assert to_ui_content(raw) == {
"kind": "fields",
"fields": (
{"key": "zeta", "value": "plain"},
{"key": "alpha", "value": '{"nested": [1, 2]}'},
{"key": "count", "value": "3"},
{"key": "missing", "value": "null"},
),
}
def test_object_with_role_but_no_content_is_fields():
assert to_ui_content('{"role": "admin", "user_id": "u1"}')["kind"] == "fields"
def test_json_string_becomes_its_text():
assert to_ui_content(json.dumps('line one\n"quoted"')) == {"kind": "text", "text": 'line one\n"quoted"'}
@pytest.mark.parametrize(
"raw",
['[{"role": "user", "content": "cut of', "plain words", "42", "[1, 2]", "[]"],
)
def test_non_message_non_object_payloads_keep_the_raw_string(raw: str):
assert to_ui_content(raw) == {"kind": "text", "text": raw}
def test_empty_is_empty_text():
assert to_ui_content("") == {"kind": "text", "text": ""}

View file

@ -312,3 +312,10 @@ def test_prompts_only_payload_returns_copy_with_response_cleared():
assert stripped["messages"] == TEST_MESSAGES
assert stripped is not payload
assert payload == snapshot
def test_legacy_s3_logger_ignores_partition_granularity_and_keeps_daily_folder():
mock_s3_client = _run_log_event({"s3_bucket_name": "b", "s3_path": "logs", "s3_partition_granularity": "hour"})
key = mock_s3_client.put_object.call_args.kwargs["Key"]
assert key.startswith("logs/2026-07-30/time-12-00-00-")

View file

@ -2552,6 +2552,253 @@ def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callba
assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name)
_PARTITION_START: Final = datetime(2026, 9, 29, 14, 5, 9, 123456)
_PARTITION_ID: Final = "chatcmpl-partition"
def _partition_payload(response_id: str = _PARTITION_ID) -> StandardLoggingPayload:
return StandardLoggingPayload(
id=response_id,
metadata={"user_api_key_team_alias": "team-a", "user_api_key_alias": "key-a"},
messages=[],
)
def _partition_logger(
monkeypatch: pytest.MonkeyPatch, callback_params: dict[str, object], **kwargs: object
) -> S3Logger:
import litellm
monkeypatch.setattr(
litellm,
"s3_callback_params",
{"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_path": "logs", **callback_params},
)
return S3Logger(
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_use_team_prefix=True,
s3_use_key_prefix=True,
**kwargs,
)
_DAILY_KEY: Final = f"logs/team-a/key-a/2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json"
_HOURLY_KEY: Final = f"logs/team-a/key-a/2026-09-29/14/time-14-05-09-123456_{_PARTITION_ID}.json"
@pytest.mark.parametrize(
("callback_params", "expected_key"),
[
({}, _DAILY_KEY),
({"s3_partition_granularity": None}, _DAILY_KEY),
({"s3_partition_granularity": "day"}, _DAILY_KEY),
({"s3_partition_granularity": "hour"}, _HOURLY_KEY),
],
)
def test_partition_granularity_sets_request_log_folder(
monkeypatch: pytest.MonkeyPatch, callback_params: dict[str, object], expected_key: str
) -> None:
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
logger = _partition_logger(monkeypatch, callback_params)
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
assert element is not None
assert element.s3_object_key == expected_key
@pytest.mark.parametrize("invalid", ["hourly", "HOUR", "1", 1, True])
def test_invalid_partition_granularity_warns_and_keeps_daily_folder(
monkeypatch: pytest.MonkeyPatch, invalid: object
) -> None:
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
with patch("litellm.integrations.s3.verbose_logger") as mock_logger:
logger = _partition_logger(monkeypatch, {"s3_partition_granularity": invalid})
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
second = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
assert element is not None
assert second is not None
assert element.s3_object_key == second.s3_object_key == _DAILY_KEY
mock_logger.warning.assert_called_once()
assert mock_logger.warning.call_args.args[1:] == (invalid,)
def test_partition_granularity_reads_admin_ui_env_var_below_callback_params(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("S3_PARTITION_GRANULARITY", "hour")
from_env = _partition_logger(monkeypatch, {}).create_s3_batch_logging_element(
_PARTITION_START, _partition_payload()
)
from_params = _partition_logger(monkeypatch, {"s3_partition_granularity": "day"}).create_s3_batch_logging_element(
_PARTITION_START, _partition_payload()
)
assert from_env is not None and from_env.s3_object_key == _HOURLY_KEY
assert from_params is not None and from_params.s3_object_key == _DAILY_KEY
def test_partition_granularity_constructor_argument_and_os_environ_reference(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
monkeypatch.setenv("MY_S3_PARTITION", "hour")
from_ctor = _partition_logger(monkeypatch, {}, s3_partition_granularity="hour")
from_secret = _partition_logger(monkeypatch, {"s3_partition_granularity": "os.environ/MY_S3_PARTITION"})
for logger in (from_ctor, from_secret):
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload())
assert element is not None and element.s3_object_key == _HOURLY_KEY
def test_hourly_partition_long_key_keeps_hour_folder_within_s3_limit(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.constants import MAX_S3_OBJECT_KEY_BYTES
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
logger = _partition_logger(monkeypatch, {"s3_partition_granularity": "hour", "s3_path": "p" * 1100})
element = logger.create_s3_batch_logging_element(_PARTITION_START, _partition_payload("r" * 600))
assert element is not None
assert len(element.s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES
assert re.search(r"/2026-09-29/14/[0-9a-f]{64}\.json$", element.s3_object_key)
@pytest.mark.asyncio
@pytest.mark.parametrize(("granularity", "hour_folder"), [("hour", True), ("day", False), (None, False)])
async def test_audit_log_key_follows_audit_callback_params_partition_granularity(
monkeypatch: pytest.MonkeyPatch, granularity: str | None, hour_folder: bool
) -> None:
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
logger = S3Logger(
s3_callback_params_override={
"s3_bucket_name": "audit-bucket",
"s3_path": "audit",
"s3_partition_granularity": granularity,
}
)
await logger.async_log_audit_log_event({"id": "audit-1"})
(element,) = logger.log_queue
match = re.fullmatch(
r"audit/audit_logs/\d{4}-\d{2}-\d{2}/(?:(\d{2})/)?(\d{2})-\d{2}-\d{2}_audit-1\.json", element.s3_object_key
)
assert match is not None, element.s3_object_key
assert (match.group(1) is not None) is hour_folder
if hour_folder:
assert match.group(1) == match.group(2)
@pytest.mark.asyncio
async def test_hourly_batch_file_upload_writes_one_file_per_hour_folder(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
logger = _partition_logger(monkeypatch, {"s3_partition_granularity": "hour"}, s3_batch_file_upload=True)
put = _RecordingPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
before = logger.create_s3_batch_logging_element(datetime(2026, 9, 29, 13, 59, 59), _partition_payload("before"))
after = logger.create_s3_batch_logging_element(datetime(2026, 9, 29, 14, 0, 1), _partition_payload("after"))
assert before is not None and after is not None
logger.log_queue = [before, after]
await logger.async_send_batch()
by_folder = {
re.sub(r"/batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", "", url.split(".com/", 1)[-1]): data
for url, data, _headers in put.calls
}
assert sorted(by_folder) == ["logs/team-a/key-a/2026-09-29/13", "logs/team-a/key-a/2026-09-29/14"]
assert [json.loads(line)["id"] for line in (by_folder["logs/team-a/key-a/2026-09-29/13"] or "").splitlines()] == [
"before"
]
assert [json.loads(line)["id"] for line in (by_folder["logs/team-a/key-a/2026-09-29/14"] or "").splitlines()] == [
"after"
]
@pytest.mark.parametrize("granularity", [None, "day", "hour"])
def test_cold_storage_object_key_matches_the_uploaded_request_log_key(
monkeypatch: pytest.MonkeyPatch, granularity: str | None
) -> None:
import litellm
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
monkeypatch.setattr(
litellm,
"s3_callback_params",
{"s3_bucket_name": "test-bucket", "s3_path": "coldlogs", "s3_partition_granularity": granularity},
)
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2")
logger = S3Logger()
uploaded = logger.create_s3_batch_logging_element(
_PARTITION_START, StandardLoggingPayload(id=_PARTITION_ID, metadata={}, messages=[])
)
monkeypatch.setattr(litellm, "callbacks", [logger])
cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=_PARTITION_START, response_id=_PARTITION_ID
)
assert uploaded is not None
assert cold_key == uploaded.s3_object_key
assert ("/2026-09-29/14/" in cold_key) is (granularity == "hour")
def test_cold_storage_key_matches_upload_when_env_var_changes_mid_request(monkeypatch: pytest.MonkeyPatch) -> None:
import litellm
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_path": "coldlogs"})
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2")
logger = S3Logger()
monkeypatch.setattr(litellm, "callbacks", [logger])
cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=_PARTITION_START, response_id=_PARTITION_ID
)
monkeypatch.setenv("S3_PARTITION_GRANULARITY", "hour")
uploaded = logger.create_s3_batch_logging_element(
_PARTITION_START,
StandardLoggingPayload(id=_PARTITION_ID, metadata={"cold_storage_object_key": cold_key}, messages=[]),
)
assert uploaded is not None
assert cold_key == uploaded.s3_object_key == f"coldlogs/2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json"
def test_hour_upload_ignores_a_cold_storage_key_owned_by_another_logger(monkeypatch: pytest.MonkeyPatch) -> None:
import litellm
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
monkeypatch.delenv("S3_PARTITION_GRANULARITY", raising=False)
monkeypatch.setattr(
litellm, "s3_callback_params", {"s3_bucket_name": "test-bucket", "s3_partition_granularity": "hour"}
)
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "gcs_bucket")
logger = S3Logger()
cold_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=_PARTITION_START, response_id=_PARTITION_ID
)
uploaded = logger.create_s3_batch_logging_element(
_PARTITION_START,
StandardLoggingPayload(id=_PARTITION_ID, metadata={"cold_storage_object_key": cold_key}, messages=[]),
)
assert cold_key == f"2026-09-29/time-14-05-09-123456_{_PARTITION_ID}.json"
assert uploaded is not None
assert uploaded.s3_object_key == f"2026-09-29/14/time-14-05-09-123456_{_PARTITION_ID}.json"
@pytest.mark.parametrize("callback_name", ["s3", "s3_v2"])
def test_partition_granularity_is_exposed_to_admin_ui(callback_name: str) -> None:
from litellm.integrations.custom_logger import CustomLogger
assert "S3_PARTITION_GRANULARITY" in CustomLogger.get_callback_env_vars(callback_name)
def _element(payload: dict[str, object], key_suffix: str) -> s3BatchLoggingElement:
return s3BatchLoggingElement(
s3_object_key=f"2025-09-14/test-{key_suffix}.json",

View file

@ -3227,6 +3227,7 @@ async def test_e2e_generate_cold_storage_object_key_successful():
prefix="", # No prefix for cold storage
start_time=start_time,
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345",
partition_granularity="day",
)
# Verify the result
@ -3276,6 +3277,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
prefix="",
start_time=start_time,
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345",
partition_granularity="day",
)
# Verify the result
@ -3320,6 +3322,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
prefix="",
start_time=start_time,
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345",
partition_granularity="day",
)
# Verify the result

View file

@ -150,7 +150,7 @@ class _AgentPersistence:
return self.row
async def update(self, *, data: Mapping[str, object], **kwargs: object) -> LiteLLM_AgentsTable:
from tests.test_litellm.proxy.agent_endpoints.test_agent_registry import _stored_agent_row
from tests.unit.proxy.agent_endpoints.test_agent_registry import _stored_agent_row
self.row = _stored_agent_row({**self.row.model_dump(), **data})
return self.row
@ -163,7 +163,7 @@ def test_identity_settings_edit_preserves_runtime_configuration_on_readback(
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
from tests.test_litellm.proxy.agent_endpoints.test_agent_registry import _stored_agent_row
from tests.unit.proxy.agent_endpoints.test_agent_registry import _stored_agent_row
runtime: Final = {
"agent_card_params": {} if cardless else _sample_agent_card_params(),
@ -1330,7 +1330,7 @@ def test_identity_providers_honor_issuer_specific_audiences_and_global_fallback(
def test_mode_only_edit_requires_the_existing_identity_sso_tenant(
monkeypatch: pytest.MonkeyPatch, change: PatchAgentRequest
) -> None:
from tests.test_litellm.proxy.agent_endpoints.test_managed_identity import BINDING, TENANT, managed_agent
from tests.unit.proxy.agent_endpoints.test_managed_identity import BINDING, TENANT, managed_agent
monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,))
monkeypatch.delenv("MICROSOFT_TENANT", raising=False)
@ -1342,7 +1342,7 @@ def test_mode_only_edit_requires_the_existing_identity_sso_tenant(
def test_identity_only_edit_preserves_delegated_mode_validation(monkeypatch: pytest.MonkeyPatch) -> None:
from tests.test_litellm.proxy.agent_endpoints.test_managed_identity import BINDING, managed_agent
from tests.unit.proxy.agent_endpoints.test_managed_identity import BINDING, managed_agent
monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,))
monkeypatch.delenv("MICROSOFT_TENANT", raising=False)
@ -1637,7 +1637,7 @@ def test_agent_detail_cache_miss_preserves_admin_identity_visibility(role, monke
def test_invalid_identity_and_untrusted_tenant_cannot_be_registered(
monkeypatch: pytest.MonkeyPatch, trusted: bool
) -> None:
from tests.test_litellm.proxy.agent_endpoints.test_managed_identity import BINDING
from tests.unit.proxy.agent_endpoints.test_managed_identity import BINDING
configuration: Final = BINDING.model_dump(
exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"}

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