diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 3a207ca1778..dc6c1f2ffaa 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -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 diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index e425c313d6a..792da5ae09c 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -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( diff --git a/.github/workflows/lens-worker.yml b/.github/workflows/lens-worker.yml index 53334abaf88..54ec2593ed8 100644 --- a/.github/workflows/lens-worker.yml +++ b/.github/workflows/lens-worker.yml @@ -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 diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 4d5b9074f03..e8483771569 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -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 diff --git a/Makefile b/Makefile index b6ce1ebabaf..ed1f9e3e1e5 100644 --- a/Makefile +++ b/Makefile @@ -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 diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 51a4d8f716c..80ca0ef22bb 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -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", diff --git a/deploy/lens/Dockerfile b/deploy/lens/Dockerfile index 360211194e2..bab5cba94ac 100644 --- a/deploy/lens/Dockerfile +++ b/deploy/lens/Dockerfile @@ -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"] diff --git a/deploy/lens/Dockerfile.dockerignore b/deploy/lens/Dockerfile.dockerignore index 8478be71be7..6db1cbdb50a 100644 --- a/deploy/lens/Dockerfile.dockerignore +++ b/deploy/lens/Dockerfile.dockerignore @@ -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 diff --git a/deploy/lens/README.md b/deploy/lens/README.md index d4bcddf8613..7a80bafa59e 100644 --- a/deploy/lens/README.md +++ b/deploy/lens/README.md @@ -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 diff --git a/deploy/lens/compose.yaml b/deploy/lens/compose.yaml index 0af04814c1e..d41cb8eb203 100644 --- a/deploy/lens/compose.yaml +++ b/deploy/lens/compose.yaml @@ -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} diff --git a/deploy/lens/screenshots/after.png b/deploy/lens/screenshots/after.png deleted file mode 100644 index 983625e2f42..00000000000 Binary files a/deploy/lens/screenshots/after.png and /dev/null differ diff --git a/deploy/lens/screenshots/before.png b/deploy/lens/screenshots/before.png deleted file mode 100644 index 5022cd2bb18..00000000000 Binary files a/deploy/lens/screenshots/before.png and /dev/null differ diff --git a/deploy/lens/screenshots/finding.png b/deploy/lens/screenshots/finding.png deleted file mode 100644 index dc8250f976e..00000000000 Binary files a/deploy/lens/screenshots/finding.png and /dev/null differ diff --git a/deploy/lens/screenshots/progress.png b/deploy/lens/screenshots/progress.png deleted file mode 100644 index f69ff1c45b3..00000000000 Binary files a/deploy/lens/screenshots/progress.png and /dev/null differ diff --git a/deploy/lens/screenshots/setup.png b/deploy/lens/screenshots/setup.png deleted file mode 100644 index 731fa012dbe..00000000000 Binary files a/deploy/lens/screenshots/setup.png and /dev/null differ diff --git a/deploy/lens/screenshots/trace.png b/deploy/lens/screenshots/trace.png deleted file mode 100644 index ef0178376d5..00000000000 Binary files a/deploy/lens/screenshots/trace.png and /dev/null differ diff --git a/deploy/lens/screenshots/worker-billing-after.png b/deploy/lens/screenshots/worker-billing-after.png deleted file mode 100644 index cb8b6991036..00000000000 Binary files a/deploy/lens/screenshots/worker-billing-after.png and /dev/null differ diff --git a/deploy/lens/screenshots/worker-billing-before.png b/deploy/lens/screenshots/worker-billing-before.png deleted file mode 100644 index 093305fb9e7..00000000000 Binary files a/deploy/lens/screenshots/worker-billing-before.png and /dev/null differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql new file mode 100644 index 00000000000..8be0ef0d031 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql @@ -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 $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 75dc7ddde9d..6f285e9dc39 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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 diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 8a83c786e02..2f74df63367 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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(), diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index e92af4b0861..2e2f3f2ce5a 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -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 = [ diff --git a/litellm/constants.py b/litellm/constants.py index b3f5b0471f4..7e1e63a112b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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)) diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 190c283d087..38928f67f42 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -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" diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index f330ca8e0ac..129fceb40bf 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -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: diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 88d7906cc4b..ea9e6a84c93 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -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) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 154893b6c21..097ca2bcd54 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7b735152065..c52d3748ce2 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 649381f24a2..4081443fda9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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", ], ) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 0524d015047..afaa41e97a5 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -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( [ diff --git a/litellm/proxy/engine/__init__.py b/litellm/proxy/lens/__init__.py similarity index 100% rename from litellm/proxy/engine/__init__.py rename to litellm/proxy/lens/__init__.py diff --git a/litellm/proxy/engine/analysis.py b/litellm/proxy/lens/analysis.py similarity index 100% rename from litellm/proxy/engine/analysis.py rename to litellm/proxy/lens/analysis.py diff --git a/litellm/proxy/engine/billing.py b/litellm/proxy/lens/billing.py similarity index 100% rename from litellm/proxy/engine/billing.py rename to litellm/proxy/lens/billing.py diff --git a/litellm/proxy/engine/endpoints.py b/litellm/proxy/lens/endpoints.py similarity index 70% rename from litellm/proxy/engine/endpoints.py rename to litellm/proxy/lens/endpoints.py index d085fe7b289..fbe6e4d9b94 100644 --- a/litellm/proxy/engine/endpoints.py +++ b/litellm/proxy/lens/endpoints.py @@ -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) diff --git a/litellm/proxy/engine/inference.py b/litellm/proxy/lens/inference.py similarity index 91% rename from litellm/proxy/engine/inference.py rename to litellm/proxy/lens/inference.py index 687f3832a1b..28b6babf177 100644 --- a/litellm/proxy/engine/inference.py +++ b/litellm/proxy/lens/inference.py @@ -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) diff --git a/litellm/proxy/engine/models.py b/litellm/proxy/lens/models.py similarity index 95% rename from litellm/proxy/engine/models.py rename to litellm/proxy/lens/models.py index 33e70ff3eca..eb88801d065 100644 --- a/litellm/proxy/engine/models.py +++ b/litellm/proxy/lens/models.py @@ -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, ...] diff --git a/litellm/proxy/engine/repository.py b/litellm/proxy/lens/repository.py similarity index 66% rename from litellm/proxy/engine/repository.py rename to litellm/proxy/lens/repository.py index 7e3c2f27282..4aa840e181b 100644 --- a/litellm/proxy/engine/repository.py +++ b/litellm/proxy/lens/repository.py @@ -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, ) diff --git a/litellm/proxy/engine/sources.py b/litellm/proxy/lens/sources.py similarity index 98% rename from litellm/proxy/engine/sources.py rename to litellm/proxy/lens/sources.py index d9d50a0b91e..12d26cd4974 100644 --- a/litellm/proxy/engine/sources.py +++ b/litellm/proxy/lens/sources.py @@ -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, diff --git a/litellm/proxy/engine/state.py b/litellm/proxy/lens/state.py similarity index 66% rename from litellm/proxy/engine/state.py rename to litellm/proxy/lens/state.py index 3ca6e881234..5fc0a88aa3a 100644 --- a/litellm/proxy/engine/state.py +++ b/litellm/proxy/lens/state.py @@ -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( { diff --git a/litellm/proxy/engine/trace_store.py b/litellm/proxy/lens/trace_store.py similarity index 100% rename from litellm/proxy/engine/trace_store.py rename to litellm/proxy/lens/trace_store.py diff --git a/litellm/proxy/engine/worker.py b/litellm/proxy/lens/worker.py similarity index 93% rename from litellm/proxy/engine/worker.py rename to litellm/proxy/lens/worker.py index e71c57ce143..2980f62deed 100644 --- a/litellm/proxy/engine/worker.py +++ b/litellm/proxy/lens/worker.py @@ -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() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7c1ab0711ea..085c38ad259 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 75dc7ddde9d..6f285e9dc39 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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 diff --git a/litellm/tracing/decode.py b/litellm/tracing/decode.py index c310a339593..b1bd951220c 100644 --- a/litellm/tracing/decode.py +++ b/litellm/tracing/decode.py @@ -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) diff --git a/litellm/tracing/normalizers/__init__.py b/litellm/tracing/normalizers/__init__.py new file mode 100644 index 00000000000..2f861330a36 --- /dev/null +++ b/litellm/tracing/normalizers/__init__.py @@ -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", +) diff --git a/litellm/tracing/normalizers/base.py b/litellm/tracing/normalizers/base.py new file mode 100644 index 00000000000..37735113ce2 --- /dev/null +++ b/litellm/tracing/normalizers/base.py @@ -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 diff --git a/litellm/tracing/normalizers/genai.py b/litellm/tracing/normalizers/genai.py new file mode 100644 index 00000000000..16986607396 --- /dev/null +++ b/litellm/tracing/normalizers/genai.py @@ -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", "") diff --git a/litellm/tracing/normalizers/langsmith.py b/litellm/tracing/normalizers/langsmith.py new file mode 100644 index 00000000000..daca932d57d --- /dev/null +++ b/litellm/tracing/normalizers/langsmith.py @@ -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) diff --git a/litellm/tracing/normalizers/messages.py b/litellm/tracing/normalizers/messages.py new file mode 100644 index 00000000000..0d552d05b82 --- /dev/null +++ b/litellm/tracing/normalizers/messages.py @@ -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 diff --git a/litellm/tracing/normalizers/openinference.py b/litellm/tracing/normalizers/openinference.py new file mode 100644 index 00000000000..f9e1295148c --- /dev/null +++ b/litellm/tracing/normalizers/openinference.py @@ -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")) diff --git a/litellm/tracing/store.py b/litellm/tracing/store.py index 806757306c0..fdb1c7820f2 100644 --- a/litellm/tracing/store.py +++ b/litellm/tracing/store.py @@ -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"], ) diff --git a/litellm/tracing/types.py b/litellm/tracing/types.py index 6cdfcd84da7..fcf0d83fb8f 100644 --- a/litellm/tracing/types.py +++ b/litellm/tracing/types.py @@ -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]] diff --git a/litellm/tracing/ui_format.py b/litellm/tracing/ui_format.py new file mode 100644 index 00000000000..d7ecf48078f --- /dev/null +++ b/litellm/tracing/ui_format.py @@ -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) diff --git a/litellm/types/integrations/s3_v2.py b/litellm/types/integrations/s3_v2.py index 3b0dad97e8c..e8ad28f1a3b 100644 --- a/litellm/types/integrations/s3_v2.py +++ b/litellm/types/integrations/s3_v2.py @@ -1,5 +1,9 @@ +from typing import Literal + from pydantic import BaseModel +S3PartitionGranularity = Literal["day", "hour"] + class s3BatchLoggingElement(BaseModel): """ diff --git a/schema.prisma b/schema.prisma index 75dc7ddde9d..6f285e9dc39 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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 diff --git a/tests/code_coverage_tests/ensure_async_clients_test.py b/tests/code_coverage_tests/ensure_async_clients_test.py index 285a5700a1c..7519c2aebb3 100644 --- a/tests/code_coverage_tests/ensure_async_clients_test.py +++ b/tests/code_coverage_tests/ensure_async_clients_test.py @@ -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", diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index 7c83e4d3aea..5424fb2d12f 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -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"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index b2682c04841..3fa9f534ffd 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -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", "") diff --git a/tests/e2e/logging/test_s3_log_e2e.py b/tests/e2e/logging/test_s3_log_e2e.py index 7a1ee1e6536..1612fa315a6 100644 --- a/tests/e2e/logging/test_s3_log_e2e.py +++ b/tests/e2e/logging/test_s3_log_e2e.py @@ -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\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\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 diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index bb0243226a8..1b3bc0183a5 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -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) diff --git a/tests/integration/database/test_engine_repository.py b/tests/integration/database/test_engine_repository.py deleted file mode 100644 index 89e019c8e1a..00000000000 --- a/tests/integration/database/test_engine_repository.py +++ /dev/null @@ -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) diff --git a/tests/integration/database/test_lens_repository.py b/tests/integration/database/test_lens_repository.py new file mode 100644 index 00000000000..29c6ad13825 --- /dev/null +++ b/tests/integration/database/test_lens_repository.py @@ -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))) diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py new file mode 100644 index 00000000000..dc0a7184cf8 --- /dev/null +++ b/tests/integration/observability/test_s3_v2_partition_granularity.py @@ -0,0 +1,1241 @@ +import json +import re +import threading +import uuid +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from pathlib import Path +from typing import Final +from urllib.parse import quote, unquote + +import httpx +import openai +import psutil +import pytest +import yaml +from _s3_v2_support import ( + BUCKET, + PREFIX, + SURFACES, + RecordingS3Sink, + call_surface, + collect_payloads, + matched_ids, + mixed_burst, + s3_config, + surface_reply, +) +from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.database_relay import database_relay +from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server + +FLUSH: Final = {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"} +HOUR: Final = {"s3_partition_granularity": "hour"} +ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-4-5-20250929" +WARNING: Final = "s3 logging: s3_partition_granularity=" +SINK_CREDENTIALS: Final = { + "s3_bucket_name": BUCKET, + "s3_region_name": "us-east-1", + "s3_path": PREFIX, + "s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", +} + + +@dataclass(slots=True) +class CountingUpstream: + """Scripted provider that answers every surface and fails any prompt ending in -fail with a 401.""" + + lock: threading.Lock = field(default_factory=threading.Lock) + prompts: list[str] = field(default_factory=list) # mutable-ok: appended per upstream request under lock + + def respond(self, request: Request) -> Reply: + if request.method != "POST" or not request.body: + return Reply(status=404) + body: Final = json.loads(request.body) + prompt: Final = str(body["input"] if "input" in body else body["messages"][0]["content"]) + with self.lock: + self.prompts.append(prompt) + if prompt.endswith("-fail"): + return Reply(status=401, body=b'{"error": {"message": "synthetic upstream rejection", "code": "401"}}') + return surface_reply(request) + + def received(self) -> tuple[str, ...]: + with self.lock: + return tuple(self.prompts) + + +def _prompt(payload: Mapping[str, JsonValue]) -> str: + messages: Final = payload["messages"] + if isinstance(messages, str): + return messages + assert isinstance(messages, list) and len(messages) == 1, payload + first: Final = messages[0] + return first if isinstance(first, str) else str(object_value(first)["content"]) + + +def _start(payload: Mapping[str, JsonValue]) -> datetime: + return datetime.fromtimestamp(float(str(payload["startTime"]))) + + +def _folder(payload: Mapping[str, JsonValue], granularity: str, prefix: str = "") -> str: + start: Final = _start(payload) + hour: Final = f"{start:%H}/" if granularity == "hour" else "" + return f"/{BUCKET}/{PREFIX}/{prefix}{start:%Y-%m-%d}/{hour}" + + +def _object_pattern(payload: Mapping[str, JsonValue], granularity: str, prefix: str = "") -> re.Pattern[str]: + return re.compile( + re.escape(_folder(payload, granularity, prefix)) + rf"time-{_start(payload):%H-%M-%S}-\d{{6}}_[^/]+\.json" + ) + + +def _outside_layout(objects: Mapping[str, bytes], granularity: str, prefix: str = "") -> tuple[str, ...]: + return tuple( + target + for target, body in objects.items() + if not _object_pattern(object_value(json.loads(body)), granularity, prefix).fullmatch(unquote(target)) + ) + + +def _batches_outside_layout(objects: Mapping[str, bytes], granularity: str) -> tuple[str, ...]: + def folders(body: bytes) -> frozenset[str]: + return frozenset(_folder(object_value(json.loads(line)), granularity) for line in body.splitlines()) + + return tuple( + target + for target, body in objects.items() + if len(folders(body)) != 1 + or not re.fullmatch( + re.escape(next(iter(folders(body)))) + r"batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl", unquote(target) + ) + ) + + +@contextmanager +def _s3_proxy( + gateway: Gateway, + tmp_path: Path, + sink_url: str, + extra: Mapping[str, JsonValue], + settings: Mapping[str, JsonValue] | None = None, + environment: Mapping[str, str] | None = None, + workers: int = 2, + models: tuple[Mapping[str, JsonValue], ...] = (), +) -> Iterator[OwnedProxy]: + config: Final = s3_config(tmp_path, sink_url, extra, settings) + if models: + declared: Final = yaml.safe_load(config.read_text()) + config.write_text(yaml.safe_dump({**declared, "model_list": [*declared["model_list"], *models]})) + with owned_proxy_process( + gateway, tmp_path, {**FLUSH, **(environment or {})}, config=config, workers=workers + ) as owned: + yield owned + + +def _models(scenario: Scenario, provider_url: str, **key_fields: JsonValue) -> tuple[str, str, str]: + openai_model: Final = scenario.model(api_base=provider_url + "/v1", api_key="synthetic-provider-key") + anthropic_model: Final = scenario.model( + model=ANTHROPIC_MODEL, api_base=provider_url, api_key="synthetic-provider-key" + ) + return openai_model, anthropic_model, scenario.key(models=[openai_model, anthropic_model], **key_fields) + + +def _config_model(name: str, model: str, api_base: str) -> Mapping[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "synthetic-provider-key"}, + } + + +def _sdk_chats(candidate: Gateway, model: str, key: str, prompts: tuple[str, ...]) -> tuple[str, ...]: + client: Final = openai.OpenAI(base_url=f"{str(candidate.client.base_url).rstrip('/')}/v1", api_key=key) + + def send(prompt: str) -> str: + reply: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": prompt}], extra_body={"cache": {"no-cache": True}} + ) + assert reply.choices[0].finish_reason == "stop", reply.model_dump_json() + return reply.id + + with ThreadPoolExecutor(max_workers=16) as pool: + return tuple(pool.map(send, prompts)) + + +def _surface_prompts(marker: str, per_surface: int) -> frozenset[str]: + return frozenset(f"{marker}-{surface}-{index}" for surface in SURFACES for index in range(per_surface)) + + +def _cold_storage_key(request_id: str, database_url: str | None = None) -> str: + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,), database_url=database_url + ), + lambda values: len(values) == 1, + seconds=60, + ) + metadata: Final = rows[0]["metadata"] + return str(object_value(json.loads(metadata) if isinstance(metadata, str) else metadata)["cold_storage_object_key"]) + + +def _update_environment(candidate: Gateway, values: Mapping[str, JsonValue]) -> None: + candidate.post( + "/config/update", + {"environment_variables": dict(values), "litellm_settings": {"success_callback": ["s3_v2"]}}, + ) + + +def _keys_on_fresh_connections(candidate: Gateway, aliases: tuple[str, ...]) -> tuple[tuple[str, str], ...]: + def generate(alias: str) -> tuple[str, str]: + with httpx.Client(base_url=candidate.client.base_url, timeout=30, trust_env=False) as fresh: + response: Final = fresh.post( + "/key/generate", + json={"key_alias": alias}, + headers={"Authorization": f"Bearer {candidate.key}", "Connection": "close"}, + ) + assert response.status_code == 200, response.text + return str(response.json()["key"]), str(response.json()["token_id"]) + + with ThreadPoolExecutor(max_workers=len(aliases)) as pool: + return tuple(pool.map(generate, aliases)) + + +def _created_key_hashes(sink: RecordingS3Sink, audit_prefix: str) -> frozenset[str]: + created: Final = ( + object_value(json.loads(body)) for target, body in sink.objects().items() if target.startswith(audit_prefix) + ) + return frozenset( + str(audit["object_id"]) + for audit in created + if audit["action"] == "created" and audit["table_name"] == "LiteLLM_VerificationToken" + ) + + +def test_s3_v2_hour_granularity_files_every_surface_under_its_hour_folder(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hour" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=2) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + log: Final = owned.log.read_text() + sent: Final = _surface_prompts(marker, 2) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert len(objects) == len(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert all(payload["status"] == "success" for payload in payloads), payloads + assert _outside_layout(objects, "hour") == (), "every object must sit in YYYY-MM-DD/HH/ of its start time" + assert WARNING not in log + + +@pytest.mark.parametrize( + "extra", + [ + pytest.param({}, id="missing"), + pytest.param({"s3_partition_granularity": "day"}, id="day"), + pytest.param({"s3_partition_granularity": ""}, id="empty"), + pytest.param({"s3_partition_granularity": None}, id="null"), + ], +) +def test_s3_v2_missing_day_empty_or_null_granularity_keeps_the_daily_layout( + gateway: Gateway, tmp_path: Path, extra: Mapping[str, JsonValue] +) -> None: + marker: Final = "s3day" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=1) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + log: Final = owned.log.read_text() + sent: Final = _surface_prompts(marker, 1) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert len(objects) == len(sent) + assert _outside_layout(objects, "day") == () + assert WARNING not in log + + +@pytest.mark.parametrize( + ("extra", "environment", "shown"), + [ + pytest.param({"s3_partition_granularity": "hourly"}, {}, "'hourly'", id="unknown_word"), + pytest.param({"s3_partition_granularity": "HOUR"}, {}, "'HOUR'", id="wrong_case"), + pytest.param({"s3_partition_granularity": 1}, {}, "1", id="integer"), + pytest.param({"s3_partition_granularity": ["hour"]}, {}, "['hour']", id="list"), + pytest.param({"s3_partition_granularity": "h" * 5120}, {}, "'[base64_data truncated: 3.8KB]'", id="five_kb"), + pytest.param({}, {"S3_PARTITION_GRANULARITY": "weekly"}, "'weekly'", id="env_unknown_word"), + ], +) +def test_s3_v2_unrecognized_granularity_warns_once_per_worker_and_keeps_the_daily_layout( + gateway: Gateway, tmp_path: Path, extra: Mapping[str, JsonValue], environment: Mapping[str, str], shown: str +) -> None: + marker: Final = "s3bad" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra, environment=environment) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=2) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + warning: Final = f"{WARNING}{shown} is not one of day, hour, using day" + log: Final = eventually(owned.log.read_text, lambda text: warning in text, seconds=15) + sent: Final = _surface_prompts(marker, 2) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert _outside_layout(objects, "day") == () + assert 1 <= log.count(warning) <= 2, "the warning is memoized per distinct value in each of the two workers" + + +@pytest.mark.parametrize( + ("extra", "environment", "granularity"), + [ + pytest.param({}, {"S3_PARTITION_GRANULARITY": "hour"}, "hour", id="env_hour_applies"), + pytest.param({"s3_partition_granularity": "day"}, {"S3_PARTITION_GRANULARITY": "hour"}, "day", id="yaml_wins"), + ], +) +def test_s3_v2_env_granularity_applies_only_when_callback_params_leave_it_unset( + gateway: Gateway, tmp_path: Path, extra: Mapping[str, JsonValue], environment: Mapping[str, str], granularity: str +) -> None: + marker: Final = "s3env" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra, environment=environment) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(8)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert frozenset(str(payload["id"]) for payload in payloads) == frozenset(prompts) + assert _outside_layout(objects, granularity) == () + + +def test_s3_v2_hour_batch_files_group_lines_under_the_hour_folder(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hbat" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, {**HOUR, "s3_batch_file_upload": True}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=4) + payloads: Final = collect_payloads(sink, len(answered)) + objects: Final = sink.objects() + sent: Final = _surface_prompts(marker, 4) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent) + assert _batches_outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_folder_sits_below_the_team_and_key_prefix(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hpre" + uuid.uuid4().hex[:8] + team_alias: Final = f"alpha-{uuid.uuid4().hex[:8]}" + key_alias: Final = f"beta-{uuid.uuid4().hex[:8]}" + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + extra: Final = {**HOUR, "s3_use_team_prefix": True, "s3_use_key_prefix": True} + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, extra) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + team: Final = scenario.team(team_alias=team_alias, models=[openai_model]) + key: Final = scenario.key(team_id=team, key_alias=key_alias, models=[openai_model]) + prompts: Final = tuple(f"{marker}-{index}" for index in range(6)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert frozenset(str(payload["id"]) for payload in payloads) == frozenset(prompts) + assert _outside_layout(objects, "hour", f"{team_alias}/{key_alias}/") == () + + +def _payload_values(payloads: tuple[dict[str, JsonValue], ...], status: str, field: str) -> frozenset[str]: + return frozenset(str(payload[field]) for payload in payloads if payload["status"] == status) + + +def test_s3_v2_hour_failure_and_rejected_requests_keep_the_hour_layout(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hfail" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + + def send(prompt: str, model: str = openai_model, caller: str = key) -> httpx.Response: + return owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": prompt}], "cache": {"no-cache": True}}, + key=caller, + ) + + successes: Final = tuple(f"{marker}-{index}" for index in range(4)) + failures: Final = tuple(f"{marker}-{index}-fail" for index in range(3)) + with ThreadPoolExecutor(max_workers=8) as pool: + responses: Final = tuple(pool.map(send, (*successes, *failures))) + ghost: Final = send(f"{marker}-ghost", model=f"ghost-{uuid.uuid4().hex}") + unauthenticated: Final = send(f"{marker}-anon", caller="sk-not-a-real-key") + after: Final = send(f"{marker}-after") + rejected_call_ids: Final = frozenset(response.headers["x-litellm-call-id"] for response in responses[4:]) + payloads: Final = eventually( + sink.payloads, + lambda stored: ( + _payload_values(stored, "success", "id") >= frozenset((*successes, f"{marker}-after")) + and _payload_values(stored, "failure", "litellm_call_id") >= rejected_call_ids + ), + seconds=60, + ) + objects: Final = sink.objects() + assert [response.status_code for response in responses[:4]] == [200] * 4, [r.text for r in responses] + assert tuple(response.json()["id"] for response in responses[:4]) == successes + assert all(response.status_code == 401 for response in responses[4:]), [r.text for r in responses[4:]] + assert all("synthetic upstream rejection" in response.text for response in responses[4:]) + assert ghost.status_code == 403 and "key_model_access_denied" in ghost.text, ghost.text + assert unauthenticated.status_code == 401 and "error" in unauthenticated.json(), unauthenticated.text + assert after.status_code == 200 and after.json()["id"] == f"{marker}-after", after.text + assert sorted(upstream.received()) == sorted((*successes, *failures, f"{marker}-after")) + assert _payload_values(payloads, "success", "id") == frozenset((*successes, f"{marker}-after")) + assert _payload_values(payloads, "failure", "litellm_call_id") >= rejected_call_ids + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_cache_hit_twins_land_one_object_each_under_the_hour_folder( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3hcache" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, anthropic_model, key = _models(scenario, provider.url) + first: Final = tuple( + call_surface(owned.gateway, surface, openai_model, anthropic_model, key, f"{marker}-{surface}", False) + for surface in ("chat", "responses") + ) + eventually(lambda: len(sink.objects()), lambda count: count >= 2, seconds=30) + repeated: Final = tuple( + call_surface(owned.gateway, surface, openai_model, anthropic_model, key, f"{marker}-{surface}", False) + for surface in ("chat", "responses") + ) + payloads: Final = collect_payloads(sink, 4) + objects: Final = sink.objects() + assert first[0][0] == f"{marker}-chat" and repeated[0][0] == first[0][0] + assert matched_ids(payloads, first + repeated) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(_prompt(payload) for payload in payloads) == sorted((f"{marker}-chat", f"{marker}-responses") * 2) + assert sorted(upstream.received()) == sorted((f"{marker}-chat", f"{marker}-responses")) + assert len(objects) == 4, list(objects) + assert sum(1 for payload in payloads if payload["cache_hit"] is True) == 2 + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_cold_storage_key_names_the_uploaded_object_and_reads_back(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hcold" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR, {"cold_storage_custom_logger": "s3_v2"}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = (f"{marker}-kept", f"{marker}-missing") + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + keys: Final = {prompt: _cold_storage_key(prompt) for prompt in prompts} + with sink.lock: + sink.store.pop(f"/{BUCKET}/{quote(keys[prompts[1]], safe='/')}") + kept: Final = eventually( + lambda: owned.gateway.request("GET", f"/spend/logs/ui/{prompts[0]}"), + lambda reply: reply.status_code == 200 and bool((reply.json() or {}).get("messages")), + seconds=30, + ) + missing: Final = owned.gateway.request("GET", f"/spend/logs/ui/{prompts[1]}") + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert frozenset(f"/{BUCKET}/{quote(key, safe='/')}" for key in keys.values()) == frozenset(objects) + assert _outside_layout(objects, "hour") == () + assert kept.json()["messages"] == [{"role": "user", "content": prompts[0]}], kept.text + assert prompts[0] in json.dumps(kept.json()["response"]), kept.text + assert missing.status_code == 200, missing.text + assert prompts[1] not in json.dumps(missing.json()["response"]), missing.text + + +def test_s3_v2_hour_layout_holds_when_another_logger_owns_cold_storage(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hgcs" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + lock: Final = threading.Lock() + puts: Final[dict[str, bytes]] = {} # mutable-ok: filled per PUT by the bucket thread under lock + + def bucket_reply(request: Request) -> Reply: + assert request.method == "PUT", request.method + with lock: + puts[unquote(request.target)] = request.body + return Reply(status=200) + + def uploaded() -> Mapping[str, bytes]: + with lock: + return dict(puts) + + with ( + wire_server(upstream.respond) as provider, + wire_server(bucket_reply) as bucket, + _s3_proxy( + gateway, tmp_path, bucket.url, {**HOUR, "s3_path": ""}, {"cold_storage_custom_logger": "gcs_bucket"} + ) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(3)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + objects: Final = eventually(uploaded, lambda values: len(values) >= len(prompts), seconds=60) + cold_keys: Final = tuple(_cold_storage_key(prompt) for prompt in prompts) + hour_object: Final = re.compile(rf"/{BUCKET}/\d{{4}}-\d{{2}}-\d{{2}}/(\d{{2}})/time-(\d{{2}})-[^/]+\.json") + matches: Final = tuple(hour_object.fullmatch(target) for target in objects) + assert returned == prompts + assert sorted(str(object_value(json.loads(body))["id"]) for body in objects.values()) == sorted(prompts) + assert all(re.fullmatch(r"\d{4}-\d{2}-\d{2}/time-[^/]+\.json", cold_key) for cold_key in cold_keys), cold_keys + assert all(match is not None and match.group(1) == match.group(2) for match in matches), sorted(objects) + + +def test_s3_v2_hour_cold_storage_rebuilds_previous_response_id_history_from_the_hour_object( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3hsess" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + histories: Final[list[str]] = [] # mutable-ok: appended per upstream request by the scripted provider thread + reads: Final[list[str]] = [] # mutable-ok: appended per sink GET by the recording sink thread + sink: Final = RecordingS3Sink(delay_seconds=0.05) + + def provider_reply(request: Request) -> Reply: + histories.append(request.body.decode()) + return upstream.respond(request) + + def bucket_reply(request: Request) -> Reply: + if request.method == "GET": + reads.append(unquote(request.target)) + return sink.respond(request) + + with ( + wire_server(provider_reply) as provider, + wire_server(bucket_reply) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR, {"cold_storage_custom_logger": "s3_v2"}) as owned, + owned.gateway.scenario() as scenario, + ): + _, anthropic_model, key = _models(scenario, provider.url) + first: Final = owned.gateway.request( + "POST", "/v1/responses", {"model": anthropic_model, "input": f"{marker}-first"}, key=key + ) + assert first.status_code == 200, first.text + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (anthropic_model,) + ), + lambda values: len(values) == 1, + seconds=60, + ) + metadata: Final = rows[0]["metadata"] + cold_key: Final = str( + object_value(json.loads(metadata) if isinstance(metadata, str) else metadata)["cold_storage_object_key"] + ) + eventually(sink.objects, lambda objects: f"/{BUCKET}/{quote(cold_key, safe='/')}" in objects, seconds=30) + second: Final = owned.gateway.request( + "POST", + "/v1/responses", + {"model": anthropic_model, "input": f"{marker}-second", "previous_response_id": first.json()["id"]}, + key=key, + ) + objects: Final = sink.objects() + assert second.status_code == 200, second.text + assert second.json()["id"] != first.json()["id"], second.text + assert re.fullmatch(rf"{re.escape(PREFIX)}/\d{{4}}-\d{{2}}-\d{{2}}/\d{{2}}/time-[^/]+\.json", cold_key), cold_key + assert _outside_layout(objects, "hour") == () + assert f"/{BUCKET}/{cold_key}" in reads, reads + assert len(histories) == 2, histories + assert f"{marker}-first" in histories[0] and f"{marker}-second" not in histories[0], histories[0] + assert f"{marker}-first" in histories[1] and f"{marker}-second" in histories[1], histories[1] + + +def test_s3_v2_audit_logs_follow_the_audit_params_granularity_not_the_request_logs( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3haudit" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with wire_server(upstream.respond) as provider, wire_server(sink.respond) as bucket: + settings: Final = { + "store_audit_logs": True, + "audit_log_callbacks": ["s3_v2"], + "s3_audit_callback_params": {**SINK_CREDENTIALS, "s3_endpoint_url": bucket.url, **HOUR}, + } + with ( + _s3_proxy(gateway, tmp_path, bucket.url, {}, settings) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url, key_alias=marker) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, (marker,)) + aliases: Final = tuple(f"{marker}-fresh{index}" for index in range(16)) + fresh_keys: Final = _keys_on_fresh_connections(owned.gateway, aliases) + audit_prefix: Final = f"/{BUCKET}/{PREFIX}/audit_logs/" + eventually( + lambda: _created_key_hashes(sink, audit_prefix), + lambda created: frozenset(token for _, token in fresh_keys) <= created, + seconds=30, + ) + owned.gateway.post("/key/delete", {"keys": [key for key, _ in fresh_keys]}) + collect_payloads(sink, 2) + objects: Final = sink.objects() + audits: Final = { + target: object_value(json.loads(body)) for target, body in objects.items() if target.startswith(audit_prefix) + } + requests: Final = {target: body for target, body in objects.items() if not target.startswith(audit_prefix)} + assert returned == (marker,) + assert upstream.received() == (marker,) + assert _outside_layout(requests, "day") == () + created: Final = tuple(audit for audit in audits.values() if audit["action"] == "created") + assert "LiteLLM_VerificationToken" in frozenset(str(audit["table_name"]) for audit in created), audits + for target, audit in audits.items(): + located: Final = re.fullmatch( + re.escape(audit_prefix) + + rf"(\d{{4}}-\d{{2}}-\d{{2}})/(\d{{2}})/(\d{{2}})-\d{{2}}-\d{{2}}_{re.escape(str(audit['id']))}\.json", + unquote(target), + ) + assert located and located[2] == located[3], (target, audit["updated_at"]) + folder: Final = datetime.fromisoformat(f"{located[1]}T{located[2]}:00:00+00:00") + updated: Final = datetime.fromisoformat(str(audit["updated_at"])) + assert timedelta(0) < folder + timedelta(hours=1) - updated <= timedelta(hours=1, minutes=1), ( + target, + audit["updated_at"], + ) + + +@pytest.mark.parametrize("level", ["key", "team"]) +def test_s3_v2_key_and_team_logging_callback_vars_cannot_change_the_proxy_hour_layout( + gateway: Gateway, tmp_path: Path, level: str +) -> None: + marker: Final = f"s3h{level}vars" + uuid.uuid4().hex[:8] + logging: Final[list[JsonValue]] = [ + {"callback_name": "s3_v2", "callback_type": "success", "callback_vars": {"s3_partition_granularity": "day"}} + ] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, _ = _models(scenario, provider.url) + key: Final = ( + scenario.key(models=[openai_model], metadata={"logging": logging}) + if level == "key" + else scenario.key(models=[openai_model], team_id=scenario.team(metadata={"logging": logging})) + ) + prompts: Final = tuple(f"{marker}-{index}" for index in range(8)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert sorted(str(object_value(json.loads(body))["id"]) for body in objects.values()) == sorted(prompts), ( + f"{level}-level s3_v2 logging must land exactly one object per request" + ) + assert _outside_layout(objects, "hour") == (), f"{level}-level callback_vars must not change the proxy granularity" + + +def test_s3_v2_admin_ui_granularity_update_moves_live_traffic_on_both_workers(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hui" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + scratch_database() as database_url, + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, {}, environment={"DATABASE_URL": database_url}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + before: Final = _sdk_chats(owned.gateway, openai_model, key, (f"{marker}-before",)) + eventually(lambda: len(sink.objects()), lambda count: count >= 1, seconds=30) + listed: Final = owned.gateway.get("/get/config/callbacks") + _update_environment(owned.gateway, {"callback": "s3_v2", "s3_partition_granularity": "hour"}) + probe_round: Final = iter(range(1000)) + + def probe() -> Mapping[str, bytes]: + round_id: Final = next(probe_round) + prompts: Final = tuple(f"{marker}-probe{round_id}-{index}" for index in range(8)) + _sdk_chats(owned.gateway, openai_model, key, prompts) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(prompts) <= landed, + seconds=20, + ) + return {target: body for target, body in sink.objects().items() if f"-probe{round_id}-" in target} + + eventually(probe, lambda probed: len(probed) == 8 and _outside_layout(probed, "hour") == (), seconds=60) + prompts: Final = tuple(f"{marker}-after-{index}" for index in range(16)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(prompts) <= landed, + seconds=30, + ) + after: Final = {target: body for target, body in sink.objects().items() if f"{marker}-after-" in target} + before_objects: Final = { + target: body for target, body in sink.objects().items() if f"{marker}-before" in target + } + readback: Final = owned.gateway.get("/get/config/callbacks") + s3_rows: Final = tuple(row for row in listed["callbacks"] if object_value(row)["name"] in ("s3", "s3_v2")) + assert s3_rows and all( + "S3_PARTITION_GRANULARITY" in object_value(object_value(row)["variables"]) for row in s3_rows + ), listed + after_rows: Final = tuple(row for row in readback["callbacks"] if object_value(row)["name"] in ("s3", "s3_v2")) + assert all( + object_value(object_value(row)["variables"])["S3_PARTITION_GRANULARITY"] == "hour" for row in after_rows + ), readback + assert before == (f"{marker}-before",) + assert returned == prompts + assert _outside_layout(before_objects, "day") == () + assert len(after) == len(prompts) + assert _outside_layout(after, "hour") == () + + +def test_s3_v2_granularity_toggles_mid_burst_keep_every_cold_storage_key_on_its_object( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3htog" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + scratch_database() as database_url, + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy( + gateway, + tmp_path, + bucket.url, + {}, + {"cold_storage_custom_logger": "s3_v2"}, + environment={"DATABASE_URL": database_url}, + ) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(32)) + with ThreadPoolExecutor(max_workers=1) as burst: + pending: Final = burst.submit(_sdk_chats, owned.gateway, openai_model, key, prompts) + for value in ("hour", "day", "hour", "day", "hour", "day"): + _update_environment(owned.gateway, {"s3_partition_granularity": value}) + returned: Final = pending.result() + collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + keys: Final = {prompt: _cold_storage_key(prompt, database_url) for prompt in prompts} + assert returned == prompts + assert sorted(upstream.received()) == sorted(prompts) + assert len(objects) == len(prompts) + assert frozenset(f"/{BUCKET}/{quote(key, safe='/')}" for key in keys.values()) == frozenset(objects), ( + "every spend log cold_storage_object_key must name the object the logger uploaded" + ) + assert all( + _object_pattern(object_value(json.loads(body)), "hour").fullmatch(unquote(target)) + or _object_pattern(object_value(json.loads(body)), "day").fullmatch(unquote(target)) + for target, body in objects.items() + ) + + +def test_s3_v2_in_flight_request_keeps_its_cold_storage_key_on_its_object_across_owner_and_granularity_switches( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3hflight" + uuid.uuid4().hex[:8] + held_prompt: Final = f"{marker}-held" + upstream: Final = CountingUpstream() + arrived: Final = threading.Event() + release: Final = threading.Event() + + def held(request: Request) -> Reply: + if held_prompt.encode() in request.body: + arrived.set() + assert release.wait(90), "held request was never released" + return upstream.respond(request) + + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + scratch_database() as database_url, + wire_server(held) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy( + gateway, + tmp_path, + bucket.url, + {}, + {"cold_storage_custom_logger": "s3_v2"}, + environment={"DATABASE_URL": database_url}, + ) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + with ThreadPoolExecutor(max_workers=1) as flight: + pending: Final = flight.submit(_sdk_chats, owned.gateway, openai_model, key, (held_prompt,)) + assert arrived.wait(60), "held request never reached the upstream" + owner_switch: Final = owned.gateway.request( + "POST", "/config/update", {"litellm_settings": {"cold_storage_custom_logger": "gcs_bucket"}} + ) + _update_environment(owned.gateway, HOUR) + probe_round: Final = iter(range(1000)) + + def probe() -> Mapping[str, bytes]: + round_id: Final = next(probe_round) + prompts: Final = tuple(f"{marker}-probe{round_id}-{index}" for index in range(8)) + _sdk_chats(owned.gateway, openai_model, key, prompts) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(prompts) <= landed, + seconds=20, + ) + return {target: body for target, body in sink.objects().items() if f"-probe{round_id}-" in target} + + eventually(probe, lambda probed: len(probed) == 8 and _outside_layout(probed, "hour") == (), seconds=60) + release.set() + returned: Final = pending.result() + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: held_prompt in landed, + seconds=30, + ) + held_objects: Final = {target: body for target, body in sink.objects().items() if held_prompt in target} + cold_key: Final = _cold_storage_key(held_prompt, database_url) + assert owner_switch.status_code == 400, owner_switch.text + assert "cold_storage_custom_logger" in owner_switch.text and "config file" in owner_switch.text, owner_switch.text + assert returned == (held_prompt,) + assert upstream.received().count(held_prompt) == 1 + assert frozenset(held_objects) == frozenset({f"/{BUCKET}/{quote(cold_key, safe='/')}"}), ( + "the in-flight request's cold_storage_object_key must name the one object the logger uploaded", + cold_key, + tuple(held_objects), + ) + assert _outside_layout(held_objects, "hour") == () + + +def test_s3_v2_cold_storage_owner_saved_through_config_update_is_not_applied_to_a_running_proxy( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3howner" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink() + with ( + scratch_database() as database_url, + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR, environment={"DATABASE_URL": database_url}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + saved: Final = owned.gateway.request( + "POST", "/config/update", {"litellm_settings": {"cold_storage_custom_logger": "s3_v2"}} + ) + prompts: Final = tuple(f"{marker}-{index}" for index in range(8)) + answered: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + landed: Final = collect_payloads(sink, len(prompts)) + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', + (list(answered),), + database_url=database_url, + ), + lambda values: len(values) == len(prompts), + seconds=60, + ) + objects: Final = sink.objects() + stored: Final = read_rows( + 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', + ("litellm_settings",), + database_url=database_url, + ) + cold_keys: Final = { + str(row["request_id"]): object_value( + json.loads(row["metadata"]) if isinstance(row["metadata"], str) else row["metadata"] + ).get("cold_storage_object_key") + for row in rows + } + assert saved.status_code == 200, saved.text + assert [ + object_value(json.loads(row["param_value"]) if isinstance(row["param_value"], str) else row["param_value"]).get( + "cold_storage_custom_logger" + ) + for row in stored + ] == ["s3_v2"], "the owner switch must be persisted, so the unchanged live keys are not a rejected write" + assert sorted(upstream.received()) == sorted(prompts) + assert sorted(_prompt(payload) for payload in landed) == sorted(prompts) + assert cold_keys == dict.fromkeys(answered), "a DB-saved cold storage owner must not change a live request" + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_postgres_outage_mid_mixed_burst_lands_every_id_exactly_once_and_recovers( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3hpg" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + sent: Final = _surface_prompts(marker, 5) + openai_model: Final = f"{marker}openai" + anthropic_model: Final = f"{marker}anthropic" + with ( + scratch_database() as database_url, + database_relay(database_url, f"{marker}-".encode()) as (relay, relayed_url), + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy( + gateway, + tmp_path, + bucket.url, + HOUR, + {"cold_storage_custom_logger": "s3_v2"}, + environment={"DATABASE_URL": relayed_url}, + models=( + _config_model(openai_model, "openai/gpt-4o-mini", provider.url + "/v1"), + _config_model(anthropic_model, ANTHROPIC_MODEL, provider.url), + ), + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key(models=[openai_model, anthropic_model]) + warm: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, f"{marker}warm", per_surface=2) + eventually( + lambda: frozenset(_prompt(payload) for payload in sink.payloads()), + lambda landed: _surface_prompts(f"{marker}warm", 2) <= landed, + seconds=60, + ) + relay.arm() + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=5) + assert relay.tripped.wait(90), "no spend log write reached the database during the burst" + eventually(lambda: relay.refused, lambda count: count >= 1, seconds=30) + assert relay.reconnected.wait(60), "the proxy never reconnected to the database after the outage" + burst_payloads: Final = eventually( + lambda: tuple(payload for payload in sink.payloads() if _prompt(payload) in sent), + lambda landed: frozenset(_prompt(payload) for payload in landed) == frozenset(sent), + seconds=60, + ) + recovered_prompt: Final = f"{marker}-recovered" + recovered: Final = _sdk_chats(owned.gateway, openai_model, key, (recovered_prompt,)) + recovered_key: Final = _cold_storage_key(recovered_prompt, database_url) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: recovered_prompt in landed, + seconds=30, + ) + objects: Final = sink.objects() + uploads: Final = sink.attempts + burst: Final = burst_payloads + assert len(warm) == len(_surface_prompts(f"{marker}warm", 2)) + assert len(answered) == len(sent) == 30 + assert sorted(prompt for prompt in upstream.received() if prompt.startswith(f"{marker}-")) == sorted( + (*sent, recovered_prompt) + ) + assert matched_ids(burst, answered) == frozenset(str(payload["id"]) for payload in burst) + assert sorted(_prompt(payload) for payload in burst) == sorted(sent), "every burst id lands exactly once" + assert uploads == len(objects), "no object is uploaded twice" + assert _outside_layout(objects, "hour") == () + assert recovered == (recovered_prompt,) + assert f"/{BUCKET}/{quote(recovered_key, safe='/')}" in objects, "cold key written after recovery names its object" + + +def test_legacy_s3_callback_ignores_hour_granularity(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3v1hour" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR, {"callbacks": [], "success_callback": ["s3"]}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(3)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts)) + objects: Final = sink.objects() + assert returned == prompts + assert frozenset(str(payload["id"]) for payload in payloads) == frozenset(prompts) + assert _outside_layout(objects, "day") == (), "legacy s3 keeps the daily layout, the setting is s3_v2 only" + + +def test_s3_v2_hour_sink_outage_mid_mixed_burst_lands_every_id_exactly_once(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hout" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05, fail_until=float("inf"), fail_status=503) + openai_model: Final = f"{marker}openai" + anthropic_model: Final = f"{marker}anthropic" + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy( + gateway, + tmp_path, + bucket.url, + HOUR, + models=( + _config_model(openai_model, "openai/gpt-4o-mini", provider.url + "/v1"), + _config_model(anthropic_model, ANTHROPIC_MODEL, provider.url), + ), + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key(models=[openai_model, anthropic_model]) + answered: Final = mixed_burst(owned.gateway, openai_model, anthropic_model, key, marker, per_surface=6) + eventually(lambda: sink.attempts, lambda attempts: attempts >= 1, seconds=30) + during: Final = owned.gateway.client.get("/health/readiness") + rejected: Final = sink.attempts + sink.fail_until = 0.0 + payloads: Final = collect_payloads(sink, len(answered), seconds=60) + objects: Final = sink.objects() + sent: Final = _surface_prompts(marker, 6) + assert len(answered) == len(sent) and len(payloads) == len(sent), payloads + assert matched_ids(payloads, answered) == frozenset(str(payload["id"]) for payload in payloads) + assert sorted(upstream.received()) == sorted(sent) + assert during.status_code == 200, during.text + assert rejected >= 1 and sink.attempts > len(objects) + assert sorted(_prompt(payload) for payload in payloads) == sorted(sent), "every burst id lands exactly once" + assert len(objects) == len(sent) + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_coded_403_retries_reuse_the_same_hour_key(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3h403" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05, fail_attempts=10, fail_status=403, fail_code="AccessDenied") + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(16)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + payloads: Final = collect_payloads(sink, len(prompts), seconds=60) + objects: Final = sink.objects() + attempted: Final = dict(sink.attempt_counts) + assert returned == prompts + assert sorted(str(payload["id"]) for payload in payloads) == sorted(prompts) + assert frozenset(attempted) == frozenset(objects), "a retried upload must reuse the key of its first attempt" + assert sum(attempted.values()) == len(objects) + 10 + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_slow_sink_batches_never_duplicate_an_upload(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hslow" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=1.5) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, {**HOUR, "s3_batch_file_upload": True}) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + prompts: Final = tuple(f"{marker}-{index}" for index in range(32)) + returned: Final = _sdk_chats(owned.gateway, openai_model, key, prompts) + + def delivered() -> int: + readiness: Final = owned.gateway.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + return sum(len(body.splitlines()) for body in sink.objects().values()) + + eventually(delivered, lambda total: total >= len(prompts), seconds=60) + payloads: Final = sink.payloads() + objects: Final = sink.objects() + targets: Final = tuple(put.target for put in bucket.drain()) + assert returned == prompts + assert len(set(targets)) == len(targets), "the same batch object was PUT more than once" + assert sorted(str(payload["id"]) for payload in payloads) == sorted(prompts) + assert _batches_outside_layout(objects, "hour") == () + + +def _worker_processes(owned: OwnedProxy) -> tuple[int, ...]: + return tuple( + process.pid + for process in group_members(owned.process.pid) + if process.pid != owned.process.pid and "spawn_main" in " ".join(process.cmdline()) + ) + + +def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3hkill" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with ( + wire_server(upstream.respond) as provider, + wire_server(sink.respond) as bucket, + _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as owned, + owned.gateway.scenario() as scenario, + ): + openai_model, _, key = _models(scenario, provider.url) + workers: Final = _worker_processes(owned) + sent: Final = tuple(f"{marker}-{index}" for index in range(40)) + + def send(prompt: str) -> tuple[str, bool]: + try: + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": openai_model, + "messages": [{"role": "user", "content": prompt}], + "cache": {"no-cache": True}, + }, + key=key, + ) + except httpx.HTTPError: + return prompt, False + return prompt, response.status_code == 200 and response.json()["id"] == prompt + + with ThreadPoolExecutor(max_workers=16) as pool: + futures: Final = tuple(pool.submit(send, prompt) for prompt in sent) + eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) + psutil.Process(workers[0]).kill() + results: Final = tuple(future.result() for future in futures) + later: Final = tuple(f"{marker}-later-{index}" for index in range(8)) + later_results: Final = tuple(send(prompt) for prompt in later) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(later) <= landed, + seconds=45, + ) + payloads: Final = sink.payloads() + objects: Final = sink.objects() + assert len(workers) == 2, workers + assert all(ok for _, ok in later_results), "the surviving worker must keep serving after the kill" + landed: Final = tuple(str(payload["id"]) for payload in payloads) + assert frozenset(landed) <= frozenset((*sent, *later)), "only ids this test sent may land" + assert len(results) == len(sent), results + assert len(landed) == len(set(landed)), "no id may land twice" + assert _outside_layout(objects, "hour") == () + + +def test_s3_v2_hour_proxy_restart_mid_burst_keeps_the_layout_without_duplicates( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "s3hterm" + uuid.uuid4().hex[:8] + upstream: Final = CountingUpstream() + sink: Final = RecordingS3Sink(delay_seconds=0.05) + with wire_server(upstream.respond) as provider, wire_server(sink.respond) as bucket: + model_name: Final = f"integration-{marker}" + + def register(candidate: Gateway) -> str: + return str( + candidate.post( + "/model/new", + { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "synthetic-provider-key", + "api_base": provider.url + "/v1", + }, + "model_info": {}, + }, + )["model_info"]["id"] + ) + + def send(candidate: Gateway, key: str, prompt: str) -> tuple[str, bool]: + try: + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model_name, + "messages": [{"role": "user", "content": prompt}], + "cache": {"no-cache": True}, + }, + key=key, + ) + except httpx.HTTPError: + return prompt, False + return prompt, response.status_code == 200 + + sent: Final = tuple(f"{marker}-{index}" for index in range(40)) + with _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as first: + model_id: Final = register(first.gateway) + first_key: Final = str(first.gateway.post("/key/generate", {"models": [model_name]})["key"]) + with ThreadPoolExecutor(max_workers=16) as pool: + futures: Final = tuple(pool.submit(send, first.gateway, first_key, prompt) for prompt in sent) + eventually(lambda: len(upstream.received()), lambda count: count >= 8, seconds=30) + first.process.terminate() + results: Final = tuple(future.result() for future in futures) + first.process.wait(timeout=30) + answered: Final = frozenset(prompt for prompt, ok in results if ok) + landed_before_restart: Final = frozenset(str(payload["id"]) for payload in sink.payloads()) + with _s3_proxy(gateway, tmp_path, bucket.url, HOUR) as second: + restarted: Final = tuple(f"{marker}-restart-{index}" for index in range(8)) + second_key: Final = second.gateway.post("/key/generate", {"models": [model_name]})["key"] + restart_results: Final = tuple(send(second.gateway, str(second_key), prompt) for prompt in restarted) + eventually( + lambda: frozenset(str(payload["id"]) for payload in sink.payloads()), + lambda landed: frozenset(restarted) <= landed, + seconds=30, + ) + second.gateway.post("/model/delete", {"id": model_id}) + payloads: Final = sink.payloads() + objects: Final = sink.objects() + assert all(ok for _, ok in restart_results) + assert landed_before_restart <= answered, "a delivered object has no answered request" + landed: Final = tuple(str(payload["id"]) for payload in payloads) + assert len(landed) == len(set(landed)), "no id may land twice across the restart" + assert frozenset(restarted) <= frozenset(landed) + targets: Final = tuple(put.target for put in bucket.drain()) + assert len(set(targets)) == len(targets) + assert _outside_layout(objects, "hour") == () diff --git a/tests/integration/spend/test_lens_billing.py b/tests/integration/spend/test_lens_billing.py index bd9afdfd954..bedcf6c5380 100644 --- a/tests/integration/spend/test_lens_billing.py +++ b/tests/integration/spend/test_lens_billing.py @@ -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", {}) diff --git a/tests/proxy_behavior/lens/evaluate.py b/tests/proxy_behavior/lens/evaluate.py index 99c15203c85..b9b30bdfa53 100644 --- a/tests/proxy_behavior/lens/evaluate.py +++ b/tests/proxy_behavior/lens/evaluate.py @@ -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), ) diff --git a/tests/proxy_behavior/lens/test_lifecycle.py b/tests/proxy_behavior/lens/test_lifecycle.py index 8a9d3873a29..3196b68bd83 100644 --- a/tests/proxy_behavior/lens/test_lifecycle.py +++ b/tests/proxy_behavior/lens/test_lifecycle.py @@ -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) diff --git a/tests/proxy_behavior/lens/worker_storage_smoke.py b/tests/proxy_behavior/lens/worker_storage_smoke.py index 8c80915f978..dca5b928321 100644 --- a/tests/proxy_behavior/lens/worker_storage_smoke.py +++ b/tests/proxy_behavior/lens/worker_storage_smoke.py @@ -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.") diff --git a/tests/proxy_behavior/spend/test_cache_activity.py b/tests/proxy_behavior/spend/test_cache_activity.py index f4a7e8eb2b2..528764fbf45 100644 --- a/tests/proxy_behavior/spend/test_cache_activity.py +++ b/tests/proxy_behavior/spend/test_cache_activity.py @@ -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 diff --git a/tests/test_litellm/proxy/test_tracing_endpoints.py b/tests/test_litellm/proxy/test_tracing_endpoints.py index 4c7c70a39f3..6391c1577f3 100644 --- a/tests/test_litellm/proxy/test_tracing_endpoints.py +++ b/tests/test_litellm/proxy/test_tracing_endpoints.py @@ -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 diff --git a/tests/test_litellm/tracing/normalizers/test_registry.py b/tests/test_litellm/tracing/normalizers/test_registry.py new file mode 100644 index 00000000000..4c2fd051d6c --- /dev/null +++ b/tests/test_litellm/tracing/normalizers/test_registry.py @@ -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) diff --git a/tests/test_litellm/tracing/test_decode.py b/tests/test_litellm/tracing/test_decode.py index 168ff2bf7fb..7185341d038 100644 --- a/tests/test_litellm/tracing/test_decode.py +++ b/tests/test_litellm/tracing/test_decode.py @@ -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 diff --git a/tests/test_litellm/tracing/test_store.py b/tests/test_litellm/tracing/test_store.py index 7ee772e078c..2e80594ff9b 100644 --- a/tests/test_litellm/tracing/test_store.py +++ b/tests/test_litellm/tracing/test_store.py @@ -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"}, } diff --git a/tests/test_litellm/tracing/test_ui_format.py b/tests/test_litellm/tracing/test_ui_format.py new file mode 100644 index 00000000000..27c554041ef --- /dev/null +++ b/tests/test_litellm/tracing/test_ui_format.py @@ -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": ""} diff --git a/tests/unit/integrations/test_s3.py b/tests/unit/integrations/test_s3.py index fd677b9dfdf..c9a53a43d34 100644 --- a/tests/unit/integrations/test_s3.py +++ b/tests/unit/integrations/test_s3.py @@ -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-") diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py index e9f5e667421..963586d9532 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -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", diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 761fa38e73f..e07ffe00d4c 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/test_litellm/proxy/a2a/__init__.py b/tests/unit/proxy/a2a/__init__.py similarity index 100% rename from tests/test_litellm/proxy/a2a/__init__.py rename to tests/unit/proxy/a2a/__init__.py diff --git a/tests/test_litellm/proxy/a2a/test_agent_card.py b/tests/unit/proxy/a2a/test_agent_card.py similarity index 100% rename from tests/test_litellm/proxy/a2a/test_agent_card.py rename to tests/unit/proxy/a2a/test_agent_card.py diff --git a/tests/test_litellm/proxy/a2a/test_discovery.py b/tests/unit/proxy/a2a/test_discovery.py similarity index 100% rename from tests/test_litellm/proxy/a2a/test_discovery.py rename to tests/unit/proxy/a2a/test_discovery.py diff --git a/tests/test_litellm/proxy/a2a/test_version_convert.py b/tests/unit/proxy/a2a/test_version_convert.py similarity index 100% rename from tests/test_litellm/proxy/a2a/test_version_convert.py rename to tests/unit/proxy/a2a/test_version_convert.py diff --git a/tests/test_litellm/proxy/agent_endpoints/__init__.py b/tests/unit/proxy/agent_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/__init__.py rename to tests/unit/proxy/agent_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/__init__.py b/tests/unit/proxy/agent_endpoints/auth/__init__.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/auth/__init__.py rename to tests/unit/proxy/agent_endpoints/auth/__init__.py diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py rename to tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_caller.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py rename to tests/unit/proxy/agent_endpoints/auth/test_agent_caller.py diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py rename to tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py rename to tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py rename to tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_version_e2e.py b/tests/unit/proxy/agent_endpoints/test_a2a_version_e2e.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_a2a_version_e2e.py rename to tests/unit/proxy/agent_endpoints/test_a2a_version_e2e.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py b/tests/unit/proxy/agent_endpoints/test_agent_header_isolation.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py rename to tests/unit/proxy/agent_endpoints/test_agent_header_isolation.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py b/tests/unit/proxy/agent_endpoints/test_agent_headers.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py rename to tests/unit/proxy/agent_endpoints/test_agent_headers.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/unit/proxy/agent_endpoints/test_agent_rbac.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_agent_rbac.py rename to tests/unit/proxy/agent_endpoints/test_agent_rbac.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/unit/proxy/agent_endpoints/test_agent_registry.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py rename to tests/unit/proxy/agent_endpoints/test_agent_registry.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/unit/proxy/agent_endpoints/test_agent_search.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_agent_search.py rename to tests/unit/proxy/agent_endpoints/test_agent_search.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_databricks_oauth.py b/tests/unit/proxy/agent_endpoints/test_databricks_oauth.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_databricks_oauth.py rename to tests/unit/proxy/agent_endpoints/test_databricks_oauth.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/unit/proxy/agent_endpoints/test_endpoints.py similarity index 99% rename from tests/test_litellm/proxy/agent_endpoints/test_endpoints.py rename to tests/unit/proxy/agent_endpoints/test_endpoints.py index 2cf81892db7..81b6c09ca12 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_endpoints.py @@ -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"} diff --git a/tests/test_litellm/proxy/agent_endpoints/test_identity.py b/tests/unit/proxy/agent_endpoints/test_identity.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_identity.py rename to tests/unit/proxy/agent_endpoints/test_identity.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_identity_store.py b/tests/unit/proxy/agent_endpoints/test_identity_store.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_identity_store.py rename to tests/unit/proxy/agent_endpoints/test_identity_store.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py b/tests/unit/proxy/agent_endpoints/test_kill_switch.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py rename to tests/unit/proxy/agent_endpoints/test_kill_switch.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py b/tests/unit/proxy/agent_endpoints/test_managed_identity.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py rename to tests/unit/proxy/agent_endpoints/test_managed_identity.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/unit/proxy/agent_endpoints/test_model_list_helpers.py similarity index 100% rename from tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py rename to tests/unit/proxy/agent_endpoints/test_model_list_helpers.py diff --git a/tests/test_litellm/proxy/analytics_endpoints/__init__.py b/tests/unit/proxy/analytics_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/analytics_endpoints/__init__.py rename to tests/unit/proxy/analytics_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py rename to tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py diff --git a/tests/test_litellm/proxy/anthropic_endpoints/__init__.py b/tests/unit/proxy/anthropic_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/anthropic_endpoints/__init__.py rename to tests/unit/proxy/anthropic_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py similarity index 100% rename from tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py rename to tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py b/tests/unit/proxy/anthropic_endpoints/test_claude_code_skill_access.py similarity index 100% rename from tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_skill_access.py rename to tests/unit/proxy/anthropic_endpoints/test_claude_code_skill_access.py diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py rename to tests/unit/proxy/anthropic_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py rename to tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_skills_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/anthropic_endpoints/test_skills_endpoints.py rename to tests/unit/proxy/anthropic_endpoints/test_skills_endpoints.py diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py similarity index 100% rename from tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py rename to tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py diff --git a/tests/test_litellm/proxy/batches_endpoints/__init__.py b/tests/unit/proxy/batches_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/batches_endpoints/__init__.py rename to tests/unit/proxy/batches_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/batches_endpoints/test_endpoints.py rename to tests/unit/proxy/batches_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py similarity index 100% rename from tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py rename to tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py diff --git a/tests/test_litellm/proxy/container_endpoints/__init__.py b/tests/unit/proxy/config_resolvers/__init__.py similarity index 100% rename from tests/test_litellm/proxy/container_endpoints/__init__.py rename to tests/unit/proxy/config_resolvers/__init__.py diff --git a/tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py b/tests/unit/proxy/config_resolvers/test_config_resolvers.py similarity index 100% rename from tests/test_litellm/proxy/config_resolvers/test_config_resolvers.py rename to tests/unit/proxy/config_resolvers/test_config_resolvers.py diff --git a/tests/test_litellm/proxy/config_resolvers/test_settings_rules.py b/tests/unit/proxy/config_resolvers/test_settings_rules.py similarity index 100% rename from tests/test_litellm/proxy/config_resolvers/test_settings_rules.py rename to tests/unit/proxy/config_resolvers/test_settings_rules.py diff --git a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py b/tests/unit/proxy/config_resolvers/test_settings_store.py similarity index 100% rename from tests/test_litellm/proxy/config_resolvers/test_settings_store.py rename to tests/unit/proxy/config_resolvers/test_settings_store.py diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py index 1d0a7475db6..0e8acf3b905 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -8,6 +8,7 @@ from collections.abc import Iterator from typing import Dict import pytest +from prisma.errors import ClientNotConnectedError import litellm @@ -15,6 +16,24 @@ import litellm.proxy.proxy_server from tests.unit.litellm_core_utils.fake_secret_vault import FakeSecretVault +class StubClientNotConnectedError(ClientNotConnectedError): + pass + + +class DisconnectedPrisma: + def is_connected(self) -> bool: + return False + + @property + def _engine(self) -> None: + raise StubClientNotConnectedError() + + +@pytest.fixture +def disconnected_prisma() -> DisconnectedPrisma: + return DisconnectedPrisma() + + # Top-level assignments of these types are the ones importlib.reload(litellm) # would have effectively reset. We snapshot them at conftest import time and # deep-copy the snapshot back before every test. diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py b/tests/unit/proxy/container_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py rename to tests/unit/proxy/container_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/container_endpoints/test_endpoints.py b/tests/unit/proxy/container_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/container_endpoints/test_endpoints.py rename to tests/unit/proxy/container_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py b/tests/unit/proxy/container_endpoints/test_handler_factory.py similarity index 100% rename from tests/test_litellm/proxy/container_endpoints/test_handler_factory.py rename to tests/unit/proxy/container_endpoints/test_handler_factory.py diff --git a/tests/test_litellm/proxy/google_endpoints/__init__.py b/tests/unit/proxy/credential_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/google_endpoints/__init__.py rename to tests/unit/proxy/credential_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/credential_endpoints/test_endpoints.py rename to tests/unit/proxy/credential_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/memory/__init__.py b/tests/unit/proxy/discovery_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/memory/__init__.py rename to tests/unit/proxy/discovery_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py b/tests/unit/proxy/discovery_endpoints/test_agent_skills_archive.py similarity index 100% rename from tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py rename to tests/unit/proxy/discovery_endpoints/test_agent_skills_archive.py diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_agent_skills_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py rename to tests/unit/proxy/discovery_endpoints/test_agent_skills_endpoints.py diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py rename to tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py diff --git a/tests/test_litellm/proxy/ocr_endpoints/__init__.py b/tests/unit/proxy/fine_tuning_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/ocr_endpoints/__init__.py rename to tests/unit/proxy/fine_tuning_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py rename to tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/google_endpoints/test_endpoints.py rename to tests/unit/proxy/google_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py rename to tests/unit/proxy/google_endpoints/test_google_api_endpoints.py diff --git a/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py b/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py similarity index 100% rename from tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py rename to tests/unit/proxy/google_endpoints/test_interactions_agent_param.py diff --git a/tests/test_litellm/proxy/google_endpoints/test_managed_agents_model_param.py b/tests/unit/proxy/google_endpoints/test_managed_agents_model_param.py similarity index 100% rename from tests/test_litellm/proxy/google_endpoints/test_managed_agents_model_param.py rename to tests/unit/proxy/google_endpoints/test_managed_agents_model_param.py diff --git a/tests/test_litellm/proxy/rag_endpoints/__init__.py b/tests/unit/proxy/health_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/rag_endpoints/__init__.py rename to tests/unit/proxy/health_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/health_endpoints/test_graceful_shutdown_endpoints.py b/tests/unit/proxy/health_endpoints/test_graceful_shutdown_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/health_endpoints/test_graceful_shutdown_endpoints.py rename to tests/unit/proxy/health_endpoints/test_graceful_shutdown_endpoints.py diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py rename to tests/unit/proxy/health_endpoints/test_health_endpoints.py diff --git a/tests/test_litellm/proxy/image_endpoints/__init__.py b/tests/unit/proxy/image_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/image_endpoints/__init__.py rename to tests/unit/proxy/image_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py b/tests/unit/proxy/image_endpoints/test_azure_routes.py similarity index 100% rename from tests/test_litellm/proxy/image_endpoints/test_azure_routes.py rename to tests/unit/proxy/image_endpoints/test_azure_routes.py diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/image_endpoints/test_endpoints.py rename to tests/unit/proxy/image_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/rerank_endpoints/__init__.py b/tests/unit/proxy/lens/__init__.py similarity index 100% rename from tests/test_litellm/proxy/rerank_endpoints/__init__.py rename to tests/unit/proxy/lens/__init__.py diff --git a/tests/unit/proxy/engine/test_analysis.py b/tests/unit/proxy/lens/test_analysis.py similarity index 91% rename from tests/unit/proxy/engine/test_analysis.py rename to tests/unit/proxy/lens/test_analysis.py index bc688d37f99..97b0c5ab022 100644 --- a/tests/unit/proxy/engine/test_analysis.py +++ b/tests/unit/proxy/lens/test_analysis.py @@ -6,8 +6,8 @@ from typing import Final import pytest -from litellm.proxy.engine.analysis import Candidate, Examined, evidence_valid, extract, investigate, partition_content -from litellm.proxy.engine.models import ( +from litellm.proxy.lens.analysis import Candidate, Examined, evidence_valid, extract, investigate, partition_content +from litellm.proxy.lens.models import ( Claim, Coverage, Evidence, @@ -18,14 +18,14 @@ from litellm.proxy.engine.models import ( Sample, TracePart, ) -from litellm.proxy.engine.state import queue_job -from tests.unit.proxy.engine.test_state import NOW, engine, finding +from litellm.proxy.lens.state import queue_job +from tests.unit.proxy.lens.test_state import NOW, lens, finding @pytest.mark.asyncio @pytest.mark.parametrize("outcome", ("complete", "cancel", "failure")) async def test_parallel_review_shares_one_model_limit_and_cleans_up(outcome: str) -> None: - from litellm.proxy.engine.analysis import ANALYSIS_CONCURRENCY, analyze_sample + from litellm.proxy.lens.analysis import ANALYSIS_CONCURRENCY, analyze_sample executions: Final = tuple( Execution(id=str(i), source="traces", trace_id=str(i), team_id="alpha", name="run", start_time="", span_count=6) @@ -70,7 +70,7 @@ async def test_parallel_review_shares_one_model_limit_and_cleans_up(outcome: str if stage == "Reading executions": counts.put(coverage.screened) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) task: Final = asyncio.create_task( analyze_sample(claim, Sample(executions=executions, eligible=len(executions)), read, model, progress) ) @@ -101,7 +101,7 @@ async def test_parallel_review_shares_one_model_limit_and_cleans_up(outcome: str @pytest.mark.asyncio async def test_independent_investigations_overlap_and_report_completions() -> None: - from litellm.proxy.engine.analysis import investigate_candidates + from litellm.proxy.lens.analysis import investigate_candidates arrived: Final = SimpleQueue[str]() progress_counts: Final = SimpleQueue[int]() @@ -124,7 +124,7 @@ async def test_independent_investigations_overlap_and_report_completions() -> No candidates: Final = tuple( Candidate(check_id="retries", title=str(i), hypothesis="Investigate", execution_ids=()) for i in range(2) ) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) results: Final = tuple( [ result @@ -185,7 +185,7 @@ async def test_reviewer_sees_final_outcome_and_catalog_across_pages() -> None: assert pages.qsize() == 2 return ModelResult(content='{"observations":[],"cannot_assess":false}', cost=0) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await extract(claim, execution, read, model) assert root in result.parts assert not result.cannot_assess @@ -193,7 +193,7 @@ async def test_reviewer_sees_final_outcome_and_catalog_across_pages() -> None: @pytest.mark.asyncio async def test_reviewer_fetches_targeted_evidence_and_rejects_outside_catalog_reads() -> None: - from litellm.proxy.engine.analysis import Observation, SpanRead, TraceReview + from litellm.proxy.lens.analysis import Observation, SpanRead, TraceReview execution: Final = Execution( id="run", source="traces", trace_id="t", team_id="", name="run", start_time="", span_count=2 @@ -246,7 +246,7 @@ async def test_reviewer_fetches_targeted_evidence_and_rejects_outside_catalog_re cost=0, ) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await extract(claim, execution, read, model) assert len(result.observations) == 1 assert result.observations[0].evidence[0].quote == "Verified result: failed" @@ -255,7 +255,7 @@ async def test_reviewer_fetches_targeted_evidence_and_rejects_outside_catalog_re @pytest.mark.asyncio async def test_reviewer_stops_repeated_read_requests() -> None: - from litellm.proxy.engine.analysis import SpanRead, TraceReview + from litellm.proxy.lens.analysis import SpanRead, TraceReview execution: Final = Execution( id="run", source="traces", trace_id="t", team_id="", name="run", start_time="", span_count=1 @@ -276,7 +276,7 @@ async def test_reviewer_stops_repeated_read_requests() -> None: content=TraceReview(reads=(SpanRead(span_id="01"),), cannot_assess=True).model_dump_json(), cost=0 ) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await extract(claim, execution, read, model) assert result.cannot_assess assert reads.qsize() == 2 @@ -311,7 +311,7 @@ async def test_investigator_rejects_a_fabricated_quote() -> None: async def read(_execution_id: str, _cursor: str, _offset: int) -> ExecutionContent: return ExecutionContent(execution=execution, parts=examined.parts) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await investigate( claim, Candidate(check_id="retries", title="Retries", hypothesis="Unrecovered", execution_ids=("run1",)), @@ -353,7 +353,7 @@ async def test_assessable_content_is_not_overridden_by_unknown_chunks(paginated: unavailable: Final = "false" if "verified result" in request.prompt else "true" return ModelResult(content='{"observations":[],"cannot_assess":' + unavailable + "}", cost=0) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await extract(claim, execution, read, model) assert result.cannot_assess is not assessable @@ -382,7 +382,7 @@ async def test_investigator_keeps_final_outcome_ahead_of_repeated_model_history( async def read(_execution_id: str, _cursor: str, _offset: int) -> ExecutionContent: return ExecutionContent(execution=execution, parts=examined.parts) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await investigate( claim, Candidate(check_id="retries", title="Retries", hypothesis="Unrecovered", execution_ids=("run1",)), @@ -425,7 +425,7 @@ async def test_oversized_model_evidence_is_retried_and_quotes_still_verified( cost=0, ) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await extract(claim, execution, read, model) assert len(result.observations) == int(accepted) assert result.cannot_assess is not accepted @@ -436,7 +436,7 @@ async def test_oversized_model_evidence_is_retried_and_quotes_still_verified( async def test_invalid_model_output_has_only_one_repair_attempt() -> None: from pydantic import ValidationError - from litellm.proxy.engine.analysis import Extraction, structured_response + from litellm.proxy.lens.analysis import Extraction, structured_response attempts: Final = iter((1, 2)) @@ -451,8 +451,8 @@ async def test_invalid_model_output_has_only_one_repair_attempt() -> None: @pytest.mark.asyncio async def test_grouping_consolidates_prior_batches_and_reports_real_progress() -> None: - from litellm.proxy.engine.analysis import Clusters, Observation, cluster_batches - from litellm.proxy.engine.models import Coverage + from litellm.proxy.lens.analysis import Clusters, Observation, cluster_batches + from litellm.proxy.lens.models import Coverage candidate: Final = Candidate( check_id="retries", title="Outage", hypothesis="Tool unavailable", execution_ids=("run1",) @@ -517,7 +517,7 @@ async def test_investigator_can_cite_a_later_page_or_offset(later_span: str) -> assert execution_id == "run1" and offset == 8000 return ExecutionContent(execution=execution, parts=(later,)) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await investigate( claim, Candidate(check_id="retries", title="Retries", hypothesis="Unrecovered", execution_ids=("run1",)), @@ -530,7 +530,7 @@ async def test_investigator_can_cite_a_later_page_or_offset(later_span: str) -> @pytest.mark.asyncio async def test_thousands_of_matching_runs_keep_all_members_without_a_growing_model_prompt() -> None: - from litellm.proxy.engine.analysis import Clusters, Observation, cluster_batches, observation_batches + from litellm.proxy.lens.analysis import Clusters, Observation, cluster_batches, observation_batches observations: Final = tuple( Observation( @@ -571,7 +571,7 @@ async def test_thousands_of_matching_runs_keep_all_members_without_a_growing_mod @pytest.mark.asyncio async def test_grouping_preserves_observations_omitted_by_model() -> None: - from litellm.proxy.engine.analysis import merge_candidates + from litellm.proxy.lens.analysis import merge_candidates original: Final = Candidate( check_id="retries", title="Unrecovered failure", hypothesis="Timeout", execution_ids=("run",) @@ -587,7 +587,7 @@ async def test_grouping_preserves_observations_omitted_by_model() -> None: @pytest.mark.asyncio async def test_grouping_repairs_duplicate_members_before_creating_findings() -> None: - from litellm.proxy.engine.analysis import Clusters, merge_candidates + from litellm.proxy.lens.analysis import Clusters, merge_candidates original: Final = Candidate( check_id="retries", title="Unrecovered failure", hypothesis="Timeout", execution_ids=("run",) @@ -609,7 +609,7 @@ async def test_grouping_repairs_duplicate_members_before_creating_findings() -> @pytest.mark.asyncio async def test_review_keeps_original_ids_in_per_run_assessments() -> None: - from litellm.proxy.engine.analysis import analyze_sample + from litellm.proxy.lens.analysis import analyze_sample execution: Final = Execution( id="opaque-original-id", @@ -636,7 +636,7 @@ async def test_review_keeps_original_ids_in_per_run_assessments() -> None: async def progress(_stage: str, _coverage: Coverage) -> None: pass - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await analyze_sample(claim, Sample(executions=(execution,), eligible=1), read, model, progress) assert result.assessments[0].execution_id == execution.id assert not result.assessments[0].cannot_assess @@ -678,7 +678,7 @@ async def test_investigation_context_accounts_for_metadata_on_thousands_of_short async def read(_identity: str, _cursor: str, _offset: int) -> ExecutionContent: pytest.fail("No read was requested") - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await investigate( claim, Candidate( @@ -696,7 +696,7 @@ async def test_investigation_context_accounts_for_metadata_on_thousands_of_short @pytest.mark.asyncio async def test_completed_read_does_not_make_supported_review_unknown() -> None: - from litellm.proxy.engine.analysis import Observation, SpanRead, TraceReview + from litellm.proxy.lens.analysis import Observation, SpanRead, TraceReview execution: Final = Execution( id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 @@ -720,7 +720,7 @@ async def test_completed_read_does_not_make_supported_review_unknown() -> None: content=TraceReview(reads=(SpanRead(span_id="s"),), observations=(observation,)).model_dump_json(), cost=0 ) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await extract(claim, execution, read, model) assert result.observations == (observation,) assert not result.cannot_assess and not result.partial @@ -770,7 +770,7 @@ async def test_echoed_feedback_page_does_not_skip_requested_evidence() -> None: cost=0, ) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await extract(claim, execution, read, model) assert tuple(requests.get_nowait() for _ in range(requests.qsize())) == (0, 1) assert len(result.observations) == 1 @@ -797,7 +797,7 @@ async def test_empty_navigation_requires_a_final_decision(action: str) -> None: return ModelResult(content='{"action":"inconclusive"}', cost=0) return ModelResult(content=json.dumps({"action": action, "page": 999, "execution_id": "run"}), cost=0) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) result: Final = await investigate( claim, Candidate(check_id="retries", title="Timeout", hypothesis="Failed", execution_ids=("run",)), @@ -812,20 +812,20 @@ async def test_empty_navigation_requires_a_final_decision(action: str) -> None: @pytest.mark.asyncio @pytest.mark.parametrize("phase", ("extract", "investigate")) async def test_large_feedback_history_is_accessible_without_overflowing_context(phase: str) -> None: - from litellm.proxy.engine.state import merge_finding + from litellm.proxy.lens.state import merge_finding execution: Final = Execution( id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 ) part: Final = TracePart(execution_id="run", span_id="span", name="task", kind="agent", content="timeout") - accepted: Final = merge_finding(engine(), finding("run"), 1, NOW) + accepted: Final = merge_finding(lens(), finding("run"), 1, NOW) prior: Final = tuple( accepted.model_copy( update=MappingProxyType({"id": str(i), "status": "dismissed", "reason": f"Accepted-{i}: " + "x" * 1900}) ) for i in range(60) ) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=prior) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=prior) pages: Final = SimpleQueue[int]() async def read(_identity: str, _cursor: str, _offset: int) -> ExecutionContent: @@ -865,7 +865,7 @@ async def test_large_feedback_history_is_accessible_without_overflowing_context( @pytest.mark.asyncio async def test_final_registry_reconciles_patterns_split_across_pages() -> None: - from litellm.proxy.engine.analysis import Clusters, Observation, cluster_batches + from litellm.proxy.lens.analysis import Clusters, Observation, cluster_batches observations: Final = tuple( Observation( @@ -902,7 +902,7 @@ async def test_final_registry_reconciles_patterns_split_across_pages() -> None: @pytest.mark.asyncio async def test_distinct_patterns_are_consolidated_in_batches_without_losing_runs() -> None: - from litellm.proxy.engine.analysis import Observation, cluster_batches, observation_batches + from litellm.proxy.lens.analysis import Observation, cluster_batches, observation_batches observations: Final = tuple( Observation( @@ -930,7 +930,7 @@ async def test_distinct_patterns_are_consolidated_in_batches_without_losing_runs @pytest.mark.asyncio async def test_invalid_candidate_response_preserves_other_findings_and_reports_inconclusive() -> None: - from litellm.proxy.engine.analysis import investigate_candidates + from litellm.proxy.lens.analysis import investigate_candidates execution: Final = Execution( id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 @@ -954,7 +954,7 @@ async def test_invalid_candidate_response_preserves_other_findings_and_reports_i async def progress(_stage: str, coverage: Coverage) -> None: counts.put(coverage.inconclusive) - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) results: Final = tuple( [ result diff --git a/tests/unit/proxy/engine/test_endpoints.py b/tests/unit/proxy/lens/test_endpoints.py similarity index 83% rename from tests/unit/proxy/engine/test_endpoints.py rename to tests/unit/proxy/lens/test_endpoints.py index e8d0095754f..97bb7759a02 100644 --- a/tests/unit/proxy/engine/test_endpoints.py +++ b/tests/unit/proxy/lens/test_endpoints.py @@ -4,7 +4,7 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.engine.endpoints import user_scope +from litellm.proxy.lens.endpoints import user_scope @pytest.mark.parametrize( @@ -27,10 +27,10 @@ def test_admin_can_configure_lens_and_viewer_can_only_read() -> None: @pytest.mark.parametrize("identity", ("not-an-execution", "W10=", "WyJvdGhlciIsICIiLCAiaWQiXQ==")) def test_invalid_explicit_execution_ids_are_rejected(identity: str) -> None: - from litellm.proxy.engine.endpoints import validate_selection - from tests.unit.proxy.engine.test_state import engine + from litellm.proxy.lens.endpoints import validate_selection + from tests.unit.proxy.lens.test_state import lens - settings: Final = engine().settings.model_copy(update={"execution_ids": (identity,)}) + settings: Final = lens().settings.model_copy(update={"execution_ids": (identity,)}) with pytest.raises(HTTPException) as error: validate_selection(settings) assert error.value.status_code == 422 @@ -38,8 +38,8 @@ def test_invalid_explicit_execution_ids_are_rejected(identity: str) -> None: @pytest.mark.asyncio async def test_incompatible_worker_is_rejected_before_claiming_work() -> None: - from litellm.proxy.engine.endpoints import claim - from tests.unit.proxy.engine.test_state import worker + from litellm.proxy.lens.endpoints import claim + from tests.unit.proxy.lens.test_state import worker with pytest.raises(HTTPException) as error: await claim(worker(), protocol_version=1) diff --git a/tests/unit/proxy/engine/test_inference.py b/tests/unit/proxy/lens/test_inference.py similarity index 58% rename from tests/unit/proxy/engine/test_inference.py rename to tests/unit/proxy/lens/test_inference.py index 90efa5cdf0f..2243759b773 100644 --- a/tests/unit/proxy/engine/test_inference.py +++ b/tests/unit/proxy/lens/test_inference.py @@ -2,18 +2,18 @@ from typing import Final import pytest -from litellm.proxy.engine.inference import Deployment, DeploymentParams, completion_charge, quote +from litellm.proxy.lens.inference import Deployment, DeploymentParams, completion_charge, quote from litellm.types.utils import ModelResponse def test_custom_priced_model_charges_reported_tokens() -> None: deployment: Final = Deployment( litellm_params=DeploymentParams( - model="openai/engine-test", input_cost_per_token=0.001, output_cost_per_token=0.002 + model="openai/lens-test", input_cost_per_token=0.001, output_cost_per_token=0.002 ) ) response: Final = ModelResponse( - model="engine-test", usage={"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30} + model="lens-test", usage={"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30} ) assert completion_charge((deployment,), response, 10) == pytest.approx(0.04) assert quote((deployment,), "hello") > 0.04 diff --git a/tests/unit/proxy/engine/test_sources.py b/tests/unit/proxy/lens/test_sources.py similarity index 87% rename from tests/unit/proxy/engine/test_sources.py rename to tests/unit/proxy/lens/test_sources.py index 7ec46922f21..5dc6e2652f0 100644 --- a/tests/unit/proxy/engine/test_sources.py +++ b/tests/unit/proxy/lens/test_sources.py @@ -4,11 +4,11 @@ from typing import Final import pytest -from litellm.proxy.engine.models import Scope, MetadataFilter -from litellm.proxy.engine.sources import SourceReader -from tests.unit.proxy.engine.test_state import engine +from litellm.proxy.lens.models import Scope, MetadataFilter +from litellm.proxy.lens.sources import SourceReader +from tests.unit.proxy.lens.test_state import lens -from litellm.proxy.engine.sources import execution_id, parse_execution +from litellm.proxy.lens.sources import execution_id, parse_execution def test_same_trace_id_from_different_keys_is_a_distinct_execution() -> None: @@ -57,7 +57,7 @@ async def test_sample_never_returns_authentication_attributes() -> None: ] reader: Final = SourceReader(StorageResponse()) - sample: Final = await reader.sample(Scope(team_id="alpha"), engine().settings, 1, 2) + sample: Final = await reader.sample(Scope(team_id="alpha"), lens().settings, 1, 2) assert sample.executions[0].metadata == (MetadataFilter(key="environment", value="production"),) assert "opaque-oauth-bearer" not in sample.model_dump_json() assert sample.eligible == 1 diff --git a/tests/unit/proxy/engine/test_state.py b/tests/unit/proxy/lens/test_state.py similarity index 87% rename from tests/unit/proxy/engine/test_state.py rename to tests/unit/proxy/lens/test_state.py index 8d56f4595da..ac70a22077e 100644 --- a/tests/unit/proxy/engine/test_state.py +++ b/tests/unit/proxy/lens/test_state.py @@ -3,17 +3,17 @@ from typing import Final import pytest -from litellm.proxy.engine.models import Check, Engine, EngineSettings, Evidence, FindingDraft, Scope, Worker -from litellm.proxy.engine.state import can_access, claim_job, current_job, merge_finding, queue_job, renew_budget +from litellm.proxy.lens.models import Check, Lens, LensSettings, Evidence, FindingDraft, Scope, Worker +from litellm.proxy.lens.state import can_access, claim_job, current_job, merge_finding, queue_job, renew_budget NOW: Final = datetime(2026, 1, 15, tzinfo=timezone.utc) -def engine() -> Engine: - return Engine( - id="engine", +def lens() -> Lens: + return Lens( + id="lens", scope=Scope(team_id="alpha"), - settings=EngineSettings( + settings=LensSettings( name="Research", model="analysis", checks=(Check(id="retries", instruction="Find unrecovered retries"),) ), created_at=NOW, @@ -50,7 +50,7 @@ def test_scope_never_crosses_another_team_or_key(viewer: Scope, target: Scope, a def test_queue_is_idempotent_and_settings_are_frozen() -> None: - original: Final = engine() + original: Final = lens() queued: Final = queue_job(original, NOW, "job") edited: Final = queued.model_copy( update={"settings": original.settings.model_copy(update={"model": "replacement"})} @@ -65,7 +65,7 @@ def test_queue_is_idempotent_and_settings_are_frozen() -> None: def test_one_off_overrides_do_not_change_saved_monitoring_settings() -> None: - original: Final = engine() + original: Final = lens() override: Final = original.settings.model_copy( update={"sample_percent": 10, "sample_size": None, "concurrency": 3, "lookback_hours": 72} ) @@ -79,7 +79,7 @@ def test_one_off_overrides_do_not_change_saved_monitoring_settings() -> None: def test_behavior_description_is_sufficient_without_separate_checks() -> None: - settings: Final = EngineSettings(name="Behavior", model="analysis", context="Answer using cited sources") + settings: Final = LensSettings(name="Behavior", model="analysis", context="Answer using cited sources") assert tuple(c.id for c in settings.analysis_checks) == ("expected_behavior",) assert settings.sample_size is None assert settings.sample_percent == 100 @@ -92,11 +92,11 @@ def test_invalid_selection_and_parallelism_are_rejected(field: str, value: int) from pydantic import ValidationError with pytest.raises(ValidationError): - EngineSettings.model_validate({**engine().settings.model_dump(), field: value}) + LensSettings.model_validate({**lens().settings.model_dump(), field: value}) def test_lease_prevents_double_claim_and_expires_with_bounded_retries() -> None: - queued: Final = queue_job(engine(), NOW, "job") + queued: Final = queue_job(lens(), NOW, "job") first: Final = claim_job(queued, worker(), NOW) assert claim_job(first, worker(identity="second"), NOW) is first assert claim_job(first, worker(team="beta"), NOW + timedelta(minutes=6)) is first @@ -110,9 +110,9 @@ def test_lease_prevents_double_claim_and_expires_with_bounded_retries() -> None: def test_replaying_evidence_does_not_reopen_but_new_occurrence_does() -> None: - from litellm.proxy.engine.state import snapshot_finding + from litellm.proxy.lens.state import snapshot_finding - original: Final = engine() + original: Final = lens() resolved: Final = merge_finding(original, finding("run1"), 1, NOW).model_copy(update={"status": "resolved"}) reviewed: Final = original.model_copy(update={"findings": (resolved,)}) assert merge_finding(reviewed, finding("run1"), 1, NOW).status == "resolved" @@ -138,7 +138,7 @@ def test_replaying_evidence_does_not_reopen_but_new_occurrence_does() -> None: def test_monthly_budget_renews_without_erasing_job_costs() -> None: - spent: Final = queue_job(engine(), NOW, "job").model_copy(update={"spent": 12}) + spent: Final = queue_job(lens(), NOW, "job").model_copy(update={"spent": 12}) renewed: Final = renew_budget(spent, datetime(2026, 2, 1, tzinfo=timezone.utc)) assert renewed.spent == 0 assert renewed.jobs == spent.jobs @@ -147,7 +147,7 @@ def test_monthly_budget_renews_without_erasing_job_costs() -> None: @pytest.mark.parametrize("hours", (24, 168, 720)) def test_every_scan_uses_the_configured_lookback_window(hours: int) -> None: - original: Final = engine() + original: Final = lens() configured: Final = original.model_copy( update={"settings": original.settings.model_copy(update={"lookback_hours": hours})} ) @@ -159,15 +159,15 @@ def test_every_scan_uses_the_configured_lookback_window(hours: int) -> None: def test_finding_keeps_uncertainty_separate_from_the_main_summary() -> None: draft: Final = finding("run1").model_copy(update={"limitation": "The final response was not recorded."}) - saved: Final = merge_finding(engine(), draft, 1, NOW) + saved: Final = merge_finding(lens(), draft, 1, NOW) assert saved.limitation == draft.limitation assert saved.description == draft.description @pytest.mark.parametrize("interval", (1, 2, 37, 90, 10080)) def test_custom_schedule_does_not_overlap_an_active_scan(interval: int) -> None: - original: Final = engine() - settings: Final = EngineSettings.model_validate({**original.settings.model_dump(), "interval_minutes": interval}) + original: Final = lens() + settings: Final = LensSettings.model_validate({**original.settings.model_dump(), "interval_minutes": interval}) configured: Final = original.model_copy(update={"settings": settings}) running: Final = claim_job(queue_job(configured, NOW, "first"), worker(), NOW) assert queue_job(running, NOW + timedelta(minutes=interval), "second") is running @@ -178,13 +178,13 @@ def test_invalid_schedule_is_rejected(interval: float) -> None: from pydantic import ValidationError with pytest.raises(ValidationError): - EngineSettings.model_validate({**engine().settings.model_dump(), "interval_minutes": interval}) + LensSettings.model_validate({**lens().settings.model_dump(), "interval_minutes": interval}) def test_batch_snapshot_keeps_feedback_identity_and_only_current_evidence() -> None: - from litellm.proxy.engine.state import snapshot_finding + from litellm.proxy.lens.state import snapshot_finding - original: Final = engine() + original: Final = lens() dismissed: Final = merge_finding(original, finding("old-run"), 1, NOW).model_copy( update={"status": "dismissed", "reason": "Expected recovery"} ) @@ -204,9 +204,9 @@ def test_batch_snapshot_keeps_feedback_identity_and_only_current_evidence() -> N @pytest.mark.parametrize("explicit_reference", (False, True)) def test_issue_and_pattern_with_same_title_keep_independent_feedback(explicit_reference: bool) -> None: - from litellm.proxy.engine.state import snapshot_finding + from litellm.proxy.lens.state import snapshot_finding - original: Final = engine() + original: Final = lens() issue: Final = merge_finding(original, finding("old"), 1, NOW).model_copy( update={"status": "dismissed", "reason": "Expected retry"} ) @@ -228,7 +228,7 @@ def test_issue_and_pattern_with_same_title_keep_independent_feedback(explicit_re def test_legacy_finding_identity_preserves_feedback_only_for_same_kind_and_check() -> None: import hashlib - original: Final = engine() + original: Final = lens() draft: Final = finding("old") legacy_id: Final = hashlib.sha256(f"{original.id}:{draft.check_id}:{draft.title.lower()}".encode()).hexdigest()[:24] legacy: Final = merge_finding(original, draft, 1, NOW).model_copy( diff --git a/tests/unit/proxy/engine/test_trace_store.py b/tests/unit/proxy/lens/test_trace_store.py similarity index 93% rename from tests/unit/proxy/engine/test_trace_store.py rename to tests/unit/proxy/lens/test_trace_store.py index f80d4348864..03667c81d3a 100644 --- a/tests/unit/proxy/engine/test_trace_store.py +++ b/tests/unit/proxy/lens/test_trace_store.py @@ -1,8 +1,8 @@ import json from typing import Final -from litellm.proxy.engine.models import Evidence, TracePart -from litellm.proxy.engine.trace_store import trace_store +from litellm.proxy.lens.models import Evidence, TracePart +from litellm.proxy.lens.trace_store import trace_store def test_trace_store_pages_large_payloads_and_recovers_exact_evidence() -> None: diff --git a/tests/unit/proxy/engine/test_worker.py b/tests/unit/proxy/lens/test_worker.py similarity index 83% rename from tests/unit/proxy/engine/test_worker.py rename to tests/unit/proxy/lens/test_worker.py index e244eff08ec..a0212e03319 100644 --- a/tests/unit/proxy/engine/test_worker.py +++ b/tests/unit/proxy/lens/test_worker.py @@ -4,7 +4,7 @@ from typing import Final import httpx import pytest -from litellm.proxy.engine.models import ( +from litellm.proxy.lens.models import ( Claim, Execution, ExecutionContent, @@ -14,9 +14,9 @@ from litellm.proxy.engine.models import ( Sample, TracePart, ) -from litellm.proxy.engine.state import queue_job -from litellm.proxy.engine.worker import EngineWorker -from tests.unit.proxy.engine.test_state import NOW, engine +from litellm.proxy.lens.state import queue_job +from litellm.proxy.lens.worker import LensWorker +from tests.unit.proxy.lens.test_state import NOW, lens @pytest.mark.asyncio @@ -39,7 +39,7 @@ async def test_model_retries_transient_failures_but_not_budget_or_revocation(fai delays.put(delay) async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: - worker: Final = EngineWorker(client, sleep=sleep) + worker: Final = LensWorker(client, sleep=sleep) if failure in (402, 409, 401): with pytest.raises(httpx.HTTPStatusError): await worker.model_request("/model", ModelRequest(purpose="extract", prompt="review")) @@ -64,7 +64,7 @@ async def test_transient_retries_are_bounded() -> None: async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: with pytest.raises(httpx.HTTPStatusError): - await EngineWorker(client, sleep=sleep).model_request( + await LensWorker(client, sleep=sleep).model_request( "/model", ModelRequest(purpose="extract", prompt="review") ) assert attempts.qsize() == 3 @@ -74,17 +74,17 @@ async def test_transient_retries_are_bounded() -> None: @pytest.mark.asyncio async def test_idle_worker_does_not_start_an_analysis() -> None: def handle(request: httpx.Request) -> httpx.Response: - assert request.url.path == "/engine/worker/claim" + assert request.url.path == "/lens/worker/claim" return httpx.Response(200, content="null") async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: - assert await EngineWorker(client).run_once() is False + assert await LensWorker(client).run_once() is False @pytest.mark.asyncio @pytest.mark.parametrize("model_status", (200, 402, 503)) async def test_worker_reads_claimed_activity_and_reports_analysis_or_failure(model_status: int) -> None: - claim: Final = Claim(engine_id="engine", job=queue_job(engine(), NOW, "job").jobs[0], findings=()) + claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) execution: Final = Execution( id="run", source="traces", trace_id="trace", team_id="alpha", name="review", start_time="", span_count=1 ) @@ -97,28 +97,28 @@ async def test_worker_reads_claimed_activity_and_reports_analysis_or_failure(mod def handle(request: httpx.Request) -> httpx.Response: match request.url.path: - case "/engine/worker/claim": + case "/lens/worker/claim": return httpx.Response(200, json=claim.model_dump(mode="json")) - case "/engine/worker/engine/job/sample": + case "/lens/worker/lens/job/sample": return httpx.Response(200, json=sample.model_dump(mode="json")) - case "/engine/worker/engine/job/content": + case "/lens/worker/lens/job/content": assert request.url.params["execution_id"] == execution.id return httpx.Response(200, json=content.model_dump(mode="json")) - case "/engine/worker/engine/job/model": + case "/lens/worker/lens/job/model": return httpx.Response( model_status, json=ModelResult(content='{"observations":[],"cannot_assess":false}', cost=0.01).model_dump(), ) - case "/engine/worker/engine/job/progress": + case "/lens/worker/lens/job/progress": return httpx.Response(200, json=True) - case "/engine/worker/engine/job/result": + case "/lens/worker/lens/job/result": saved.put(Result.model_validate_json(request.content)) return httpx.Response(200, json=True) case _: pytest.fail(f"Unexpected analyzer request: {request.url.path}") async with httpx.AsyncClient(base_url="https://proxy.test", transport=httpx.MockTransport(handle)) as client: - assert await EngineWorker(client).run_once() is True + assert await LensWorker(client).run_once() is True result: Final = saved.get_nowait() assert saved.empty() if model_status == 200: diff --git a/tests/test_litellm/proxy/response_api_endpoints/__init__.py b/tests/unit/proxy/list_api/__init__.py similarity index 100% rename from tests/test_litellm/proxy/response_api_endpoints/__init__.py rename to tests/unit/proxy/list_api/__init__.py diff --git a/tests/test_litellm/proxy/list_api/test_common.py b/tests/unit/proxy/list_api/test_common.py similarity index 100% rename from tests/test_litellm/proxy/list_api/test_common.py rename to tests/unit/proxy/list_api/test_common.py diff --git a/tests/test_litellm/proxy/list_api/test_in_memory.py b/tests/unit/proxy/list_api/test_in_memory.py similarity index 100% rename from tests/test_litellm/proxy/list_api/test_in_memory.py rename to tests/unit/proxy/list_api/test_in_memory.py diff --git a/tests/test_litellm/proxy/list_api/test_list_framework.py b/tests/unit/proxy/list_api/test_list_framework.py similarity index 100% rename from tests/test_litellm/proxy/list_api/test_list_framework.py rename to tests/unit/proxy/list_api/test_list_framework.py diff --git a/tests/test_litellm/proxy/utils/__init__.py b/tests/unit/proxy/memory/__init__.py similarity index 100% rename from tests/test_litellm/proxy/utils/__init__.py rename to tests/unit/proxy/memory/__init__.py diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/unit/proxy/memory/test_memory_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/memory/test_memory_endpoints.py rename to tests/unit/proxy/memory/test_memory_endpoints.py diff --git a/tests/test_litellm/proxy/utils/helpers/__init__.py b/tests/unit/proxy/ocr_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/__init__.py rename to tests/unit/proxy/ocr_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py b/tests/unit/proxy/ocr_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py rename to tests/unit/proxy/ocr_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/__init__.py b/tests/unit/proxy/openai_files_endpoint/__init__.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/__init__.py rename to tests/unit/proxy/openai_files_endpoint/__init__.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/unit/proxy/openai_files_endpoint/test_batch_guardrails.py similarity index 100% rename from tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py rename to tests/unit/proxy/openai_files_endpoint/test_batch_guardrails.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py b/tests/unit/proxy/openai_files_endpoint/test_files_batch_file_validation.py similarity index 100% rename from tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py rename to tests/unit/proxy/openai_files_endpoint/test_files_batch_file_validation.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py similarity index 100% rename from tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py rename to tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py similarity index 100% rename from tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py rename to tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py b/tests/unit/proxy/openai_files_endpoint/test_general_upload_validation.py similarity index 100% rename from tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py rename to tests/unit/proxy/openai_files_endpoint/test_general_upload_validation.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py b/tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py similarity index 100% rename from tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py rename to tests/unit/proxy/openai_files_endpoint/test_storage_backend_service.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/__init__.py b/tests/unit/proxy/prompts/__init__.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/__init__.py rename to tests/unit/proxy/prompts/__init__.py diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py b/tests/unit/proxy/prompts/test_prompt_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/prompts/test_prompt_endpoints.py rename to tests/unit/proxy/prompts/test_prompt_endpoints.py diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/unit/proxy/prompts/test_prompt_endpoints_crud.py similarity index 100% rename from tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py rename to tests/unit/proxy/prompts/test_prompt_endpoints_crud.py diff --git a/tests/test_litellm/proxy/prompts/test_prompt_environment.py b/tests/unit/proxy/prompts/test_prompt_environment.py similarity index 100% rename from tests/test_litellm/proxy/prompts/test_prompt_environment.py rename to tests/unit/proxy/prompts/test_prompt_environment.py diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/unit/proxy/prompts/test_prompt_registry.py similarity index 100% rename from tests/test_litellm/proxy/prompts/test_prompt_registry.py rename to tests/unit/proxy/prompts/test_prompt_registry.py diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py b/tests/unit/proxy/public_endpoints/public_v1/__init__.py similarity index 100% rename from tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py rename to tests/unit/proxy/public_endpoints/public_v1/__init__.py diff --git a/tests/test_litellm/proxy/public_endpoints/public_v1/test_model_hub.py b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py similarity index 100% rename from tests/test_litellm/proxy/public_endpoints/public_v1/test_model_hub.py rename to tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py rename to tests/unit/proxy/public_endpoints/test_public_endpoints.py diff --git a/tests/test_litellm/proxy/video_endpoints/__init__.py b/tests/unit/proxy/rag_endpoints/__init__.py similarity index 100% rename from tests/test_litellm/proxy/video_endpoints/__init__.py rename to tests/unit/proxy/rag_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py rename to tests/unit/proxy/rag_endpoints/test_rag_endpoints.py diff --git a/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py b/tests/unit/proxy/rag_endpoints/test_upload_security.py similarity index 100% rename from tests/test_litellm/proxy/rag_endpoints/test_upload_security.py rename to tests/unit/proxy/rag_endpoints/test_upload_security.py diff --git a/tests/unit/proxy/engine/__init__.py b/tests/unit/proxy/realtime_endpoints/__init__.py similarity index 100% rename from tests/unit/proxy/engine/__init__.py rename to tests/unit/proxy/realtime_endpoints/__init__.py diff --git a/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py b/tests/unit/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py rename to tests/unit/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py diff --git a/tests/unit/proxy/rerank_endpoints/__init__.py b/tests/unit/proxy/rerank_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/unit/proxy/rerank_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py rename to tests/unit/proxy/rerank_endpoints/test_endpoints.py diff --git a/tests/unit/proxy/response_api_endpoints/__init__.py b/tests/unit/proxy/response_api_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py rename to tests/unit/proxy/response_api_endpoints/test_endpoints.py diff --git a/tests/unit/proxy/shutdown/__init__.py b/tests/unit/proxy/shutdown/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/shutdown/test_graceful_shutdown_manager.py b/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py similarity index 100% rename from tests/test_litellm/proxy/shutdown/test_graceful_shutdown_manager.py rename to tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/unit/proxy/shutdown/test_scheduled_jobs.py similarity index 100% rename from tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py rename to tests/unit/proxy/shutdown/test_scheduled_jobs.py diff --git a/tests/unit/proxy/ui_crud_endpoints/__init__.py b/tests/unit/proxy/ui_crud_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py rename to tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py rename to tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_user_banner_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/ui_crud_endpoints/test_user_banner_endpoints.py rename to tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py diff --git a/tests/unit/proxy/utils/__init__.py b/tests/unit/proxy/utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/utils/helpers/__init__.py b/tests/unit/proxy/utils/helpers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/unit/proxy/utils/helpers/test_error_helpers.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_error_helpers.py rename to tests/unit/proxy/utils/helpers/test_error_helpers.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_guardrail_merge.py b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_guardrail_merge.py rename to tests/unit/proxy/utils/helpers/test_guardrail_merge.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_misc_helpers.py b/tests/unit/proxy/utils/helpers/test_misc_helpers.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_misc_helpers.py rename to tests/unit/proxy/utils/helpers/test_misc_helpers.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/unit/proxy/utils/helpers/test_model_access.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_model_access.py rename to tests/unit/proxy/utils/helpers/test_model_access.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_month_end_projection.py b/tests/unit/proxy/utils/helpers/test_month_end_projection.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_month_end_projection.py rename to tests/unit/proxy/utils/helpers/test_month_end_projection.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_premium_user_check.py b/tests/unit/proxy/utils/helpers/test_premium_user_check.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_premium_user_check.py rename to tests/unit/proxy/utils/helpers/test_premium_user_check.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_team_configs.py b/tests/unit/proxy/utils/helpers/test_team_configs.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_team_configs.py rename to tests/unit/proxy/utils/helpers/test_team_configs.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_to_ns.py b/tests/unit/proxy/utils/helpers/test_to_ns.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_to_ns.py rename to tests/unit/proxy/utils/helpers/test_to_ns.py diff --git a/tests/test_litellm/proxy/utils/helpers/test_url_helpers.py b/tests/unit/proxy/utils/helpers/test_url_helpers.py similarity index 100% rename from tests/test_litellm/proxy/utils/helpers/test_url_helpers.py rename to tests/unit/proxy/utils/helpers/test_url_helpers.py diff --git a/tests/unit/proxy/utils/prisma_and_spend/__init__.py b/tests/unit/proxy/utils/prisma_and_spend/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/_harness_smoke_test.py b/tests/unit/proxy/utils/prisma_and_spend/_harness_smoke_test.py similarity index 93% rename from tests/test_litellm/proxy/utils/prisma_and_spend/_harness_smoke_test.py rename to tests/unit/proxy/utils/prisma_and_spend/_harness_smoke_test.py index 2243d46ae7f..dd4bd0f1f72 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/_harness_smoke_test.py +++ b/tests/unit/proxy/utils/prisma_and_spend/_harness_smoke_test.py @@ -15,14 +15,14 @@ from litellm.proxy.utils import PrismaClient def test_normalize_scrubs_volatile_keys() -> None: - from tests.test_litellm.proxy.utils.prisma_and_spend.conftest import normalize + from tests.unit.proxy.utils.prisma_and_spend.conftest import normalize out = normalize({"id": 1, "spend": 2.0, "team_id": "t1"}) assert out == {"id": "", "spend": "", "team_id": "t1"} def test_normalize_recurses_into_lists() -> None: - from tests.test_litellm.proxy.utils.prisma_and_spend.conftest import normalize + from tests.unit.proxy.utils.prisma_and_spend.conftest import normalize out = normalize([{"id": "x"}, {"team_id": "t"}]) assert out == [{"id": ""}, {"team_id": "t"}] diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/unit/proxy/utils/prisma_and_spend/conftest.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py rename to tests/unit/proxy/utils/prisma_and_spend/conftest.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_cache_user_row.py b/tests/unit/proxy/utils/prisma_and_spend/test_cache_user_row.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_cache_user_row.py rename to tests/unit/proxy/utils/prisma_and_spend/test_cache_user_row.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_config_param_cache.py b/tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_config_param_cache.py rename to tests/unit/proxy/utils/prisma_and_spend/test_config_param_cache.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_password_helpers.py b/tests/unit/proxy/utils/prisma_and_spend/test_password_helpers.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_password_helpers.py rename to tests/unit/proxy/utils/prisma_and_spend/test_password_helpers.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py rename to tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py rename to tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_health.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py rename to tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_health.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py rename to tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py rename to tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py rename to tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py rename to tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py b/tests/unit/proxy/utils/prisma_and_spend/test_send_email.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py rename to tests/unit/proxy/utils/prisma_and_spend/test_send_email.py diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py similarity index 100% rename from tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py rename to tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py diff --git a/tests/unit/proxy/utils/proxy_logging/__init__.py b/tests/unit/proxy/utils/proxy_logging/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/utils/proxy_logging/_harness_smoke_test.py b/tests/unit/proxy/utils/proxy_logging/_harness_smoke_test.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/_harness_smoke_test.py rename to tests/unit/proxy/utils/proxy_logging/_harness_smoke_test.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/conftest.py b/tests/unit/proxy/utils/proxy_logging/conftest.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/conftest.py rename to tests/unit/proxy/utils/proxy_logging/conftest.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py b/tests/unit/proxy/utils/proxy_logging/test_alerting.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py rename to tests/unit/proxy/utils/proxy_logging/test_alerting.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py b/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py rename to tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_dataclass.py b/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_dataclass.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_dataclass.py rename to tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_dataclass.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_during_call_hook.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py rename to tests/unit/proxy/utils/proxy_logging/test_during_call_hook.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py rename to tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_internal_usage_cache.py b/tests/unit/proxy/utils/proxy_logging/test_internal_usage_cache.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_internal_usage_cache.py rename to tests/unit/proxy/utils/proxy_logging/test_internal_usage_cache.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py b/tests/unit/proxy/utils/proxy_logging/test_lifecycle.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py rename to tests/unit/proxy/utils/proxy_logging/test_lifecycle.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/unit/proxy/utils/proxy_logging/test_mcp_bridging.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py rename to tests/unit/proxy/utils/proxy_logging/test_mcp_bridging.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_module_helpers.py b/tests/unit/proxy/utils/proxy_logging/test_module_helpers.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_module_helpers.py rename to tests/unit/proxy/utils/proxy_logging/test_module_helpers.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py rename to tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py rename to tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py rename to tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py similarity index 100% rename from tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py rename to tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py diff --git a/tests/unit/proxy/vector_store_endpoints/__init__.py b/tests/unit/proxy/vector_store_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_access_control.py similarity index 100% rename from tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py rename to tests/unit/proxy/vector_store_endpoints/test_vector_store_access_control.py diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py rename to tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_rbac.py similarity index 100% rename from tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py rename to tests/unit/proxy/vector_store_endpoints/test_vector_store_rbac.py diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py similarity index 100% rename from tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py rename to tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py diff --git a/tests/unit/proxy/vector_store_files_endpoints/__init__.py b/tests/unit/proxy/vector_store_files_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py rename to tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py diff --git a/tests/unit/proxy/video_endpoints/__init__.py b/tests/unit/proxy/video_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py b/tests/unit/proxy/video_endpoints/test_endpoints.py similarity index 100% rename from tests/test_litellm/proxy/video_endpoints/test_endpoints.py rename to tests/unit/proxy/video_endpoints/test_endpoints.py diff --git a/tests/test_litellm/proxy/video_endpoints/test_utils.py b/tests/unit/proxy/video_endpoints/test_utils.py similarity index 100% rename from tests/test_litellm/proxy/video_endpoints/test_utils.py rename to tests/unit/proxy/video_endpoints/test_utils.py diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/ActivityScope.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/ActivityScope.tsx index 912e7686972..188c1e6db92 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/ActivityScope.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/ActivityScope.tsx @@ -7,7 +7,7 @@ import { apiClient } from "@/components/networking"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { TracePanel } from "./TracePanel"; -import { type Sample, type Settings, runTime, durationLabel } from "./engineData"; +import { type Sample, type Settings, runTime, durationLabel } from "./lensData"; import { DurationInput } from "./DurationInput"; @@ -72,7 +72,7 @@ export function ActivityScope({ const valid = validWindow && validSampling && validFilters; const load = (selection: ActivitySelection, pageOffset = 0) => { const { lookback_hours, ...selectionSettings } = selection; - return apiClient.post("/engine/preview/sample", { + return apiClient.post("/lens/preview/sample", { accessToken, body: { offset: pageOffset, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineProgress.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensProgress.tsx similarity index 93% rename from ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineProgress.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensProgress.tsx index b479fe287e8..9413fbb5157 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineProgress.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensProgress.tsx @@ -3,11 +3,11 @@ import { useEffect, useState } from "react"; import { Check, Loader2 } from "lucide-react"; import { Button } from "@/components/ui/button"; -import { analysisElapsed, analysisProgress, nextCheckStatus, type Engine, type Job } from "./engineData"; +import { analysisElapsed, analysisProgress, nextCheckStatus, type Lens, type Job } from "./lensData"; const steps = ["Review runs", "Find patterns", "Check evidence"]; -export function EngineProgress({ job, onCancel }: { job: Job; onCancel?: () => void }) { +export function LensProgress({ job, onCancel }: { job: Job; onCancel?: () => void }) { const [now, setNow] = useState(Date.now); useEffect(() => { const timer = window.setInterval(() => setNow(Date.now()), 1000); @@ -69,13 +69,13 @@ export function EngineProgress({ job, onCancel }: { job: Job; onCancel?: () => v ); } -export function NextCheck({ engine }: { engine: Engine }) { +export function NextCheck({ lens }: { lens: Lens }) { const [now, setNow] = useState(Date.now); useEffect(() => { const timer = window.setInterval(() => setNow(Date.now()), 15000); return () => window.clearInterval(timer); }, []); - const label = nextCheckStatus(engine, now); + const label = nextCheckStatus(lens, now); if (!label) return null; return

{label}

; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensRuns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensRuns.tsx index fc664ed5998..8c4534d8745 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensRuns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensRuns.tsx @@ -2,7 +2,7 @@ import { useState } from "react"; import { ArrowUpRight } from "lucide-react"; import { Button } from "@/components/ui/button"; import { RunList } from "./ActivityScope"; -import type { Job } from "./engineData"; +import type { Job } from "./lensData"; function assessmentLabel(assessment: Job["assessments"][number] | undefined): string { if (!assessment) return "Not reviewed"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineSetup.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensSetup.integration.test.tsx similarity index 92% rename from ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineSetup.integration.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensSetup.integration.test.tsx index dfc95369e3c..41b898a2bdc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineSetup.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensSetup.integration.test.tsx @@ -2,9 +2,9 @@ import { fireEvent, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "@/../tests/test-utils"; -import { EngineSetup } from "./EngineSetup"; +import { LensSetup } from "./LensSetup"; import { apiClient } from "@/components/networking"; -import type { Settings } from "./engineData"; +import type { Settings } from "./lensData"; vi.mock("@/components/networking", () => ({ apiClient: { post: vi.fn() } })); @@ -30,7 +30,7 @@ const settings: Settings = { ], }; -describe("Engine setup", () => { +describe("Lens setup", () => { beforeEach(() => { vi.mocked(apiClient.post).mockReset(); vi.mocked(apiClient.post).mockResolvedValue({ eligible: 0, executions: [] }); @@ -39,7 +39,7 @@ describe("Engine setup", () => { const save = vi.fn().mockResolvedValue(undefined); const user = userEvent.setup(); renderWithProviders( - , + , ); fireEvent.change(screen.getByRole("textbox", { name: "Specific checks (optional)" }), { target: { value: "Find incomplete reports\nFind repeated searches" }, @@ -52,7 +52,7 @@ describe("Engine setup", () => { it("rejects invalid metadata before reviewing the selection", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); fireEvent.change(screen.getByRole("textbox", { name: "Name" }), { target: { value: "Research" } }); await user.click(screen.getByRole("button", { name: "Continue" })); await user.click(screen.getByRole("button", { name: "Add condition" })); @@ -82,7 +82,7 @@ describe("Engine setup", () => { } : { eligible: 0, executions: [] }; }); - renderWithProviders(); + renderWithProviders(); fireEvent.change(screen.getByRole("textbox", { name: "Name" }), { target: { value: "Research" } }); await user.click(screen.getByRole("button", { name: "Continue" })); await user.click(screen.getByRole("button", { name: "Add condition" })); @@ -106,7 +106,7 @@ it("searches providers and saves custom history and schedule values", async () = const user = userEvent.setup(); const save = vi.fn().mockResolvedValue(undefined); renderWithProviders( - ({ apiClient: { get: vi.fn(), post: vi.fn() }, proxyBaseUrl: "" })); @@ -35,7 +35,7 @@ const issue: Finding = { kind: "issue", priority: "high", }; -const engine: Engine = { +const lens: Lens = { version: 0, spent: 0, id: "lens", @@ -134,15 +134,15 @@ describe("Lens findings and runs", () => { beforeEach(() => { vi.mocked(apiClient.get).mockReset(); vi.mocked(apiClient.get).mockImplementation(async (path) => { - if (path === "/engine") return { engines: [engine], workers: [], tracing_enabled: true }; - if (path === "/engine/lens/runs") return engine.jobs; + if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true }; + if (path === "/lens/lens/runs") return lens.jobs; return { data: [] }; }); }); it("separates patterns from issues and reveals original evidence only when requested", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); expect(await screen.findByText("Review used the wrong defect rate")).toBeInTheDocument(); expect(screen.queryByText(pattern.title)).not.toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Patterns (1)" })); @@ -159,7 +159,7 @@ describe("Lens findings and runs", () => { it("shows the actual frozen run selection in the Runs tab", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); await user.click(await screen.findByRole("tab", { name: "Runs" })); expect(screen.getByText("Release-42")).toBeInTheDocument(); expect(screen.getByText("trace-42")).toBeInTheDocument(); @@ -170,27 +170,27 @@ describe("Lens findings and runs", () => { it("shows the actual next schedule and avoids a stale countdown during active scans", () => { const now = Date.parse("2026-09-30T10:00:00Z"); const monitoring = { - ...engine, - settings: { ...engine.settings, enabled: true }, + ...lens, + settings: { ...lens.settings, enabled: true }, next_run_at: "2026-09-30T10:12:00Z", }; expect(nextCheckStatus(monitoring, now)).toContain("in 12 minutes"); expect(nextCheckStatus(monitoring, now + 12 * 60000)).toBe("Due now · waiting for an analyzer"); - expect(nextCheckStatus({ ...monitoring, jobs: [{ ...engine.jobs[0], status: "running" }] }, now)).toBe( + expect(nextCheckStatus({ ...monitoring, jobs: [{ ...lens.jobs[0], status: "running" }] }, now)).toBe( "Next check scheduled after this scan finishes", ); - expect(nextCheckStatus({ ...monitoring, jobs: [{ ...engine.jobs[0], status: "queued" }] }, now)).toBe( + expect(nextCheckStatus({ ...monitoring, jobs: [{ ...lens.jobs[0], status: "queued" }] }, now)).toBe( "Waiting for an analyzer", ); - expect(nextCheckStatus(engine, now)).toBeNull(); + expect(nextCheckStatus(lens, now)).toBeNull(); }); it("runs saved settings immediately without opening setup", async () => { testQueryClient.clear(); vi.mocked(apiClient.get).mockImplementation(async (path) => { - if (path === "/engine") + if (path === "/lens") return { - engines: [engine], + lenses: [lens], tracing_enabled: true, workers: [ { @@ -198,29 +198,29 @@ it("runs saved settings immediately without opening setup", async () => { name: "Worker", revoked: false, analysis_key_id: "a".repeat(64), - scope: engine.scope, + scope: lens.scope, last_seen: new Date().toISOString(), }, ], }; - if (path === "/engine/lens/runs") return engine.jobs; + if (path === "/lens/lens/runs") return lens.jobs; return { data: [] }; }); - vi.mocked(apiClient.post).mockResolvedValue(engine); + vi.mocked(apiClient.post).mockResolvedValue(lens); const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); await user.click(await screen.findByRole("button", { name: "Run now" })); - expect(apiClient.post).toHaveBeenCalledWith("/engine/lens/runs", { accessToken: "test", body: {} }); + expect(apiClient.post).toHaveBeenCalledWith("/lens/lens/runs", { accessToken: "test", body: {} }); expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); }); it("guides a first-time administrator into worker connection and lens setup", async () => { testQueryClient.clear(); vi.mocked(apiClient.get).mockImplementation(async (path) => - path === "/engine" ? { engines: [], workers: [], tracing_enabled: true } : { data: [] }, + path === "/lens" ? { lenses: [], workers: [], tracing_enabled: true } : { data: [] }, ); const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); const guide = within(await screen.findByRole("region", { name: "Understand what your agents are doing" })); expect(guide.getByRole("link", { name: "View logs" })).toHaveAttribute("href", "/ui/logs/"); await user.click(guide.getByRole("button", { name: "Connect analyzer" })); @@ -234,20 +234,20 @@ it("guides a first-time administrator into worker connection and lens setup", as it("opens the saved results of an older batch", async () => { testQueryClient.clear(); const older = { - ...engine.jobs[0], + ...lens.jobs[0], id: "older", created_at: "2026-09-29T10:00:00Z", finished_at: "2026-09-29T10:02:13Z", findings: [{ ...issue, title: "Earlier batch finding" }], }; vi.mocked(apiClient.get).mockImplementation(async (path) => { - if (path === "/engine") return { engines: [engine], workers: [], tracing_enabled: true }; - if (path === "/engine/lens/runs") return [engine.jobs[0], older]; - if (path === "/engine/lens/runs/older") return older; + if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true }; + if (path === "/lens/lens/runs") return [lens.jobs[0], older]; + if (path === "/lens/lens/runs/older") return older; return { data: [] }; }); const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); await screen.findByRole("option", { name: `${new Date(older.created_at).toLocaleString()} · completed` }); await user.selectOptions(screen.getByRole("combobox", { name: "Investigation batch" }), "older"); expect(await screen.findByText("Earlier batch finding")).toBeVisible(); @@ -264,15 +264,15 @@ it("reads request content from the beginning after its abbreviated preview", asy testQueryClient.clear(); const requestId = btoa(JSON.stringify(["requests", "", "request-1"])); const job = { - ...engine.jobs[0], + ...lens.jobs[0], sample: { eligible: 1, - executions: [{ ...engine.jobs[0].sample!.executions[0], id: requestId, source: "requests" as const }], + executions: [{ ...lens.jobs[0].sample!.executions[0], id: requestId, source: "requests" as const }], }, }; vi.mocked(apiClient.get).mockImplementation(async (path, options) => { - if (path === "/engine") return { engines: [{ ...engine, jobs: [job] }], workers: [], tracing_enabled: true }; - if (path === "/engine/lens/runs") return [job]; + if (path === "/lens") return { lenses: [{ ...lens, jobs: [job] }], workers: [], tracing_enabled: true }; + if (path === "/lens/lens/runs") return [job]; const offset = options?.query?.offset ?? 0; return { parts: [ @@ -285,7 +285,7 @@ it("reads request content from the beginning after its abbreviated preview", asy }; }); const user = userEvent.setup(); - renderWithProviders(); + renderWithProviders(); await user.click(await screen.findByRole("tab", { name: "Runs" })); await user.click(screen.getByRole("button", { name: "Open request" })); expect(await screen.findByText("Abbreviated preview")).toBeVisible(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensView.tsx similarity index 87% rename from ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineView.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensView.tsx index 1ad36c17299..dacd93310fd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/EngineView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/LensView.tsx @@ -22,22 +22,22 @@ import { Popover, PopoverContent, PopoverTitle, PopoverTrigger } from "@/compone import { Textarea } from "@/components/ui/textarea"; import { apiClient } from "@/components/networking"; import { TracePanel } from "./TracePanel"; -import { EngineSetup } from "./EngineSetup"; +import { LensSetup } from "./LensSetup"; import { LensRuns } from "./LensRuns"; -import { EngineProgress, NextCheck, ScanDuration } from "./EngineProgress"; +import { LensProgress, NextCheck, ScanDuration } from "./LensProgress"; import { WorkerSetup } from "./WorkerSetup"; import { LensWelcome } from "./LensWelcome"; import { - engineStatus, + lensStatus, evidenceTarget, sortedFindings, runTime, - type Engine, - type EngineList, + type Lens, + type LensList, type Finding, type Settings, type Job, -} from "./engineData"; +} from "./lensData"; const money = (n: number) => new Intl.NumberFormat("en-US", { style: "currency", currency: "USD", maximumFractionDigits: 3 }).format(n); @@ -50,22 +50,22 @@ function emptyFindingTitle(active: boolean, scanned: boolean) { return scanned ? "No matching findings" : "Ready for the first analysis"; } -export function EngineView({ accessToken, readOnly = false }: { accessToken: string; readOnly?: boolean }) { +export function LensView({ accessToken, readOnly = false }: { accessToken: string; readOnly?: boolean }) { const client = useQueryClient(); - const key = ["engines", accessToken]; + const key = ["lenses", accessToken]; const query = useQuery({ queryKey: key, - queryFn: () => apiClient.get("/engine", { accessToken }), + queryFn: () => apiClient.get("/lens", { accessToken }), refetchInterval: 10000, }); const models = useQuery({ - queryKey: ["engine-models", accessToken], + queryKey: ["lens-models", accessToken], queryFn: () => apiClient.get<{ data: { id: string }[] }>("/models", { accessToken }), }); const modelDetails = useQuery({ queryKey: ["lens-model-details", accessToken], queryFn: () => - apiClient.get<{ data: import("./engineData").AnalysisModelInfo[] }>("/model_group/info", { accessToken }), + apiClient.get<{ data: import("./lensData").AnalysisModelInfo[] }>("/model_group/info", { accessToken }), }); const [selected, setSelected] = useState(() => typeof window === "undefined" ? null : new URLSearchParams(window.location.search).get("lens"), @@ -91,32 +91,31 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str const [error, setError] = useState(""); const [busy, setBusy] = useState(false); const [evidence, setEvidence] = useState<{ id: string; span: string } | null>(null); - const engines = [...(query.data?.engines ?? [])].sort((a, b) => Date.parse(b.created_at) - Date.parse(a.created_at)); - const showEmpty = !query.isLoading && !query.error && engines.length === 0; - const engine = engines.find((e) => e.id === selected) ?? engines[0]; + const lenses = [...(query.data?.lenses ?? [])].sort((a, b) => Date.parse(b.created_at) - Date.parse(a.created_at)); + const showEmpty = !query.isLoading && !query.error && lenses.length === 0; + const lens = lenses.find((e) => e.id === selected) ?? lenses[0]; const connected = query.data?.workers?.some( (w) => !w.revoked && w.analysis_key_id && query.dataUpdatedAt - Date.parse(w.last_seen) < 120000, ) ?? false; const historyQuery = { - queryKey: ["lens-history", engine?.id, historyOffset, accessToken], - enabled: !!engine, - queryFn: () => - apiClient.get(`/engine/${engine?.id}/runs`, { accessToken, query: { offset: historyOffset } }), + queryKey: ["lens-history", lens?.id, historyOffset, accessToken], + enabled: !!lens, + queryFn: () => apiClient.get(`/lens/${lens?.id}/runs`, { accessToken, query: { offset: historyOffset } }), refetchInterval: 10000, }; const history = useQuery(historyQuery); const historical = useQuery({ - queryKey: ["lens-batch", engine?.id, batchId, accessToken], - enabled: !!engine && !["latest", "all"].includes(batchId), - queryFn: () => apiClient.get(`/engine/${engine?.id}/runs/${batchId}`, { accessToken }), + queryKey: ["lens-batch", lens?.id, batchId, accessToken], + enabled: !!lens && !["latest", "all"].includes(batchId), + queryFn: () => apiClient.get(`/lens/${lens?.id}/runs/${batchId}`, { accessToken }), }); - const job = ["latest", "all"].includes(batchId) ? engine?.jobs?.[0] : historical.data; + const job = ["latest", "all"].includes(batchId) ? lens?.jobs?.[0] : historical.data; const missingSnapshot = job?.status === "completed" && job.findings == null && batchId !== "all"; const selectedOutsideHistory = !["latest", "all"].includes(batchId) && !history.data?.some((j) => j.id === batchId); - const batchSettings = job?.settings ?? engine?.settings; - const batchFindings = (batchId === "all" ? engine?.findings ?? [] : job?.findings ?? []).map((f) => { - const feedback = engine?.findings?.find((current) => current.id === f.id); + const batchSettings = job?.settings ?? lens?.settings; + const batchFindings = (batchId === "all" ? lens?.findings ?? [] : job?.findings ?? []).map((f) => { + const feedback = lens?.findings?.find((current) => current.id === f.id); return feedback ? { ...f, status: feedback.status, reason: feedback.reason } : f; }); const finding = batchFindings.find((f) => f.id === findingId); @@ -127,12 +126,12 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str }; const setupSettings = () => { if (editing === "new") return undefined; - if (editing === "duplicate" && engine) - return { ...engine.settings, name: `${engine.settings.name} copy`, enabled: false }; - return engine?.settings; + if (editing === "duplicate" && lens) + return { ...lens.settings, name: `${lens.settings.name} copy`, enabled: false }; + return lens?.settings; }; - const lastCompleted = engine?.jobs?.find((j) => j.status === "completed"); - const active = engine?.jobs?.find((j) => j.status === "queued" || j.status === "running"); + const lastCompleted = lens?.jobs?.find((j) => j.status === "completed"); + const active = lens?.jobs?.find((j) => j.status === "queued" || j.status === "running"); const visibleFindings = sortedFindings( batchFindings.filter((f) => (filter === "all" || f.status === filter) && f.kind === kind), ); @@ -147,11 +146,11 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str const target = evidence ? evidenceTarget(evidence.id) : null; const [requestOffset, setRequestOffset] = useState(0); const requestEvidence = useQuery({ - queryKey: ["engine-evidence", engine?.id, evidence?.id, requestOffset, accessToken], - enabled: !!engine && target?.source === "requests", + queryKey: ["lens-evidence", lens?.id, evidence?.id, requestOffset, accessToken], + enabled: !!lens && target?.source === "requests", queryFn: () => apiClient.get( - `/engine/${engine?.id}/executions/${encodeURIComponent(evidence?.id ?? "")}`, + `/lens/${lens?.id}/executions/${encodeURIComponent(evidence?.id ?? "")}`, { accessToken, query: { offset: requestOffset } }, ), }); @@ -172,9 +171,9 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str } }; const save = async (settings: Settings) => { - const saved = await apiClient.request( + const saved = await apiClient.request( editing === "edit" ? "PUT" : "POST", - editing === "edit" ? `/engine/${engine.id}` : "/engine", + editing === "edit" ? `/lens/${lens.id}` : "/lens", { accessToken, body: settings }, ); selectLens(saved.id); @@ -182,8 +181,8 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str refresh(); }; const changeFinding = async (status: Finding["status"]) => { - if (!engine || !finding) return; - await update(`/engine/${engine.id}/findings/${finding.id}`, { status, reason }, "patch"); + if (!lens || !finding) return; + await update(`/lens/${lens.id}/findings/${finding.id}`, { status, reason }, "patch"); }; return ( @@ -206,7 +205,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str /> {connected ? "Analyzer connected" : "Set up analysis"} - {engines.length > 0 && ( + {lenses.length > 0 && ( ))}
-

{engine.settings.name}

+

{lens.settings.name}

- {sourceLabels[engine.settings.source ?? "traces"]} ·{" "} - {engine.settings.service || "All accessible activity"} - {engine.settings.filters?.length ? ` · ${engine.settings.filters.length} filters` : ""} + {sourceLabels[lens.settings.source ?? "traces"]} ·{" "} + {lens.settings.service || "All accessible activity"} + {lens.settings.filters?.length ? ` · ${lens.settings.filters.length} filters` : ""}

{!readOnly && ( @@ -273,16 +272,13 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str variant="outline" disabled={busy} onClick={() => - update(`/engine/${engine.id}`, { ...engine.settings, enabled: !engine.settings.enabled }, "put") + update(`/lens/${lens.id}`, { ...lens.settings, enabled: !lens.settings.enabled }, "put") } > - {engine.settings.enabled ? : } - {engine.settings.enabled ? "Pause" : "Resume"} + {lens.settings.enabled ? : } + {lens.settings.enabled ? "Pause" : "Resume"} - @@ -298,18 +294,18 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str

Status

- {engineStatus(engine, connected)} + {lensStatus(lens, connected)}

- {engine.settings.enabled - ? `Checks every ${engine.settings.interval_minutes} minutes` + {lens.settings.enabled + ? `Checks every ${lens.settings.interval_minutes} minutes` : "Manual analysis available"}

- +

Last successful scan

-

{when(lastCompleted?.finished_at ?? engine.last_scan_at)}

+

{when(lastCompleted?.finished_at ?? lens.last_scan_at)}

{lastCompleted && (

{lastCompleted.coverage?.screened ?? 0} of {lastCompleted.coverage?.eligible ?? 0} eligible runs @@ -320,21 +316,21 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str

Analysis spend this month

- {money(engine.budget_month === new Date().toISOString().slice(0, 7) ? engine.spent ?? 0 : 0)}{" "} - / {money(engine.settings.monthly_budget ?? 20)} + {money(lens.budget_month === new Date().toISOString().slice(0, 7) ? lens.spent ?? 0 : 0)}{" "} + / {money(lens.settings.monthly_budget ?? 20)}

Includes reservations for pending calls

{active && ( - { - void update(`/engine/${engine.id}/cancel`, {}); + void update(`/lens/${lens.id}/cancel`, {}); } } /> @@ -344,7 +340,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str {job.error}

)} - +
Findings @@ -369,7 +365,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str {when(job.created_at)} · {job.status} )} - {(history.data ?? engine.jobs)?.map((j) => ( + {(history.data ?? lens.jobs)?.map((j) => ( @@ -482,7 +478,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str {visibleFindings.length === 0 && (
-

{emptyFindingTitle(!!active, !!engine.last_scan_at)}

+

{emptyFindingTitle(!!active, !!lens.last_scan_at)}

{active ? "Lens is reviewing the selected activity." @@ -517,10 +513,10 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str variant="ghost" onClick={() => update( - `/engine/${engine.id}`, + `/lens/${lens.id}`, { - ...engine.settings, - checks: engine.settings.checks.map((q) => + ...lens.settings, + checks: lens.settings.checks.map((q) => q.id === c.id ? { ...q, enabled: !q.enabled } : q, ), }, @@ -537,7 +533,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str @@ -590,7 +586,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str

{history.error &&

{history.error.message}

} - {(history.data ?? engine.jobs)?.map((j) => ( + {(history.data ?? lens.jobs)?.map((j) => (
{j.stage} @@ -621,7 +617,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str
)} {editing && ( - m.id) ?? []} @@ -750,7 +746,7 @@ export function EngineView({ accessToken, readOnly = false }: { accessToken: str )} - {engine && target?.source === "traces" && ( + {lens && target?.source === "traces" && ( { await user.click(screen.getByRole("combobox", { name: "Charge analysis to" })); await user.click(await screen.findByRole("option", { name: "Analysis" })); await user.click(screen.getByRole("button", { name: "Generate setup command" })); - expect(apiClient.post).toHaveBeenCalledWith("/engine/workers/register", { + expect(apiClient.post).toHaveBeenCalledWith("/lens/workers/register", { accessToken: "admin", body: { name: "Lens analyzer", analysis_key_id: "b".repeat(64) }, }); @@ -71,7 +71,7 @@ describe("Worker setup", () => { await user.click(screen.getByRole("combobox", { name: "Charge analysis to" })); await user.click(await screen.findByRole("option", { name: "Analysis" })); await user.click(screen.getByRole("button", { name: "Save billing key" })); - expect(apiClient.put).toHaveBeenCalledWith("/engine/workers/worker/billing-key", { + expect(apiClient.put).toHaveBeenCalledWith("/lens/workers/worker/billing-key", { accessToken: "admin", body: { analysis_key_id: "b".repeat(64) }, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/WorkerSetup.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/WorkerSetup.tsx index 7800f445a17..e0037330dd5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/WorkerSetup.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/WorkerSetup.tsx @@ -7,10 +7,10 @@ import { Input } from "@/components/ui/input"; import { serverRootPath } from "@/lib/serverRootPath"; import { apiClient, proxyBaseUrl } from "@/components/networking"; import { AnalysisKey } from "./AnalysisKey"; -import type { EngineList, WorkerCreated } from "./engineData"; +import type { LensList, WorkerCreated } from "./lensData"; export const LENS_WORKER_IMAGE = - "ghcr.io/berriai/litellm-lens-worker@sha256:c41e932eaf3e4efbcaf8cc5027c7e93021e5b2823f21cb8785cd107e37b91c9a"; + "ghcr.io/berriai/litellm-lens-worker@sha256:a8e8731d954916594eea462969946b9292fb771681ff515a9fd296b53f856c77"; function initialProxyAddress(): string { const url = new URL(proxyBaseUrl || serverRootPath, window.location.origin); @@ -30,7 +30,7 @@ export function workerSetupCommand(address: string, token: string): string { ].join(" \\\n"); } -function workerStatus(worker: EngineList["workers"][number], now: number): string { +function workerStatus(worker: LensList["workers"][number], now: number): string { if (!worker.analysis_key_id) return "Billing key required"; return now - Date.parse(worker.last_seen) < 120000 ? "Connected · ready to analyze" : "Not connected"; } @@ -42,7 +42,7 @@ export function WorkerSetup({ onChanged, }: { accessToken: string; - workers: EngineList["workers"]; + workers: LensList["workers"]; onClose: () => void; onChanged: () => void; }) { @@ -59,7 +59,7 @@ export function WorkerSetup({ const [error, setError] = useState(""); const [busy, setBusy] = useState(false); const actionLabel = editingWorker ? "Save billing key" : "Generate setup command"; - const editBilling = (worker: EngineList["workers"][number]) => { + const editBilling = (worker: LensList["workers"][number]) => { setCreated(null); setEditingWorker(worker.id); setAnalysisKey(worker.analysis_key_id ?? null); @@ -69,7 +69,7 @@ export function WorkerSetup({ setError(""); try { if (editingWorker) { - await apiClient.put(`/engine/workers/${editingWorker}/billing-key`, { + await apiClient.put(`/lens/workers/${editingWorker}/billing-key`, { accessToken, body: { analysis_key_id: analysisKey }, }); @@ -79,7 +79,7 @@ export function WorkerSetup({ return; } setCreated( - await apiClient.post("/engine/workers/register", { + await apiClient.post("/lens/workers/register", { accessToken, body: { name: "Lens analyzer", analysis_key_id: analysisKey }, }), @@ -189,7 +189,7 @@ export function WorkerSetup({ size="sm" onClick={async () => { try { - await apiClient.delete(`/engine/workers/${worker.id}`, { accessToken }); + await apiClient.delete(`/lens/workers/${worker.id}`, { accessToken }); onChanged(); } catch (e) { setError(e instanceof Error ? e.message : "Could not revoke worker"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/engineData.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/lensData.test.ts similarity index 99% rename from ui/litellm-dashboard/src/app/(dashboard)/lens/_components/engineData.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/lens/_components/lensData.test.ts index 6d546261e27..90eccf4bcd4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/engineData.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/lensData.test.ts @@ -6,7 +6,7 @@ import { sortedFindings, type Finding, type Job, -} from "./engineData"; +} from "./lensData"; const coverage: Job["coverage"] = { eligible: 0, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/engineData.ts b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/lensData.ts similarity index 86% rename from ui/litellm-dashboard/src/app/(dashboard)/lens/_components/engineData.ts rename to ui/litellm-dashboard/src/app/(dashboard)/lens/_components/lensData.ts index f7ebe086148..043fcb67845 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/engineData.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/_components/lensData.ts @@ -1,8 +1,8 @@ import type { components } from "@/lib/http/schema"; -export type Engine = components["schemas"]["Engine"]; -export type Settings = components["schemas"]["EngineSettings"]; -export type EngineList = components["schemas"]["EngineList"]; +export type Lens = components["schemas"]["Lens"]; +export type Settings = components["schemas"]["LensSettings"]; +export type LensList = components["schemas"]["LensList"]; export type Finding = components["schemas"]["Finding"]; export type Sample = components["schemas"]["Sample"]; export type WorkerCreated = components["schemas"]["WorkerCreated"]; @@ -33,12 +33,12 @@ export function sortedFindings(findings: Finding[]): Finding[] { ); } -export function engineStatus(engine: Engine, connected: boolean): string { - const active = engine.jobs?.find((job) => ["queued", "running"].includes(job.status ?? "")); +export function lensStatus(lens: Lens, connected: boolean): string { + const active = lens.jobs?.find((job) => ["queued", "running"].includes(job.status ?? "")); if (active) return connected ? active.stage ?? "Queued" : "Waiting for analyzer"; - const spent = engine.budget_month === new Date().toISOString().slice(0, 7) ? engine.spent ?? 0 : 0; - if (spent >= (engine.settings.monthly_budget ?? 20)) return "Budget reached"; - if (!engine.settings.enabled) return "Paused"; + const spent = lens.budget_month === new Date().toISOString().slice(0, 7) ? lens.spent ?? 0 : 0; + if (spent >= (lens.settings.monthly_budget ?? 20)) return "Budget reached"; + if (!lens.settings.enabled) return "Paused"; return connected ? "Monitoring" : "Analyzer disconnected"; } @@ -152,12 +152,12 @@ const nextCheckTimeFormat: Intl.DateTimeFormatOptions = { minute: "2-digit", }; -export function nextCheckStatus(engine: Engine, now: number): string | null { - if (!engine.settings.enabled) return null; - const active = engine.jobs.find((job) => job.status === "queued" || job.status === "running"); +export function nextCheckStatus(lens: Lens, now: number): string | null { + if (!lens.settings.enabled) return null; + const active = lens.jobs.find((job) => job.status === "queued" || job.status === "running"); if (active?.status === "running") return "Next check scheduled after this scan finishes"; if (active?.status === "queued") return "Waiting for an analyzer"; - const next = new Date(engine.next_run_at); + const next = new Date(lens.next_run_at); const remaining = next.getTime() - now; if (remaining <= 0) return "Due now · waiting for an analyzer"; const minutes = Math.ceil(remaining / 60000); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx index f833188bb21..dd0dc1518a7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx @@ -2,13 +2,13 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole, isProxyAdminTierRole } from "@/utils/roles"; -import { EngineView } from "./_components/EngineView"; +import { LensView } from "./_components/LensView"; -export default function EnginePage() { +export default function LensPage() { const { accessToken, userRole } = useAuthorized(); if (!accessToken) return null; if (!isProxyAdminTierRole(userRole ?? "")) { return

Lens requires proxy administrator access.

; } - return ; + return ; } diff --git a/ui/litellm-dashboard/src/app/globals.css b/ui/litellm-dashboard/src/app/globals.css index 87959ca0139..e1ff921648f 100644 --- a/ui/litellm-dashboard/src/app/globals.css +++ b/ui/litellm-dashboard/src/app/globals.css @@ -64,6 +64,61 @@ } @theme inline { + --animate-slide-left: slide-left 200ms cubic-bezier(0, 0, 0.2, 1) both; + --animate-view-fade-in: view-fade-in 100ms cubic-bezier(0.4, 0, 0.2, 1) both; + --animate-slot-slide-in: slot-slide-in 150ms cubic-bezier(0, 0, 0.2, 1) both; + --animate-trace-drawer-in: trace-drawer-in 200ms cubic-bezier(0.25, 1, 0.5, 1) both; + --animate-trace-drawer-out: trace-drawer-out 200ms cubic-bezier(0.4, 0, 1, 1) both; + + @keyframes trace-drawer-in { + from { + opacity: 0; + transform: translateX(2rem) scaleX(0.98); + } + to { + opacity: 1; + transform: none; + } + } + @keyframes trace-drawer-out { + from { + opacity: 1; + transform: none; + } + to { + opacity: 0; + transform: translateX(2rem) scaleX(0.98); + } + } + + @keyframes slide-left { + from { + opacity: 0; + transform: translateX(2rem) scaleX(0.98); + } + to { + opacity: 1; + transform: none; + } + } + @keyframes view-fade-in { + from { + opacity: 0; + } + to { + opacity: 1; + } + } + @keyframes slot-slide-in { + from { + opacity: 0; + transform: translateY(4px); + } + to { + opacity: 1; + transform: none; + } + } @keyframes scroll-fade-reveal-e { from { --scroll-fade-e: var(--_scroll-fade-size-e, var(--scroll-fade-size, min(12%, calc(var(--spacing) * 10)))); @@ -74,6 +129,22 @@ } } +@layer utilities { + .animate-trace-drawer-in, + .animate-trace-drawer-out { + transform-origin: right center; + } + @media (prefers-reduced-motion: reduce) { + .animate-trace-drawer-in, + .animate-trace-drawer-out, + .animate-slide-left, + .animate-view-fade-in, + .animate-slot-slide-in { + animation: none !important; + } + } +} + @utility scroll-fade-e { --_scroll-fade-size-e: var(--scroll-fade-e-size, var(--scroll-fade-size, min(12%, calc(var(--spacing) * 10)))); --scroll-fade-mask: linear-gradient(to right, #000 0, #000 calc(100% - var(--scroll-fade-e, 0px)), transparent 100%); @@ -141,6 +212,38 @@ --sidebar-ring: oklch(0.707 0.022 261.325); --neutral-border: #dcddeb; --logo-surface: oklch(1 0 0); + --trace-text: oklch(0.21 0.03 256); + --trace-text-2: oklch(0.35 0.03 256); + --trace-text-secondary: oklch(0.35 0.03 256); + --trace-duration: oklch(0.48 0.03 230); + --trace-key: oklch(0.55 0.03 240); + --trace-placeholder: oklch(0.7 0.02 240); + --trace-surface: oklch(1 0 0); + --trace-chip: oklch(0.975 0.006 220); + --trace-row-hover: oklch(0.975 0.008 215); + --trace-row-selected: oklch(0.95 0.035 200); + --trace-brand: oklch(0.6 0.13 195); + --trace-border: oklch(0.92 0.01 230); + --trace-line: oklch(0.88 0.03 205); + --trace-card-border: oklch(0.93 0.01 230); + --trace-dot: oklch(0.86 0.05 190); + --trace-tab-active: oklch(0.95 0.025 205); + --trace-tab-hover: oklch(0.93 0.02 215); + --trace-tag: oklch(0.95 0.02 205); + --trace-chain: oklch(0.56 0.17 255); + --trace-llm: oklch(0.6 0.13 215); + --trace-tool: oklch(0.64 0.14 165); + --trace-glyph: oklch(0.99 0 0); + --trace-human: oklch(0.5 0.15 260); + --trace-human-glyph: oklch(0.95 0.04 210); + --trace-turn: oklch(0.96 0.03 200); + --trace-turn-border: oklch(0.75 0.1 200); + --trace-ok: oklch(0.92 0.08 160); + --trace-ok-glyph: oklch(0.55 0.15 155); + --trace-called: oklch(0.93 0.06 185); + --trace-called-text: oklch(0.38 0.08 195); + --trace-shadow-md: 0 4px 6px -1px #0b1b2e1a, 0 2px 4px -1px #0b1b2e0f; + --trace-shadow-xs: 0 1px 2px 0 #0b1b2e0d; } .dark { @@ -183,9 +286,73 @@ --sidebar-border: oklch(0.187 0 0); --sidebar-ring: oklch(0.569 0 0); --neutral-border: var(--border); + --trace-text: oklch(0.96 0.005 220); + --trace-text-2: oklch(0.88 0.01 220); + --trace-text-secondary: oklch(0.88 0.01 220); + --trace-duration: oklch(0.78 0.03 200); + --trace-key: oklch(0.68 0.03 220); + --trace-placeholder: oklch(0.5 0.02 230); + --trace-surface: oklch(0.19 0.012 240); + --trace-chip: oklch(0.23 0.015 235); + --trace-row-hover: oklch(0.23 0.018 230); + --trace-row-selected: oklch(0.29 0.05 210); + --trace-brand: oklch(0.78 0.13 190); + --trace-border: oklch(0.3 0.02 235); + --trace-line: oklch(0.36 0.04 210); + --trace-card-border: oklch(0.27 0.02 235); + --trace-dot: oklch(0.45 0.06 195); + --trace-tab-active: oklch(0.28 0.03 215); + --trace-tab-hover: oklch(0.32 0.03 220); + --trace-tag: oklch(0.28 0.03 215); + --trace-chain: oklch(0.6 0.17 255); + --trace-llm: oklch(0.64 0.13 215); + --trace-tool: oklch(0.68 0.14 165); + --trace-glyph: oklch(0.99 0 0); + --trace-human: oklch(0.56 0.15 260); + --trace-human-glyph: oklch(0.95 0.04 210); + --trace-turn: oklch(0.29 0.05 210); + --trace-turn-border: oklch(0.5 0.09 200); + --trace-ok: oklch(0.35 0.07 160); + --trace-ok-glyph: oklch(0.82 0.15 155); + --trace-called: oklch(0.32 0.06 190); + --trace-called-text: oklch(0.88 0.08 185); + --trace-shadow-md: 0 4px 6px -1px #00000080, 0 2px 4px -1px #00000066; + --trace-shadow-xs: 0 1px 2px 0 #0000004d; } @theme inline { + --color-trace-text: var(--trace-text); + --color-trace-text-2: var(--trace-text-2); + --color-trace-text-secondary: var(--trace-text-secondary); + --color-trace-duration: var(--trace-duration); + --color-trace-key: var(--trace-key); + --color-trace-placeholder: var(--trace-placeholder); + --color-trace-surface: var(--trace-surface); + --color-trace-chip: var(--trace-chip); + --color-trace-row-hover: var(--trace-row-hover); + --color-trace-row-selected: var(--trace-row-selected); + --color-trace-brand: var(--trace-brand); + --color-trace-border: var(--trace-border); + --color-trace-line: var(--trace-line); + --color-trace-card-border: var(--trace-card-border); + --color-trace-dot: var(--trace-dot); + --color-trace-tab-active: var(--trace-tab-active); + --color-trace-tab-hover: var(--trace-tab-hover); + --color-trace-tag: var(--trace-tag); + --color-trace-chain: var(--trace-chain); + --color-trace-llm: var(--trace-llm); + --color-trace-tool: var(--trace-tool); + --color-trace-glyph: var(--trace-glyph); + --color-trace-human: var(--trace-human); + --color-trace-human-glyph: var(--trace-human-glyph); + --color-trace-turn: var(--trace-turn); + --color-trace-turn-border: var(--trace-turn-border); + --color-trace-ok: var(--trace-ok); + --color-trace-ok-glyph: var(--trace-ok-glyph); + --color-trace-called: var(--trace-called); + --color-trace-called-text: var(--trace-called-text); + --shadow-trace-md: var(--trace-shadow-md); + --shadow-trace-xs: var(--trace-shadow-xs); --radius-sm: calc(var(--radius) - 4px); --radius-md: calc(var(--radius) - 2px); --radius-lg: var(--radius); diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index 4ba5dd23fd1..9b67657dbca 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -313,6 +313,7 @@ describe("Settings", () => { "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME", "S3_LOG_PROMPTS_ONLY", + "S3_PARTITION_GRANULARITY", ], ui_callback_name: "s3 Bucket (AWS)", }, @@ -326,6 +327,12 @@ describe("Settings", () => { dynamic_params: { s3_bucket_name: { type: "text", ui_name: "S3 Bucket Name", required: false }, s3_log_prompts_only: { type: "boolean", ui_name: "Log Prompts Only", required: false }, + s3_partition_granularity: { + type: "select", + ui_name: "Folder Partitioning", + options: ["day", "hour"], + required: false, + }, }, }, ]); @@ -409,6 +416,46 @@ describe("Settings", () => { }); }); + it("should show the saved s3_v2 folder partitioning and post the newly selected value", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: null, S3_PARTITION_GRANULARITY: "hour" }, "s3_v2"); + const user = await openS3EditModal("s3_v2"); + + const dialog = screen.getByRole("dialog"); + const partitioning = await within(dialog).findByRole("combobox", { name: "Folder Partitioning" }); + expect(partitioning).toHaveTextContent("hour"); + + await user.click(partitioning); + await user.click(await screen.findByRole("option", { name: "day" })); + await user.click(within(dialog).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + environment_variables: expect.objectContaining({ callback: "s3_v2", s3_partition_granularity: "day" }), + litellm_settings: { success_callback: ["s3_v2"] }, + }), + ); + }); + }); + + it("should not offer folder partitioning for the legacy s3 callback, which cannot honour it", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: null, S3_PARTITION_GRANULARITY: null }); + const user = await openS3EditModal(); + + const dialog = screen.getByRole("dialog"); + await within(dialog).findByRole("switch", { name: "Log Prompts Only" }); + expect(within(dialog).queryByRole("combobox", { name: "Folder Partitioning" })).not.toBeInTheDocument(); + + await user.click(within(dialog).getByRole("button", { name: "Save Changes" })); + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledTimes(1); + }); + const [, payload] = vi.mocked(setCallbacksCall).mock.calls[0]; + expect(payload.environment_variables).not.toHaveProperty("s3_partition_granularity"); + expect(payload.environment_variables).not.toHaveProperty("S3_PARTITION_GRANULARITY"); + }); + it("should send the typed webhook url for an alert type when the alerting tab is saved", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 9247f22ec28..e376d858df8 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -238,6 +238,7 @@ export const CallbackSelector: React.FC = ({ }; const CALLBACK_CONFIG_ALIASES: Record = { s3_v2: "s3" }; +const CALLBACK_UNSUPPORTED_PARAMS: Record = { s3: ["s3_partition_granularity"] }; interface DynamicParamConfig { type?: string; @@ -274,7 +275,8 @@ const getDynamicParamsForCallback = ( const callbackConfig = findCallbackConfig(callbackConfigs, callbackName); if (callbackConfig?.dynamic_params) { - return Object.keys(callbackConfig.dynamic_params); + const unsupportedParams = CALLBACK_UNSUPPORTED_PARAMS[callbackName] ?? []; + return Object.keys(callbackConfig.dynamic_params).filter((param) => !unsupportedParams.includes(param)); } return fallbackVariables ? Object.keys(fallbackVariables) : []; diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx index 126542e848f..666859b1cc5 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.test.tsx @@ -163,17 +163,51 @@ describe("AgentTracesSection", () => { expect(failed.length + ok.length).toBe(runs.length); }); - it("opens the run in place and goes back to the list", async () => { + it("opens a run in a side drawer over the list and swaps runs without closing it", async () => { vi.mocked(agentTraceListCall).mockResolvedValue(traceList as TracePage); renderSection(); const rows = await screen.findAllByTestId("agent-trace-row"); fireEvent.click(rows[0]); - expect(screen.getByTestId("run-view")).toHaveTextContent(`run ${runs[0].trace_id}`); - expect(screen.queryByTestId("runs-table")).not.toBeInTheDocument(); - - fireEvent.click(screen.getByText("back")); + const drawer = screen.getByRole("complementary", { name: "Trace details" }); + expect(within(drawer).getByTestId("run-view")).toHaveTextContent(`run ${runs[0].trace_id}`); expect(screen.getByTestId("runs-table")).toBeInTheDocument(); + expect(rows[0]).toHaveAttribute("aria-selected", "true"); + + fireEvent.click(rows[1]); + expect(screen.getByRole("complementary", { name: "Trace details" })).toBe(drawer); + expect(within(drawer).getByTestId("run-view")).toHaveTextContent(`run ${runs[1].trace_id}`); + expect(rows[1]).toHaveAttribute("aria-selected", "true"); + expect(rows[0]).toHaveAttribute("aria-selected", "false"); + }); + + it("closes the drawer when the open row is clicked again or Escape is pressed", async () => { + vi.mocked(agentTraceListCall).mockResolvedValue(traceList as TracePage); + renderSection(); + const rows = await screen.findAllByTestId("agent-trace-row"); + + fireEvent.click(rows[0]); + fireEvent.click(rows[0]); + expect(rows[0]).toHaveAttribute("aria-selected", "false"); + + fireEvent.click(rows[1]); + fireEvent.keyDown(window, { key: "Escape" }); + expect(rows[1]).toHaveAttribute("aria-selected", "false"); + }); + + it("moves to the next and previous run with j / k and the header arrows", async () => { + vi.mocked(agentTraceListCall).mockResolvedValue(traceList as TracePage); + renderSection(); + const rows = await screen.findAllByTestId("agent-trace-row"); + + fireEvent.click(rows[0]); + fireEvent.keyDown(window, { key: "j" }); + expect(screen.getByTestId("run-view")).toHaveTextContent(`run ${runs[1].trace_id}`); + fireEvent.keyDown(window, { key: "k" }); + expect(screen.getByTestId("run-view")).toHaveTextContent(`run ${runs[0].trace_id}`); + expect(screen.getByRole("button", { name: "Previous trace (K)" })).toBeDisabled(); + fireEvent.click(screen.getByRole("button", { name: "Next trace (J)" })); + expect(screen.getByTestId("run-view")).toHaveTextContent(`run ${runs[1].trace_id}`); }); it("plots every loaded run on the timeline", async () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx index 61e4c737b24..8f95573d0dd 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/AgentTracesSection.tsx @@ -1,11 +1,14 @@ "use client"; +import { Plug } from "lucide-react"; import moment from "moment"; import { useMemo, useState } from "react"; +import { Button } from "@/components/ui/button"; + import { AgentTracesTable } from "./AgentTracesTable"; +import { RunDrawer } from "./RunDrawer"; import { ALL_SERVICES, RunsToolbar, type RunStatusFilter } from "./RunsToolbar"; -import { RunView } from "./TraceDrawer"; import type { TraceSummary } from "./traceTypes"; import { previewText } from "./traceUtils"; import { TimeRangeControls } from "./TimeRangeControls"; @@ -31,6 +34,8 @@ export function filterRuns( }); } +const runKey = (run: TraceSummary): string => run.trace_ref || run.trace_id; + const filterByWindow = (runs: TraceSummary[], range: TimeWindow): TraceSummary[] => runs.filter((run) => { const t = moment(run.start_time).valueOf(); @@ -120,19 +125,12 @@ export function AgentTracesSection({ ); } - if (openTrace !== null) { - return ( - openRun(null)} - /> - ); - } + const toggleRun = (trace: TraceSummary | null) => + openRun(trace !== null && openTrace !== null && runKey(trace) === runKey(openTrace) ? null : trace); return (
+ - + {timeControls && (
void; onOpenTrace: (trace: TraceSummary) => void; + selectedKey?: string | null; } const SECOND_MS = 1000; @@ -53,6 +55,7 @@ export function AgentTracesTable({ hasMore, onLoadMore, onOpenTrace, + selectedKey = null, }: AgentTracesTableProps) { const isEmpty = !isLoading && !error && traces.length === 0; return ( @@ -81,7 +84,13 @@ export function AgentTracesTable({ key={run.trace_ref || run.trace_id} data-testid="agent-trace-row" onClick={() => onOpenTrace(run)} - className="h-9 cursor-pointer border-b border-border/60 text-[12px] hover:bg-accent/50" + aria-selected={selectedKey === (run.trace_ref || run.trace_id)} + className={cn( + "h-9 cursor-pointer border-b border-border/60 text-[12px] transition-colors duration-150 motion-reduce:transition-none", + selectedKey === (run.trace_ref || run.trace_id) + ? "bg-trace-row-selected shadow-[inset_2px_0_0_var(--trace-brand)]" + : "hover:bg-trace-row-hover", + )} > +

{title}

+ {children} +
+ ); +} + +/** Ids, then the raw OTEL attributes, as dot-bulleted key / value rows. */ export function AttributesDetail({ traceId, span, attributes, isLoading }: AttributesDetailProps) { - const entries: [string, string][] = [ + const ids: KeyValue[] = [ ["trace_id", traceId], ["span_id", span.span_id], ["parent_span_id", span.parent_span_id ?? "—"], - ...Object.entries(attributes ?? {}).sort(([a], [b]) => a.localeCompare(b)), ]; + const attributeEntries: KeyValue[] = Object.entries(attributes ?? {}).sort(([a], [b]) => a.localeCompare(b)); return ( -
-
- {entries.map(([key, value]) => ( -
-
{key}
-
{value}
-
- ))} -
- {isLoading &&
Loading attributes…
} +
+ + + + {attributeEntries.length > 0 && ( + + + + )} + {isLoading &&
Loading attributes…
}
); } diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/Collapse.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/Collapse.tsx new file mode 100644 index 00000000000..fbbd6f71bf2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/Collapse.tsx @@ -0,0 +1,35 @@ +"use client"; + +import { ChevronRight } from "lucide-react"; + +import { cn } from "@/lib/cva.config"; + +/** Snaps between 0 and auto height; children stay mounted but inert while closed. */ +export function Collapse({ + open, + children, + className, +}: { + open: boolean; + children: React.ReactNode; + className?: string; +}) { + return ( +
+
{children}
+
+ ); +} + +/** Right-pointing chevron that rotates to point down when open. */ +export function FoldChevron({ open, className }: { open: boolean; className?: string }) { + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailContent.tsx index c95a4318391..9c156d88bd5 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailContent.tsx @@ -1,17 +1,24 @@ "use client"; import { useQuery, type UseQueryOptions } from "@tanstack/react-query"; -import { AlertTriangle, Bot, CornerDownRight, Wrench } from "lucide-react"; +import { AlertTriangle } from "lucide-react"; + +import { cn } from "@/lib/cva.config"; import { agentTraceSpanCall } from "../../networking"; +import { type KeyValue, KeyValueRows, objectEntries } from "./KeyValueRows"; +import { Card, MessageCard, Section, ToolResultCard } from "./MessageCard"; import type { ErrorSource } from "./traceTree"; -import type { Span, SpanDetail, TraceMessage } from "./traceTypes"; -import { errorSource, parseMessages, prettyPayload } from "./traceUtils"; +import type { Span, SpanDetail, TraceMessage, UIContent, UIMessage } from "./traceTypes"; +import { errorSource, parseJson, parseMessages, prettyPayload } from "./traceUtils"; const ERROR_SOURCE_LABEL: Record = { tool: "Tool", model: "Model", litellm: "LiteLLM" }; const TRACEBACK_MARKER = "Traceback (most recent call last):"; +const STATUS_TEXT = "px-5 py-2 text-[13px] tracking-[-0.26px] text-trace-duration"; +const PAYLOAD_PRE = + "font-mono text-[13px] leading-[1.5] tracking-[-0.26px] break-words whitespace-pre-wrap text-trace-text"; -/** LangSmith records `repr(exc)` + traceback with no separator; keep the exception line. */ +/** Exporters record `repr(exc)` + traceback with no separator; keep the exception line. */ export const errorHeadline = (error: string): string => (error.split(TRACEBACK_MARKER, 1)[0].split("\n")[0] ?? "").trim() || error.trim(); @@ -29,93 +36,104 @@ export function useSpanDetail(accessToken: string, traceId: string, spanId: stri return useQuery(queryOptions); } -export function SectionLabel({ children }: { children: React.ReactNode }) { - return ( -
- {children} -
- ); -} - -export function TextBlock({ label, value, mono = false }: { label: string; value: string; mono?: boolean }) { - return ( -
- {label} -
- {value} -
-
- ); -} - -function RoleIcon({ role }: { role: string }) { - if (role === "assistant") return ; - if (role === "tool") return ; - return ; -} - -export function MessageBlock({ message }: { message: TraceMessage }) { - return ( -
-
- - {message.role} - {message.name ? · {message.name} : null} -
- {(message.tool_calls ?? []).map((call, i) => ( -
- {call.name} - ( - {JSON.stringify(call.args)} - ) -
- ))} - {message.content && ( -
- {message.content} -
- )} -
- ); -} - export function ErrorBlock({ span }: { span: Span }) { const source = errorSource(span); if (!source) return null; const headline = errorHeadline(span.error ?? "") || "Span reported an error status."; return ( -
-
- +
+
+ {ERROR_SOURCE_LABEL[source]} · {errorReason(headline)}
- 
-        {headline}
-
+ 
{headline}
); } -function Payload({ label, value, mono }: { label: string; value: string; mono: boolean }) { - const messages = parseMessages(value); - if (messages) { - return ( - <> - {`${label}${messages.length > 1 ? ` · ${messages.length} messages` : ""}`} - {messages.map((message, i) => ( - - ))} - - ); +function FieldsCard({ entries }: { entries: readonly KeyValue[] }) { + return ( + + + + ); +} + +function TextCard({ text }: { text: string }) { + return ( + + 
{text}
+ + ); +} + +function PlainPayload({ value }: { value: string }) { + const entries = objectEntries(parseJson(value)); + if (entries && entries.length > 0) return ; + return ; +} + +function Messages({ messages, model }: { messages: TraceMessage[]; model: string | null }) { + return ( + <> + {messages.map((message, i) => ( + + ))} + + ); +} + +const toTraceMessage = (message: UIMessage): TraceMessage => ({ + ...message, + tool_calls: message.tool_calls?.map((call) => ({ + name: call.name, + args: parseJson(call.arguments) ?? call.arguments, + })), +}); + +interface PayloadProps { + value: string; + span: Span; + role: "input" | "output"; +} + +const isToolResult = ({ span, role }: Omit): boolean => + span.type === "tool" && role === "output"; + +function ToolResult({ value, span }: Omit) { + return ; +} + +const singleText = (content: UIContent): string | null => + content.kind === "messages" && content.messages.length === 1 && !content.messages[0].tool_calls?.length + ? content.messages[0].content + : null; + +function UIPayload({ content, ...props }: PayloadProps & { content: UIContent }) { + const toolText = isToolResult(props) ? singleText(content) : null; + if (toolText !== null) return ; + if (content.kind === "messages") { + return ; } - return ; + if (isToolResult(props)) return ; + if (content.kind === "fields" && content.fields.length > 0) { + return [field.key, field.value])} />; + } + return ; +} + +function Payload(props: PayloadProps) { + const messages = parseMessages(props.value); + if (messages) return ; + if (isToolResult(props)) return ; + return ; +} + +function SpanPayload({ content, ...props }: PayloadProps & { content: UIContent | undefined }) { + return content ? : ; } interface DetailContentProps { @@ -125,26 +143,29 @@ interface DetailContentProps { span: Span; } -/** Content tab: the error first (if any), then what went in and what came out. */ +/** Content tab: the error first (if any), then collapsible Input and Output rendered as chat cards. */ export function DetailContent({ accessToken, traceId, traceRef, span }: DetailContentProps) { const detailQuery = useSpanDetail(accessToken, traceId, span.span_id, traceRef); const detail = detailQuery.data; - const isTool = span.type === "tool"; const empty = detail && !detail.input && !detail.output; return ( -
+
- {detailQuery.isLoading &&
Loading span…
} - {detailQuery.isError && ( -
- Could not load span: {detailQuery.error.message} -
- )} - {detail?.input ? : null} - {detail?.output ? : null} + {detailQuery.isLoading &&
Loading span…
} + {detailQuery.isError &&
Could not load span: {detailQuery.error.message}
} + {detail?.input ? ( +
+ +
+ ) : null} + {detail?.output ? ( +
+ +
+ ) : null} {empty && span.status !== "error" && ( -
+
No content recorded for this span.
)} diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.test.tsx index 9a541d26966..1adcda0f6cc 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.test.tsx @@ -1,9 +1,10 @@ -import { screen, waitFor } from "@testing-library/react"; +import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, testQueryClient } from "../../../../tests/test-utils"; import { DetailPane } from "./DetailPane"; +import { absoluteTime, SpanHoverCard, spanFacts } from "./SpanHoverCard"; import type { GroupRowData, SpanRowData } from "./traceTree"; import type { Span, SpanDetail, Trace } from "./traceTypes"; @@ -14,7 +15,9 @@ vi.mock("../../networking", () => ({ import { agentTraceSpanCall } from "../../networking"; -const span = (overrides: Partial & Pick): Span => ({ +type SpanFields = Partial & Pick; + +const span = (overrides: SpanFields): Span => ({ parent_span_id: "root", name: overrides.span_id, type: "chain", @@ -76,6 +79,8 @@ const trace: Trace = { spans: [root, llm, failedTool], }; +const LONG_NOTE = "Escalated twice already. ".repeat(6).trim(); + const details: Record = { llm1: { span_id: "llm1", @@ -86,7 +91,7 @@ const details: Record = { output: JSON.stringify({ role: "assistant", content: "", - tool_calls: [{ name: "get_customer_plan", args: { customer_id: "acme-404" } }], + tool_calls: [{ name: "get_customer_plan", args: { customer_id: "acme-404", note: LONG_NOTE } }], }), attributes: { "gen_ai.request.model": "claude-sonnet-4-5" }, }, @@ -99,6 +104,40 @@ const details: Record = { }, }; +const standardDetail: SpanDetail = { + span_id: "llm1", + input: "raw input left unparsed", + output: "raw output left unparsed", + input_ui: { kind: "fields", fields: [{ key: "ticket_id", value: "T-981" }] }, + output_ui: { + kind: "messages", + messages: [ + { + role: "assistant", + content: "Refund approved for T-981.", + tool_calls: [{ name: "issue_refund", arguments: '{"amount_usd": 40}' }], + }, + ], + }, + attributes: {}, +}; + +const textDetail: SpanDetail = { + span_id: "llm1", + input: "", + output: '{"answer": "all done"}', + output_ui: { kind: "text", text: "all done" }, + attributes: {}, +}; + +const failedToolMessageDetail: SpanDetail = { + span_id: "tool1", + input: '{"customer_id":"acme-404"}', + output: "raw tool output", + output_ui: { kind: "messages", messages: [{ role: "tool", content: "permission denied: /etc/shadow" }] }, + attributes: {}, +}; + const spanRow = (s: Span): SpanRowData => ({ kind: "span", id: s.span_id, @@ -124,7 +163,7 @@ describe("DetailPane", () => { expect(screen.getByRole("tab", { name: "Request" })).toBeInTheDocument(); expect(screen.getByRole("tab", { name: "Attributes" })).toBeInTheDocument(); expect(await screen.findByText("You are a LiteLLM support agent.")).toBeInTheDocument(); - expect(screen.getByText("get_customer_plan")).toBeInTheDocument(); + expect(screen.getAllByText("get_customer_plan").length).toBeGreaterThan(0); expect(vi.mocked(agentTraceSpanCall)).toHaveBeenCalledWith("sk-test", "t1", "llm1", undefined); }); @@ -134,8 +173,9 @@ describe("DetailPane", () => { expect(error).toHaveTextContent("Tool · ValueError"); expect(error).toHaveTextContent("ValueError('customer acme-404 not found in billing DB')"); expect(error).not.toHaveTextContent("Traceback"); - // tool args render as pretty JSON under "Input" - expect(await screen.findByText(/"customer_id": "acme-404"/)).toBeInTheDocument(); + const input = await screen.findByRole("region", { name: "Input" }); + expect(input).toHaveTextContent("customer_id"); + expect(input).toHaveTextContent("acme-404"); }); it("shows the LiteLLM request facts on the Request tab", async () => { @@ -188,4 +228,97 @@ describe("DetailPane", () => { await waitFor(() => expect(writeText).toHaveBeenCalled()); expect(writeText.mock.calls[0][0]).toContain("http://proxy.test/v1/traces/t1?format=md&span_id=llm1"); }); + + it("renders the AI tool call as a card and expands a long argument on click", async () => { + const user = userEvent.setup(); + renderPane(spanRow(llm)); + const output = await screen.findByRole("region", { name: "Output" }); + expect(output).toHaveTextContent("AI"); + expect(output).toHaveTextContent("get_customer_plan"); + const expand = within(output).getAllByRole("button", { name: "Expand note" })[0]; + expect(within(output).queryAllByText(LONG_NOTE, { selector: "pre", ignore: "[inert] *" })).toHaveLength(0); + await user.click(expand); + expect(within(output).getAllByRole("button", { name: "Collapse note" })[0]).toHaveAttribute( + "aria-expanded", + "true", + ); + expect(within(output).getAllByText(LONG_NOTE, { selector: "pre", ignore: "[inert] *" })).not.toHaveLength(0); + }); + + it("collapses the Input section without touching Output", async () => { + const user = userEvent.setup(); + renderPane(spanRow(llm)); + const input = await screen.findByRole("region", { name: "Input" }); + const systemText = "You are a LiteLLM support agent."; + expect(within(input).getByText(systemText, { ignore: "[inert] *" })).toBeInTheDocument(); + await user.click(within(input).getByRole("button", { name: "Input" })); + expect(within(input).getByRole("button", { name: "Input" })).toHaveAttribute("aria-expanded", "false"); + expect(within(input).queryByText(systemText, { ignore: "[inert] *" })).not.toBeInTheDocument(); + const output = screen.getByRole("region", { name: "Output" }); + expect(within(output).getAllByText("get_customer_plan", { ignore: "[inert] *" })).not.toHaveLength(0); + }); + + it("renders the standard input_ui / output_ui instead of re-parsing the raw payload", async () => { + vi.mocked(agentTraceSpanCall).mockResolvedValue(standardDetail); + renderPane(spanRow(llm)); + const input = await screen.findByRole("region", { name: "Input" }); + expect(input).toHaveTextContent("ticket_id"); + expect(input).toHaveTextContent("T-981"); + expect(input).not.toHaveTextContent("raw input left unparsed"); + const output = screen.getByRole("region", { name: "Output" }); + expect(output).toHaveTextContent("AI"); + expect(output).toHaveTextContent("Refund approved for T-981."); + expect(output).toHaveTextContent("issue_refund"); + expect(output).toHaveTextContent("amount_usd"); + expect(output).not.toHaveTextContent("raw output left unparsed"); + }); + + it("keeps the failed-tool styling when a tool's output arrives as a single message", async () => { + vi.mocked(agentTraceSpanCall).mockResolvedValue(failedToolMessageDetail); + renderPane(spanRow(failedTool)); + const output = await screen.findByRole("region", { name: "Output" }); + const result = within(output).getByText("permission denied: /etc/shadow"); + expect(result).toHaveClass("text-destructive"); + expect(output).not.toHaveTextContent("AI"); + }); + + it("shows a text output_ui as its plain text", async () => { + vi.mocked(agentTraceSpanCall).mockResolvedValue(textDetail); + renderPane(spanRow(llm)); + const output = await screen.findByRole("region", { name: "Output" }); + expect(within(output).getByText("all done", { selector: "pre" })).toBeInTheDocument(); + expect(output).not.toHaveTextContent("answer"); + }); + + it("groups ids and OTEL attributes into separate key / value sections on the Attributes tab", async () => { + const user = userEvent.setup(); + renderPane(spanRow(llm)); + await user.click(screen.getByRole("tab", { name: "Attributes" })); + const ids = screen.getByRole("region", { name: "Identifiers" }); + expect(ids).toHaveTextContent("span_idllm1"); + expect(ids).toHaveTextContent("parent_span_idroot"); + const attributes = await screen.findByRole("region", { name: "Attributes" }); + expect(attributes).toHaveTextContent("gen_ai.request.model"); + expect(attributes).not.toHaveTextContent("span_id"); + }); +}); + +describe("SpanHoverCard", () => { + it("shows absolute Start / End times and the agent tag after hovering the row", async () => { + const user = userEvent.setup(); + const traceStartMs = Date.parse(trace.summary.start_time); + const timed = span({ span_id: "timed", start_offset_ms: 2000, duration_ms: 3000 }); + renderWithProviders( + + + , + ); + expect(screen.queryByTestId("span-hover-card")).not.toBeInTheDocument(); + await user.hover(screen.getByRole("button", { name: "row" })); + const card = await screen.findByTestId("span-hover-card", {}, { timeout: 2000 }); + const time = within(card).getByRole("region", { name: "Time" }); + expect(time).toHaveTextContent(`Start${absoluteTime(traceStartMs, 2000)}`); + expect(time).toHaveTextContent(`End${absoluteTime(traceStartMs, 5000)}`); + expect(within(card).getByRole("region", { name: "Tags" })).toHaveTextContent("agent:support_triage_agent"); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.tsx index 4424b95cd52..985de97cdc9 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/DetailPane.tsx @@ -9,10 +9,12 @@ import { cn } from "@/lib/cva.config"; import { AttributesDetail } from "./AttributesDetail"; import { CopyButton } from "./CopyButton"; import { DetailContent, errorHeadline, useSpanDetail } from "./DetailContent"; +import { IdChip } from "./IdChip"; import { RequestDetail } from "./RequestDetail"; +import { SpanIcon } from "./SpanIcon"; import { agentHandoffText } from "./TraceDrawer"; import type { GroupRowData, TreeRow } from "./traceTree"; -import type { Span, Trace } from "./traceTypes"; +import type { Span, SpanType, Trace } from "./traceTypes"; import { fmtMs, fmtTok } from "./traceUtils"; interface DetailPaneProps { @@ -30,25 +32,61 @@ const TABS: { id: Tab; label: string }[] = [ { id: "attributes", label: "Attributes" }, ]; -function PaneHeader({ children, onClose }: { children: React.ReactNode; onClose: () => void }) { +function PaneHeader({ + type, + model, + failed, + title, + idValue, + onClose, +}: { + type: SpanType; + model: string | null; + failed: boolean; + title: React.ReactNode; + idValue?: string; + onClose: () => void; +}) { return ( -
- {children} -
); } function PaneFooter({ children }: { children: React.ReactNode }) { - return
{children}
; + return ( +
+ {children} +
+ ); } function Meta({ label, value }: { label: string; value: string }) { return ( - {label}= + {label} {value} ); @@ -75,20 +113,19 @@ function SpanPane({ ); const tokens = span.input_tokens + span.output_tokens; return ( -