From fbb6facc0f64e08464e3a0f211dc4a6e2d6a5972 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 4 Oct 2026 17:19:50 -0700 Subject: [PATCH] wip --- .../crates/gateway-traces/src/runs.rs | 11 +- .../crates/gateway-traces/tests/routes.rs | 4 +- .../crates/python-bridge/src/cache/mod.rs | 9 +- .../src/cache/native/activation.rs | 7 +- .../python-bridge/src/cache/native/backend.rs | 2 +- .../src/cache/native/semantic.rs | 3 +- .../python-bridge/src/cache/native/v2.rs | 8 +- .../crates/python-bridge/src/cache/runtime.rs | 13 +- .../python-bridge/src/cache/selection.rs | 6 +- .../crates/python-bridge/src/callable.rs | 8 +- .../crates/python-bridge/src/credentials.rs | 3 +- litellm-rust/crates/python-bridge/src/lib.rs | 7 +- .../crates/python-bridge/src/logger/mod.rs | 6 +- .../crates/python-bridge/src/logger/tests.rs | 1 - .../crates/python-bridge/src/preflight.rs | 6 +- .../src/routes/audio_transcription.rs | 2 +- .../src/routes/chat_completions.rs | 12 +- .../src/routes/chat_completions/host.rs | 3 +- .../python-bridge/src/routes/messages/host.rs | 8 +- .../crates/python-bridge/src/routes/mod.rs | 3 +- .../python-bridge/src/routes/ocr/host.rs | 5 +- .../python-bridge/src/routes/responses.rs | 3 +- .../src/routes/responses/host.rs | 3 +- .../python-bridge/src/routes/token_counter.rs | 7 +- .../crates/python-bridge/src/routes/traces.rs | 125 ++-- .../python-bridge/src/secrets/callback.rs | 3 +- .../python-bridge/src/secrets/config.rs | 3 +- .../python-bridge/src/secrets/mutation.rs | 3 +- .../python-bridge/src/secrets/operations.rs | 4 +- .../python-bridge/src/secrets/provider.rs | 3 +- .../python-bridge/src/secrets/python.rs | 2 +- .../python-bridge/src/secrets/resolved.rs | 13 +- .../crates/python-bridge/src/secrets/vault.rs | 3 +- .../crates/python-bridge/src/tokenizer.rs | 21 +- .../crates/traces-cache/src/cursor.rs | 80 ++- litellm-rust/crates/traces-cache/src/lib.rs | 2 +- litellm-rust/crates/traces-cache/src/list.rs | 1 + .../crates/traces-cache/src/reader.rs | 215 +++++-- litellm-rust/crates/traces-cache/src/store.rs | 3 +- .../crates/traces-cache/tests/read.rs | 117 ++-- .../0016_trace_rollup_agent_labels.sql | 4 +- .../0017_trace_rollup_agent_labels_mv.sql | 4 +- .../traces-clickhouse/query/lens_agents.sql | 6 - .../query/lens_availability.sql | 8 - .../traces-clickhouse/query/lens_content.sql | 35 -- .../traces-clickhouse/query/lens_evidence.sql | 14 - .../traces-clickhouse/query/lens_sample.sql | 67 --- .../traces-clickhouse/query/matching_runs.sql | 22 +- .../query/run_attribute_counts.sql | 19 + .../traces-clickhouse/query/run_counts.sql | 2 + .../traces-clickhouse/query/run_spans.sql | 1 + .../crates/traces-clickhouse/query/runs.sql | 26 +- .../traces-clickhouse/query/runs_page.sql | 16 + .../traces-clickhouse/query/span_text.sql | 16 +- .../crates/traces-clickhouse/src/access.rs | 11 + .../crates/traces-clickhouse/src/error.rs | 2 - .../crates/traces-clickhouse/src/lib.rs | 4 +- .../crates/traces-clickhouse/src/query.rs | 1 - .../traces-clickhouse/src/query/lens.rs | 268 --------- .../traces-clickhouse/src/query/named.rs | 162 +++-- .../traces-clickhouse/src/query/number.rs | 60 +- .../crates/traces-clickhouse/src/reads.rs | 27 +- .../crates/traces-clickhouse/src/sql.rs | 74 --- .../traces-clickhouse/src/wire_schema.rs | 117 +--- .../traces-clickhouse/tests/migrations.rs | 499 +++------------- .../crates/traces-clickhouse/tests/queries.rs | 6 +- .../crates/traces-clickhouse/tests/reads.rs | 93 ++- .../crates/traces-clickhouse/tests/search.rs | 553 +++++++++++++++++- litellm-rust/crates/traces/src/error.rs | 4 - litellm-rust/crates/traces/src/lib.rs | 3 +- litellm-rust/crates/traces/src/query.rs | 16 - litellm-rust/crates/traces/src/schema.rs | 2 + litellm-rust/crates/traces/src/search.rs | 46 +- litellm-rust/crates/traces/src/store.rs | 136 ++++- litellm-rust/crates/traces/tests/query.rs | 25 - litellm-rust/crates/traces/tests/search.rs | 8 +- litellm/proxy/lens/AGENTS.md | 6 + litellm/proxy/lens/analysis.py | 2 +- litellm/proxy/lens/endpoints.py | 43 +- litellm/proxy/lens/models.py | 131 ++++- litellm/proxy/lens/search.py | 11 +- litellm/proxy/lens/sources.py | 410 +++++++++---- litellm/proxy/tracing_endpoints.py | 6 +- litellm/rust_bridge/_native.pyi | 70 ++- litellm/rust_bridge/trace/generated/models.py | 285 +-------- litellm/rust_bridge/trace/generated/types.py | 32 +- litellm/rust_bridge/trace/queries.py | 52 +- litellm/rust_bridge/trace/storage.py | 114 ++-- litellm/tracing/receiver.py | 13 +- scripts/generate_trace_types.py | 9 +- .../ActivityAvailability.json | 57 -- .../schemas/traces-clickhouse/AgentRow.json | 13 - .../schemas/traces-clickhouse/CountRow.json | 29 - .../traces-clickhouse/ExecutionRow.json | 151 ----- .../traces-clickhouse/LensAccessParams.json | 26 - .../traces-clickhouse/LensContentParams.json | 62 -- .../traces-clickhouse/LensEvidenceParams.json | 59 -- .../traces-clickhouse/LensSampleParams.json | 127 ---- .../schemas/traces-clickhouse/PartRow.json | 53 -- .../traces-clickhouse/ReadQueryName.json | 12 - .../schemas/traces/RunField.json | 4 +- .../schemas/traces/RunOrder.json | 31 + .../schemas/traces/SpanText.json | 33 ++ tests/unit/proxy/lens/test_analysis.py | 89 +-- tests/unit/proxy/lens/test_endpoints.py | 43 +- tests/unit/proxy/lens/test_sources.py | 399 ++++++++++--- tests/unit/proxy/lens/test_worker.py | 17 +- tests/unit/proxy/test_tracing_endpoints.py | 43 +- tests/unit/rust_bridge/trace/test_queries.py | 124 +--- .../src/app/(dashboard)/layout.test.tsx | 6 +- .../src/app/(dashboard)/layout.tsx | 12 +- .../src/components/lens/LensModeSwitch.tsx | 11 +- .../lens/LensSetup.integration.test.tsx | 100 +--- .../lens/LensWorkspace.integration.test.tsx | 46 +- .../src/components/lens/LensWorkspace.tsx | 15 +- .../src/components/lens/data/LensServices.tsx | 2 +- .../lens/data/demo/createLensDemo.ts | 4 +- .../src/components/lens/data/demo/fixtures.ts | 20 +- .../src/components/lens/data/queries.ts | 57 +- .../src/components/lens/data/service.ts | 22 +- .../lens/investigations/Evidence.tsx | 64 +- .../FindingDetails.integration.test.tsx | 4 +- .../lens/investigations/FindingDetails.tsx | 15 +- .../InvestigationsView.integration.test.tsx | 178 ++---- .../investigations/InvestigationsView.tsx | 1 - .../detail/HistoryTimeline.test.ts | 6 +- .../detail/InvestigationDetail.tsx | 8 +- .../investigations/detail/RunNowDialog.tsx | 24 +- .../lens/investigations/detail/RunsTab.tsx | 8 +- .../investigations/investigationQuery.test.ts | 21 +- .../lens/investigations/investigationQuery.ts | 3 +- .../investigationScreen.test.ts | 6 +- .../src/components/lens/model/findings.ts | 21 +- .../src/components/lens/model/format.test.ts | 6 +- .../src/components/lens/model/format.ts | 10 +- .../src/components/lens/model/inbox.test.ts | 29 +- .../src/components/lens/model/inbox.ts | 16 +- .../components/lens/model/progress.test.ts | 6 +- .../components/lens/model/readiness.test.ts | 6 +- .../src/components/lens/model/readiness.ts | 7 +- .../src/components/lens/model/runRequest.ts | 8 + .../src/components/lens/model/status.test.ts | 6 +- .../src/components/lens/model/types.ts | 19 +- .../lens/onboarding/OnboardingSteps.tsx | 7 +- .../src/components/lens/route.ts | 3 +- .../InvestigationSetup.integration.test.tsx | 236 ++------ .../lens/setup/InvestigationSetup.tsx | 36 +- .../lens/setup/MatchingActivityPreview.tsx | 113 ++-- .../lens/setup/fields/MetadataFilters.tsx | 75 --- .../lens/setup/fields/ScopeFields.tsx | 121 +--- .../src/components/lens/setup/filters.test.ts | 13 - .../src/components/lens/setup/filters.ts | 8 - .../lens/setup/investigationSchema.test.ts | 21 +- .../lens/setup/investigationSchema.ts | 49 +- .../lens/setup/useMatchingActivity.ts | 54 +- .../lens/traces/list/AgentTracesSection.tsx | 47 +- .../lens/traces/list/AgentTracesTable.tsx | 29 +- .../src/components/lens/traces/routing.ts | 7 + .../shared/search/SearchBox.test.tsx | 12 +- .../components/shared/search/SearchBox.tsx | 7 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 157 +---- 161 files changed, 3140 insertions(+), 4273 deletions(-) delete mode 100644 litellm-rust/crates/traces-clickhouse/query/lens_agents.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/lens_availability.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/lens_content.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/lens_evidence.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/lens_sample.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/run_attribute_counts.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/runs_page.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/src/query/lens.rs delete mode 100644 litellm-rust/crates/traces-clickhouse/src/sql.rs delete mode 100644 litellm-rust/crates/traces/tests/query.rs create mode 100644 litellm/proxy/lens/AGENTS.md delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/ActivityAvailability.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/AgentRow.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/CountRow.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/ExecutionRow.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/LensAccessParams.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/LensContentParams.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/LensEvidenceParams.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/LensSampleParams.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/PartRow.json delete mode 100644 scripts/trace_codegen/schemas/traces-clickhouse/ReadQueryName.json create mode 100644 scripts/trace_codegen/schemas/traces/RunOrder.json create mode 100644 scripts/trace_codegen/schemas/traces/SpanText.json delete mode 100644 ui/litellm-dashboard/src/components/lens/setup/fields/MetadataFilters.tsx delete mode 100644 ui/litellm-dashboard/src/components/lens/setup/filters.test.ts delete mode 100644 ui/litellm-dashboard/src/components/lens/setup/filters.ts diff --git a/litellm-rust/crates/gateway-traces/src/runs.rs b/litellm-rust/crates/gateway-traces/src/runs.rs index b5ee4b92641..030a7c2a5a5 100644 --- a/litellm-rust/crates/gateway-traces/src/runs.rs +++ b/litellm-rust/crates/gateway-traces/src/runs.rs @@ -10,6 +10,7 @@ use axum::{ use litellm_traces::{ QueryScope, TracePage, search::{RunField, RunFilter, RunSearch, RunValues, TraceHistogram}, + store::RunOrder, }; use litellm_traces_cache::{PageRequest, TraceStore}; use serde::Deserialize; @@ -34,6 +35,7 @@ impl Runs { start_ms: self.start_ms.unwrap_or(now_ms - DAY_MS), end_ms: self.end_ms.unwrap_or(now_ms), search: RunSearch::parse(&self.q), + trace_refs: Vec::new(), } } } @@ -52,11 +54,18 @@ pub(crate) async fn list( let page = PageRequest { cursor, limit: PAGE_SIZE, + ..PageRequest::default() }; Ok(Json( traces .reader - .list_traces(&traces.store, &access, &runs.filter(), &page) + .list_traces( + &traces.store, + &access, + &runs.filter(), + RunOrder::NEWEST, + &page, + ) .await?, )) } diff --git a/litellm-rust/crates/gateway-traces/tests/routes.rs b/litellm-rust/crates/gateway-traces/tests/routes.rs index 444c5bf4bb8..de4b00449d3 100644 --- a/litellm-rust/crates/gateway-traces/tests/routes.rs +++ b/litellm-rust/crates/gateway-traces/tests/routes.rs @@ -102,8 +102,8 @@ impl TraceStore for FakeStore { &self, _: &QueryScope, _: &SpanTextQuery, - ) -> StoreResult, FakeError> { - Ok(None) + ) -> StoreResult, FakeError> { + Ok(Vec::new()) } async fn calls(&self, _: &QueryScope, _: &CallQuery) -> StoreResult, FakeError> { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 179f16c4a1f..09653f8f68a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -4,16 +4,15 @@ mod python; mod runtime; mod selection; -pub(crate) use native::NativeCacheHandle; -pub(crate) use python::{CacheCall, PythonCache}; -pub(crate) use runtime::ResolvedCache; -pub(crate) use selection::{Cached, Selection, admit_native, configure, configured_native}; - use litellm_cache::Error; +pub(crate) use native::NativeCacheHandle; use pyo3::{ exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; +pub(crate) use python::{CacheCall, PythonCache}; +pub(crate) use runtime::ResolvedCache; +pub(crate) use selection::{Cached, Selection, admit_native, configure, configured_native}; fn cache_error(error: Error) -> PyErr { match error { diff --git a/litellm-rust/crates/python-bridge/src/cache/native/activation.rs b/litellm-rust/crates/python-bridge/src/cache/native/activation.rs index 3ac038d2c39..774802ae719 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/activation.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/activation.rs @@ -1,5 +1,3 @@ -use crate::cache::cache_error; -use crate::execution::run_sync_value; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_host_python::release_gil; @@ -11,8 +9,9 @@ use super::{ config::{CacheBackendConfig, NativeCacheConfig, UnsupportedCacheConfig}, embedder::PythonEmbedder, }; -use crate::errors::RustBridgeDeclined; -use crate::http::host_client; +use crate::{ + cache::cache_error, errors::RustBridgeDeclined, execution::run_sync_value, http::host_client, +}; fn declined(reason: UnsupportedCacheConfig) -> PyErr { RustBridgeDeclined::new_err(reason.message()) diff --git a/litellm-rust/crates/python-bridge/src/cache/native/backend.rs b/litellm-rust/crates/python-bridge/src/cache/native/backend.rs index 1151ed5cc9d..1d0c08d16db 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/backend.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/backend.rs @@ -1,4 +1,3 @@ -use crate::cache::cache_error; use std::{sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error, semantic::SemanticLookup}; @@ -25,6 +24,7 @@ use super::{ request::{NativeRequest, now}, semantic::{EmbeddingFailure, SemanticExecution, SemanticOperation, drive}, }; +use crate::cache::cache_error; /// What the Python embedder receives for one semantic request. pub(in crate::cache) struct EmbeddingInput { diff --git a/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs index 8d9bf270be0..2489f0b9add 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs @@ -1,5 +1,3 @@ -use crate::cache::cache_error; -use crate::execution::run_async; use std::{collections::VecDeque, time::Duration}; use litellm_cache::Error; @@ -16,6 +14,7 @@ use super::{ embedder::{PythonEmbedder, with_prepared_embedding}, request::{NativeRequest, now}, }; +use crate::{cache::cache_error, execution::run_async}; pub(super) enum SemanticOperation { Lookup(NativeRequest), diff --git a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs index 0dd70a042e9..01bbf1cc834 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs @@ -1,17 +1,17 @@ -use crate::cache::cache_error; use std::{sync::Arc, time::Duration}; use litellm_cache::{DeleteCache, DisconnectCache, PingCache}; -use litellm_host_python::{from_py, release_gil, to_py}; -use serde_json::Value; - use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_response::{ CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, }; +use litellm_host_python::{from_py, release_gil, to_py}; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use serde_json::Value; + +use crate::cache::cache_error; #[pyclass( frozen, diff --git a/litellm-rust/crates/python-bridge/src/cache/runtime.rs b/litellm-rust/crates/python-bridge/src/cache/runtime.rs index eec82f2ba4c..deaf86bd341 100644 --- a/litellm-rust/crates/python-bridge/src/cache/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/cache/runtime.rs @@ -1,4 +1,3 @@ -use crate::execution::run_async; use litellm_cache_response::PartialHits; use litellm_host_python::{ExecutionStep, from_py, release_gil, to_py}; use pyo3::{ @@ -12,13 +11,15 @@ use serde_json::Value; use super::{ cache_error, future::{ready_none, ready_value}, - native::activation::activate, - native::backend::{NativeResponseCache, SemanticReply}, - native::config::{CacheConfigProjection, NativeCacheConfig}, - native::request::{now, request, requests}, + native::{ + activation::activate, + backend::{NativeResponseCache, SemanticReply}, + config::{CacheConfigProjection, NativeCacheConfig}, + request::{now, request, requests}, + }, python::PythonCallback, }; -use crate::errors::RustBridgeDeclined; +use crate::{errors::RustBridgeDeclined, execution::run_async}; pub(super) enum CacheBinding { Disabled, diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index e6e9f8d2d4f..10f641b1037 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -1,4 +1,5 @@ -use super::{native, python}; +use std::sync::Arc; + use litellm_cache_response::{ CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache, }; @@ -7,7 +8,8 @@ use litellm_host::{ protocol::Protocol, }; use pyo3::{prelude::*, types::PyDict}; -use std::sync::Arc; + +use super::{native, python}; pub(crate) struct Cached

(std::marker::PhantomData

); diff --git a/litellm-rust/crates/python-bridge/src/callable.rs b/litellm-rust/crates/python-bridge/src/callable.rs index d0ed76a3e8b..0686d763bb1 100644 --- a/litellm-rust/crates/python-bridge/src/callable.rs +++ b/litellm-rust/crates/python-bridge/src/callable.rs @@ -1,8 +1,10 @@ //! Failures raised by a caller-supplied Python callable. -use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError}; -use pyo3::prelude::*; -use pyo3::types::PyString; +use pyo3::{ + exceptions::{PyException, PyRuntimeError, PyTypeError}, + prelude::*, + types::PyString, +}; /// Reports a caller-supplied callable's failure under `template`, a Python format string /// with one field for the original exception, while leaving alone the failures a caller diff --git a/litellm-rust/crates/python-bridge/src/credentials.rs b/litellm-rust/crates/python-bridge/src/credentials.rs index 04d8e94c8fe..a01983d9c98 100644 --- a/litellm-rust/crates/python-bridge/src/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/credentials.rs @@ -1,7 +1,6 @@ //! Credentials the caller supplies as Python callables, projected out of a route's //! keyword arguments and acquired on the host's own thread when the call asks for one. -use crate::callable::wrap_failure; use litellm_auth::{ResolvedCredential, SecretValue}; use pyo3::{ exceptions::PyTypeError, @@ -10,6 +9,8 @@ use pyo3::{ types::{PyDict, PyString}, }; +use crate::callable::wrap_failure; + const NOT_CALLABLE: &str = "Azure AD token provider must be callable"; const NOT_A_STRING: &str = "Azure AD token must be a string, got {}"; const FAILED: &str = "Failed to get Azure AD token: {}"; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 6659be5160e..b06229f9c08 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -17,6 +17,10 @@ mod tokenizer; #[pymodule(gil_used = true)] mod _native { + #[pymodule_export] + use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; + use pyo3::{prelude::*, types::PyModule}; + use crate::cache::ResolvedCache; #[cfg(feature = "panic-test")] #[pymodule_export] @@ -52,9 +56,6 @@ mod _native { use crate::tokenizer::HuggingFaceEncoding; #[pymodule_export] use crate::tokenizer::Tokenizer; - #[pymodule_export] - use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; - use pyo3::{prelude::*, types::PyModule}; #[pymodule_init] fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/litellm-rust/crates/python-bridge/src/logger/mod.rs b/litellm-rust/crates/python-bridge/src/logger/mod.rs index bf5c735360b..d02263948b7 100644 --- a/litellm-rust/crates/python-bridge/src/logger/mod.rs +++ b/litellm-rust/crates/python-bridge/src/logger/mod.rs @@ -1,11 +1,9 @@ mod machine; -pub(crate) use machine::LoggedMachine; - use litellm_host_python::Pythonized; use litellm_tracing::{DiagnosticInput, Level, Logger, Metadata, Policy, Processor, Record, Sink}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::*; +pub(crate) use machine::LoggedMachine; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; const MODULE: &str = "litellm.rust_bridge.logger"; type NativeDiagnosticOutput = (String, Option, Option, Vec, bool); diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs index 1fca3be720e..26b6c74c617 100644 --- a/litellm-rust/crates/python-bridge/src/logger/tests.rs +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -4,7 +4,6 @@ use litellm_host::{ machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, protocol::Protocol, }; - use pyo3::{prelude::*, types::PyDict}; struct DiagnosticMachine; diff --git a/litellm-rust/crates/python-bridge/src/preflight.rs b/litellm-rust/crates/python-bridge/src/preflight.rs index bc692498880..307056cd5cb 100644 --- a/litellm-rust/crates/python-bridge/src/preflight.rs +++ b/litellm-rust/crates/python-bridge/src/preflight.rs @@ -105,11 +105,11 @@ fn inherit_credentials<'py>( #[cfg(test)] mod tests { - use std::collections::BTreeSet; - use std::sync::Mutex; + use std::{collections::BTreeSet, sync::Mutex}; + + use strum::VariantArray; use super::*; - use strum::VariantArray; /// Tests share one interpreter, and the stub module below is global state, so the /// tests that install it run one at a time. diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 8d434dbbc74..b29f5e092a1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -1,4 +1,3 @@ -use crate::execution::{run_async, run_sync}; use litellm_core::audio_transcription::{ AudioTranscriptionRoute, Error, types::AudioTranscriptionRequest, }; @@ -8,6 +7,7 @@ use serde_json::{Map, Value}; use crate::{ errors::route_error_to_pyerr, + execution::{run_async, run_sync}, marshal::{RouteOptions, extra_headers_argument, optional_params_argument, optional_timeout}, }; diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 5955729d6e9..c00471d12e4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -1,15 +1,16 @@ mod host; -use pyo3::types::{PyDict, PyTuple}; - -use crate::execution::{run_async, run_sync}; use litellm_core::chat_completions::{ChatCompletionsRoute, Error, types::ChatCompletionsRequest}; use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse; -use pyo3::prelude::*; +use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, +}; use serde_json::{Map, Value}; use crate::{ errors::route_error_to_pyerr, + execution::{run_async, run_sync}, marshal::{ RouteOptions, extra_headers_argument, messages_argument, optional_params_argument, optional_timeout, @@ -136,8 +137,9 @@ fn run_public( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { - use super::inference::InferenceHost; use litellm_callbacks_legacy_python::LoggingOperation; + + use super::inference::InferenceHost; let host = InferenceHost::new( request.clone().unbind(), "litellm.rust_bridge.chat_completions.route_host", diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 5011e151b48..c45b1d7ac6c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -1,6 +1,5 @@ use std::convert::Infallible; -use super::super::inference::InferenceHost; use litellm_core::chat_completions::{Error, route::ChatCompletions, types::ChatCompletionsCall}; use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; use pyo3::{ @@ -9,6 +8,8 @@ use pyo3::{ types::PyDict, }; +use super::super::inference::InferenceHost; + pub(super) struct ChatCompletionsPythonHost(pub InferenceHost); pub(super) fn project( diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 2f67151374d..da9623d478c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -1,12 +1,11 @@ -use crate::cache::{CacheCall, Cached, PythonCache, Selection}; -use litellm_host_python::{PythonHostCalls, PythonOwned}; - use bytes::Bytes; use litellm_core::messages::{ Error, MessagesCall, MessagesShaping, messages_body, route::{Messages, MessagesStreamHead}, }; -use litellm_host_python::{InvokeError, PythonBinding, from_py, lookup, to_py}; +use litellm_host_python::{ + InvokeError, PythonBinding, PythonHostCalls, PythonOwned, from_py, lookup, to_py, +}; use litellm_http::transport::Error as TransportError; use litellm_llms_types::headers::ProviderSpecificHeaders; use pyo3::{ @@ -18,6 +17,7 @@ use pyo3::{ use serde_json::{Map, Value}; use crate::{ + cache::{CacheCall, Cached, PythonCache, Selection}, errors::{RustUpstreamError, route_error_to_pyerr}, marshal::{optional_timeout, python_timeout_seconds}, }; diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 2380274001e..75d6d2934fc 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -8,8 +8,7 @@ pub(crate) mod responses; pub(crate) mod token_counter; pub(crate) mod traces; -use litellm_callbacks_legacy_python::LoggingOperation; -use litellm_callbacks_legacy_python::{LegacyLogging, PublicCall}; +use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCall}; use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol}; use litellm_host_python::{HookChain, PythonBinding, PythonCallHooks, PythonHostCalls}; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 28317317544..f011e87bbba 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,7 +1,8 @@ use litellm_auth::ResolvedCredential; use litellm_core::ocr::route::{Ocr, OcrCall, OcrOp}; -use litellm_host_python::{InvokeError, PythonBinding, missing_state, to_py}; -use litellm_host_python::{PythonHostCalls, PythonOwned}; +use litellm_host_python::{ + InvokeError, PythonBinding, PythonHostCalls, PythonOwned, missing_state, to_py, +}; use litellm_llms::base_llm::ocr::error::Error; use litellm_llms_types::formats::ocr::LiteLLMOcrResponse; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/routes/responses.rs b/litellm-rust/crates/python-bridge/src/routes/responses.rs index 4e1426cd298..da46ab334a3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses.rs @@ -19,8 +19,9 @@ fn run_public( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { - use super::inference::InferenceHost; use litellm_callbacks_legacy_python::LoggingOperation; + + use super::inference::InferenceHost; let host = InferenceHost::new( request.clone().unbind(), "litellm.rust_bridge.responses.route_host", diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index b054cbe5f38..c04e88849f2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -1,6 +1,5 @@ use std::convert::Infallible; -use super::super::inference::InferenceHost; use litellm_core::responses::{Error, route::Responses, types::ResponsesCall}; use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; use pyo3::{ @@ -9,6 +8,8 @@ use pyo3::{ types::PyDict, }; +use super::super::inference::InferenceHost; + pub(super) struct ResponsesPythonHost(pub InferenceHost); pub(super) fn project( diff --git a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs index 2c26311231f..5f731f8864c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs @@ -1,6 +1,4 @@ -use crate::execution::run_async; -use std::sync::Arc; -use std::{num::NonZero, thread::available_parallelism}; +use std::{num::NonZero, sync::Arc, thread::available_parallelism}; use litellm_host_python::enter_native; use litellm_token_counter::{ @@ -13,8 +11,7 @@ use pyo3::{ }; use tokio::sync::Semaphore; -use crate::errors::RustBridgeDeclined; -use crate::tokenizer::Tokenizer; +use crate::{errors::RustBridgeDeclined, execution::run_async, tokenizer::Tokenizer}; /// Counts the input tokens of a raw request body off the Python event loop with /// the GIL released. Python owns which requests get here and what to do with diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index cf8d9c46fa8..1892bda37bc 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -2,13 +2,12 @@ use std::{collections::BTreeMap, sync::Arc}; use litellm_http::ClientVariant; use litellm_traces::{ - QueryScope, ReadQuery, Tenant, + QueryScope, Tenant, search::{RunField, RunFilter, RunSearch}, + store::{RunOrder, SpanPart, TextRange}, }; use litellm_traces_cache::{PageRequest, ReadError, TraceReader}; -use litellm_traces_clickhouse::{ - ClickHouseTraces, Config, Error, InsertTable, Parameter, QueryReaders, -}; +use litellm_traces_clickhouse::{ClickHouseTraces, Config, Error, InsertTable, QueryReaders}; use prost::Message; use pyo3::{ exceptions::{PyOverflowError, PyRuntimeError, PyValueError}, @@ -51,7 +50,6 @@ fn map_error_ref(error: &Error) -> PyErr { | Error::InvalidTable | Error::Decode(_) | Error::InvalidSchema - | Error::InvalidQuery | Error::InvalidParameters | Error::InvalidScope => PyValueError::new_err(error.to_string()), Error::Task @@ -95,14 +93,21 @@ fn map_read_error(error: ReadError) -> PyErr { } } -fn run_filter(start_ms: i64, end_ms: i64, q: &str) -> RunFilter { +fn run_filter(start_ms: i64, end_ms: i64, q: &str, trace_refs: Vec) -> RunFilter { RunFilter { start_ms, end_ms, search: RunSearch::parse(q), + trace_refs, } } +fn parsed(kind: &str, value: &str) -> PyResult { + value + .parse() + .map_err(|_| PyValueError::new_err(format!("unknown {kind} {value}"))) +} + fn map_sql_error(error: Error) -> PyErr { match error { Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => { @@ -235,7 +240,7 @@ impl NativeTraceStorage { ) } - #[pyo3(signature = (scope, start_ms, end_ms, q, cursor, limit))] + #[pyo3(signature = (scope, start_ms, end_ms, q, cursor, limit, order, trace_refs=Vec::new()))] #[expect( clippy::too_many_arguments, reason = "one parameter per Python argument" @@ -249,8 +254,10 @@ impl NativeTraceStorage { q: &str, cursor: Option, limit: u32, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] order: RunOrder, + trace_refs: Vec, ) -> PyResult> { - let filter = run_filter(start_ms, end_ms, q); + let filter = run_filter(start_ms, end_ms, q, trace_refs); let page = PageRequest { cursor, limit }; let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); @@ -259,7 +266,75 @@ impl NativeTraceStorage { py, async move { let store = ClickHouseTraces::new(client, connection); - reader.list_traces(&store, &scope, &filter, &page).await + reader + .list_traces(&store, &scope, &filter, order, &page) + .await + }, + map_read_error, + ) + } + + #[pyo3(signature = (scope, start_ms, end_ms, q, trace_refs=Vec::new()))] + fn count_traces<'py>( + &self, + py: Python<'py>, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, + start_ms: i64, + end_ms: i64, + q: &str, + trace_refs: Vec, + ) -> PyResult> { + let filter = run_filter(start_ms, end_ms, q, trace_refs); + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); + crate::execution::run_async( + py, + async move { + let store = ClickHouseTraces::new(client, connection); + reader.count_traces(&store, &scope, &filter).await + }, + map_read_error, + ) + } + + /// `tail` reads the last `max_chars` characters instead of starting at `offset`. + #[pyo3(signature = (trace_id, trace_ref, span_ids, part, scope, offset=0, max_chars=None, tail=false, contains=None))] + #[expect( + clippy::too_many_arguments, + reason = "one parameter per Python argument" + )] + fn span_text<'py>( + &self, + py: Python<'py>, + trace_id: String, + trace_ref: String, + span_ids: Vec, + part: &str, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, + offset: u64, + max_chars: Option, + tail: bool, + contains: Option, + ) -> PyResult> { + let part: SpanPart = parsed("span part", part)?; + let range = match (tail, max_chars) { + (true, Some(chars)) => TextRange::Last { chars }, + (true, None) => return Err(PyValueError::new_err("tail reads need max_chars")), + (false, max_chars) => TextRange::From { offset, max_chars }, + }; + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); + crate::execution::run_async( + py, + async move { + let store = ClickHouseTraces::new(client, connection); + reader + .span_text( + &store, &scope, &trace_id, &trace_ref, span_ids, part, range, contains, + ) + .await }, map_read_error, ) @@ -274,7 +349,7 @@ impl NativeTraceStorage { q: &str, buckets: u32, ) -> PyResult> { - let filter = run_filter(start_ms, end_ms, q); + let filter = run_filter(start_ms, end_ms, q, Vec::new()); let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); let reader = Arc::clone(&self.reader); @@ -306,7 +381,7 @@ impl NativeTraceStorage { let field = field .parse::() .map_err(|_| PyValueError::new_err(format!("unknown run field {field}")))?; - let filter = run_filter(start_ms, end_ms, q); + let filter = run_filter(start_ms, end_ms, q, Vec::new()); let contains = contains.to_owned(); let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); @@ -461,34 +536,6 @@ impl NativeTraceStorage { map_sql_error, ) } - - fn query<'py>( - &self, - py: Python<'py>, - query: &str, - #[pyo3(from_py_with = litellm_host_python::from_py_argument)] parameters: BTreeMap< - String, - Parameter, - >, - ) -> PyResult> { - let query = - ReadQuery::parse(query).map_err(|error| PyValueError::new_err(error.to_string()))?; - let connection = self.config.storage().reader().clone(); - let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; - crate::execution::run_async( - py, - async move { - litellm_traces_clickhouse::execute_named_read( - &client, - &connection, - query, - ¶meters, - ) - .await - }, - map_error, - ) - } } /// The `otel_traces` rows an export would be stored as, without writing them. diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 6b61ca22dcd..7d44fb8e241 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -130,6 +130,7 @@ fn log_environment_fallback(py: Python<'_>, name: &str, error: &PyErr) -> PyResu mod tests { use std::sync::{Arc, Mutex, MutexGuard}; + use litellm_host_python::PythonContext; use litellm_secrets::{ FailurePolicy, KeyManagementSettings, KeyManagementSystem, OidcResolver, SecretManager, SecretManagerState, SecretResolver, @@ -137,8 +138,6 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; use rstest::rstest; - use litellm_host_python::PythonContext; - use super::{HANDLER_MODULE, PythonSecretManager, python_name}; use crate::secrets::python_error; diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index d1a1eafd468..8d7d8e6429d 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -1,12 +1,11 @@ use std::sync::Arc; +use litellm_host_python::PythonContext; use litellm_secrets::{SecretManager, SecretManagerState}; use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue}; use pyo3::prelude::*; use serde_json::Value; -use litellm_host_python::PythonContext; - use super::callback::PythonSecretManager; use crate::{ coercion::{Field, FieldSpec, ProjectionError}, diff --git a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs index eaab06a5b12..33490db0bd1 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs @@ -1,8 +1,9 @@ -use super::operations::{PythonMutationError, PythonMutationResponse}; use litellm_host_python::{json_loads, to_py}; use litellm_secrets::cyberark; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use super::operations::{PythonMutationError, PythonMutationResponse}; + pub(super) fn mutation_value( result: Result, context: &super::vault::ErrorContext, diff --git a/litellm-rust/crates/python-bridge/src/secrets/operations.rs b/litellm-rust/crates/python-bridge/src/secrets/operations.rs index 8e338741aff..8b7d8578616 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/operations.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/operations.rs @@ -1,7 +1,5 @@ use litellm_core_utils::settings::Lookup; -use litellm_secrets::Secret; -use litellm_secrets::cyberark::AuthenticationRetry; -use litellm_secrets::{Error, SecretManager}; +use litellm_secrets::{Error, Secret, SecretManager, cyberark::AuthenticationRetry}; use litellm_secrets_types::{PythonSecretRead, SecretOperationContext}; pub(super) struct PythonReadRequest { diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index 568a0cd2228..b19b33243e0 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -1,6 +1,5 @@ use std::time::Duration; -use super::operations::PythonReadRequest; use litellm_secrets::{KeyManagementSystem, SecretValue}; use litellm_secrets_types::{ AwsOperationContext, CyberarkOperationContext, GoogleOperationContext, @@ -8,6 +7,8 @@ use litellm_secrets_types::{ }; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use super::operations::PythonReadRequest; + pub(super) fn read_request( system: KeyManagementSystem, secret_name: String, diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index fb26b1e80fc..8be7793d2f9 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -60,13 +60,13 @@ impl SecretSource for PythonSecrets { #[cfg(test)] mod tests { + use litellm_host_python::PythonContext; use litellm_secrets::source::SecretSource; use pyo3::{prelude::*, types::PyDict}; use rstest::{fixture, rstest}; use super::PythonSecrets; use crate::secrets::python_error; - use litellm_host_python::PythonContext; #[fixture] fn namespace() -> Py { diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index 5a606ab1039..3376e4fbc59 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -4,9 +4,9 @@ use futures_util::future::BoxFuture; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::PythonContext; use litellm_http::Client; -use litellm_secrets::source::SecretSource; use litellm_secrets::{ Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver, SecretValue, + source::SecretSource, }; use super::config::SecretManagerSnapshot; @@ -49,11 +49,13 @@ impl SecretSource for ResolvedSecrets { mod tests { use std::sync::Arc; - use aws_sdk_secretsmanager::Client; - use aws_sdk_secretsmanager::config::{ - BehaviorVersion, Credentials, Region, retry::RetryConfig, + use aws_sdk_secretsmanager::{ + Client, + config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, + }; + use litellm_secrets::{ + AccessMode, KeyManagementSettings, SecretManager, SecretManagerState, source::SecretSource, }; - use litellm_secrets::{AccessMode, KeyManagementSettings, SecretManager, SecretManagerState}; use litellm_secrets_aws::AwsSecretsManagerV2; use serde_json::json; use wiremock::{ @@ -62,7 +64,6 @@ mod tests { }; use super::ResolvedSecrets; - use litellm_secrets::source::SecretSource; fn state(server: &MockServer, settings: KeyManagementSettings) -> Arc { let client = Client::from_conf( diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault.rs b/litellm-rust/crates/python-bridge/src/secrets/vault.rs index de0fa95bcef..7de23254970 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault.rs @@ -1,8 +1,7 @@ mod operation; -pub(super) use operation::{Failure, FailureKind, FailureStage, delete, rotate, write}; - use litellm_secrets::hashicorp::{Error, RawOperationError}; +pub(super) use operation::{Failure, FailureKind, FailureStage, delete, rotate, write}; use pyo3::prelude::*; use super::mutation::{error_value, http_message, json_value}; diff --git a/litellm-rust/crates/python-bridge/src/tokenizer.rs b/litellm-rust/crates/python-bridge/src/tokenizer.rs index df219d55eb6..bab3c6a7b48 100644 --- a/litellm-rust/crates/python-bridge/src/tokenizer.rs +++ b/litellm-rust/crates/python-bridge/src/tokenizer.rs @@ -1,23 +1,28 @@ //! The Python face of the text codecs: one `Tokenizer` class over the tiktoken and Hugging //! Face backends, carrying the read-only surface of `tiktoken.Encoding` and //! `tokenizers.Tokenizer` that `litellm/litellm_core_utils/tokenizer.py` wraps. -use std::borrow::Cow; #[cfg(any(feature = "tiktoken", feature = "huggingface"))] use std::collections::HashMap; -use std::sync::Arc; #[cfg(feature = "fast")] use std::sync::OnceLock; +use std::{borrow::Cow, sync::Arc}; use litellm_host_python::{enter_native, release_gil}; #[cfg(feature = "fast")] use litellm_token_counter::fast::{FastCounter, FastTokenizer}; +#[cfg(feature = "huggingface")] +use litellm_token_counter::huggingface::{ + EncodeInput, Encoding, HuggingFaceTokenizer, InputSequence, PaddingDirection, PaddingStrategy, + TruncationDirection, encoding_from_json, encoding_to_json, +}; +#[cfg(feature = "tiktoken")] +use litellm_token_counter::tiktoken::{TiktokenTokenizer, Vocabulary}; use litellm_token_counter::{Error, TextCodec}; -use pyo3::{exceptions::PyUnicodeEncodeError, prelude::*, types::PyString}; - #[cfg(any(feature = "tiktoken", feature = "huggingface"))] use pyo3::exceptions::PyValueError; #[cfg(feature = "huggingface")] use pyo3::{exceptions::PyIOError, types::PyDict}; +use pyo3::{exceptions::PyUnicodeEncodeError, prelude::*, types::PyString}; #[cfg(feature = "tiktoken")] use pyo3::{ exceptions::{PyKeyError, PyRuntimeError}, @@ -28,14 +33,6 @@ use pyo3::{ use crate::errors::RustBridgeDeclined; use crate::routes::token_counter::token_count_error_to_pyerr; -#[cfg(feature = "huggingface")] -use litellm_token_counter::huggingface::{ - EncodeInput, Encoding, HuggingFaceTokenizer, InputSequence, PaddingDirection, PaddingStrategy, - TruncationDirection, encoding_from_json, encoding_to_json, -}; -#[cfg(feature = "tiktoken")] -use litellm_token_counter::tiktoken::{TiktokenTokenizer, Vocabulary}; - #[cfg(feature = "tiktoken")] pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { enter_native()?; diff --git a/litellm-rust/crates/traces-cache/src/cursor.rs b/litellm-rust/crates/traces-cache/src/cursor.rs index 436fbc3049d..82d5a61e282 100644 --- a/litellm-rust/crates/traces-cache/src/cursor.rs +++ b/litellm-rust/crates/traces-cache/src/cursor.rs @@ -1,5 +1,5 @@ use base64::{Engine, engine::general_purpose::URL_SAFE}; -use litellm_traces::store::{RunCursor, SpanPart}; +use litellm_traces::store::{RunCursor, RunOrder, RunRow, SpanPart}; use serde::{Deserialize, Serialize}; use crate::ReadError; @@ -12,7 +12,7 @@ use crate::ReadError; deny_unknown_fields )] pub(super) enum Cursor { - Run(RunCursor), + Run(RunPosition), Span(SpanPosition), Text(TextPosition), } @@ -40,6 +40,25 @@ impl Cursor { } } +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub(super) struct RunPosition { + order: RunOrder, + value: i64, + trace_ref: String, +} + +impl RunPosition { + pub(super) fn after(order: RunOrder, row: &RunRow) -> Self { + let RunCursor { value, trace_ref } = order.cursor(row); + Self { + order, + value, + trace_ref, + } + } +} + #[derive(Deserialize, Serialize)] #[serde(deny_unknown_fields)] pub(super) struct SpanPosition { @@ -57,13 +76,19 @@ pub(super) struct TextPosition { pub(super) version: String, } -pub(super) fn run_position(cursor: Option<&str>) -> Result, ReadError> { +pub(super) fn run_position( + cursor: Option<&str>, + order: RunOrder, +) -> Result, ReadError> { let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else { return Ok(None); }; match Cursor::decode(cursor, "trace")? { - Cursor::Run(position) if position.start_ms > 0 && !position.trace_ref.is_empty() => { - Ok(Some(position)) + Cursor::Run(position) if position.order == order && !position.trace_ref.is_empty() => { + Ok(Some(RunCursor { + value: position.value, + trace_ref: position.trace_ref, + })) } _ => Err(ReadError::InvalidCursor("trace")), } @@ -99,13 +124,15 @@ pub(super) fn text_position( #[cfg(test)] mod tests { + use litellm_traces::store::RunSortKey; use rstest::rstest; use super::*; - fn run(start_ms: i64, trace_ref: &str) -> String { - Cursor::Run(RunCursor { - start_ms, + fn run(order: RunOrder, value: i64, trace_ref: &str) -> String { + Cursor::Run(RunPosition { + order, + value, trace_ref: trace_ref.into(), }) .encode() @@ -134,36 +161,49 @@ mod tests { URL_SAFE.encode(value.to_string()) } + const BY_ERRORS: RunOrder = RunOrder { + key: RunSortKey::ErrorCount, + descending: false, + }; + #[rstest] - fn run_cursor_round_trips_the_last_listed_run() { - let position = run_position::(Some(&run(1_790_742_989_377, "4BAD"))) + #[case::newest(RunOrder::NEWEST, 1_790_742_989_377)] + #[case::zero_value(BY_ERRORS, 0)] + fn run_cursor_round_trips_under_its_own_order(#[case] order: RunOrder, #[case] value: i64) { + let position = run_position::(Some(&run(order, value, "4BAD")), order) .unwrap() .unwrap(); assert_eq!( - (position.start_ms, position.trace_ref.as_str()), - (1_790_742_989_377, "4BAD") + (position.value, position.trace_ref.as_str()), + (value, "4BAD") ); } #[rstest] #[case::absent(None)] #[case::empty(Some(""))] - fn missing_run_cursor_starts_from_the_newest(#[case] cursor: Option<&str>) { - assert!(run_position::(cursor).unwrap().is_none()); + fn missing_run_cursor_starts_from_the_first_page(#[case] cursor: Option<&str>) { + assert!( + run_position::(cursor, RunOrder::NEWEST) + .unwrap() + .is_none() + ); } #[rstest] #[case::not_base64("abc".into())] #[case::not_json(URL_SAFE.encode("not-json"))] #[case::untagged_tuple(json(serde_json::json!([1, "ref"])))] - #[case::zero_start(run(0, "ref"))] - #[case::empty_ref(run(1, ""))] + #[case::other_key(run(BY_ERRORS, 1, "ref"))] + #[case::other_direction(run(RunOrder { descending: false, ..RunOrder::NEWEST }, 1, "ref"))] + #[case::empty_ref(run(RunOrder::NEWEST, 1, ""))] #[case::span_cursor(span())] #[case::text_cursor(text(SpanPart::Error, 0, "A".repeat(64)))] - #[case::unknown_field(json(serde_json::json!({"kind": "run", "position": {"start_ms": 1, "trace_ref": "r", "extra": 1}})))] - fn malformed_run_cursors_are_rejected(#[case] cursor: String) { + #[case::without_order(json(serde_json::json!({"kind": "run", "position": {"value": 1, "trace_ref": "r"}})))] + #[case::unknown_field(json(serde_json::json!({"kind": "run", "position": {"order": {"key": "start_ms", "descending": true}, "value": 1, "trace_ref": "r", "extra": 1}})))] + fn run_cursors_not_minted_under_the_requested_order_are_rejected(#[case] cursor: String) { assert!(matches!( - run_position::(Some(&cursor)), + run_position::(Some(&cursor), RunOrder::NEWEST), Err(ReadError::InvalidCursor("trace")) )); } @@ -182,7 +222,7 @@ mod tests { } #[rstest] - #[case::run_cursor(run(1, "ref"))] + #[case::run_cursor(run(RunOrder::NEWEST, 1, "ref"))] #[case::text_cursor(text(SpanPart::Error, 0, "a".repeat(64)))] fn other_kinds_are_not_span_cursors(#[case] cursor: String) { assert!(matches!( diff --git a/litellm-rust/crates/traces-cache/src/lib.rs b/litellm-rust/crates/traces-cache/src/lib.rs index 8e8863cc232..ece59fec654 100644 --- a/litellm-rust/crates/traces-cache/src/lib.rs +++ b/litellm-rust/crates/traces-cache/src/lib.rs @@ -9,5 +9,5 @@ mod store; pub use cache::{Freshness, LIVE_TTL, SETTLED_TTL, Snapshot, SnapshotCache, SnapshotKey}; pub use error::{Error, ReadError}; -pub use reader::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, PageRequest, TraceReader}; +pub use reader::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, MAX_TEXT_SPANS, PageRequest, TraceReader}; pub use store::{StoreError, StoreResult, TraceStore}; diff --git a/litellm-rust/crates/traces-cache/src/list.rs b/litellm-rust/crates/traces-cache/src/list.rs index c00ece2d1bd..d70a6339721 100644 --- a/litellm-rust/crates/traces-cache/src/list.rs +++ b/litellm-rust/crates/traces-cache/src/list.rs @@ -104,6 +104,7 @@ async fn resolve_runs( return Ok(Vec::new()); }; let selection = SpanSelection::Runs { + trace_ids: runs.iter().map(|row| row.trace_id.clone()).collect(), trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(), window: start_ms..end_ms.saturating_add(1), }; diff --git a/litellm-rust/crates/traces-cache/src/reader.rs b/litellm-rust/crates/traces-cache/src/reader.rs index 78d58a80f1b..9c94c63b60c 100644 --- a/litellm-rust/crates/traces-cache/src/reader.rs +++ b/litellm-rust/crates/traces-cache/src/reader.rs @@ -7,8 +7,8 @@ use litellm_traces::{ histogram, }, store::{ - CountBy, CountValue, RunCountQuery, RunQuery, RunSelection, SpanPart, SpanQuery, SpanRow, - SpanSelection, SpanText, SpanTextQuery, + CountBy, CountValue, RunCountQuery, RunOrder, RunQuery, RunSelection, SpanPart, SpanQuery, + SpanRow, SpanSelection, SpanText, SpanTextQuery, TextRange, }, to_ui_content, }; @@ -16,7 +16,9 @@ use litellm_traces::{ use crate::{ ReadError, Snapshot, SnapshotCache, SnapshotKey, StoreError, TraceStore, cache::{Freshness, ListCache}, - cursor::{Cursor, SpanPosition, TextPosition, run_position, span_position, text_position}, + cursor::{ + Cursor, RunPosition, SpanPosition, TextPosition, run_position, span_position, text_position, + }, list::{list_summaries, run_batches}, pages::read_all, spend::spend, @@ -27,6 +29,7 @@ pub const MAX_GRAPH_SPANS: usize = 100_000; const SNAPSHOT_IDLE: Duration = Duration::from_secs(120); const ERROR_PAGE_CHARS: u64 = 16_384; +pub const MAX_TEXT_SPANS: usize = 100; #[derive(Clone, Debug, Default)] pub struct PageRequest { @@ -77,33 +80,38 @@ impl TraceReader { store: &S, access: &QueryScope, filter: &RunFilter, + order: RunOrder, page: &PageRequest, ) -> Result> { if page.limit == 0 || filter.start_ms >= filter.end_ms { return Err(ReadError::InvalidParameters); } - let after = run_position(page.cursor.as_deref())?; + let after = run_position(page.cursor.as_deref(), order)?; let scope = SnapshotKey::scope(store.source(), access)?; let accepted = self.lists.limits.get(&scope).await.unwrap_or(u32::MAX); - let mut query = RunQuery { - selection: RunSelection::Matching(filter.clone()), - after, - limit: page.limit.min(500).min(accepted), - }; - let rows = loop { + let mut page_size = page.limit.min(500).min(accepted); + let mut rows = loop { + let query = RunQuery { + selection: RunSelection::Matching(filter.clone()), + order, + after: after.clone(), + limit: page_size + 1, + }; match store.runs(access, &query).await { - Err(StoreError::TooLarge) if query.limit > 1 => { - query.limit /= 2; - self.lists.limits.insert(scope.clone(), query.limit).await; + Err(StoreError::TooLarge) if page_size > 1 => { + page_size /= 2; + self.lists.limits.insert(scope.clone(), page_size).await; } Err(StoreError::TooLarge) => return Err(ReadError::TooLarge), result => break result.map_err(map_store_error)?, } }; + let more = rows.len() > page_size as usize; + rows.truncate(page_size as usize); let next_cursor = rows .last() - .filter(|_| rows.len() == query.limit as usize) - .map(|last| Cursor::Run(last.cursor()).encode()); + .filter(|_| more) + .map(|last| Cursor::Run(RunPosition::after(order, last)).encode()); let data = { let mut summaries = Vec::with_capacity(rows.len()); for batch in run_batches(&rows) { @@ -171,6 +179,61 @@ impl TraceReader { }) } + pub async fn count_traces( + &self, + store: &S, + access: &QueryScope, + filter: &RunFilter, + ) -> Result> { + if filter.start_ms >= filter.end_ms { + return Err(ReadError::InvalidParameters); + } + let query = RunCountQuery { + filter: filter.clone(), + by: CountBy::default(), + contains: String::new(), + limit: None, + }; + let counts = store + .run_counts(access, &query) + .await + .map_err(map_store_error)?; + Ok(counts.iter().map(|count| count.runs).sum()) + } + + /// One part of each listed span, for readers that page or search a run's text themselves. + #[expect(clippy::too_many_arguments, reason = "one argument per read dimension")] + pub async fn span_text( + &self, + store: &S, + access: &QueryScope, + trace_id: &str, + trace_ref: &str, + span_ids: Vec, + part: SpanPart, + range: TextRange, + contains: Option, + ) -> Result, ReadError> { + if span_ids.len() > MAX_TEXT_SPANS || trace_ref.is_empty() { + return Err(ReadError::InvalidParameters); + } + if span_ids.is_empty() { + return Ok(Vec::new()); + } + let query = SpanTextQuery { + trace_id: trace_id.to_owned(), + trace_ref: trace_ref.to_owned(), + span_ids, + part, + range, + contains, + }; + store + .span_text(access, &query) + .await + .map_err(map_store_error) + } + pub async fn get_trace( &self, store: &S, @@ -296,21 +359,21 @@ impl TraceReader { return Ok(None); }; let read = |part| { - let query = SpanTextQuery { - trace_id: trace_id.to_owned(), - trace_ref: trace_ref.clone(), - span_id: span_id.to_owned(), + span_texts( + store, + access, + trace_id, + &trace_ref, + vec![span_id.to_owned()], part, - offset: 0, - max_chars: None, - }; - async move { store.span_text(access, &query).await } + TextRange::ALL, + ) }; - let Some(input) = read(SpanPart::Input).await.map_err(map_store_error)? else { + let Some(input) = read(SpanPart::Input).await?.pop() else { return Ok(None); }; - let output = text_of(read(SpanPart::Output).await)?; - let attributes = text_of(read(SpanPart::Attributes).await)?; + let output = text_of(read(SpanPart::Output).await?); + let attributes = text_of(read(SpanPart::Attributes).await?); let output = if output.is_empty() { self.agent_answer(store, access, trace_id, &trace_ref, span_id) .await? @@ -345,28 +408,33 @@ impl TraceReader { { return Ok(String::new()); } - let mut calls: Vec<_> = spans + let calls: Vec<_> = spans .iter() .filter(|span| { span.kind == ObservationType::Llm && span.parent_span_id.as_deref() == Some(span_id) }) .collect(); - calls.sort_by(|left, right| right.start_offset_ms.total_cmp(&left.start_offset_ms)); - for call in calls { - let query = SpanTextQuery { - trace_id: trace_id.to_owned(), - trace_ref: trace_ref.to_owned(), - span_id: call.span_id.clone(), - part: SpanPart::Output, - offset: 0, - max_chars: None, - }; - let output = text_of(store.span_text(access, &query).await)?; - if !output.is_empty() { - return Ok(output); - } - } - Ok(String::new()) + let outputs = span_texts( + store, + access, + trace_id, + trace_ref, + calls.iter().map(|call| call.span_id.clone()).collect(), + SpanPart::Output, + TextRange::ALL, + ) + .await?; + Ok(calls + .iter() + .filter_map(|call| { + outputs + .iter() + .find(|output| output.span_id == call.span_id && !output.text.is_empty()) + .map(|output| (call.start_offset_ms, &output.text)) + }) + .max_by(|left, right| left.0.total_cmp(&right.0)) + .map(|(_, text)| text.clone()) + .unwrap_or_default()) } pub async fn get_span_error( @@ -383,19 +451,20 @@ impl TraceReader { return Ok(None); }; let offset = position.as_ref().map_or(0, |position| position.offset); - let query = SpanTextQuery { - trace_id: trace_id.to_owned(), - trace_ref, - span_id: span_id.to_owned(), - part: SpanPart::Error, - offset, - max_chars: Some(ERROR_PAGE_CHARS), - }; - let Some(text) = store - .span_text(access, &query) - .await - .map_err(map_store_error)? - else { + let Some(text) = span_texts( + store, + access, + trace_id, + &trace_ref, + vec![span_id.to_owned()], + SpanPart::Error, + TextRange::From { + offset, + max_chars: Some(ERROR_PAGE_CHARS), + }, + ) + .await? + .pop() else { return Ok(None); }; if position.is_some_and(|position| position.version != text.version) { @@ -419,11 +488,34 @@ impl TraceReader { } } -fn text_of(result: Result, StoreError>) -> Result> { - Ok(result - .map_err(map_store_error)? - .map(|text| text.text) - .unwrap_or_default()) +fn text_of(mut texts: Vec) -> String { + texts.pop().map(|text| text.text).unwrap_or_default() +} + +pub(super) async fn span_texts( + store: &S, + access: &QueryScope, + trace_id: &str, + trace_ref: &str, + span_ids: Vec, + part: SpanPart, + range: TextRange, +) -> Result, ReadError> { + if span_ids.is_empty() { + return Ok(Vec::new()); + } + let query = SpanTextQuery { + trace_id: trace_id.to_owned(), + trace_ref: trace_ref.to_owned(), + span_ids, + part, + range, + contains: None, + }; + store + .span_text(access, &query) + .await + .map_err(map_store_error) } fn parse_attributes(json: &str) -> Result, ReadError> { @@ -464,6 +556,7 @@ async fn reference( } let query = RunQuery { selection: RunSelection::TraceId(trace_id.to_owned()), + order: RunOrder::NEWEST, after: None, limit: 2, }; diff --git a/litellm-rust/crates/traces-cache/src/store.rs b/litellm-rust/crates/traces-cache/src/store.rs index 9fb7e96bf14..f24ba4f6da5 100644 --- a/litellm-rust/crates/traces-cache/src/store.rs +++ b/litellm-rust/crates/traces-cache/src/store.rs @@ -45,12 +45,11 @@ pub trait TraceStore: Sync { query: &SpanQuery, ) -> impl Future, Self::Error>> + Send; - /// `None` when the span is not visible to `access`. fn span_text( &self, access: &QueryScope, query: &SpanTextQuery, - ) -> impl Future, Self::Error>> + Send; + ) -> impl Future, Self::Error>> + Send; fn calls( &self, diff --git a/litellm-rust/crates/traces-cache/tests/read.rs b/litellm-rust/crates/traces-cache/tests/read.rs index 02c54e45049..39c0a103993 100644 --- a/litellm-rust/crates/traces-cache/tests/read.rs +++ b/litellm-rust/crates/traces-cache/tests/read.rs @@ -11,8 +11,9 @@ use litellm_traces::{ CallEvidenceKind, CallKey, ObservationType, QueryScope, SpanStatus, search::{AgentRuns, HistogramBucket, RunField, RunFilter, RunSearch}, store::{ - CallQuery, CallRow, CountBy, CountValue, RunCount, RunCountQuery, RunQuery, RunRow, - RunSelection, SpanPart, SpanQuery, SpanRow, SpanSelection, SpanText, SpanTextQuery, + CallQuery, CallRow, CountBy, CountValue, RunCount, RunCountQuery, RunOrder, RunQuery, + RunRow, RunSelection, SpanPart, SpanQuery, SpanRow, SpanSelection, SpanText, SpanTextQuery, + TextRange, }, }; use litellm_traces_cache::{ @@ -244,22 +245,38 @@ impl TraceStore for FakeStore { &self, _: &QueryScope, query: &SpanTextQuery, - ) -> StoreResult, FakeError> { + ) -> StoreResult, FakeError> { self.record(Operation::SpanText); let state = self.state.lock().unwrap(); Self::failure(&state, Operation::SpanText)?; - let Some(text) = state.texts.get(&(query.span_id.clone(), query.part)) else { - return Ok(None); - }; - let rest = text.chars().skip(query.offset as usize); - Ok(Some(SpanText { - text: match query.max_chars { - Some(max) => rest.take(max as usize).collect(), - None => rest.collect(), - }, - total_chars: text.chars().count() as u64, - version: format!("{:0>64}", text.len()), - })) + Ok(query + .span_ids + .iter() + .filter_map(|span_id| { + let text = state.texts.get(&(span_id.clone(), query.part))?; + let total = text.chars().count() as u64; + let (skip, take) = match query.range { + TextRange::From { offset, max_chars } => { + (offset, max_chars.unwrap_or(u64::MAX)) + } + TextRange::Last { chars } => (total.saturating_sub(chars), chars), + }; + Some(SpanText { + span_id: span_id.clone(), + text: text + .chars() + .skip(skip as usize) + .take(take as usize) + .collect(), + total_chars: total, + version: format!("{:0>64}", text.len()), + contains: query + .contains + .as_ref() + .is_some_and(|needle| text.contains(needle.as_str())), + }) + }) + .collect()) } async fn calls( @@ -294,6 +311,7 @@ fn everything() -> RunFilter { start_ms: 0, end_ms: i64::MAX, search: RunSearch::default(), + ..Default::default() } } @@ -302,6 +320,7 @@ fn window(start_ms: i64, end_ms: i64, q: &str) -> RunFilter { start_ms, end_ms, search: RunSearch::parse(q), + ..Default::default() } } @@ -309,6 +328,7 @@ fn newest(limit: u32) -> PageRequest { PageRequest { cursor: None, limit, + ..Default::default() } } @@ -508,34 +528,37 @@ async fn response_size_splits_pages_and_rejects_a_single_oversized_span() { #[rstest] #[tokio::test] -async fn list_run_budget_halves_the_limit_and_cursor_requires_a_full_page() { - let store = FakeStore::default(); - store.set_list_runs( - (0..3) +async fn list_run_budget_halves_the_limit_and_cursor_requires_a_run_past_the_page() { + let runs = |count: usize| { + (0..count) .map(|index| run(&format!("trace-{index}"), &format!("ref-{index}"))) - .collect(), - ); - store.set_list_runs_too_large_above(2); + .collect() + }; + let store = FakeStore::default(); + store.set_list_runs(runs(3)); + store.set_list_runs_too_large_above(3); let reader = TraceReader::new(usize::MAX); let access = access(); let page = reader - .list_traces(&store, &access, &everything(), &newest(8)) + .list_traces(&store, &access, &everything(), RunOrder::NEWEST, &newest(8)) .await .unwrap(); assert_eq!(page.data.len(), 2); assert!(page.next_cursor.is_some()); assert_eq!(store.calls(Operation::ListRuns), 3); - let shorter = FakeStore::default(); - shorter.set_list_runs(vec![run("only", "ref-only")]); - shorter.set_list_runs_too_large_above(2); - let page = reader - .list_traces(&shorter, &access, &everything(), &newest(8)) - .await - .unwrap(); - assert_eq!(page.data.len(), 1); - assert!(page.next_cursor.is_none()); - assert_eq!(shorter.calls(Operation::ListRuns), 1); + for (remaining, listed) in [(1, 1), (2, 2)] { + let rest = FakeStore::default(); + rest.set_list_runs(runs(remaining)); + rest.set_list_runs_too_large_above(3); + let page = reader + .list_traces(&rest, &access, &everything(), RunOrder::NEWEST, &newest(8)) + .await + .unwrap(); + assert_eq!(page.data.len(), listed); + assert!(page.next_cursor.is_none()); + assert_eq!(rest.calls(Operation::ListRuns), 1); + } } #[rstest] @@ -551,7 +574,13 @@ async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() store.set_failure(Operation::RunSpans, Failure::TooLarge); let reader = TraceReader::new(usize::MAX); let page = reader - .list_traces(&store, &access(), &everything(), &newest(2)) + .list_traces( + &store, + &access(), + &everything(), + RunOrder::NEWEST, + &newest(2), + ) .await .unwrap(); assert_eq!(page.data.len(), 2); @@ -565,7 +594,13 @@ async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() ); let again = reader - .list_traces(&store, &access(), &everything(), &newest(2)) + .list_traces( + &store, + &access(), + &everything(), + RunOrder::NEWEST, + &newest(2), + ) .await .unwrap(); assert_eq!(again.data, page.data); @@ -624,7 +659,13 @@ async fn failed_batch_spend_lookup_falls_back_to_each_run_instead_of_losing_ever store.set_spend_fails_above_response_ids(1); let page = TraceReader::new(usize::MAX) - .list_traces(&store, &access(), &everything(), &newest(8)) + .list_traces( + &store, + &access(), + &everything(), + RunOrder::NEWEST, + &newest(8), + ) .await .unwrap(); @@ -715,7 +756,7 @@ async fn invalid_list_reads_are_rejected_before_storage( let store = FakeStore::default(); assert!(matches!( TraceReader::new(usize::MAX) - .list_traces(&store, &access(), &filter, &newest(limit)) + .list_traces(&store, &access(), &filter, RunOrder::NEWEST, &newest(limit)) .await, Err(ReadError::InvalidParameters) )); @@ -936,7 +977,7 @@ async fn listed_runs_are_read_once_until_a_live_run_expires() { let access = access(); let filter = everything(); let page = newest(2); - let list = || reader.list_traces(&store, &access, &filter, &page); + let list = || reader.list_traces(&store, &access, &filter, RunOrder::NEWEST, &page); let first = list().await.unwrap(); assert!(first.data.iter().all(|summary| summary.name == "agent")); diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql b/litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql index b59a21f2d4f..25eacc740ac 100644 --- a/litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql +++ b/litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql @@ -1,2 +1,4 @@ ALTER TABLE {database}.agent_traces_by_key - ADD COLUMN IF NOT EXISTS AgentLabels SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [] + ADD COLUMN IF NOT EXISTS AgentLabels SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [], + ADD COLUMN IF NOT EXISTS AgentIdentities SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [], + ADD COLUMN IF NOT EXISTS Frameworks SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [] diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql b/litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql index c20373ef706..1fd7b1295a4 100644 --- a/litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql +++ b/litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql @@ -17,7 +17,9 @@ SELECT sum(OutputTokens) AS OutputTokens, groupUniqArrayIf(toString(Model), Model != '') AS Models, groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames, - groupUniqArrayIf(if(AgentName != '', AgentName, SpanName), ObservationType = 'agent' OR AgentName != '') AS AgentLabels, + groupUniqArrayIf(AgentName, AgentName != '') AS AgentLabels, + groupUniqArrayIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS AgentIdentities, + groupUniqArrayIf(toString(Framework), Framework != '') AS Frameworks, groupArrayIf(LiteLLMRequestId, ObservationType = 'llm' OR LiteLLMRequestId != '') AS RequestIds FROM {database}.otel_traces GROUP BY TeamId, ApiKeyHash, TraceId diff --git a/litellm-rust/crates/traces-clickhouse/query/lens_agents.sql b/litellm-rust/crates/traces-clickhouse/query/lens_agents.sql deleted file mode 100644 index fbdd578f8e7..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/lens_agents.sql +++ /dev/null @@ -1,6 +0,0 @@ -SELECT DISTINCT AgentName AS agent_name -FROM otel_traces -WHERE AgentName != '' - AND ({all_teams:UInt8}=1 OR TeamId={team:String}) - AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) -ORDER BY agent_name diff --git a/litellm-rust/crates/traces-clickhouse/query/lens_availability.sql b/litellm-rust/crates/traces-clickhouse/query/lens_availability.sql deleted file mode 100644 index 8d350dd1779..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/lens_availability.sql +++ /dev/null @@ -1,8 +0,0 @@ -SELECT - EXISTS(SELECT 1 FROM otel_traces - WHERE ({all_teams:UInt8}=1 OR TeamId={team:String}) - AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})) AS traces, - EXISTS(SELECT 1 FROM spend_logs - WHERE ({all_teams:UInt8}=1 OR team_id={team:String}) - AND ({key_hash:String}='' OR api_key={key_hash:String}) - AND NOT JSONExtractBool(metadata,'litellm_lens_internal')) AS requests diff --git a/litellm-rust/crates/traces-clickhouse/query/lens_content.sql b/litellm-rust/crates/traces-clickhouse/query/lens_content.sql deleted file mode 100644 index f0572796bd5..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/lens_content.sql +++ /dev/null @@ -1,35 +0,0 @@ -WITH greatest(toInt64({offset:UInt32})-1,1) AS content_offset, -(value, budget) -> if(lengthUTF8(value) <= budget, value, - concat(substringUTF8(value, 1, intDiv(budget, 3)), '\n[... content omitted ...]\n', - substringUTF8(value, -(budget - intDiv(budget, 3))))) AS excerpt -SELECT * FROM ( - SELECT SpanId AS span_id, ParentSpanId AS parent_span_id, SpanName AS name, - ObservationType AS kind, - if({offset:UInt32}=1 AND lengthUTF8(concat('Input: ',Input,'\nOutput: ',Output,'\nStatus: ',StatusCode,' ',StatusMessage))>8000, - concat('Input: ',excerpt(Input,2000),'\nOutput: ',excerpt(Output,5000), - '\nStatus: ',StatusCode,' ',excerpt(StatusMessage,500)), - substringUTF8(concat('Input: ',Input,'\nOutput: ',Output,'\nStatus: ',StatusCode,' ',StatusMessage), - content_offset,8000)) AS content, - lengthUTF8(concat('Input: ',Input,'\nOutput: ',Output,'\nStatus: ',StatusCode,' ',StatusMessage)) - >= content_offset+8000 AS truncated - FROM otel_traces WHERE {source:String}='traces' - AND ({all_teams:UInt8}=1 OR TeamId={team:String}) - AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) - AND ({trace_ref:String}='' OR hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId)))={trace_ref:String}) - AND TraceId={id:String} AND TeamId={record_team:String} AND SpanId > {cursor:String} - ORDER BY SpanId LIMIT 1 BY SpanId LIMIT 40 -) -UNION ALL -SELECT * FROM ( - SELECT request_id AS span_id, '' AS parent_span_id, model AS name, 'llm' AS kind, - if({offset:UInt32}=1 AND lengthUTF8(concat('Input: ',messages,'\nOutput: ',response,'\nError: ',error_str))>8000, - concat('Input: ',excerpt(messages,2000),'\nOutput: ',excerpt(response,5000),'\nError: ',excerpt(error_str,500)), - substringUTF8(concat('Input: ',messages,'\nOutput: ',response,'\nError: ',error_str), - content_offset,8000)) AS content, - lengthUTF8(concat('Input: ',messages,'\nOutput: ',response,'\nError: ',error_str)) - >= content_offset+8000 AS truncated - FROM spend_logs FINAL WHERE {source:String}='requests' - AND ({all_teams:UInt8}=1 OR team_id={team:String}) - AND ({key_hash:String}='' OR api_key={key_hash:String}) - AND request_id={id:String} AND team_id={record_team:String} LIMIT 1 -) diff --git a/litellm-rust/crates/traces-clickhouse/query/lens_evidence.sql b/litellm-rust/crates/traces-clickhouse/query/lens_evidence.sql deleted file mode 100644 index a0d600cdfde..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/lens_evidence.sql +++ /dev/null @@ -1,14 +0,0 @@ -SELECT sum(matches) AS count FROM ( - SELECT count() AS matches FROM otel_traces WHERE {source:String}='traces' - AND ({all_teams:UInt8}=1 OR TeamId={team:String}) - AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) - AND ({trace_ref:String}='' OR hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId)))={trace_ref:String}) - AND TraceId={id:String} AND TeamId={record_team:String} AND SpanId={span:String} - AND position(concat('Input: ',Input,'\nOutput: ',Output,'\nStatus: ',StatusCode,' ',StatusMessage),{quote:String})>0 - UNION ALL - SELECT count() AS matches FROM spend_logs FINAL WHERE {source:String}='requests' - AND ({all_teams:UInt8}=1 OR team_id={team:String}) - AND ({key_hash:String}='' OR api_key={key_hash:String}) - AND request_id={id:String} AND team_id={record_team:String} AND request_id={span:String} - AND position(concat('Input: ',messages,'\nOutput: ',response,'\nError: ',error_str),{quote:String})>0 -) diff --git a/litellm-rust/crates/traces-clickhouse/query/lens_sample.sql b/litellm-rust/crates/traces-clickhouse/query/lens_sample.sql deleted file mode 100644 index 92086c33c13..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/lens_sample.sql +++ /dev/null @@ -1,67 +0,0 @@ -WITH concat(leftPad(toString(cityHash64(concat(source,team_id,trace_ref,trace_id))),20,'0'), - hex(concat(source,char(0),team_id,char(0),trace_ref,char(0),trace_id))) AS selection_key -SELECT *, selection_key FROM ( - SELECT *, if({sample_cap:UInt64}=0, ceiling(eligible*{sample_percent:Float64}/100), - least(toFloat64({sample_cap:UInt64}),ceiling(eligible*{sample_percent:Float64}/100))) AS selected - FROM ( - SELECT *, count() OVER () AS eligible, - row_number() OVER (ORDER BY selection_key) AS position - FROM ( - SELECT 'traces' AS source, TraceId AS trace_id, TeamId AS team_id, hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref, - coalesce(nullIf(argMin(ResourceAttributes['run.name'], Timestamp), ''), - argMin(SpanName, Timestamp)) AS name, toString(min(Timestamp)) AS start_time, - uniqExact(SpanId) AS span_count, countIf(ParentSpanId='') > 0 AS root_seen, - argMin(ServiceName, Timestamp) AS service, - arrayZip(mapKeys(argMin(mapConcat(ResourceAttributes, SpanAttributes), tuple(ParentSpanId!='',Timestamp))), - mapValues(argMin(mapConcat(ResourceAttributes, SpanAttributes), tuple(ParentSpanId!='',Timestamp)))) AS attributes - FROM otel_traces - WHERE {source:String} IN ('traces','both') - AND ({all_teams:UInt8}=1 OR TeamId={team:String}) - AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) - AND (TeamId,ApiKeyHash,TraceId) IN ( - SELECT TeamId,ApiKeyHash,TraceId FROM otel_traces - WHERE ({all_teams:UInt8}=1 OR TeamId={team:String}) - AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) - AND if(EngineReceivedMs>0,toInt64(EngineReceivedMs), - toUnixTimestamp64Milli(Timestamp)+toInt64(intDiv(Duration,1000000))) >= {start:UInt64} - ) - GROUP BY TeamId,ApiKeyHash,TraceId - HAVING max(EngineReceivedMs) < {end:UInt64} - AND max(toUnixTimestamp64Milli(Timestamp)+toInt64(intDiv(Duration,1000000))) < {end:UInt64} - AND ({agent_name:String}='' OR countIf(AgentName={agent_name:String}) > 0) - AND countIf(arrayAll((k,v) -> ResourceAttributes[k]=v OR SpanAttributes[k]=v, - {filter_keys:Array(String)},{filter_values:Array(String)}) - AND ({service:String}='' OR ServiceName={service:String})) > 0 - UNION ALL - SELECT 'requests' AS source, request_id AS trace_id, team_id, '' AS trace_ref, model AS name, - toString(start_time) AS start_time, toUInt64(1) AS span_count, toUInt8(1) AS root_seen, - model_group AS service, - arrayConcat(JSONExtractKeysAndValues(metadata, 'requester_metadata', 'String'), - arrayMap(t -> tuple('tag', t), request_tags)) AS attributes - FROM spend_logs FINAL - WHERE {source:String} IN ('requests','both') - AND ({all_teams:UInt8}=1 OR team_id={team:String}) - AND ({key_hash:String}='' OR api_key={key_hash:String}) - AND if(EngineReceivedMs>0,toInt64(EngineReceivedMs),toUnixTimestamp64Milli(end_time)) >= {start:UInt64} - AND EngineReceivedMs < {end:UInt64} - AND toUnixTimestamp64Milli(end_time) < {end:UInt64} - AND arrayAll((k,v) -> JSONExtractString(metadata,k)=v - OR JSONExtractString(metadata,'requester_metadata',k)=v OR (k='tag' AND has(request_tags,v)), - {filter_keys:Array(String)},{filter_values:Array(String)}) - AND ({service:String}='' OR model_group={service:String}) - AND {agent_name:String}='' - AND NOT JSONExtractBool(metadata,'litellm_lens_internal') - AND ({source:String}!='both' OR (team_id,api_key,response_id) NOT IN ( - SELECT TeamId,ApiKeyHash,LiteLLMRequestId FROM otel_traces - WHERE ({all_teams:UInt8}=1 OR TeamId={team:String}) - AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String}) AND LiteLLMRequestId!='' - )) -) -WHERE ({selected_team:String}='' OR team_id={selected_team:String}) - AND (empty({execution_ids:Array(String)}) OR has({execution_ids:Array(String)}, - concat(source,char(0),team_id,char(0),if(trace_ref='',trace_id,trace_ref)))) -) -) -WHERE ({preview:UInt8}=1 OR position <= selected) - AND selection_key > {after:String} -ORDER BY selection_key LIMIT {limit:UInt32} OFFSET {offset:UInt64} diff --git a/litellm-rust/crates/traces-clickhouse/query/matching_runs.sql b/litellm-rust/crates/traces-clickhouse/query/matching_runs.sql index c033e7f1f60..14964d09495 100644 --- a/litellm-rust/crates/traces-clickhouse/query/matching_runs.sql +++ b/litellm-rust/crates/traces-clickhouse/query/matching_runs.sql @@ -4,7 +4,6 @@ SELECT TraceId AS trace_id, ifNull(any(RootName), '') AS name, any(ServiceName) AS service, ifNull(any(RootInput), '') AS input_preview, ifNull(any(RootStatus), '') AS status, toUnixTimestamp64Milli(min(StartTs)) AS start_ms, - min(StartTs) AS trace_start, max(EndTs) AS trace_end, dateDiff('millisecond', min(StartTs), max(EndTs)) AS duration_ms, sum(SpanCount) AS span_count, sum(AgentCount) AS agent_invocations, @@ -14,8 +13,20 @@ SELECT TraceId AS trace_id, arraySort(if(empty(groupUniqArrayArray(AgentLabels)), groupUniqArrayArray(AgentNames), groupUniqArrayArray(AgentLabels))) AS search_agents, - if(error_count > 0, 'error', 'ok') AS search_status + if(error_count > 0, 'error', 'ok') AS search_status, + length(groupUniqArrayArray(AgentIdentities)) AS agent_count, + arraySort(groupUniqArrayArray(Frameworks)) AS frameworks FROM owned_runs +LEFT JOIN ( + SELECT TeamId, ApiKeyHash, TraceId, + groupUniqArrayArray(arrayFilter(i -> ResourceAttributes[{attribute_keys:Array(String)}[i]] ILIKE {attribute_patterns:Array(String)}[i] + OR SpanAttributes[{attribute_keys:Array(String)}[i]] ILIKE {attribute_patterns:Array(String)}[i], + arrayEnumerate({attribute_keys:Array(String)}))) AS matched_attributes + FROM owned_spans + WHERE notEmpty({attribute_keys:Array(String)}) + AND Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64}) + GROUP BY TeamId, ApiKeyHash, TraceId +) AS attributes USING (TeamId, ApiKeyHash, TraceId) WHERE {trace_id:String} = '' OR TraceId = {trace_id:String} GROUP BY TeamId, ApiKeyHash, TraceId HAVING {trace_id:String} != '' @@ -29,5 +40,10 @@ HAVING {trace_id:String} != '' f = 'model', arrayExists(x -> x ILIKE p, models), f = 'input', input_preview ILIKE p, f = 'trace_id', trace_id ILIKE p, + f = 'service', service ILIKE p, + f = 'team', team_id ILIKE p, false), - {filter_fields:Array(String)}, {filter_patterns:Array(String)}, {filter_modes:Array(String)})) + {filter_fields:Array(String)}, {filter_patterns:Array(String)}, {filter_modes:Array(String)}) + AND arrayAll((i, m) -> (m = 'exclude') != has(any(matched_attributes), i), + arrayEnumerate({attribute_keys:Array(String)}), {attribute_modes:Array(String)}) + AND (empty({trace_refs:Array(String)}) OR trace_ref IN {trace_refs:Array(String)})) diff --git a/litellm-rust/crates/traces-clickhouse/query/run_attribute_counts.sql b/litellm-rust/crates/traces-clickhouse/query/run_attribute_counts.sql new file mode 100644 index 00000000000..f2c3d7969ae --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/run_attribute_counts.sql @@ -0,0 +1,19 @@ +SELECT bucket, failed, value, uniqExact(team_id, api_key_hash, trace_id) AS runs +FROM ( + SELECT runs.team_id AS team_id, runs.api_key_hash AS api_key_hash, runs.trace_id AS trace_id, + if({buckets:UInt32} = 0, toUInt32(0), + toUInt32(intDiv((runs.start_ms - {start_ms:Int64}) * {buckets:UInt32}, {end_ms:Int64} - {start_ms:Int64}))) AS bucket, + toUInt8({by_failed:UInt8} = 1 AND runs.error_count > 0) AS failed, + value + FROM owned_spans AS spans + INNER JOIN runs ON spans.TeamId = runs.team_id AND spans.ApiKeyHash = runs.api_key_hash + AND spans.TraceId = runs.trace_id + ARRAY JOIN if({attribute_key:String} = '', + arrayConcat(mapKeys(spans.ResourceAttributes), mapKeys(spans.SpanAttributes)), + [spans.ResourceAttributes[{attribute_key:String}], spans.SpanAttributes[{attribute_key:String}]]) AS value + WHERE spans.Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64}) + AND value != '' AND value ILIKE {contains:String} +) +GROUP BY bucket, failed, value +ORDER BY runs DESC, bucket, failed, value +LIMIT {limit:UInt64} diff --git a/litellm-rust/crates/traces-clickhouse/query/run_counts.sql b/litellm-rust/crates/traces-clickhouse/query/run_counts.sql index 91ff6974a53..00ad7b25d45 100644 --- a/litellm-rust/crates/traces-clickhouse/query/run_counts.sql +++ b/litellm-rust/crates/traces-clickhouse/query/run_counts.sql @@ -13,6 +13,8 @@ ARRAY JOIN multiIf( {value:String} = 'model', models, {value:String} = 'input', [input_preview], {value:String} = 'trace_id', [trace_id], + {value:String} = 'service', [service], + {value:String} = 'team', [team_id], []) AS value WHERE ({value:String} = '' OR value != '') AND value ILIKE {contains:String} GROUP BY bucket, failed, value diff --git a/litellm-rust/crates/traces-clickhouse/query/run_spans.sql b/litellm-rust/crates/traces-clickhouse/query/run_spans.sql index 65b06d9e3b8..993c9ecfb95 100644 --- a/litellm-rust/crates/traces-clickhouse/query/run_spans.sql +++ b/litellm-rust/crates/traces-clickhouse/query/run_spans.sql @@ -1,3 +1,4 @@ WHERE Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64}) AND Timestamp < fromUnixTimestamp64Milli({end_ms:Int64}) + AND TraceId IN {trace_ids:Array(String)} AND hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) IN {trace_refs:Array(String)} diff --git a/litellm-rust/crates/traces-clickhouse/query/runs.sql b/litellm-rust/crates/traces-clickhouse/query/runs.sql index 6cf5d5c9f70..b70d3561896 100644 --- a/litellm-rust/crates/traces-clickhouse/query/runs.sql +++ b/litellm-rust/crates/traces-clickhouse/query/runs.sql @@ -1,26 +1,2 @@ -page AS ( -SELECT * EXCEPT (search_agents, search_status) -FROM runs -WHERE {has_cursor:UInt8} = 0 OR (start_ms, trace_ref) < ({cursor_ms:Int64}, {cursor_ref:String}) -ORDER BY start_ms DESC, trace_ref DESC -LIMIT {limit:UInt32} -) -SELECT page.* EXCEPT (trace_start, trace_end), - identities.agent_names AS agent_names, identities.agent_count AS agent_count, - identities.frameworks AS frameworks +SELECT * EXCEPT (search_agents, sort_value), search_agents AS agent_names FROM page -LEFT JOIN ( - SELECT TeamId, ApiKeyHash, TraceId, - arraySort(groupUniqArrayIf(AgentName, AgentName != '')) AS agent_names, - arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks, - uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count - FROM owned_spans - WHERE Timestamp >= (SELECT min(trace_start) FROM page) - AND Timestamp <= (SELECT max(trace_end) FROM page) - AND TraceId IN (SELECT trace_id FROM page) - AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page) - GROUP BY TeamId, ApiKeyHash, TraceId -) AS identities -ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash - AND page.trace_id = identities.TraceId -ORDER BY page.start_ms DESC, page.trace_ref DESC diff --git a/litellm-rust/crates/traces-clickhouse/query/runs_page.sql b/litellm-rust/crates/traces-clickhouse/query/runs_page.sql new file mode 100644 index 00000000000..cd6e9d46192 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/runs_page.sql @@ -0,0 +1,16 @@ +page AS ( +SELECT * EXCEPT (search_status), + multiIf({sort_key:String} = 'duration_ms', duration_ms, + {sort_key:String} = 'span_count', toInt64(span_count), + {sort_key:String} = 'error_count', toInt64(error_count), + {sort_key:String} = 'trace_ref', toInt64(0), + start_ms) AS sort_value +FROM runs +WHERE {has_cursor:UInt8} = 0 + OR if({descending:UInt8} = 1, + (sort_value, trace_ref) < ({cursor_value:Int64}, {cursor_ref:String}), + (sort_value, trace_ref) > ({cursor_value:Int64}, {cursor_ref:String})) +ORDER BY if({descending:UInt8} = 1, sort_value, 0) DESC, if({descending:UInt8} = 1, trace_ref, '') DESC, + sort_value, trace_ref +LIMIT {limit:UInt32} +) diff --git a/litellm-rust/crates/traces-clickhouse/query/span_text.sql b/litellm-rust/crates/traces-clickhouse/query/span_text.sql index ebec7b82fab..824c4fbdf91 100644 --- a/litellm-rust/crates/traces-clickhouse/query/span_text.sql +++ b/litellm-rust/crates/traces-clickhouse/query/span_text.sql @@ -1,15 +1,19 @@ -SELECT if({bounded:UInt8} = 1, substringUTF8(part, {offset:UInt64} + 1, {max_chars:UInt64}), - substringUTF8(part, {offset:UInt64} + 1)) AS text, +SELECT span_id, + multiIf({range:String} = 'last', substringUTF8(part, toUInt64(greatest(toInt64(lengthUTF8(part)) - toInt64({chars:UInt64}), 0)) + 1), + {bounded:UInt8} = 1, substringUTF8(part, {offset:UInt64} + 1, {chars:UInt64}), + substringUTF8(part, {offset:UInt64} + 1)) AS text, lengthUTF8(part) AS total_chars, - hex(SHA256(part)) AS version + hex(SHA256(part)) AS version, + {needle:String} != '' AND position(part, {needle:String}) > 0 AS contains FROM ( - SELECT multiIf({part:String} = 'input', Input, + SELECT SpanId AS span_id, + multiIf({part:String} = 'input', Input, {part:String} = 'output', Output, {part:String} = 'error', StatusMessage, toJSONString(SpanAttributes)) AS part FROM owned_spans - WHERE TraceId = {trace_id:String} AND SpanId = {span_id:String} + WHERE TraceId = {trace_id:String} AND SpanId IN {span_ids:Array(String)} AND hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String} ORDER BY Timestamp, EngineReceivedMs, StatusMessage - LIMIT 1 + LIMIT 1 BY SpanId ) diff --git a/litellm-rust/crates/traces-clickhouse/src/access.rs b/litellm-rust/crates/traces-clickhouse/src/access.rs index 941b1358dd6..4320fe4e125 100644 --- a/litellm-rust/crates/traces-clickhouse/src/access.rs +++ b/litellm-rust/crates/traces-clickhouse/src/access.rs @@ -15,6 +15,14 @@ macro_rules! owned_by { }; } +/// Rollup rows of one run merge in the background, so a user-owned row can belong to a run that +/// other users also wrote to. Trusted reads drop such runs; a row policy cannot express this. +macro_rules! whole_runs { + () => { + "({access_all:UInt8} = 1 OR has({access_teams:Array(String)}, TeamId) OR (TeamId, ApiKeyHash, TraceId) NOT IN (SELECT TeamId, ApiKeyHash, TraceId FROM agent_traces_by_key WHERE {access_user:String} != '' AND UserIds != [{access_user:String}]))" + }; +} + /// Prefixes a read with the rows its caller may see: `owned_spans`, `owned_runs` and /// `owned_calls`. Trusted SQL reads only these, never the tables. macro_rules! owned { @@ -24,6 +32,8 @@ macro_rules! owned { $crate::access::owned_by!(otel_traces), "),\nowned_runs AS (SELECT * FROM agent_traces_by_key WHERE ", $crate::access::owned_by!(agent_traces_by_key), + " AND ", + $crate::access::whole_runs!(), "),\nowned_calls AS (SELECT * FROM spend_logs FINAL WHERE ", $crate::access::owned_by!(spend_logs), ")", @@ -34,6 +44,7 @@ macro_rules! owned { pub(crate) use owned; pub(crate) use owned_by; +pub(crate) use whole_runs; #[derive(Debug, Serialize)] pub(crate) struct AccessParams { diff --git a/litellm-rust/crates/traces-clickhouse/src/error.rs b/litellm-rust/crates/traces-clickhouse/src/error.rs index 1d15c556316..3aee1f336a8 100644 --- a/litellm-rust/crates/traces-clickhouse/src/error.rs +++ b/litellm-rust/crates/traces-clickhouse/src/error.rs @@ -8,8 +8,6 @@ pub enum Error { InvalidTable, #[error("database must be a nonempty SQL identifier and retention must be positive")] InvalidSchema, - #[error("unknown ClickHouse read query")] - InvalidQuery, #[error("invalid ClickHouse query parameters")] InvalidParameters, #[error("ClickHouse returned an invalid or failed JSON query response")] diff --git a/litellm-rust/crates/traces-clickhouse/src/lib.rs b/litellm-rust/crates/traces-clickhouse/src/lib.rs index f44116b8c86..1eec998b0a1 100644 --- a/litellm-rust/crates/traces-clickhouse/src/lib.rs +++ b/litellm-rust/crates/traces-clickhouse/src/lib.rs @@ -19,7 +19,6 @@ mod query_access; mod reads; mod schema; mod span_row; -mod sql; mod table; #[cfg(feature = "schema")] pub mod wire_schema; @@ -28,7 +27,7 @@ pub use config::Config; pub use error::Error; pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows}; pub use litellm_storage_clickhouse::{Connection, Parameter}; -pub use litellm_traces::{QueryScope, ReadQuery}; +pub use litellm_traces::QueryScope; pub use query::{QueryHelp, execute_read, query_help, query_sql}; pub use query_access::QueryReaders; pub use reads::ClickHouseTraces; @@ -36,5 +35,4 @@ pub use schema::{ NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements, }; pub use span_row::span_rows; -pub use sql::execute_named_read; pub use table::TraceTable; diff --git a/litellm-rust/crates/traces-clickhouse/src/query.rs b/litellm-rust/crates/traces-clickhouse/src/query.rs index 98603ca142c..5c027e8ef6f 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query.rs @@ -17,7 +17,6 @@ use super::{ use crate::TraceTable; mod guide; -pub mod lens; pub mod named; mod number; diff --git a/litellm-rust/crates/traces-clickhouse/src/query/lens.rs b/litellm-rust/crates/traces-clickhouse/src/query/lens.rs deleted file mode 100644 index ff30f127000..00000000000 --- a/litellm-rust/crates/traces-clickhouse/src/query/lens.rs +++ /dev/null @@ -1,268 +0,0 @@ -use litellm_storage_clickhouse::Query; - -pub const LENS_QUERIES: [litellm_traces::ReadQuery; 5] = [ - litellm_traces::ReadQuery::Availability, - litellm_traces::ReadQuery::Agents, - litellm_traces::ReadQuery::Sample, - litellm_traces::ReadQuery::Content, - litellm_traces::ReadQuery::Evidence, -]; - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[serde(rename_all = "lowercase")] -pub enum ExecutionSource { - Traces, - Requests, - Both, -} - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[serde(rename_all = "lowercase")] -pub enum ContentSource { - Traces, - Requests, -} - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))] -pub struct LensAccessParams { - #[serde( - deserialize_with = "super::number::boolean", - serialize_with = "litellm_traces::wire::serialize_flag" - )] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "litellm_traces::schema::flag") - )] - pub all_teams: bool, - pub team: String, - pub key_hash: String, -} - -pub struct LensAvailability; - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[serde(deny_unknown_fields)] -pub struct LensAvailabilityParams { - #[serde(flatten)] - pub access: LensAccessParams, -} - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[cfg_attr(feature = "schema", schemars(rename = "ActivityAvailability"))] -pub struct LensAvailabilityRow { - #[serde(default, deserialize_with = "super::number::flag")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::boolean_flag") - )] - pub traces: u8, - #[serde(default, deserialize_with = "super::number::flag")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::boolean_flag") - )] - pub requests: u8, -} - -impl Query for LensAvailability { - type Params = LensAvailabilityParams; - type Row = LensAvailabilityRow; - - const SQL: &'static str = include_str!("../../query/lens_availability.sql"); -} - -pub struct LensAgents; - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[serde(deny_unknown_fields)] -pub struct LensAgentsParams { - #[serde(flatten)] - pub access: LensAccessParams, -} - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[cfg_attr(feature = "schema", schemars(rename = "AgentRow"))] -pub struct LensAgentsRow { - pub agent_name: String, -} - -impl Query for LensAgents { - type Params = LensAgentsParams; - type Row = LensAgentsRow; - - const SQL: &'static str = include_str!("../../query/lens_agents.sql"); -} - -pub struct LensSample; - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[serde(deny_unknown_fields)] -pub struct LensSampleParams { - #[serde(flatten)] - pub access: LensAccessParams, - pub source: ExecutionSource, - #[serde(deserialize_with = "super::number::deserialize")] - pub start: u64, - #[serde(deserialize_with = "super::number::deserialize")] - pub end: u64, - pub agent_name: String, - pub service: String, - pub filter_keys: Vec, - pub filter_values: Vec, - pub selected_team: String, - pub execution_ids: Vec, - #[serde(deserialize_with = "super::number::deserialize")] - pub sample_cap: u64, - #[serde(deserialize_with = "super::number::percent")] - #[cfg_attr(feature = "schema", schemars(range(min = 0, max = 100)))] - pub sample_percent: f64, - #[serde(deserialize_with = "super::number::flag")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "litellm_traces::schema::flag") - )] - pub preview: u8, - pub after: String, - #[serde(deserialize_with = "super::number::deserialize")] - pub limit: u32, - #[serde(deserialize_with = "super::number::deserialize")] - pub offset: u64, -} - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[cfg_attr(feature = "schema", schemars(rename = "ExecutionRow"))] -pub struct LensSampleRow { - pub source: ContentSource, - pub trace_id: String, - pub team_id: String, - #[serde(default)] - pub trace_ref: String, - pub name: String, - pub start_time: String, - #[serde(deserialize_with = "super::number::deserialize")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::u64_number") - )] - pub span_count: u64, - #[serde(deserialize_with = "super::number::flag")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::flag_number") - )] - pub root_seen: u8, - #[serde(default)] - pub service: String, - #[serde(default)] - pub attributes: Vec<(String, String)>, - #[serde(deserialize_with = "super::number::deserialize")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::u64_number") - )] - pub eligible: u64, - #[serde(deserialize_with = "super::number::deserialize")] - #[cfg_attr(feature = "schema", schemars(skip))] - pub position: u64, - #[serde(default, deserialize_with = "super::number::deserialize")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::selected") - )] - pub selected: f64, - #[serde(default)] - pub selection_key: String, -} - -impl Query for LensSample { - type Params = LensSampleParams; - type Row = LensSampleRow; - - const SQL: &'static str = include_str!("../../query/lens_sample.sql"); -} - -pub struct LensContent; - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[serde(deny_unknown_fields)] -pub struct LensContentParams { - #[serde(flatten)] - pub access: LensAccessParams, - pub source: ContentSource, - pub id: String, - pub record_team: String, - pub trace_ref: String, - pub cursor: String, - #[serde(deserialize_with = "super::number::deserialize")] - pub offset: u32, -} - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[cfg_attr(feature = "schema", schemars(rename = "PartRow"))] -pub struct LensContentRow { - pub span_id: String, - pub parent_span_id: String, - pub name: String, - pub kind: String, - pub content: String, - #[serde(deserialize_with = "super::number::flag")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::flag_number") - )] - pub truncated: u8, -} - -impl Query for LensContent { - type Params = LensContentParams; - type Row = LensContentRow; - - const SQL: &'static str = include_str!("../../query/lens_content.sql"); -} - -pub struct LensEvidence; - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[serde(deny_unknown_fields)] -pub struct LensEvidenceParams { - #[serde(flatten)] - pub access: LensAccessParams, - pub source: ContentSource, - pub id: String, - pub record_team: String, - pub trace_ref: String, - pub span: String, - pub quote: String, -} - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Debug)] -#[cfg_attr(feature = "schema", schemars(rename = "CountRow"))] -pub struct LensEvidenceRow { - #[serde(deserialize_with = "super::number::deserialize")] - #[cfg_attr( - feature = "schema", - schemars(schema_with = "crate::wire_schema::u64_number") - )] - pub count: u64, -} - -impl Query for LensEvidence { - type Params = LensEvidenceParams; - type Row = LensEvidenceRow; - - const SQL: &'static str = include_str!("../../query/lens_evidence.sql"); -} diff --git a/litellm-rust/crates/traces-clickhouse/src/query/named.rs b/litellm-rust/crates/traces-clickhouse/src/query/named.rs index 5929c9765c2..4b94365f181 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/named.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/named.rs @@ -1,10 +1,10 @@ use litellm_storage_clickhouse::Query; use litellm_traces::{ QueryScope, - search::{RunFilter, RunSearch}, + search::{FieldFilter, RunFilter, RunSearch, SearchKey}, store::{ CallQuery, CallRow, CountValue, RunCount, RunCountQuery, RunQuery, RunRow, RunSelection, - SpanQuery, SpanRow, SpanSelection, SpanText, SpanTextQuery, + RunSortKey, SpanQuery, SpanRow, SpanSelection, SpanText, SpanTextQuery, TextRange, }, }; use serde::{Deserialize, Serialize}; @@ -26,6 +26,14 @@ fn contains(value: &str) -> String { format!("%{}%", like_literal(value)) } +fn pattern(filter: &FieldFilter) -> String { + like_literal(&filter.pattern).replace('*', "%") +} + +fn mode(filter: &FieldFilter) -> &'static str { + if filter.exclude { "exclude" } else { "include" } +} + /// Query parameters only carry string arrays, so filters travel as parallel columns. #[derive(Debug, Default, Serialize)] struct SearchColumns { @@ -33,27 +41,40 @@ struct SearchColumns { filter_fields: Vec<&'static str>, filter_patterns: Vec, filter_modes: Vec<&'static str>, + attribute_keys: Vec, + attribute_patterns: Vec, + attribute_modes: Vec<&'static str>, } impl From<&RunSearch> for SearchColumns { fn from(search: &RunSearch) -> Self { + let fields: Vec<_> = search + .filters + .iter() + .filter_map(|filter| match &filter.key { + SearchKey::Field(field) => Some(((*field).into(), filter)), + SearchKey::Attribute(_) => None, + }) + .collect(); + let attributes: Vec<_> = search + .filters + .iter() + .filter_map(|filter| match &filter.key { + SearchKey::Attribute(key) => Some((key.clone(), filter)), + SearchKey::Field(_) => None, + }) + .collect(); Self { text: search.text.iter().map(|term| contains(term)).collect(), - filter_fields: search - .filters + filter_fields: fields.iter().map(|(field, _)| *field).collect(), + filter_patterns: fields.iter().map(|(_, filter)| pattern(filter)).collect(), + filter_modes: fields.iter().map(|(_, filter)| mode(filter)).collect(), + attribute_patterns: attributes .iter() - .map(|filter| filter.field.into()) - .collect(), - filter_patterns: search - .filters - .iter() - .map(|filter| like_literal(&filter.pattern).replace('*', "%")) - .collect(), - filter_modes: search - .filters - .iter() - .map(|filter| if filter.exclude { "exclude" } else { "include" }) + .map(|(_, filter)| pattern(filter)) .collect(), + attribute_modes: attributes.iter().map(|(_, filter)| mode(filter)).collect(), + attribute_keys: attributes.into_iter().map(|(key, _)| key).collect(), } } } @@ -65,6 +86,7 @@ struct RunsFilter { end_ms: i64, #[serde(flatten)] search: SearchColumns, + trace_refs: Vec, } impl From<&RunFilter> for RunsFilter { @@ -74,6 +96,7 @@ impl From<&RunFilter> for RunsFilter { start_ms: filter.start_ms, end_ms: filter.end_ms, search: (&filter.search).into(), + trace_refs: filter.trace_refs.clone(), } } } @@ -96,8 +119,10 @@ pub(crate) struct RunsParams { access: AccessParams, #[serde(flatten)] filter: RunsFilter, + sort_key: RunSortKey, + descending: u8, has_cursor: u8, - cursor_ms: i64, + cursor_value: i64, cursor_ref: String, limit: u32, } @@ -108,8 +133,10 @@ impl RunsParams { Self { access: access.into(), filter: (&query.selection).into(), + sort_key: query.order.key, + descending: query.order.descending.into(), has_cursor: query.after.is_some().into(), - cursor_ms: after.start_ms, + cursor_value: after.value, cursor_ref: after.trace_ref, limit: query.limit, } @@ -170,13 +197,17 @@ macro_rules! over_matching_runs { }; } -pub(crate) struct Runs; +pub(crate) struct RunsPage; -impl Query for Runs { +impl Query for RunsPage { type Params = RunsParams; type Row = RunRowWire; - const SQL: &'static str = over_matching_runs!(",\n", include_str!("../../query/runs.sql")); + const SQL: &'static str = over_matching_runs!( + ",\n", + include_str!("../../query/runs_page.sql"), + include_str!("../../query/runs.sql") + ); } #[derive(Debug, Serialize)] @@ -188,6 +219,7 @@ pub(crate) struct RunCountsParams { buckets: u32, by_failed: u8, value: &'static str, + attribute_key: String, contains: String, limit: u64, } @@ -199,10 +231,15 @@ impl RunCountsParams { filter: (&query.filter).into(), buckets: query.by.buckets.unwrap_or(0), by_failed: query.by.failed.into(), - value: match query.by.value { + value: match &query.by.value { None => "", Some(CountValue::PrimaryAgent) => "primary_agent", - Some(CountValue::Field(field)) => field.into(), + Some(CountValue::Field(field)) => (*field).into(), + Some(CountValue::AttributeKey | CountValue::Attribute(_)) => "attribute", + }, + attribute_key: match &query.by.value { + Some(CountValue::Attribute(key)) => key.clone(), + _ => String::new(), }, contains: contains(&query.contains), limit: query.limit.map_or(u64::MAX, u64::from), @@ -228,6 +265,12 @@ struct RunCountEncoding { #[derive(Debug, Deserialize, Serialize)] pub(crate) struct RunCountRow(#[serde(with = "RunCountEncoding")] pub RunCount); +impl RunCountsParams { + pub(crate) fn counts_attributes(&self) -> bool { + self.value == "attribute" + } +} + pub(crate) struct RunCounts; impl Query for RunCounts { @@ -237,6 +280,16 @@ impl Query for RunCounts { const SQL: &'static str = over_matching_runs!("\n", include_str!("../../query/run_counts.sql")); } +pub(crate) struct RunAttributeCounts; + +impl Query for RunAttributeCounts { + type Params = RunCountsParams; + type Row = RunCountRow; + + const SQL: &'static str = + over_matching_runs!("\n", include_str!("../../query/run_attribute_counts.sql")); +} + #[derive(Debug, Default, Serialize)] struct SpanKeyset { as_of_ms: u64, @@ -275,6 +328,7 @@ pub(crate) struct TraceSpansParams { pub(crate) struct RunSpansParams { #[serde(flatten)] access: AccessParams, + trace_ids: Vec, trace_refs: Vec, start_ms: i64, end_ms: i64, @@ -300,8 +354,13 @@ impl SpansParams { trace_ref: trace_ref.clone(), keyset, }), - SpanSelection::Runs { trace_refs, window } => Self::Runs(RunSpansParams { + SpanSelection::Runs { + trace_ids, + trace_refs, + window, + } => Self::Runs(RunSpansParams { access: access.into(), + trace_ids: trace_ids.clone(), trace_refs: trace_refs.clone(), start_ms: window.start, end_ms: window.end, @@ -403,24 +462,32 @@ pub(crate) struct SpanTextParams { access: AccessParams, trace_id: String, trace_ref: String, - span_id: String, + span_ids: Vec, part: &'static str, + range: &'static str, offset: u64, bounded: u8, - max_chars: u64, + chars: u64, + needle: String, } impl SpanTextParams { pub(crate) fn new(access: &QueryScope, query: &SpanTextQuery) -> Self { + let (range, offset, chars) = match query.range { + TextRange::From { offset, max_chars } => ("from", offset, max_chars), + TextRange::Last { chars } => ("last", 0, Some(chars)), + }; Self { access: access.into(), trace_id: query.trace_id.clone(), trace_ref: query.trace_ref.clone(), - span_id: query.span_id.clone(), + span_ids: query.span_ids.clone(), part: query.part.into(), - offset: query.offset, - bounded: query.max_chars.is_some().into(), - max_chars: query.max_chars.unwrap_or(0), + range, + offset, + bounded: chars.is_some().into(), + chars: chars.unwrap_or(0), + needle: query.contains.clone().unwrap_or_default(), } } } @@ -428,10 +495,16 @@ impl SpanTextParams { #[derive(Deserialize, Serialize)] #[serde(remote = "SpanText")] struct SpanTextEncoding { + pub span_id: String, pub text: String, #[serde(deserialize_with = "super::number::deserialize")] pub total_chars: u64, pub version: String, + #[serde( + deserialize_with = "super::number::boolean", + serialize_with = "litellm_traces::wire::serialize_flag" + )] + pub contains: bool, } #[derive(Debug, Deserialize, Serialize)] @@ -513,7 +586,7 @@ impl Query for Calls { #[cfg(test)] mod tests { - use litellm_traces::search::{FieldFilter, RunField}; + use litellm_traces::search::RunField; use rstest::rstest; use serde_json::{Value, json}; @@ -539,20 +612,37 @@ mod tests { text: Vec::new(), filters: vec![ FieldFilter { - field: RunField::Model, + key: SearchKey::Field(RunField::Model), pattern: pattern.into(), exclude: true, }, FieldFilter { - field: RunField::Agent, + key: SearchKey::Attribute("tenant.tier".into()), pattern: "x".into(), exclude: false, }, ], }); - assert_eq!(columns.filter_fields, ["model", "agent"]); - assert_eq!(columns.filter_patterns, [like, "x"]); - assert_eq!(columns.filter_modes, ["exclude", "include"]); + assert_eq!( + ( + columns.filter_fields, + columns.filter_patterns, + columns.filter_modes + ), + (vec!["model"], vec![like.to_owned()], vec!["exclude"]) + ); + assert_eq!( + ( + columns.attribute_keys, + columns.attribute_patterns, + columns.attribute_modes + ), + ( + vec!["tenant.tier".to_owned()], + vec!["x".to_owned()], + vec!["include"] + ) + ); } fn decoded(wire: Value, quoted: bool) -> Value { @@ -583,7 +673,7 @@ mod tests { assert_eq!(decoded::(span.clone(), quoted), span); let count = json!({"bucket": 2, "failed": 1, "value": "v", "runs": u64::MAX}); assert_eq!(decoded::(count.clone(), quoted), count); - let text = json!({"text": "error", "total_chars": u64::MAX, "version": "version"}); + let text = json!({"span_id": "span", "text": "error", "total_chars": u64::MAX, "version": "version", "contains": 1}); assert_eq!(decoded::(text.clone(), quoted), text); let call = json!({"request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "upstream", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}); assert_eq!(decoded::(call.clone(), quoted), call); diff --git a/litellm-rust/crates/traces-clickhouse/src/query/number.rs b/litellm-rust/crates/traces-clickhouse/src/query/number.rs index 27c88bd3768..8f7ee364821 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/number.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/number.rs @@ -40,17 +40,6 @@ pub(super) fn flag<'de, D: Deserializer<'de>>(deserializer: D) -> Result>(deserializer: D) -> Result { - let value: f64 = deserialize(deserializer)?; - if value.is_finite() && (0.0..=100.0).contains(&value) { - Ok(value) - } else { - Err(serde::de::Error::custom( - "expected a finite percentage between 0 and 100", - )) - } -} - pub(super) fn boolean<'de, D: Deserializer<'de>>(deserializer: D) -> Result { flag(deserializer).map(|value| value == 1) } @@ -61,53 +50,6 @@ mod tests { use crate::query::named::SpanTextRow; - #[rstest] - #[case::flag_zero(serde_json::json!(0), true)] - #[case::flag_one(serde_json::json!("1"), true)] - #[case::invalid_flag(serde_json::json!(2), false)] - fn access_rejects_non_boolean_flags(#[case] value: serde_json::Value, #[case] valid: bool) { - let parameters = serde_json::json!({"all_teams": value, "team": "team", "key_hash": ""}); - assert_eq!( - serde_json::from_value::(parameters).is_ok(), - valid - ); - } - - #[rstest] - #[case::zero(serde_json::json!(0), true)] - #[case::hundred(serde_json::json!("100"), true)] - #[case::negative(serde_json::json!(-0.1), false)] - #[case::too_large(serde_json::json!(100.1), false)] - #[case::nan(serde_json::json!("NaN"), false)] - fn sampling_rejects_invalid_percentages(#[case] value: serde_json::Value, #[case] valid: bool) { - let parameters = serde_json::json!({ - "all_teams": 0, "team": "team", "key_hash": "", "source": "both", "start": 0, "end": 1, - "agent_name": "", "service": "", "filter_keys": [], "filter_values": [], "selected_team": "", - "execution_ids": [], "sample_cap": 0, "sample_percent": value, "preview": 0, "after": "", - "limit": 10, "offset": 0 - }); - assert_eq!( - serde_json::from_value::(parameters).is_ok(), - valid - ); - } - - #[rstest] - #[case::trace("traces", true)] - #[case::request("requests", true)] - #[case::both("both", false)] - #[case::unknown("unknown", false)] - fn content_rejects_unsupported_sources(#[case] source: &str, #[case] valid: bool) { - let parameters = serde_json::json!({ - "all_teams": 0, "team": "team", "key_hash": "", "source": source, "id": "id", - "record_team": "team", "trace_ref": "", "cursor": "", "offset": 0 - }); - assert_eq!( - serde_json::from_value::(parameters).is_ok(), - valid - ); - } - #[rstest] #[case::quoted_max(serde_json::json!(u64::MAX.to_string()), Some(u64::MAX))] #[case::unquoted_max(serde_json::json!(u64::MAX), Some(u64::MAX))] @@ -119,7 +61,7 @@ mod tests { #[case] expected: Option, ) { let row = serde_json::from_value::(serde_json::json!({ - "text": "error", "total_chars": value, "version": "hash" + "span_id": "span", "text": "error", "total_chars": value, "version": "hash", "contains": 0 })); match expected { Some(value) => assert_eq!(row.unwrap().0.total_chars, value), diff --git a/litellm-rust/crates/traces-clickhouse/src/reads.rs b/litellm-rust/crates/traces-clickhouse/src/reads.rs index 32562972ee7..e71253ddf9d 100644 --- a/litellm-rust/crates/traces-clickhouse/src/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/src/reads.rs @@ -12,8 +12,8 @@ use litellm_traces_cache::{StoreError, StoreResult, TraceStore}; use crate::{ Connection, Error, query::named::{ - Calls, CallsParams, RunCounts, RunCountsParams, RunSpans, Runs, RunsParams, SpanTextParams, - SpanTexts, SpansParams, TraceSpans, + Calls, CallsParams, RunAttributeCounts, RunCounts, RunCountsParams, RunSpans, RunsPage, + RunsParams, SpanTextParams, SpanTexts, SpansParams, TraceSpans, }, }; @@ -45,8 +45,14 @@ impl TraceStore for ClickHouseTraces { } async fn runs(&self, access: &QueryScope, query: &RunQuery) -> StoreResult, Error> { - let rows = self.fetch::(&RunsParams::new(access, query)).await?; - Ok(rows.into_iter().map(|row| row.0).collect()) + let mut rows: Vec = self + .fetch::(&RunsParams::new(access, query)) + .await? + .into_iter() + .map(|row| row.0) + .collect(); + rows.sort_by(|left, right| query.order.compare(left, right)); + Ok(rows) } async fn run_counts( @@ -54,9 +60,12 @@ impl TraceStore for ClickHouseTraces { access: &QueryScope, query: &RunCountQuery, ) -> StoreResult, Error> { - let rows = self - .fetch::(&RunCountsParams::new(access, query)) - .await?; + let params = RunCountsParams::new(access, query); + let rows = if params.counts_attributes() { + self.fetch::(¶ms).await? + } else { + self.fetch::(¶ms).await? + }; Ok(rows.into_iter().map(|row| row.0).collect()) } @@ -76,11 +85,11 @@ impl TraceStore for ClickHouseTraces { &self, access: &QueryScope, query: &SpanTextQuery, - ) -> StoreResult, Error> { + ) -> StoreResult, Error> { let rows = self .fetch::(&SpanTextParams::new(access, query)) .await?; - Ok(rows.into_iter().next().map(|row| row.0)) + Ok(rows.into_iter().map(|row| row.0).collect()) } async fn calls( diff --git a/litellm-rust/crates/traces-clickhouse/src/sql.rs b/litellm-rust/crates/traces-clickhouse/src/sql.rs deleted file mode 100644 index 40eb278c76c..00000000000 --- a/litellm-rust/crates/traces-clickhouse/src/sql.rs +++ /dev/null @@ -1,74 +0,0 @@ -use std::collections::BTreeMap; - -use litellm_http::Client; -use litellm_storage_clickhouse::{Query, fetch_json}; -use litellm_traces::ReadQuery; - -use super::{Connection, Error, Parameter, query::lens::*}; - -pub async fn execute_named_read( - client: &Client, - connection: &Connection, - query: ReadQuery, - parameters: &BTreeMap, -) -> Result { - match query { - ReadQuery::Availability => { - named_json::(client, connection, parameters).await - } - ReadQuery::Agents => named_json::(client, connection, parameters).await, - ReadQuery::Sample => named_json::(client, connection, parameters).await, - ReadQuery::Content => named_json::(client, connection, parameters).await, - ReadQuery::Evidence => named_json::(client, connection, parameters).await, - } -} - -async fn named_json( - client: &Client, - connection: &Connection, - parameters: &BTreeMap, -) -> Result -where - Q::Params: serde::de::DeserializeOwned, -{ - let value = serde_json::to_value(parameters).map_err(|_| Error::InvalidParameters)?; - let params = - serde_json::from_value::(value).map_err(|_| Error::InvalidParameters)?; - fetch_json::(client, connection, ¶ms) - .await - .map_err(Error::from) -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - - use super::*; - - #[rstest] - #[case::missing_cursor(serde_json::json!({"offset": 1}))] - #[case::negative_offset(serde_json::json!({"cursor": "", "offset": -1}))] - #[case::overflow(serde_json::json!({"cursor": "", "offset": "4294967296"}))] - #[tokio::test] - async fn named_read_rejects_invalid_parameters_before_transport( - #[case] specific: serde_json::Value, - ) { - let common = serde_json::json!({ - "all_teams": 1, "team": "", "key_hash": "", "source": "traces", "id": "trace", - "record_team": "team", "trace_ref": "" - }); - let parameters: BTreeMap = common - .as_object() - .unwrap() - .iter() - .chain(specific.as_object().unwrap().iter()) - .map(|(name, value)| (name.clone(), serde_json::from_value(value.clone()).unwrap())) - .collect(); - let client = Client::no_redirect_for_test(); - let connection = Connection::parse("http://127.0.0.1:1").unwrap(); - assert!(matches!( - execute_named_read(&client, &connection, ReadQuery::Content, ¶meters).await, - Err(Error::InvalidParameters) - )); - } -} diff --git a/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs b/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs index 82473f1727f..c5ace2535c1 100644 --- a/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs +++ b/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs @@ -1,120 +1,7 @@ use std::collections::BTreeMap; -use schemars::{JsonSchema, Schema, SchemaGenerator, generate::SchemaSettings}; -use serde_json::json; - -use crate::query::lens; - -fn quoted_u64() -> Schema { - let upper = u64::MAX.to_string(); - let alternatives = upper - .char_indices() - .filter_map(|(index, digit)| { - let lower = if index == 0 { '1' } else { '0' }; - if digit <= lower { - return None; - } - Some(format!( - "{}[{}-{}][0-9]{{{}}}", - &upper[..index], - lower, - char::from(digit as u8 - 1), - upper.len() - index - 1 - )) - }) - .collect::>() - .join("|"); - json!({ - "type": "string", - "pattern": format!("^(?:0|[1-9][0-9]{{0,{}}}|{alternatives}|{upper})$", upper.len() - 2), - }) - .try_into() - .unwrap() -} - -fn numeric_wire(normalized: Schema, python_type: String) -> Schema { - json!({ - "anyOf": [normalized, quoted_u64()], - "x-python-normalized": {"type": python_type, "minimum": 0, "maximum": u64::MAX}, - }) - .try_into() - .unwrap() -} - -pub(crate) fn u64_number(generator: &mut SchemaGenerator) -> Schema { - numeric_wire(u64::json_schema(generator), "int".to_owned()) -} - -pub(crate) fn flag_number(_: &mut SchemaGenerator) -> Schema { - json!({ - "anyOf": [{"type": "integer", "enum": [0, 1]}, {"type": "string", "enum": ["0", "1"]}], - "x-python-normalized": {"type": "int", "minimum": 0, "maximum": 1} - }) - .try_into() - .unwrap() -} - -pub(crate) fn boolean_flag(_: &mut SchemaGenerator) -> Schema { - json!({ - "anyOf": [{"type": "boolean"}, {"type": "integer", "enum": [0, 1]}, {"type": "string", "enum": ["0", "1"]}], - "default": false, - "x-python-normalized": {"type": "bool"} - }).try_into().unwrap() -} - -pub(crate) fn selected(generator: &mut SchemaGenerator) -> Schema { - u64_number(generator) -} - -fn received() -> Schema { - SchemaSettings::draft2020_12() - .for_deserialize() - .with_transform(litellm_traces::schema::integer_bounds) - .into_generator() - .into_root_schema_for::() -} +use schemars::Schema; pub fn schemas() -> BTreeMap<&'static str, Schema> { - BTreeMap::from([ - ("ReadQueryName", json!({"$schema": "https://json-schema.org/draft/2020-12/schema", "title": "ReadQueryName", "type": "string", "enum": lens::LENS_QUERIES.map(|query| query.to_string())}).try_into().unwrap()), - ("LensAccessParams", received::()), - ("LensSampleParams", received::()), - ("LensContentParams", received::()), - ("LensEvidenceParams", received::()), - ( - "ActivityAvailability", - received::(), - ), - ("ExecutionRow", received::()), - ("PartRow", received::()), - ("CountRow", received::()), - ("AgentRow", received::()), - ("TraceQueryHelp", crate::query::help_schema()), - ]) -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - - use super::*; - - #[rstest] - #[case::zero(json!(0), true)] - #[case::quoted_zero(json!("0"), true)] - #[case::maximum(json!(u64::MAX), true)] - #[case::quoted_maximum(json!(u64::MAX.to_string()), true)] - #[case::negative(json!(-1), false)] - #[case::overflow(json!((u128::from(u64::MAX) + 1).to_string()), false)] - #[case::fraction(json!(1.5), false)] - fn count_schema_enforces_the_native_range( - #[case] value: serde_json::Value, - #[case] valid: bool, - ) { - let schema = received::(); - assert_eq!( - jsonschema::is_valid(schema.as_value(), &json!({"count": value})), - valid - ); - } + BTreeMap::from([("TraceQueryHelp", crate::query::help_schema())]) } diff --git a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs index dc806182c1e..6f154ec32f3 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs @@ -3,12 +3,14 @@ use std::{collections::BTreeMap, time::Duration}; use litellm_http::Client; use litellm_traces::{ search::RunFilter, - store::{CallQuery, RunCursor, RunQuery, RunSelection, SpanQuery, SpanSelection}, + store::{ + CallQuery, RunCursor, RunQuery, RunSelection, SpanPart, SpanQuery, SpanSelection, TextRange, + }, }; use litellm_traces_cache::{TraceReader, TraceStore}; use litellm_traces_clickhouse::{ - ClickHouseTraces, Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, - QueryScope, encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements, + ClickHouseTraces, Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, QueryScope, + encode_rows, ensure_schema, execute_read, schema_statements, }; use rstest::rstest; use sha2::{Digest, Sha256}; @@ -43,10 +45,12 @@ async fn list_runs( limit: u32, ) -> TestResult { let query = RunQuery { + order: Default::default(), selection: RunSelection::Matching(RunFilter { start_ms: window.start, end_ms: window.end, search: Default::default(), + ..Default::default() }), after, limit, @@ -481,7 +485,7 @@ async fn listed_agent_names_preserve_scope_and_cursor( let window = timestamp / 1_000_000 - 1000..timestamp / 1_000_000 + 1000; let first = list_runs(&database, &connection, &owner, window.clone(), None, 1).await?; let after = RunCursor { - start_ms: first["data"][0]["start_ms"] + value: first["data"][0]["start_ms"] .as_i64() .ok_or("missing start")?, trace_ref: first["data"][0]["trace_ref"] @@ -779,10 +783,9 @@ fn schema_rejects_invalid_configuration(#[case] database: &str, #[case] retentio #[rstest] #[tokio::test] -async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate( +async fn reused_trace_ids_stay_separate_runs_through_filters_and_span_text( #[future(awt)] database: TestResult, ) -> TestResult { - use litellm_traces_clickhouse::{Parameter, ReadQuery}; let database = database?; let writer = Connection::writer(&database.url)?; ensure_schema(&database.client, &writer, "trace_test", 7).await?; @@ -795,325 +798,79 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate( }))?]).await?; } let connection = Connection::configured(&database.url, "trace_test", "default", "")?; - let sample_parameters = BTreeMap::from([ - ("source".into(), Parameter::Text("traces".into())), - ("all_teams".into(), Parameter::Integer(1)), - ("team".into(), Parameter::Text(String::new())), - ("key_hash".into(), Parameter::Text(String::new())), - ( - "start".into(), - Parameter::Integer(timestamp / 1_000_000 - 1000), - ), - ( - "end".into(), - Parameter::Integer(timestamp / 1_000_000 + 1000), - ), - ("agent_name".into(), Parameter::Text(String::new())), - ("service".into(), Parameter::Text("review".into())), - ( - "filter_keys".into(), - Parameter::Strings(vec!["swarm".into()]), - ), - ( - "filter_values".into(), - Parameter::Strings(vec!["release".into()]), - ), - ("limit".into(), Parameter::Integer(10)), - ("offset".into(), Parameter::Integer(0)), - ("after".into(), Parameter::Text(String::new())), - ("sample_percent".into(), Parameter::Text("100".into())), - ("sample_cap".into(), Parameter::Integer(0)), - ("preview".into(), Parameter::Integer(0)), - ("selected_team".into(), Parameter::Text(String::new())), - ("execution_ids".into(), Parameter::Strings(vec![])), - ]); - let sample: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::Sample, - &sample_parameters, + let store = traces(&database, &connection); + let window = timestamp / 1_000_000 - 1000..timestamp / 1_000_000 + 1000; + let matched = list_runs( + &database, + &connection, + &QueryScope::All, + window.clone(), + None, + 10, + ) + .await?; + let filtered = store + .runs( + &QueryScope::All, + &RunQuery { + order: Default::default(), + selection: RunSelection::Matching(RunFilter { + start_ms: window.start, + end_ms: window.end, + search: litellm_traces::search::RunSearch::parse( + "service:review attr.swarm:release", + ), + ..Default::default() + }), + after: None, + limit: 10, + }, ) - .await?, - )?; - let rows = sample["data"].as_array().expect("sample rows"); - assert_eq!(rows.len(), 2); - assert_ne!(rows[0]["trace_ref"], rows[1]["trace_ref"]); + .await?; + assert_eq!(filtered.len(), 2); + assert_eq!(matched["data"].as_array().map(Vec::len), Some(2)); + assert_ne!(filtered[0].trace_ref, filtered[1].trace_ref); let by_trace_id = RunQuery { + order: Default::default(), selection: RunSelection::TraceId("shared".into()), after: None, limit: 10, }; - let store = traces(&database, &connection); - let identities = store.runs(&owned("", &["team"]), &by_trace_id).await?; - assert_eq!(identities.len(), 2); - let identity = serde_json::json!({ - "data": store.runs(&owned("one", &[]), &by_trace_id).await?, - }); - assert_eq!(identity["data"].as_array().map(Vec::len), Some(1)); - assert!( - rows.iter() - .any(|row| row["trace_ref"] == identity["data"][0]["trace_ref"]) - ); - let first_ref = rows[0]["trace_ref"].as_str().expect("reference"); - let content_parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(1)), - ("team".into(), Parameter::Text(String::new())), - ("key_hash".into(), Parameter::Text(String::new())), - ("source".into(), Parameter::Text("traces".into())), - ("id".into(), Parameter::Text("shared".into())), - ("record_team".into(), Parameter::Text("team".into())), - ("trace_ref".into(), Parameter::Text(first_ref.into())), - ("cursor".into(), Parameter::Text(String::new())), - ("offset".into(), Parameter::Integer(1)), - ]); - let content: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::Content, - &content_parameters, - ) - .await?, - )?; - assert_eq!(content["data"].as_array().map(Vec::len), Some(1)); - let text = content["data"][0]["content"].as_str().expect("content"); - let opposite = if text.contains("timeout") { - "success" - } else { - "timeout" - }; - let evidence_parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(1)), - ("team".into(), Parameter::Text(String::new())), - ("key_hash".into(), Parameter::Text(String::new())), - ("source".into(), Parameter::Text("traces".into())), - ("id".into(), Parameter::Text("shared".into())), - ("record_team".into(), Parameter::Text("team".into())), - ("trace_ref".into(), Parameter::Text(first_ref.into())), - ("span".into(), Parameter::Text("root".into())), - ("quote".into(), Parameter::Text(opposite.into())), - ]); - let evidence: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::Evidence, - &evidence_parameters, - ) - .await?, - )?; - assert_eq!(evidence["data"][0]["count"], 0); - Ok(()) -} - -#[rstest] -#[tokio::test] -async fn lens_request_sample_does_not_trust_caller_tags( - #[future(awt)] database: TestResult, -) -> TestResult { - use litellm_traces_clickhouse::{Parameter, ReadQuery}; - let database = database?; - let writer = Connection::writer(&database.url)?; - ensure_schema(&database.client, &writer, "trace_test", 7).await?; - let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64 / 1_000_000; - for (id, internal) in [("external", false), ("internal", true)] { - let row = serde_json::from_value(serde_json::json!({ - "request_id": id, "team_id": "team", "start_time": timestamp, "end_time": timestamp, - "request_tags": ["litellm-engine"], - "metadata": serde_json::json!({"litellm_lens_internal": internal}).to_string() - }))?; - insert_rows(&database, "spend_logs", vec![row]).await?; - } - let connection = Connection::configured(&database.url, "trace_test", "default", "")?; - let parameters = BTreeMap::from([ - ("source".into(), Parameter::Text("requests".into())), - ("all_teams".into(), Parameter::Integer(1)), - ("team".into(), Parameter::Text(String::new())), - ("key_hash".into(), Parameter::Text(String::new())), - ("start".into(), Parameter::Integer(timestamp - 1000)), - ("end".into(), Parameter::Integer(timestamp + 60000)), - ("agent_name".into(), Parameter::Text(String::new())), - ("service".into(), Parameter::Text(String::new())), - ("filter_keys".into(), Parameter::Strings(vec![])), - ("filter_values".into(), Parameter::Strings(vec![])), - ("limit".into(), Parameter::Integer(10)), - ("offset".into(), Parameter::Integer(0)), - ("after".into(), Parameter::Text(String::new())), - ("sample_percent".into(), Parameter::Text("100".into())), - ("sample_cap".into(), Parameter::Integer(0)), - ("preview".into(), Parameter::Integer(0)), - ("selected_team".into(), Parameter::Text(String::new())), - ("execution_ids".into(), Parameter::Strings(vec![])), - ]); - let sample: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::Sample, - ¶meters, - ) - .await?, - )?; - let rows = sample["data"].as_array().expect("sample rows"); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0]["trace_id"], "external"); - Ok(()) -} - -#[rstest] -#[case::changing("100", 0, 0, 1001, 100, true)] -#[case::all("100", 0, 0, 1001, 100, false)] -#[case::percentage("10", 0, 0, 101, 100, false)] -#[case::capped("100", 25, 0, 25, 100, false)] -#[case::preview("10", 25, 1, 1001, 100, false)] -#[tokio::test] -async fn lens_selection_pages_without_losing_or_repeating_runs( - #[future(awt)] database: TestResult, - #[case] percent: &str, - #[case] cap: i64, - #[case] preview: i64, - #[case] expected: usize, - #[case] page_size: usize, - #[case] changing: bool, -) -> TestResult { - use litellm_traces_clickhouse::ReadQuery; - let database = database?; - ensure_schema( - &database.client, - &Connection::writer(&database.url)?, - "trace_test", - 7, - ) - .await?; - execute_write(&database, "INSERT INTO trace_test.spend_logs (request_id,team_id,start_time,end_time) SELECT toString(number),'team',now64(3)-INTERVAL 5 MINUTE,now64(3)-INTERVAL 5 MINUTE FROM numbers(1001)").await?; - let connection = Connection::configured(&database.url, "trace_test", "default", "")?; - let end = time::OffsetDateTime::now_utc().unix_timestamp() * 1000 + 60000; - let mut seen = std::collections::BTreeSet::new(); - let mut cursor = String::new(); - let step = if page_size == 0 { expected } else { page_size }; - for offset in (0..expected).step_by(step) { - let parameters = BTreeMap::from([ - ("source".into(), Parameter::Text("requests".into())), - ("all_teams".into(), Parameter::Integer(0)), - ("team".into(), Parameter::Text("team".into())), - ("key_hash".into(), Parameter::Text(String::new())), - ("start".into(), Parameter::Integer(0)), - ("end".into(), Parameter::Integer(end)), - ("agent_name".into(), Parameter::Text(String::new())), - ("service".into(), Parameter::Text(String::new())), - ("filter_keys".into(), Parameter::Strings(vec![])), - ("filter_values".into(), Parameter::Strings(vec![])), - ("limit".into(), Parameter::Integer(page_size as i64)), - ( - "offset".into(), - Parameter::Integer(if changing { 0 } else { offset as i64 }), - ), - ("after".into(), Parameter::Text(cursor.clone())), - ("sample_percent".into(), Parameter::Text(percent.into())), - ("sample_cap".into(), Parameter::Integer(cap)), - ("preview".into(), Parameter::Integer(preview)), - ("selected_team".into(), Parameter::Text(String::new())), - ("execution_ids".into(), Parameter::Strings(vec![])), - ]); - let body = execute_named_read( - &database.client, - &connection, - ReadQuery::Sample, - ¶meters, - ) - .await?; - let json: serde_json::Value = serde_json::from_str(&body)?; - let rows = json["data"].as_array().expect("sample rows"); - assert_eq!(rows.len(), step.min(expected - offset)); - for row in rows { - assert_eq!( - row["eligible"], - if changing && offset > 0 { 1000 } else { 1001 } - ); - assert!(seen.insert(row["trace_id"].as_str().expect("run id").to_owned())); - } - if changing { - cursor = rows.last().expect("last run")["selection_key"] - .as_str() - .expect("selection key") - .to_owned(); - if offset == 0 { - let removed = rows[0]["trace_id"].as_str().expect("request id"); - execute_write(&database, &format!("ALTER TABLE trace_test.spend_logs DELETE WHERE request_id='{removed}' SETTINGS mutations_sync=1")).await?; - } - } - } - assert_eq!(seen.len(), expected); - Ok(()) -} - -#[rstest] -#[case::short(100)] -#[case::boundary(7970)] -#[case::long(16000)] -#[tokio::test] -async fn lens_content_keeps_output_visible_after_long_input( - #[future(awt)] database: TestResult, - #[case] input_length: usize, -) -> TestResult { - use litellm_traces_clickhouse::ReadQuery; - let database = database?; - ensure_schema( - &database.client, - &Connection::writer(&database.url)?, - "trace_test", - 7, - ) - .await?; - insert_rows(&database, "spend_logs", vec![serde_json::from_value(serde_json::json!({ - "request_id": "request", "team_id": "team", "start_time": time::OffsetDateTime::now_utc().unix_timestamp()*1000, "end_time": time::OffsetDateTime::now_utc().unix_timestamp()*1000, "messages": "x".repeat(input_length), "response": "Delivered result" - }))?]).await?; - let connection = Connection::configured(&database.url, "trace_test", "default", "")?; - let mut parameters = BTreeMap::from([ - ("source".into(), Parameter::Text("requests".into())), - ("all_teams".into(), Parameter::Integer(0)), - ("team".into(), Parameter::Text("team".into())), - ("record_team".into(), Parameter::Text("team".into())), - ("key_hash".into(), Parameter::Text(String::new())), - ("trace_ref".into(), Parameter::Text(String::new())), - ("id".into(), Parameter::Text("request".into())), - ("cursor".into(), Parameter::Text(String::new())), - ("offset".into(), Parameter::Integer(1)), - ]); - let body = execute_named_read( - &database.client, - &connection, - ReadQuery::Content, - ¶meters, - ) - .await?; - let json: serde_json::Value = serde_json::from_str(&body)?; - let text = json["data"][0]["content"].as_str().expect("content"); - assert!(text.contains("Output: Delivered result")); - assert!(text.len() <= 8000); assert_eq!( - json["data"][0]["truncated"], - u8::from(input_length + "Input: \nOutput: Delivered result\nError: ".len() > 8000) + store.runs(&owned("", &["team"]), &by_trace_id).await?.len(), + 2 ); - let original = format!( - "Input: {}\nOutput: Delivered result\nError: ", - "x".repeat(input_length) + let own = store.runs(&owned("one", &[]), &by_trace_id).await?; + assert_eq!(own.len(), 1); + let reader = TraceReader::new(usize::MAX); + let read = |trace_ref: String, contains: &'static str| { + let reader = &reader; + let store = &store; + async move { + reader + .span_text( + store, + &QueryScope::All, + "shared", + &trace_ref, + vec!["root".into()], + SpanPart::Input, + TextRange::ALL, + Some(contains.into()), + ) + .await + } + }; + let texts = read(own[0].trace_ref.clone(), "timeout").await?; + assert_eq!( + texts + .iter() + .map(|text| (text.text.as_str(), text.contains)) + .collect::>(), + [("timeout", true)] ); - let mut recovered = String::new(); - for offset in (2..original.len() + 2).step_by(8000) { - parameters.insert("offset".into(), Parameter::Integer(offset as i64)); - let body = execute_named_read( - &database.client, - &connection, - ReadQuery::Content, - ¶meters, - ) - .await?; - let page: serde_json::Value = serde_json::from_str(&body)?; - recovered.push_str(page["data"][0]["content"].as_str().expect("content")); - } - assert_eq!(recovered, original); + let other = read(own[0].trace_ref.clone(), "success").await?; + assert!(other.iter().all(|text| !text.contains)); Ok(()) } @@ -1276,116 +1033,6 @@ fn schema_includes_every_migration_file() -> TestResult { Ok(()) } -#[rstest] -#[tokio::test] -async fn lens_agent_discovery_and_selection_preserve_scope( - #[future] database: TestResult, -) -> TestResult { - use litellm_traces_clickhouse::ReadQuery; - let database = database.await?; - let writer = Connection::writer(&database.url)?; - ensure_schema(&database.client, &writer, "trace_test", 7).await?; - let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64; - for (team, key, trace, agent, span, parent) in [ - ("alpha", "one", "research", "research_agent", "root", ""), - ("alpha", "one", "research", "", "tool", "root"), - ("alpha", "one", "support", "support_agent", "root", ""), - ("alpha", "two", "hidden-key", "private_agent", "root", ""), - ("beta", "one", "hidden-team", "other_agent", "root", ""), - ] { - insert_rows( - &database, - "otel_traces", - vec![serde_json::from_value(serde_json::json!({ - "Timestamp": timestamp, "TraceId": trace, "SpanId": span, "ParentSpanId": parent, - "ServiceName": "shared-app", "SpanName": "run", "Input": "test", - "SpanAttributes": {"gen_ai.agent.name": agent}, - "ResourceAttributes": {"litellm.team_id": team, "litellm.api_key_hash": key} - }))?], - ) - .await?; - } - let connection = Connection::configured(&database.url, "trace_test", "default", "")?; - let agent_parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(0)), - ("team".into(), Parameter::Text("alpha".into())), - ("key_hash".into(), Parameter::Text("one".into())), - ]); - let agents: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::Agents, - &agent_parameters, - ) - .await?, - )?; - assert_eq!( - agents["data"], - serde_json::json!([ - {"agent_name": "research_agent"}, {"agent_name": "support_agent"} - ]) - ); - let sample_parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(0)), - ("team".into(), Parameter::Text("alpha".into())), - ("key_hash".into(), Parameter::Text("one".into())), - ("source".into(), Parameter::Text("traces".into())), - ( - "start".into(), - Parameter::Integer(timestamp / 1_000_000 - 1000), - ), - ( - "end".into(), - Parameter::Integer(timestamp / 1_000_000 + 1000), - ), - ("service".into(), Parameter::Text("shared-app".into())), - ( - "agent_name".into(), - Parameter::Text("research_agent".into()), - ), - ("filter_keys".into(), Parameter::Strings(vec![])), - ("filter_values".into(), Parameter::Strings(vec![])), - ("limit".into(), Parameter::Integer(100)), - ("offset".into(), Parameter::Integer(0)), - ("after".into(), Parameter::Text(String::new())), - ("sample_percent".into(), Parameter::Text("100".into())), - ("sample_cap".into(), Parameter::Integer(0)), - ("preview".into(), Parameter::Integer(1)), - ("selected_team".into(), Parameter::Text(String::new())), - ("execution_ids".into(), Parameter::Strings(vec![])), - ]); - let sample: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::Sample, - &sample_parameters, - ) - .await?, - )?; - assert_eq!(sample["data"].as_array().expect("rows").len(), 1); - assert_eq!(sample["data"][0]["trace_id"], "research"); - assert_eq!(sample["data"][0]["span_count"], 2); - let availability_parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(0)), - ("team".into(), Parameter::Text("alpha".into())), - ("key_hash".into(), Parameter::Text("one".into())), - ]); - let available: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::Availability, - &availability_parameters, - ) - .await?, - )?; - assert_eq!(available["data"][0]["traces"], 1); - assert_eq!(available["data"][0]["requests"], 0); - Ok(()) -} - #[rstest] #[case::empty(false)] #[case::custom_metadata(true)] diff --git a/litellm-rust/crates/traces-clickhouse/tests/queries.rs b/litellm-rust/crates/traces-clickhouse/tests/queries.rs index b9235b5c775..78df4b30d85 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/queries.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/queries.rs @@ -2,7 +2,7 @@ use std::collections::BTreeMap; use litellm_traces::{ search::RunFilter, - store::{RunQuery, RunRow, RunSelection, SpanQuery, SpanRow, SpanSelection}, + store::{RunOrder, RunQuery, RunRow, RunSelection, SpanQuery, SpanRow, SpanSelection}, }; use litellm_traces_cache::{StoreResult, TraceStore}; use litellm_traces_clickhouse::{ClickHouseTraces, Error, QueryScope, query_help, query_sql}; @@ -134,12 +134,14 @@ fn fixture_clock() -> TestResult { fn newest(limit: u32, after: Option<&RunRow>) -> RunQuery { RunQuery { + order: Default::default(), selection: RunSelection::Matching(RunFilter { start_ms: 0, end_ms: i64::MAX / 1_000_000, search: Default::default(), + ..Default::default() }), - after: after.map(RunRow::cursor), + after: after.map(|row| RunOrder::NEWEST.cursor(row)), limit, } } diff --git a/litellm-rust/crates/traces-clickhouse/tests/reads.rs b/litellm-rust/crates/traces-clickhouse/tests/reads.rs index 45c4b34745c..a1a73637564 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/reads.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use litellm_http::Client; -use litellm_traces::search::RunFilter; +use litellm_traces::{search::RunFilter, store::RunOrder}; use litellm_traces_cache::{PageRequest, ReadError, TraceReader}; use litellm_traces_clickhouse::{ ClickHouseTraces, Connection, InsertTable, QueryScope, insert_rows, @@ -14,6 +14,7 @@ fn all_runs() -> RunFilter { start_ms: 0, end_ms: 2_000_000_000_000, search: Default::default(), + ..Default::default() } } @@ -105,9 +106,11 @@ async fn list_costs_match_each_run_when_response_ids_are_reused( &store, &access, &all_runs(), + RunOrder::NEWEST, &PageRequest { cursor: None, limit: 50, + ..Default::default() }, ) .await?; @@ -241,9 +244,11 @@ async fn large_runs_remain_complete_under_default_reader_limits( &store, &access, &all_runs(), + RunOrder::NEWEST, &PageRequest { cursor: None, limit: 500, + ..Default::default() }, ) .await?; @@ -401,9 +406,11 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( &store, &access, &all_runs(), + RunOrder::NEWEST, &PageRequest { cursor: None, limit: 10, + ..Default::default() }, ) .await?; @@ -567,9 +574,11 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( &store, &access, &all_runs(), + RunOrder::NEWEST, &PageRequest { cursor: None, limit: 50, + ..Default::default() }, ) .await?; @@ -604,9 +613,11 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( &store, &access, &all_runs(), + RunOrder::NEWEST, &PageRequest { cursor: None, limit: 50, + ..Default::default() }, ) .await?; @@ -617,9 +628,11 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( &store, &access, &all_runs(), + RunOrder::NEWEST, &PageRequest { cursor: None, limit: 50, + ..Default::default() }, ) .await?; @@ -743,9 +756,11 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback &store, &access, &all_runs(), + RunOrder::NEWEST, &PageRequest { cursor: None, limit: 50, + ..Default::default() }, ) .await?; @@ -765,3 +780,79 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback } Ok(()) } + +#[rstest] +#[tokio::test] +async fn a_run_shared_with_another_user_stays_hidden_before_its_rollup_rows_merge( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let client = &fixture.database.client; + let writer = Connection::writer(&fixture.database.url)?; + let start_ms = 1_790_000_000_000_i64; + let span = |trace_id: &str, span_id: &str, user_id: &str, offset_ms: i64| { + BTreeMap::from([ + ( + "Timestamp".into(), + json!((start_ms + offset_ms) * 1_000_000), + ), + ("Duration".into(), json!(1_000_000)), + ("TraceId".into(), json!(trace_id)), + ("SpanId".into(), json!(span_id)), + ( + "ParentSpanId".into(), + json!(if offset_ms == 0 { "" } else { "root" }), + ), + ("ObservationType".into(), json!("llm")), + ("TeamId".into(), json!("team-a")), + ("ApiKeyHash".into(), json!("key-a")), + ("UserId".into(), json!(user_id)), + ]) + }; + for batch in [ + vec![ + span("shared", "root", "alice", 0), + span("shared", "alice-call", "alice", 1), + ], + vec![span("shared", "bob-call", "bob", 2)], + vec![span("solo", "root", "alice", 10)], + ] { + insert_rows(client, &writer, DATABASE, InsertTable::OtelTraces, batch).await?; + } + let connection = fixture + .readers + .connection(client, &QueryScope::All, "fixture-secret") + .await?; + let (reader, store) = make_reader(client, connection); + let page = PageRequest { + cursor: None, + limit: 50, + ..Default::default() + }; + let team = QueryScope::Owned { + user_id: String::new(), + team_ids: vec!["team-a".into()], + }; + let shared = reader + .list_traces(&store, &team, &all_runs(), RunOrder::NEWEST, &page) + .await? + .data + .into_iter() + .find(|run| run.trace_id == "shared") + .ok_or("missing shared run")?; + assert_eq!(shared.span_count, 3); + let alice = QueryScope::Owned { + user_id: "alice".into(), + team_ids: Vec::new(), + }; + let listed = reader + .list_traces(&store, &alice, &all_runs(), RunOrder::NEWEST, &page) + .await?; + let trace_ids: Vec<&str> = listed + .data + .iter() + .map(|run| run.trace_id.as_str()) + .collect(); + assert_eq!(trace_ids, ["solo"]); + Ok(()) +} diff --git a/litellm-rust/crates/traces-clickhouse/tests/search.rs b/litellm-rust/crates/traces-clickhouse/tests/search.rs index defbb87f5b1..3d9b228c7a3 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/search.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/search.rs @@ -1,7 +1,13 @@ use std::collections::{BTreeMap, BTreeSet}; -use litellm_traces::search::{AgentRuns, HistogramBucket, RunField, RunFilter, RunSearch}; -use litellm_traces_cache::{PageRequest, TraceReader}; +use litellm_traces::{ + search::{AgentRuns, HistogramBucket, RunField, RunFilter, RunSearch}, + store::{ + CountBy, CountValue, RunCountQuery, RunOrder, RunQuery, RunSelection, RunSortKey, SpanPart, + SpanQuery, SpanSelection, TextRange, + }, +}; +use litellm_traces_cache::{PageRequest, TraceReader, TraceStore}; use litellm_traces_clickhouse::{ ClickHouseTraces, Connection, InsertTable, QueryScope, insert_rows, }; @@ -27,11 +33,16 @@ fn filter(start_ms: i64, q: &str) -> RunFilter { start_ms, end_ms: WINDOW_END_MS, search: RunSearch::parse(q), + ..Default::default() } } fn page(cursor: Option, limit: u32) -> PageRequest { - PageRequest { cursor, limit } + PageRequest { + cursor, + limit, + ..Default::default() + } } fn reader() -> TraceReader { @@ -240,7 +251,13 @@ async fn list_q_selects_matching_runs_before_paging( ]; for (q, expected) in cases { let page = reader - .list_traces(&store, &team_a(), &filter(0, q), &page(None, 50)) + .list_traces( + &store, + &team_a(), + &filter(0, q), + RunOrder::NEWEST, + &page(None, 50), + ) .await?; let listed: Vec<&str> = page.data.iter().map(|run| run.trace_id.as_str()).collect(); assert_eq!(&listed, expected, "q = {q:?}"); @@ -258,13 +275,14 @@ async fn list_q_pages_through_matches_only( let reader = reader(); let filter = filter(0, "model:gpt-x"); let first = reader - .list_traces(&store, &team_a(), &filter, &page(None, 1)) + .list_traces(&store, &team_a(), &filter, RunOrder::NEWEST, &page(None, 1)) .await?; let second = reader .list_traces( &store, &team_a(), &filter, + RunOrder::NEWEST, &page(first.next_cursor.clone(), 1), ) .await?; @@ -399,13 +417,20 @@ async fn every_field_suggests_values_that_filter_back_to_their_runs( for value in &values.values { let q = format!(r#"{key}:"{value}""#); let included = reader - .list_traces(&store, &team_a(), &filter(0, &q), &page(None, 50)) + .list_traces( + &store, + &team_a(), + &filter(0, &q), + RunOrder::NEWEST, + &page(None, 50), + ) .await?; let excluded = reader .list_traces( &store, &team_a(), &filter(0, &format!("-{q}")), + RunOrder::NEWEST, &page(None, 50), ) .await?; @@ -419,3 +444,519 @@ async fn every_field_suggests_values_that_filter_back_to_their_runs( } Ok(()) } + +async fn sql(fixture: &SeededDatabase, query: String) -> TestResult { + let writer = Connection::writer(&fixture.database.url)?; + Ok(fixture + .database + .client + .post(writer.url().clone()) + .body(query) + .send() + .await? + .error_for_status()? + .text() + .await? + .trim() + .to_owned()) +} + +async fn table_rows(fixture: &SeededDatabase, table: &str) -> TestResult { + Ok( + sql(fixture, format!("SELECT count() FROM {DATABASE}.{table}")) + .await? + .parse()?, + ) +} + +async fn rows_read_by(fixture: &SeededDatabase, marker: &str) -> TestResult { + sql(fixture, "SYSTEM FLUSH LOGS".into()).await?; + let read = sql( + fixture, + format!( + "SELECT read_rows FROM system.query_log WHERE type = 'QueryFinish' \ + AND current_database = '{DATABASE}' AND position(query, '{marker}') > 0 \ + AND query NOT LIKE '%system.query_log%' \ + ORDER BY event_time_microseconds DESC LIMIT 1" + ), + ) + .await?; + if read.is_empty() { + return Err(format!("no finished query mentions {marker}").into()); + } + Ok(read.parse()?) +} + +#[rstest] +#[tokio::test] +async fn listed_runs_name_the_agent_they_matched_even_when_resolution_is_limited( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let writer = Connection::writer(&fixture.database.url)?; + insert_rows( + &fixture.database.client, + &writer, + DATABASE, + InsertTable::OtelTraces, + vec![BTreeMap::from([ + ("Timestamp".into(), json!((T0_MS + HOUR_MS + 5) * 1_000_000)), + ("TraceId".into(), json!("beta")), + ("SpanId".into(), json!("beta-oversized")), + ("ParentSpanId".into(), json!("beta-root")), + ( + "SpanName".into(), + json!("x".repeat(litellm_storage_clickhouse::READ_LIMITS.response_bytes + 1)), + ), + ("ObservationType".into(), json!("tool")), + ("TeamId".into(), json!("team-a")), + ("ApiKeyHash".into(), json!("key-a")), + ("Duration".into(), json!(1_000_000)), + ])], + ) + .await?; + let reader = reader(); + let agents = reader + .values(&store, &team_a(), &filter(0, ""), RunField::Agent, "", 10) + .await?; + let mut limited = 0; + for agent in &agents.values { + let listed = reader + .list_traces( + &store, + &team_a(), + &filter(0, &format!(r#"agent:"{agent}""#)), + RunOrder::NEWEST, + &page(None, 50), + ) + .await?; + assert!(!listed.data.is_empty(), "agent:{agent} lists nothing"); + for run in &listed.data { + limited += usize::from(run.resolution_limited); + assert!( + run.agent_names.contains(agent), + "{} matched agent:{agent} but lists {:?}", + run.trace_id, + run.agent_names + ); + } + } + assert_eq!( + limited, 1, + "the oversized span should leave exactly one run on its rollup summary" + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn listing_runs_scans_the_rollup_once( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let query = RunQuery { + selection: RunSelection::Matching(filter(0, "")), + order: RunOrder::NEWEST, + after: None, + limit: 50, + }; + let listed = store.runs(&QueryScope::All, &query).await?; + assert_eq!(listed.len(), 4); + let budget = table_rows(&fixture, "agent_traces_by_key").await? + + table_rows(&fixture, "otel_traces").await?; + let read = rows_read_by(&fixture, "FROM owned_runs").await?; + assert!( + read <= budget, + "listing 4 runs read {read} rows, more than the {budget} rollup and span rows that exist" + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn reading_the_spans_of_listed_runs_skips_other_runs_in_the_window( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let client = &fixture.database.client; + let writer = Connection::writer(&fixture.database.url)?; + let rows_for = |index: i64| -> Vec> { + let trace_id = format!("run-{index:02}"); + let start_ms = T0_MS + index * 60_000; + (0..3) + .map(|offset| { + BTreeMap::from([ + ("Timestamp".into(), json!((start_ms + offset) * 1_000_000)), + ("TraceId".into(), json!(trace_id)), + ("SpanId".into(), json!(format!("{trace_id}-{offset}"))), + ( + "ParentSpanId".into(), + json!(if offset == 0 { + String::new() + } else { + format!("{trace_id}-0") + }), + ), + ("SpanName".into(), json!("step")), + ("ServiceName".into(), json!("svc")), + ( + "ObservationType".into(), + json!(if offset == 0 { "chain" } else { "tool" }), + ), + ("TeamId".into(), json!("team-a")), + ("ApiKeyHash".into(), json!("key-a")), + ("Duration".into(), json!(1_000_000)), + ]) + }) + .collect() + }; + for index in 0..12 { + insert_rows( + client, + &writer, + DATABASE, + InsertTable::OtelTraces, + rows_for(index), + ) + .await?; + } + let connection = fixture + .readers + .connection(client, &QueryScope::All, "fixture-secret") + .await?; + let store = ClickHouseTraces::new(client.clone(), connection); + let run = store + .runs( + &QueryScope::All, + &RunQuery { + selection: RunSelection::TraceId("run-05".into()), + order: RunOrder::NEWEST, + after: None, + limit: 2, + }, + ) + .await?; + let spans = store + .spans( + &QueryScope::All, + &SpanQuery { + selection: SpanSelection::Runs { + trace_ids: vec![run[0].trace_id.clone()], + trace_refs: vec![run[0].trace_ref.clone()], + window: T0_MS..WINDOW_END_MS, + }, + as_of_ms: u64::MAX, + after: None, + limit: 256, + }, + ) + .await?; + assert_eq!(spans.len(), 3); + let read = rows_read_by(&fixture, &run[0].trace_ref).await?; + assert!( + read <= 2 * spans.len() as u64, + "reading one run's 3 spans read {read} of the 36 spans in the window" + ); + Ok(()) +} + +async fn add_attributes(fixture: &SeededDatabase) -> TestResult { + let writer = Connection::writer(&fixture.database.url)?; + let span = |trace_id: &str, start_ms: i64, column: &str, value: &str| { + BTreeMap::from([ + ("Timestamp".into(), json!((start_ms + 9) * 1_000_000)), + ("TraceId".into(), json!(trace_id)), + ("SpanId".into(), json!(format!("{trace_id}-tagged"))), + ("ParentSpanId".into(), json!(format!("{trace_id}-root"))), + ("SpanName".into(), json!("tag")), + ("ServiceName".into(), json!("svc")), + ("ObservationType".into(), json!("tool")), + ("TeamId".into(), json!("team-a")), + ("ApiKeyHash".into(), json!("key-a")), + ("Duration".into(), json!(1_000_000)), + (column.into(), json!({"tenant.tier": value})), + ]) + }; + insert_rows( + &fixture.database.client, + &writer, + DATABASE, + InsertTable::OtelTraces, + vec![ + span("alpha", T0_MS, "SpanAttributes", "gold"), + span("beta", T0_MS + HOUR_MS, "ResourceAttributes", "silver"), + ], + ) + .await?; + Ok(()) +} + +#[rstest] +#[case::service("service:svc", &["gamma", "beta", "alpha"])] +#[case::other_service("service:other", &[])] +#[case::team("team:team-a", &["gamma", "beta", "alpha"])] +#[case::excluded_team("-team:team-a", &[])] +#[case::span_attribute("attr.tenant.tier:gold", &["alpha"])] +#[case::resource_attribute("attr.tenant.tier:SILV*", &["beta"])] +#[case::excluded_attribute("-attr.tenant.tier:gold", &["gamma", "beta"])] +#[case::two_attributes("attr.tenant.tier:gold attr.tenant.tier:silver", &[])] +#[case::attribute_and_field("attr.tenant.tier:* status:error", &["beta"])] +#[case::unknown_attribute("attr.missing:gold", &[])] +#[tokio::test] +async fn service_team_and_attribute_filters_select_runs( + #[future(awt)] migrated_database: TestResult, + #[case] q: &str, + #[case] expected: &[&str], +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + add_attributes(&fixture).await?; + let page = reader() + .list_traces( + &store, + &team_a(), + &filter(0, q), + RunOrder::NEWEST, + &page(None, 50), + ) + .await?; + let listed: Vec<&str> = page.data.iter().map(|run| run.trace_id.as_str()).collect(); + assert_eq!(listed, expected); + Ok(()) +} + +#[rstest] +#[case::newest(RunOrder::NEWEST)] +#[case::oldest(RunOrder { descending: false, ..RunOrder::NEWEST })] +#[case::longest(RunOrder { key: RunSortKey::DurationMs, descending: true })] +#[case::shortest(RunOrder { key: RunSortKey::DurationMs, descending: false })] +#[case::most_spans(RunOrder { key: RunSortKey::SpanCount, descending: true })] +#[case::fewest_spans(RunOrder { key: RunSortKey::SpanCount, descending: false })] +#[case::most_errors(RunOrder { key: RunSortKey::ErrorCount, descending: true })] +#[case::fewest_errors(RunOrder { key: RunSortKey::ErrorCount, descending: false })] +#[case::by_reference(RunOrder::BY_REFERENCE)] +#[case::by_reference_descending(RunOrder { descending: true, ..RunOrder::BY_REFERENCE })] +#[tokio::test] +async fn one_run_pages_walk_every_order_without_gaps_or_repeats( + #[future(awt)] migrated_database: TestResult, + #[case] order: RunOrder, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let reader = reader(); + let all = RunQuery { + selection: RunSelection::Matching(filter(0, "")), + order: RunOrder::NEWEST, + after: None, + limit: 50, + }; + let mut expected = store.runs(&team_a(), &all).await?; + expected.sort_by(|left, right| order.compare(left, right)); + let expected: Vec = expected.into_iter().map(|run| run.trace_id).collect(); + + let mut listed = Vec::new(); + let mut cursor = None; + loop { + let request = PageRequest { + cursor: cursor.take(), + limit: 1, + }; + let page = reader + .list_traces(&store, &team_a(), &filter(0, ""), order, &request) + .await?; + listed.extend(page.data.into_iter().map(|run| run.trace_id)); + let Some(next) = page.next_cursor else { break }; + cursor = Some(next); + } + assert_eq!(listed, expected); + + let whole = reader + .list_traces( + &store, + &team_a(), + &filter(0, ""), + order, + &PageRequest { + cursor: None, + limit: 3, + }, + ) + .await?; + assert_eq!(whole.data.len(), 3); + assert!(whole.next_cursor.is_none()); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn listed_references_restrict_runs_and_order_by_reference_pages_stably( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let reader = reader(); + let by_reference = |cursor| PageRequest { cursor, limit: 2 }; + let first = reader + .list_traces( + &store, + &team_a(), + &filter(0, ""), + RunOrder::BY_REFERENCE, + &by_reference(None), + ) + .await?; + let second = reader + .list_traces( + &store, + &team_a(), + &filter(0, ""), + RunOrder::BY_REFERENCE, + &by_reference(first.next_cursor.clone()), + ) + .await?; + let refs: Vec = first + .data + .iter() + .chain(&second.data) + .map(|run| run.trace_ref.clone()) + .collect(); + let mut sorted = refs.clone(); + sorted.sort(); + assert_eq!(refs, sorted); + assert_eq!(refs.len(), 3); + assert!(second.next_cursor.is_none()); + + let picked = RunFilter { + trace_refs: vec![refs[1].clone(), "not-a-run".into()], + ..filter(0, "") + }; + let only = reader + .list_traces( + &store, + &team_a(), + &picked, + RunOrder::NEWEST, + &page(None, 50), + ) + .await?; + assert_eq!( + only.data + .iter() + .map(|run| &run.trace_ref) + .collect::>(), + [&refs[1]] + ); + assert_eq!(reader.count_traces(&store, &team_a(), &picked).await?, 1); + assert_eq!( + reader + .count_traces(&store, &team_a(), &filter(0, "model:gpt-x")) + .await?, + 2 + ); + Ok(()) +} + +#[rstest] +#[case::keys(CountValue::AttributeKey, "", &[("tenant.tier", 2)])] +#[case::values(CountValue::Attribute("tenant.tier".into()), "", &[("gold", 1), ("silver", 1)])] +#[case::values_containing(CountValue::Attribute("tenant.tier".into()), "IL", &[("silver", 1)])] +#[tokio::test] +async fn attribute_keys_and_values_count_runs_in_scope( + #[future(awt)] migrated_database: TestResult, + #[case] value: CountValue, + #[case] contains: &str, + #[case] expected: &[(&str, u64)], +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + add_attributes(&fixture).await?; + let query = RunCountQuery { + filter: filter(0, ""), + by: CountBy { + value: Some(value), + ..CountBy::default() + }, + contains: contains.into(), + limit: Some(10), + }; + let counts = store.run_counts(&team_a(), &query).await?; + let counted: Vec<(&str, u64)> = counts + .iter() + .map(|count| (count.value.as_str(), count.runs)) + .collect(); + assert_eq!(counted, expected); + Ok(()) +} + +#[rstest] +#[case::whole(TextRange::ALL, None, &[("alpha-root", "book a flight to Paris", false)])] +#[case::window(TextRange::From { offset: 5, max_chars: Some(6) }, None, &[("alpha-root", "a flig", false)])] +#[case::tail(TextRange::Last { chars: 5 }, None, &[("alpha-root", "Paris", false)])] +#[case::tail_longer_than_text(TextRange::Last { chars: 500 }, None, &[("alpha-root", "book a flight to Paris", false)])] +#[case::contains(TextRange::From { offset: 0, max_chars: Some(0) }, Some("flight to"), &[("alpha-root", "", true)])] +#[case::contains_is_case_sensitive(TextRange::From { offset: 0, max_chars: Some(0) }, Some("PARIS"), &[("alpha-root", "", false)])] +#[tokio::test] +async fn span_text_reads_ranges_of_each_listed_span( + #[future(awt)] migrated_database: TestResult, + #[case] range: TextRange, + #[case] contains: Option<&str>, + #[case] expected: &[(&str, &str, bool)], +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let runs = store + .runs( + &team_a(), + &RunQuery { + selection: RunSelection::TraceId("alpha".into()), + order: RunOrder::NEWEST, + after: None, + limit: 2, + }, + ) + .await?; + let texts = reader() + .span_text( + &store, + &team_a(), + "alpha", + &runs[0].trace_ref, + vec!["alpha-root".into(), "alpha-llm".into(), "missing".into()], + SpanPart::Input, + range, + contains.map(str::to_owned), + ) + .await?; + let inputs: Vec<(&str, &str, bool)> = texts + .iter() + .filter(|text| text.total_chars > 0) + .map(|text| (text.span_id.as_str(), text.text.as_str(), text.contains)) + .collect(); + assert_eq!(inputs, expected); + assert_eq!( + texts + .iter() + .map(|text| text.span_id.as_str()) + .collect::>(), + BTreeSet::from(["alpha-llm", "alpha-root"]) + ); + let foreign = reader() + .span_text( + &store, + &QueryScope::Owned { + user_id: String::new(), + team_ids: vec!["team-b".into()], + }, + "alpha", + &runs[0].trace_ref, + vec!["alpha-root".into()], + SpanPart::Input, + range, + None, + ) + .await?; + assert!(foreign.is_empty()); + Ok(()) +} diff --git a/litellm-rust/crates/traces/src/error.rs b/litellm-rust/crates/traces/src/error.rs index ee14abf8b8f..e01b0541216 100644 --- a/litellm-rust/crates/traces/src/error.rs +++ b/litellm-rust/crates/traces/src/error.rs @@ -14,10 +14,6 @@ pub enum Error { #[error("invalid trace query scope")] pub struct InvalidScope; -#[derive(Debug, thiserror::Error)] -#[error("unknown ClickHouse read query")] -pub struct InvalidQuery; - #[derive(Debug, thiserror::Error)] #[error("invalid trace call key")] pub struct InvalidCallKey; diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index 72c4baa2e67..4a7afa5694b 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -27,13 +27,12 @@ mod ui; mod view; pub mod wire; -pub use error::{Error, InvalidCallKey, InvalidQuery, InvalidScope}; +pub use error::{Error, InvalidCallKey, InvalidScope}; pub use normalize::{ AgentMetadata, AgentType, CallEvidence, CallEvidenceKind, CallKey, Integration, NormalizedSpan, ObservationType, }; pub use otlp::{DecodeLimits, DecodedEvent, DecodedSpan, decode_otlp, decode_otlp_with_limits}; -pub use query::ReadQuery; pub use query_access::QueryScope; pub use resolve::{SpendLookup, iso_time, listed_summary, resolve_trace}; pub use shared::{Shared, SharedIdentity}; diff --git a/litellm-rust/crates/traces/src/query.rs b/litellm-rust/crates/traces/src/query.rs index a0524dbb22b..ef9a0976845 100644 --- a/litellm-rust/crates/traces/src/query.rs +++ b/litellm-rust/crates/traces/src/query.rs @@ -1,17 +1 @@ pub mod guide; - -#[derive(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString, strum::Display, strum::AsRefStr)] -#[strum(serialize_all = "snake_case")] -pub enum ReadQuery { - Availability, - Agents, - Sample, - Content, - Evidence, -} - -impl ReadQuery { - pub fn parse(value: &str) -> Result { - value.parse().map_err(|_| crate::InvalidQuery) - } -} diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs index 1f9fa8f491c..9ef05f1ce74 100644 --- a/litellm-rust/crates/traces/src/schema.rs +++ b/litellm-rust/crates/traces/src/schema.rs @@ -49,11 +49,13 @@ pub fn schemas() -> BTreeMap<&'static str, Schema> { ("QueryScope", received::()), ("Tenant", received::()), ("TracePage", emitted::()), + ("SpanText", emitted::()), ("Trace", emitted::()), ("SpanDetail", emitted::()), ("SpanErrorPage", emitted::()), ("TraceHistogram", emitted::()), ("RunValues", emitted::()), ("RunField", received::()), + ("RunOrder", received::()), ]) } diff --git a/litellm-rust/crates/traces/src/search.rs b/litellm-rust/crates/traces/src/search.rs index b2b5ceef126..7be72d3c689 100644 --- a/litellm-rust/crates/traces/src/search.rs +++ b/litellm-rust/crates/traces/src/search.rs @@ -24,6 +24,28 @@ pub enum RunField { Model, Input, TraceId, + Service, + Team, +} + +/// What a `key:value` filter matches: a run field, or `attr.`, a span or resource +/// attribute that any span of the run carries. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum SearchKey { + Field(RunField), + Attribute(String), +} + +const ATTRIBUTE_PREFIX: &str = "attr."; + +impl SearchKey { + pub fn parse(key: &str) -> Option { + if let Some(attribute) = key.strip_prefix(ATTRIBUTE_PREFIX) { + return (!attribute.is_empty()).then(|| Self::Attribute(attribute.to_owned())); + } + let named = !key.is_empty() && key.chars().all(|c| c.is_ascii_alphabetic() || c == '_'); + named.then(|| key.parse().ok().map(Self::Field)).flatten() + } } #[derive(Clone, Debug, Default, Eq, PartialEq)] @@ -31,11 +53,13 @@ pub struct RunFilter { pub start_ms: i64, pub end_ms: i64, pub search: RunSearch, + /// When not empty, only these runs can match. + pub trace_refs: Vec, } #[derive(Clone, Debug, Eq, PartialEq)] pub struct FieldFilter { - pub field: RunField, + pub key: SearchKey, /// Matched against the whole value, ignoring case; `*` matches any run of characters. pub pattern: String, pub exclude: bool, @@ -67,11 +91,11 @@ impl RunSearch { .into_iter() .filter_map(|clause| match clause { Clause::Field { - field, + key, exclude, value, } => Some(FieldFilter { - field, + key, pattern: value, exclude, }), @@ -85,7 +109,7 @@ impl RunSearch { enum Clause { Text(String), Field { - field: RunField, + key: SearchKey, exclude: bool, value: String, }, @@ -135,16 +159,12 @@ fn clause(raw: &str) -> Clause { let (exclude, body) = raw .strip_prefix('-') .map_or((false, raw), |body| (true, body)); - let field = body.split_once(':').and_then(|(key, value)| { - let named = !key.is_empty() && key.chars().all(|c| c.is_ascii_alphabetic() || c == '_'); - named - .then(|| key.parse::().ok()) - .flatten() - .map(|field| (field, value)) - }); + let field = body + .split_once(':') + .and_then(|(key, value)| SearchKey::parse(key).map(|key| (key, value))); match field { - Some((field, value)) => Clause::Field { - field, + Some((key, value)) => Clause::Field { + key, exclude, value: unquote(value), }, diff --git a/litellm-rust/crates/traces/src/store.rs b/litellm-rust/crates/traces/src/store.rs index f9a67ab8af6..34d7ec75b6e 100644 --- a/litellm-rust/crates/traces/src/store.rs +++ b/litellm-rust/crates/traces/src/store.rs @@ -1,6 +1,6 @@ //! What trace storage must answer, independent of the engine behind it. -use std::ops::Range; +use std::{cmp::Ordering, ops::Range}; use serde::{Deserialize, Serialize}; @@ -13,17 +13,84 @@ pub enum RunSelection { TraceId(String), } +/// The last row of a page in its order: the row's sort value and its reference. #[derive(Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)] #[serde(deny_unknown_fields)] pub struct RunCursor { - pub start_ms: i64, + pub value: i64, pub trace_ref: String, } -/// Runs newest first, by `(start_ms, trace_ref)` descending. +#[macro_rules_attribute::apply(wire_type)] +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum RunSortKey { + #[default] + StartMs, + DurationMs, + SpanCount, + ErrorCount, + TraceRef, +} + +/// Runs by `key`, ties broken by `trace_ref` in the same direction. +#[macro_rules_attribute::apply(wire_type)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct RunOrder { + pub key: RunSortKey, + pub descending: bool, +} + +impl RunOrder { + pub const NEWEST: Self = Self { + key: RunSortKey::StartMs, + descending: true, + }; + pub const BY_REFERENCE: Self = Self { + key: RunSortKey::TraceRef, + descending: false, + }; + + pub fn value(self, row: &RunRow) -> i64 { + let count = |count: u64| i64::try_from(count).unwrap_or(i64::MAX); + match self.key { + RunSortKey::StartMs => row.start_ms, + RunSortKey::DurationMs => row.duration_ms, + RunSortKey::SpanCount => count(row.span_count), + RunSortKey::ErrorCount => count(row.error_count), + RunSortKey::TraceRef => 0, + } + } + + pub fn compare(self, left: &RunRow, right: &RunRow) -> Ordering { + let ascending = + (self.value(left), &left.trace_ref).cmp(&(self.value(right), &right.trace_ref)); + if self.descending { + ascending.reverse() + } else { + ascending + } + } + + pub fn cursor(self, row: &RunRow) -> RunCursor { + RunCursor { + value: self.value(row), + trace_ref: row.trace_ref.clone(), + } + } +} + +impl Default for RunOrder { + fn default() -> Self { + Self::NEWEST + } +} + #[derive(Clone, Debug, PartialEq)] pub struct RunQuery { pub selection: RunSelection, + pub order: RunOrder, pub after: Option, pub limit: u32, } @@ -57,25 +124,20 @@ pub struct RunRow { pub error_count: u64, } -impl RunRow { - pub fn cursor(&self) -> RunCursor { - RunCursor { - start_ms: self.start_ms, - trace_ref: self.trace_ref.clone(), - } - } -} - /// What a run is counted under. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[derive(Clone, Debug, Eq, PartialEq)] pub enum CountValue { Field(RunField), /// The run's alphabetically first agent label, or its service when it has none. PrimaryAgent, + /// Keys of the span and resource attributes its spans carry. + AttributeKey, + /// Values of one attribute across its spans. + Attribute(String), } /// Each dimension left unset collapses to one group: bucket 0, not failed, or an empty value. -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct CountBy { /// Equal-width slices of the filter window; run `i` lands in /// `(start_ms - window.start) * buckets / window.len()`. @@ -114,8 +176,10 @@ pub enum SpanSelection { trace_id: String, trace_ref: String, }, - /// Spans of several runs that started within `window`. + /// Spans of several runs that started within `window`. `trace_ids` are those runs' trace ids, + /// which narrow the read before references are checked. Runs { + trace_ids: Vec, trace_refs: Vec, window: Range, }, @@ -200,7 +264,18 @@ impl SpanRow { } } -#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Deserialize, Serialize, strum::IntoStaticStr)] +#[derive( + Clone, + Copy, + Debug, + Eq, + Hash, + PartialEq, + Deserialize, + Serialize, + strum::EnumString, + strum::IntoStaticStr, +)] #[serde(rename_all = "snake_case")] #[strum(serialize_all = "snake_case")] pub enum SpanPart { @@ -211,24 +286,43 @@ pub enum SpanPart { Attributes, } -/// A character range of one part of one span, read from the copy [`SpanQuery`] would return. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum TextRange { + /// Up to `max_chars` characters from `offset`; `None` reads to the end. + From { offset: u64, max_chars: Option }, + /// The last `chars` characters. + Last { chars: u64 }, +} + +impl TextRange { + pub const ALL: Self = Self::From { + offset: 0, + max_chars: None, + }; +} + +/// One part of each listed span of one run, read from the copy [`SpanQuery`] would return. +/// Spans that are not visible to the reader, or do not exist, are left out. #[derive(Clone, Debug, PartialEq)] pub struct SpanTextQuery { pub trace_id: String, pub trace_ref: String, - pub span_id: String, + pub span_ids: Vec, pub part: SpanPart, - pub offset: u64, - /// `None` reads to the end. - pub max_chars: Option, + pub range: TextRange, + /// Reports whether the whole part contains this text, case-sensitive. + pub contains: Option, } #[derive(Clone, Debug, Deserialize, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub struct SpanText { + pub span_id: String, pub text: String, pub total_chars: u64, /// Uppercase hex SHA-256 of the whole part, so a reader can tell when it changed. pub version: String, + pub contains: bool, } /// Gateway calls that can be priced against spans: those whose response id, call id or trace id diff --git a/litellm-rust/crates/traces/tests/query.rs b/litellm-rust/crates/traces/tests/query.rs deleted file mode 100644 index f0ad89bc54e..00000000000 --- a/litellm-rust/crates/traces/tests/query.rs +++ /dev/null @@ -1,25 +0,0 @@ -use litellm_traces::{InvalidQuery, ReadQuery}; -use rstest::rstest; - -#[rstest] -#[case::availability("availability", ReadQuery::Availability)] -#[case::agents("agents", ReadQuery::Agents)] -#[case::sample("sample", ReadQuery::Sample)] -#[case::content("content", ReadQuery::Content)] -#[case::evidence("evidence", ReadQuery::Evidence)] -fn names_select_the_public_query(#[case] name: &str, #[case] query: ReadQuery) { - assert_eq!(ReadQuery::parse(name).unwrap(), query); - assert_eq!(query.as_ref(), name); - assert_eq!(query.to_string(), name); -} - -#[rstest] -#[case::unknown("unknown")] -#[case::case_sensitive("Sample")] -#[case::whitespace(" sample")] -#[case::empty("")] -fn invalid_names_preserve_the_public_error(#[case] name: &str) { - let error = ReadQuery::parse(name).unwrap_err(); - assert!(matches!(error, InvalidQuery)); - assert_eq!(error.to_string(), "unknown ClickHouse read query"); -} diff --git a/litellm-rust/crates/traces/tests/search.rs b/litellm-rust/crates/traces/tests/search.rs index d2c7f4a8cd4..6644021d042 100644 --- a/litellm-rust/crates/traces/tests/search.rs +++ b/litellm-rust/crates/traces/tests/search.rs @@ -1,12 +1,12 @@ use litellm_traces::{ - search::{AgentRuns, FieldFilter, RunField, RunSearch, histogram}, + search::{AgentRuns, FieldFilter, RunField, RunSearch, SearchKey, histogram}, store::RunCount, }; use rstest::rstest; fn filter(field: RunField, pattern: &str, exclude: bool) -> FieldFilter { FieldFilter { - field, + key: SearchKey::Field(field), pattern: pattern.into(), exclude, } @@ -29,6 +29,10 @@ fn filter(field: RunField, pattern: &str, exclude: bool) -> FieldFilter { #[case::unknown_key_is_text("color:red", &["color:red"], vec![])] #[case::non_word_key_is_text("k1:v", &["k1:v"], vec![])] #[case::negated_text_stays_text("-foo", &["-foo"], vec![])] +#[case::service("service:billing", &[], vec![filter(RunField::Service, "billing", false)])] +#[case::team("-team:acme", &[], vec![filter(RunField::Team, "acme", true)])] +#[case::attribute("attr.gen_ai.system:openai", &[], vec![FieldFilter { key: SearchKey::Attribute("gen_ai.system".into()), pattern: "openai".into(), exclude: false }])] +#[case::attribute_without_a_key_is_text("attr.:x", &["attr.:x"], vec![])] fn parse_matches_the_dashboard_search_grammar( #[case] q: &str, #[case] text: &[&str], diff --git a/litellm/proxy/lens/AGENTS.md b/litellm/proxy/lens/AGENTS.md new file mode 100644 index 00000000000..cd5a58bc64a --- /dev/null +++ b/litellm/proxy/lens/AGENTS.md @@ -0,0 +1,6 @@ +- Lens product rules live here, in Python: sample percent, cap and preview, analyzer excerpt budgets and labels, evidence verification, and job and finding state +- Read traces only through the general trace reads the native bridge exposes: run listing with a sort order, run counts, the trace graph, and batched span text with ranges and substring checks +- Never add a Lens-only read to the Rust trace reader (`litellm-rust/crates/traces-cache`) or the storage port (`litellm_traces::store`). If only Lens needs it, compose it here from the general reads +- Never write SQL against trace storage from this package. Storage engines stay behind the Rust store port +- A target run is a run of the traces list, identified by its `trace_ref`, and its filter is the same `q` search the Traces tab uses +- Lens reads traces only, never the gateway request log diff --git a/litellm/proxy/lens/analysis.py b/litellm/proxy/lens/analysis.py index 7286cec9ba8..e5456fbbe08 100644 --- a/litellm/proxy/lens/analysis.py +++ b/litellm/proxy/lens/analysis.py @@ -498,7 +498,7 @@ async def investigate_stored( "workflow_outlines": tuple( { "execution_id": item.execution.id, - "recorded_span_count": item.execution.span_count, + "recorded_span_count": item.execution.summary["span_count"] if item.execution.summary else None, "partial": item.partial, "cannot_assess": item.cannot_assess, "available_unique_spans": len(frozenset(p.span_id for p in item.parts)), diff --git a/litellm/proxy/lens/endpoints.py b/litellm/proxy/lens/endpoints.py index 32339e7befe..335459d38c0 100644 --- a/litellm/proxy/lens/endpoints.py +++ b/litellm/proxy/lens/endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper from litellm.proxy.lens.billing import validate_key from litellm.proxy.lens.inference import Deployment, deployment_prices from litellm.proxy.lens.models import ( + ActivityAvailability, ActivitySelection, Claim, Execution, @@ -43,11 +44,12 @@ from litellm.proxy.lens.models import ( WatchSkipped, Worker, WorkerCreated, + parse_execution, ) from litellm.proxy.lens.release import PROTOCOL_VERSION, release_tag, worker_image from litellm.proxy.lens.repository import LensRepository, WriterDatabase -from litellm.proxy.lens.search import LensField, parse_search -from litellm.proxy.lens.sources import ActivityAvailability, SourceReader, Storage, parse_execution +from litellm.proxy.lens.search import LensField, parse_search, search_terms +from litellm.proxy.lens.sources import SourceReader, Storage from litellm.proxy.lens.state import ( add_step, can_access, @@ -134,9 +136,7 @@ def required(lens: Lens | None) -> Lens: def validate_selection(settings: ActivitySelection) -> None: for identity in settings.execution_ids: try: - source, _, _, _ = parse_execution(identity) - if source not in ("traces", "requests"): - raise ValueError("Unsupported source") + parse_execution(identity) except ValueError: raise HTTPException(422, "Choose execution IDs returned by the activity preview") @@ -241,12 +241,6 @@ async def activity_available(auth: Auth, storage: StorageDep) -> ActivityAvailab return await source_reader(storage).availability(scope) if storage is not None else ActivityAvailability() -@router.get("/agents", response_model=tuple[str, ...]) -async def list_agents(auth: Auth, storage: StorageDep) -> tuple[str, ...]: - scope: Final = user_scope(auth) - return await source_reader(storage).agents(scope) if storage is not None else () - - @router.get("/values/{field}", response_model=tuple[str, ...]) async def list_lens_values( field: LensField, @@ -324,7 +318,12 @@ def run_window(lens: Lens, body: RunRequest, now: datetime) -> tuple[datetime, d def run_settings(lens: Lens, body: RunRequest) -> LensSettings | None: if body.agent_name is None: return body.settings - return (body.settings or lens.settings).model_copy(update=MappingProxyType({"agent_name": body.agent_name})) + settings: Final = body.settings or lens.settings + agent: Final = ( + f'agent:"{body.agent_name}"' if any(c.isspace() for c in body.agent_name) else f"agent:{body.agent_name}" + ) + kept: Final = tuple(term for term in search_terms(settings.q) if not term.lower().startswith("agent:")) + return settings.model_copy(update=MappingProxyType({"q": " ".join((*kept, agent))})) @router.post("/{lens_id}/runs", response_model=Lens) @@ -414,7 +413,7 @@ async def update_finding(lens_id: str, finding_id: str, body: FindingUpdate, aut class Preview(BaseModel): as_of: AwareDatetime | None = None - offset: int = Field(default=0, ge=0) + cursor: str = "" selection: ActivitySelection lookback_hours: LookbackHours = 24 @@ -433,7 +432,7 @@ async def preview_sample(body: Preview, auth: Auth, storage: StorageDep) -> Samp body.selection, start, end, - offset=body.offset, + cursor=body.cursor, preview=True, ) @@ -748,20 +747,8 @@ async def evidence_content( ) -> ExecutionContent: lens: Final = await get_lens(lens_id, user_scope(auth)) try: - source, team, trace_id, trace_ref = parse_execution(execution_id) + trace_ref, trace_id = parse_execution(execution_id) except ValueError: raise HTTPException(404, "Execution not found") - 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, - source="traces" if source == "traces" else "requests", - trace_id=trace_id, - trace_ref=trace_ref, - team_id=team, - name=trace_id, - start_time="", - span_count=1, - root_seen=source == "requests", - ) + execution: Final = Execution(id=execution_id, trace_id=trace_id, trace_ref=trace_ref) return await source_reader(storage).content(lens.scope, execution, cursor, offset) diff --git a/litellm/proxy/lens/models.py b/litellm/proxy/lens/models.py index 90c92cc7acd..85fa7a4e039 100644 --- a/litellm/proxy/lens/models.py +++ b/litellm/proxy/lens/models.py @@ -1,7 +1,11 @@ +import base64 from datetime import datetime, timedelta, timezone from typing import Annotated, Final, Literal, TypeAlias -from pydantic import AfterValidator, BaseModel, ConfigDict, Field, model_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, model_validator +from typing_extensions import ReadOnly, TypedDict + +from litellm.rust_bridge.trace.generated.types import TraceSummary def calendar_lookback(hours: int) -> int: @@ -34,9 +38,69 @@ class Scope(Record): all_teams: bool = False -class MetadataFilter(Record): - key: str = Field(min_length=1) - value: str = Field(min_length=1) +TRACE_REF_LENGTH: Final = 64 + + +def execution_id(trace_ref: str, trace_id: str) -> str: + return f"{trace_ref}:{trace_id}" + + +def parse_execution(value: str) -> tuple[str, str]: + """`(trace_ref, trace_id)` of an execution id.""" + trace_ref, separator, trace_id = value.partition(":") + if len(trace_ref) != TRACE_REF_LENGTH or not separator or not trace_id: + raise ValueError("Not an execution ID") + return trace_ref, trace_id + + +_LEGACY_ID: Final[TypeAdapter[tuple[str, str, str] | tuple[str, str, str, str]]] = TypeAdapter( + tuple[str, str, str] | tuple[str, str, str, str] +) + + +def _legacy_execution_id(value: str) -> str | None: + """Selections saved before executions were runs named them by source, team, trace id and reference.""" + try: + parts: Final = _LEGACY_ID.validate_json(base64.urlsafe_b64decode(value)) + except (ValueError, ValidationError): + return value + trace_ref: Final = parts[3] if len(parts) == 4 else "" + return execution_id(trace_ref, parts[2]) if parts[0] == "traces" and trace_ref else None + + +def _term(key: str, value: str) -> str: + return f'{key}:"{value}"' if any(c.isspace() for c in value) else f"{key}:{value}" + + +class _LegacyFilter(BaseModel): + key: str + value: str + + +class _LegacySelection(BaseModel): + model_config = ConfigDict(extra="allow") + q: str = "" + source: str = "" + service: str = "" + agent_name: str = "" + filters: tuple[_LegacyFilter, ...] = () + team_id: str = "" + execution_ids: tuple[str, ...] = () + + def current(self) -> dict[str, object]: + terms: Final = ( + self.q, + _term("agent", self.agent_name) if self.agent_name else "", + _term("service", self.service) if self.service else "", + _term("team", self.team_id) if self.team_id else "", + *(_term(f"attr.{f.key}", f.value) for f in self.filters), + ) + ids: Final = tuple(i for i in map(_legacy_execution_id, self.execution_ids) if i is not None) + return {**(self.model_extra or {}), "q": " ".join(t for t in terms if t), "execution_ids": ids} + + +_LEGACY_KEYS: Final = frozenset({"source", "service", "agent_name", "filters", "team_id"}) +_FIELDS: Final = TypeAdapter(dict[str, object]) class Check(Record): @@ -46,15 +110,23 @@ class Check(Record): class ActivitySelection(Record): - source: Literal["traces", "requests", "both"] = "traces" - service: str = Field(default="") - agent_name: str = Field(default="") - filters: tuple[MetadataFilter, ...] = Field(default=()) + q: str = Field(default="", max_length=2000) sample_size: int | None = Field(default=None, ge=1) sample_percent: float = Field(default=100, gt=0, le=100, allow_inf_nan=False) - team_id: str = "" execution_ids: tuple[str, ...] = () + @model_validator(mode="before") + @classmethod + def from_saved_filters(cls, data: object) -> object: + """Selections saved with separate agent, service, team and attribute filters load as `q`.""" + try: + fields: Final = _FIELDS.validate_python(data) + except ValidationError: + return data + if not _LEGACY_KEYS & fields.keys(): + return data + return _LegacySelection.model_validate(fields).current() + class LensSettings(ActivitySelection): name: str = Field(min_length=1) @@ -149,16 +221,9 @@ class Coverage(Record): class Execution(Record): id: str - source: Literal["traces", "requests"] trace_id: str - trace_ref: str = "" - team_id: str - name: str - start_time: str - span_count: int - root_seen: bool = False - service: str = "" - metadata: tuple[MetadataFilter, ...] = () + trace_ref: str + summary: TraceSummary | None = None class TracePart(Record): @@ -182,10 +247,24 @@ class Sample(Record): executions: tuple[Execution, ...] eligible: int selected: int = 0 - next_offset: int | None = None next_cursor: str | None = None +class ActivityAvailability(Record): + traces: bool = False + + +class _SavedSample(TypedDict, total=False): + executions: ReadOnly[list[dict[str, object]]] + + +class _SavedJob(TypedDict, total=False): + sample: ReadOnly[_SavedSample | None] + + +_SAVED_JOB: Final = TypeAdapter(_SavedJob) + + class RunAssessment(Record): execution_id: str issue_checks: tuple[str, ...] = () @@ -208,6 +287,20 @@ class Step(Record): class Job(Record): + @model_validator(mode="before") + @classmethod + def without_legacy_sample(cls, data: object) -> object: + """Samples saved before executions were runs no longer resolve, so they load as absent.""" + try: + saved: Final = _SAVED_JOB.validate_python(data) + fields: Final = _FIELDS.validate_python(data) + except ValidationError: + return data + sample: Final = saved.get("sample") or {} + if not any("source" in e for e in sample.get("executions", [])): + return data + return {**fields, "sample": None} + id: str status: Literal["queued", "running", "completed", "failed", "cancelled"] = "queued" stage: str = "Queued" diff --git a/litellm/proxy/lens/search.py b/litellm/proxy/lens/search.py index 78701cc1517..902771e5756 100644 --- a/litellm/proxy/lens/search.py +++ b/litellm/proxy/lens/search.py @@ -9,13 +9,13 @@ _SETTINGS: Final = "data->'settings'" FIELD_VALUES: Final[MappingProxyType[LensField, str]] = MappingProxyType( { "name": f"ARRAY[{_SETTINGS}->>'name']", - "agent": f"ARRAY[{_SETTINGS}->>'agent_name', {_SETTINGS}->>'service']", + "agent": f"ARRAY[{_SETTINGS}->>'q', {_SETTINGS}->>'agent_name', {_SETTINGS}->>'service']", "status": "ARRAY[COALESCE(data->'jobs'->0->>'status', 'never')]", "schedule": f"ARRAY[CASE WHEN ({_SETTINGS}->>'enabled')::boolean THEN 'watching' ELSE 'paused' END]", } ) -_SCOPE_LABEL: Final = f"""COALESCE(NULLIF(concat_ws(' · ', +_SCOPE_LABEL: Final = f"""COALESCE(NULLIF({_SETTINGS}->>'q', ''), NULLIF(concat_ws(' · ', NULLIF({_SETTINGS}->>'agent_name', ''), NULLIF({_SETTINGS}->>'service', ''), (SELECT string_agg((f->>'key') || ': ' || (f->>'value'), ' · ') @@ -23,6 +23,13 @@ _SCOPE_LABEL: Final = f"""COALESCE(NULLIF(concat_ws(' · ', FREE_TEXT: Final = f"ARRAY[{_SETTINGS}->>'name', {_SCOPE_LABEL}]" _TOKEN: Final = re.compile(r'(?:"[^"]*"?|\S)+') + + +def search_terms(q: str) -> tuple[str, ...]: + """Whitespace-separated terms of a search, keeping quoted stretches whole.""" + return tuple(_TOKEN.findall(q)) + + _FIELD_TOKEN: Final = re.compile(r"^(-?)([A-Za-z_]+):(.*)$", re.DOTALL) _QUOTED: Final = re.compile(r'^"([^"]*)"?$') diff --git a/litellm/proxy/lens/sources.py b/litellm/proxy/lens/sources.py index e36653aa091..2cb535714c0 100644 --- a/litellm/proxy/lens/sources.py +++ b/litellm/proxy/lens/sources.py @@ -1,61 +1,143 @@ import base64 -import json -from collections.abc import Awaitable, Sequence -from typing import Final, Protocol, TypeAlias +import math +import time +from collections.abc import Awaitable, Mapping, Sequence +from itertools import accumulate, chain +from typing import Final, Protocol -from pydantic import TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.proxy.lens.models import ( + ActivityAvailability, ActivitySelection, Evidence, Execution, ExecutionContent, - MetadataFilter, Sample, Scope, TracePart, + execution_id, + parse_execution, ) -from litellm.rust_bridge.trace.generated.models import ( - ActivityAvailability, - AgentRow, - CountRow, - ExecutionRow, - LensAccessParams, - LensContentParams, - LensEvidenceParams, - LensSampleParams, - PartRow, +from litellm.rust_bridge.trace.generated.types import ( + AllQueryScope, + OwnedQueryScope, + QueryScope, + RunOrder, + Span, + SpanText, + Trace, + TracePage, + TraceSummary, ) +from litellm.rust_bridge.trace.storage import BY_REFERENCE, NEWEST, SpanPart class Storage(Protocol): - def lens_availability(self, parameters: LensAccessParams) -> Awaitable[Sequence[ActivityAvailability]]: ... - def lens_agents(self, parameters: LensAccessParams) -> Awaitable[Sequence[AgentRow]]: ... - def lens_sample(self, parameters: LensSampleParams) -> Awaitable[Sequence[ExecutionRow]]: ... - def lens_content(self, parameters: LensContentParams) -> Awaitable[Sequence[PartRow]]: ... - def lens_evidence(self, parameters: LensEvidenceParams) -> Awaitable[Sequence[CountRow]]: ... + def list_traces( + self, + scope: QueryScope, + start_ms: int, + end_ms: int, + q: str = "", + cursor: str | None = None, + limit: int = 50, + order: RunOrder = NEWEST, + trace_refs: Sequence[str] = (), + ) -> Awaitable[TracePage]: ... + + def count_traces( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str = "", trace_refs: Sequence[str] = () + ) -> Awaitable[int]: ... + + def get_trace( + self, + trace_id: str, + scope: QueryScope, + trace_ref: str = "", + cursor: str | None = None, + page_size: int | None = None, + ) -> Awaitable[Trace | None]: ... + + def span_text( + self, + trace_id: str, + trace_ref: str, + span_ids: Sequence[str], + part: SpanPart, + scope: QueryScope, + offset: int = 0, + max_chars: int | None = None, + tail: bool = False, + contains: str | None = None, + ) -> Awaitable[tuple[SpanText, ...]]: ... -ExecutionIdParts: TypeAlias = tuple[str, str, str] | tuple[str, str, str, str] -_EXECUTION_ID: Final[TypeAdapter[ExecutionIdParts]] = TypeAdapter(ExecutionIdParts) +PAGE_SPANS: Final = 40 +BUDGET: Final = 8_000 +OMITTED: Final = "\n[... content omitted ...]\n" +PARTS: Final[tuple[tuple[SpanPart, str, int], ...]] = ( + ("input", "Input: ", 2_000), + ("output", "\nOutput: ", 5_000), + ("error", "\nStatus: ", 500), +) +Texts = Mapping[tuple[str, SpanPart], SpanText] -def execution_id(source: str, team_id: str, trace_id: str, trace_ref: str = "") -> str: - return base64.urlsafe_b64encode(json.dumps((source, team_id, trace_id, trace_ref)).encode()).decode() +def lens_access(scope: Scope) -> QueryScope: + if scope.all_teams: + return AllQueryScope(kind="all") + return OwnedQueryScope(kind="owned", user_id="", team_ids=(scope.team_id,) if scope.team_id else ()) -def parse_execution(value: str) -> tuple[str, str, str, str]: - parts: Final = _EXECUTION_ID.validate_json(base64.urlsafe_b64decode(value)) - return (parts[0], parts[1], parts[2], parts[3] if len(parts) == 4 else "") +def execution_of(run: TraceSummary) -> Execution: + trace_ref: Final = run.get("trace_ref", "") + return Execution( + id=execution_id(trace_ref, run["trace_id"]), trace_id=run["trace_id"], trace_ref=trace_ref, summary=run + ) -def access_parameters(scope: Scope) -> LensAccessParams: - return LensAccessParams(all_teams=1 if scope.all_teams else 0, team=scope.team_id, key_hash=scope.api_key_hash) +class SamplePosition(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + cursor: str | None + offset: int -def selection_id(value: str) -> str: - source, team, trace_id, trace_ref = parse_execution(value) - return "\0".join((source, team, trace_ref or trace_id)) +_POSITION: Final = TypeAdapter(SamplePosition) + + +def _encode(position: SamplePosition) -> str: + return base64.urlsafe_b64encode(position.model_dump_json().encode()).decode() + + +def _decode(cursor: str) -> SamplePosition: + if not cursor: + return SamplePosition(cursor=None, offset=0) + try: + return _POSITION.validate_json(base64.urlsafe_b64decode(cursor)) + except (ValueError, ValidationError) as error: + raise ValueError("Invalid sample cursor") from error + + +def selected_count(selection: ActivitySelection, eligible: int) -> int: + share: Final = math.ceil(eligible * selection.sample_percent / 100) + return min(share, selection.sample_size) if selection.sample_size else share + + +def _status(span: Span) -> str: + return f"{span['status']} " + + +def _label(span: Span, part: SpanPart, label: str) -> str: + return label + _status(span) if part == "error" else label + + +def _pieces(span: Span, texts: Texts) -> tuple[tuple[SpanPart, str, SpanText | None], ...]: + return tuple((part, _label(span, part, label), texts.get((span["span_id"], part))) for part, label, _ in PARTS) + + +def _total(pieces: tuple[tuple[SpanPart, str, SpanText | None], ...]) -> int: + return sum(len(label) + (text["total_chars"] if text else 0) for _, label, text in pieces) class SourceReader: @@ -63,118 +145,194 @@ class SourceReader: self.storage: Final = storage async def availability(self, scope: Scope) -> ActivityAvailability: - rows: Final = await self.storage.lens_availability(access_parameters(scope)) - return rows[0] if rows else ActivityAvailability() - - async def agents(self, scope: Scope) -> tuple[str, ...]: - rows: Final = await self.storage.lens_agents(access_parameters(scope)) - return tuple(row.agent_name for row in rows) + found: Final = await self.storage.count_traces(lens_access(scope), 0, int(time.time() * 1000) + 1) + return ActivityAvailability(traces=found > 0) async def sample( self, scope: Scope, - settings: ActivitySelection, + selection: ActivitySelection, start: int, end: int, - offset: int = 0, + cursor: str = "", page_size: int = 100, preview: bool = False, - cursor: str = "", ) -> Sample: - params: Final = LensSampleParams( - all_teams=1 if scope.all_teams else 0, - team=scope.team_id, - key_hash=scope.api_key_hash, - source=settings.source, - start=start, - end=end, - service=settings.service, - agent_name=settings.agent_name, - filter_keys=tuple(f.key for f in settings.filters), - filter_values=tuple(f.value for f in settings.filters), - limit=page_size, - offset=offset, - after=cursor, - sample_percent=settings.sample_percent, - sample_cap=settings.sample_size or 0, - preview=1 if preview else 0, - selected_team=settings.team_id, - execution_ids=tuple(selection_id(value) for value in settings.execution_ids), + """The selected runs in reference order: a stable order unrelated to time, so a prefix is a fair sample.""" + access: Final = lens_access(scope) + refs: Final = tuple(parse_execution(identity)[0] for identity in selection.execution_ids) + eligible: Final = await self.storage.count_traces(access, start, end, selection.q, refs) + selected: Final = selected_count(selection, eligible) + bound: Final = eligible if preview else selected + position: Final = _decode(cursor) + limit: Final = min(page_size, bound - position.offset) + if limit <= 0: + return Sample(executions=(), eligible=eligible, selected=selected) + page: Final = await self.storage.list_traces( + access, start, end, selection.q, position.cursor, limit, BY_REFERENCE, refs ) - rows: Final = await self.storage.lens_sample(params) + executions: Final = tuple(execution_of(run) for run in page["data"]) + offset: Final = position.offset + len(executions) + next_page: Final = page["next_cursor"] return Sample( - eligible=rows[0].eligible if rows else 0, - selected=rows[0].selected if rows else 0, - next_cursor=rows[-1].selection_key if len(rows) == page_size else None, - next_offset=( - offset + len(rows) - if page_size and rows and offset + len(rows) < (rows[0].eligible if preview else rows[0].selected) - else None - ), - executions=tuple( - Execution( - id=execution_id(row.source, row.team_id, row.trace_id, row.trace_ref), - source=row.source, - trace_id=row.trace_id, - trace_ref=row.trace_ref, - team_id=row.team_id, - name=row.name, - start_time=row.start_time, - span_count=row.span_count, - root_seen=bool(row.root_seen), - service=row.service, - metadata=tuple( - MetadataFilter(key=k, value=v) - for k, v in row.attributes - if k != "litellm.api_key_hash" and k and v - ), - ) - for row in rows + executions=executions, + eligible=eligible, + selected=selected, + next_cursor=( + _encode(SamplePosition(cursor=next_page, offset=offset)) if next_page and offset < bound else None ), ) + async def _texts( + self, access: QueryScope, execution: Execution, span_ids: Sequence[str], max_chars: int, tail: bool = False + ) -> Mapping[tuple[str, SpanPart], SpanText]: + if not span_ids: + return {} + reads: Final[list[tuple[SpanPart, tuple[SpanText, ...]]]] = [ + ( + part, + await self.storage.span_text( + execution.trace_id, + execution.trace_ref, + span_ids, + part, + access, + max_chars=max_chars if not tail else budget - budget // 3, + tail=tail, + ), + ) + for part, _, budget in PARTS + ] + return { + (text["span_id"], part): text for part, texts in reads for text in texts + } # comprehension-ok: flatten one read per part + async def content(self, scope: Scope, execution: Execution, cursor: str = "", offset: int = 0) -> ExecutionContent: - params: Final = LensContentParams( - all_teams=1 if scope.all_teams else 0, - team=scope.team_id, - key_hash=scope.api_key_hash, - source=execution.source, - id=execution.trace_id, - trace_ref=execution.trace_ref, - record_team=execution.team_id, - cursor=cursor, - offset=offset + 1, + """Each span as `Input: … Output: … Status: …`. A first read keeps the start and end of parts over budget; + a later `offset` reads the budget's worth of the full text from there.""" + access: Final = lens_access(scope) + trace: Final = await self.storage.get_trace(execution.trace_id, access, execution.trace_ref) + if trace is None: + return ExecutionContent(execution=execution, parts=(), partial=True) + spans: Final = trace["spans"] + ids: Final = tuple(span["span_id"] for span in spans) + start: Final = ids.index(cursor) + 1 if cursor in ids else 0 if not cursor else len(ids) + page: Final = spans[start : start + PAGE_SPANS] + page_ids: Final = tuple(span["span_id"] for span in page) + heads: Final = await self._texts(access, execution, page_ids, BUDGET) + long: Final = tuple(span["span_id"] for span in page if offset == 0 and _total(_pieces(span, heads)) > BUDGET) + tails: Final = await self._texts(access, execution, long, 0, tail=True) + parts: Final = tuple( + [ + await self._part(access, execution, span, heads, tails, offset) + for span in page # comprehension-ok: sequential reads keep storage load bounded + ] ) - rows: Final = await self.storage.lens_content(params) + root_seen: Final = any(span.get("parent_span_id") is None for span in spans) return ExecutionContent( execution=execution, - parts=tuple( - TracePart( - execution_id=execution.id, - span_id=row.span_id, - parent_span_id=row.parent_span_id, - name=row.name, - kind=row.kind, - content=row.content, - truncated=bool(row.truncated), - ) - for row in rows - ), - next_cursor=rows[-1].span_id if len(rows) == 40 else None, - partial=not execution.root_seen or any(row.truncated for row in rows), + parts=parts, + next_cursor=page_ids[-1] if page_ids and start + len(page) < len(spans) else None, + partial=not root_seen or any(part.truncated for part in parts), ) - async def verify_evidence(self, scope: Scope, execution: Execution, evidence: Evidence) -> bool: - params: Final = LensEvidenceParams( - all_teams=1 if scope.all_teams else 0, - team=scope.team_id, - key_hash=scope.api_key_hash, - source=execution.source, - id=execution.trace_id, - trace_ref=execution.trace_ref, - record_team=execution.team_id, - span=evidence.span_id, - quote=evidence.quote, + async def _part( + self, access: QueryScope, execution: Execution, span: Span, heads: Texts, tails: Texts, offset: int + ) -> TracePart: + pieces: Final = _pieces(span, heads) + total: Final = _total(pieces) + content, truncated = ( + (_excerpt(span, pieces, tails), total > BUDGET) + if offset == 0 + else (await self._window(access, execution, span, pieces, offset), offset + BUDGET < total) ) - rows: Final = await self.storage.lens_evidence(params) - return bool(rows and rows[0].count) + return TracePart( + execution_id=execution.id, + span_id=span["span_id"], + parent_span_id=span.get("parent_span_id") or "", + name=span["name"], + kind=span["type"], + content=content, + truncated=truncated, + ) + + async def _window( + self, + access: QueryScope, + execution: Execution, + span: Span, + pieces: tuple[tuple[SpanPart, str, SpanText | None], ...], + offset: int, + ) -> str: + starts: Final = tuple( + accumulate((len(label) + (text["total_chars"] if text else 0) for _, label, text in pieces), initial=0) + ) + chunks: Final = [ + await self._window_piece(access, execution, span, piece, start, offset) + for piece, start in zip(pieces, starts) + ] + return "".join(chunks) + + async def _window_piece( + self, + access: QueryScope, + execution: Execution, + span: Span, + piece: tuple[SpanPart, str, SpanText | None], + start: int, + offset: int, + ) -> str: + part, label, text = piece + end: Final = offset + BUDGET + shown_label: Final = label[max(offset - start, 0) : max(end - start, 0)] + text_start: Final = start + len(label) + if text is None: + return shown_label + low, high = max(offset, text_start), min(end, text_start + text["total_chars"]) + if low >= high: + return shown_label + if high - text_start <= len(text["text"]): + return shown_label + text["text"][low - text_start : high - text_start] + read: Final = await self.storage.span_text( + execution.trace_id, + execution.trace_ref, + (span["span_id"],), + part, + access, + offset=low - text_start, + max_chars=high - low, + ) + return shown_label + (read[0]["text"] if read else "") + + async def verify_evidence(self, scope: Scope, execution: Execution, evidence: Evidence) -> bool: + access: Final = lens_access(scope) + found: Final = [ + await self.storage.span_text( + execution.trace_id, + execution.trace_ref, + (evidence.span_id,), + part, + access, + max_chars=0, + contains=evidence.quote, + ) + for part, _, _ in PARTS + ] + return any(text["contains"] for text in chain.from_iterable(found)) + + +def _excerpt(span: Span, pieces: tuple[tuple[SpanPart, str, SpanText | None], ...], tails: Texts) -> str: + if _total(pieces) <= BUDGET: + return "".join(label + (text["text"] if text else "") for _, label, text in pieces) + return "".join( + label + _shortened(text, budget, tails.get((span["span_id"], part))) + for (part, label, text), (_, _, budget) in zip(pieces, PARTS) + ) + + +def _shortened(text: SpanText | None, budget: int, tail: SpanText | None) -> str: + if text is None: + return "" + if text["total_chars"] <= budget: + return text["text"] + return text["text"][: budget // 3] + OMITTED + (tail["text"] if tail else "") diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index 1d91615da6e..e6392c9b2c9 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -36,6 +36,7 @@ from litellm.rust_bridge.trace.generated.types import ( OwnedQueryScope, QueryScope, RunField, + RunOrder, RunValues, SpanDetail, SpanErrorPage, @@ -206,11 +207,14 @@ async def list_agent_traces( window: Annotated[TraceWindow, Depends(trace_window)], q: RunQuery = "", cursor: Annotated[str | None, Query(max_length=512)] = None, + sort_by: Literal["start_ms", "duration_ms", "span_count", "error_count"] = "start_ms", + sort_dir: Literal["asc", "desc"] = "desc", ) -> TracePage: + order: Final = RunOrder(key=sort_by, descending=sort_dir == "desc") try: tracing, scope = context.reader() return await tracing.list_traces( - scope=scope, start_ms=window.start_ms, end_ms=window.end_ms, q=q, cursor=cursor + scope=scope, start_ms=window.start_ms, end_ms=window.end_ms, q=q, cursor=cursor, order=order ) except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: raise read_failure(error) from error diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index f3f7b4114d8..4f9e3b8f115 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -11,7 +11,7 @@ from litellm.rust_bridge.embeddings.entrypoints import LiteLLMEmbeddingRequest from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest -from litellm.rust_bridge.trace.generated.types import QueryScope, ReadQueryName +from litellm.rust_bridge.trace.generated.types import QueryScope from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -43,7 +43,30 @@ class NativeTraceStorage: def insert_rows(self, table: str, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... def ingest(self, payload: bytes, content_type: str | None, tenant: Mapping[str, str]) -> Future[int]: ... def list_traces( - self, scope: QueryScope, start_ms: int, end_ms: int, q: str, cursor: str | None, limit: int + self, + scope: QueryScope, + start_ms: int, + end_ms: int, + q: str, + cursor: str | None, + limit: int, + order: str = "newest", + trace_refs: Sequence[str] = (), + ) -> Future[JsonValue]: ... + def count_traces( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, trace_refs: Sequence[str] = () + ) -> Future[JsonValue]: ... + def span_text( + self, + trace_id: str, + trace_ref: str, + span_ids: Sequence[str], + part: str, + scope: QueryScope, + offset: int = 0, + max_chars: int | None = None, + tail: bool = False, + contains: str | None = None, ) -> Future[JsonValue]: ... def trace_histogram( self, scope: QueryScope, start_ms: int, end_ms: int, q: str, buckets: int @@ -60,7 +83,6 @@ class NativeTraceStorage: ) -> Future[JsonValue]: ... def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Future[str]: ... def query_help(self, scope: QueryScope, secret: str) -> Future[JsonValue]: ... - def query(self, query: ReadQueryName, parameters: Mapping[str, str | int | float | Sequence[str]]) -> Future[str]: ... @final class NativeDiagnosticProcessor: @@ -403,27 +425,42 @@ class _SecretManagerRuntime: def read_secret(self, name: str, settings: Mapping[str, object] | None = None) -> JsonValue: ... def read_secret_async(self, name: str, settings: Mapping[str, object] | None = None) -> Future[JsonValue]: ... def async_write_secret( - self, secret_name: str, secret_value: str, description: str | None = None, + self, + secret_name: str, + secret_value: str, + description: str | None = None, optional_params: Mapping[str, object] | None = None, - timeout: float | httpx.Timeout | None = None, tags: object = None, + timeout: float | httpx.Timeout | None = None, + tags: object = None, ) -> Future[dict[str, JsonValue]]: ... def async_delete_secret( - self, secret_name: str, recovery_window_in_days: int | None = None, + self, + secret_name: str, + recovery_window_in_days: int | None = None, optional_params: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, ) -> Future[dict[str, JsonValue]]: ... def async_rotate_secret( - self, current_secret_name: str, new_secret_name: str, new_secret_value: str, + self, + current_secret_name: str, + new_secret_name: str, + new_secret_value: str, optional_params: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, ) -> Future[dict[str, JsonValue]]: ... def sync_read_secret( - self, secret_name: str, optional_params: Mapping[str, object] | None = None, - timeout: float | httpx.Timeout | None = None, primary_secret_name: str | None = None, + self, + secret_name: str, + optional_params: Mapping[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + primary_secret_name: str | None = None, ) -> JsonValue: ... def async_read_secret( - self, secret_name: str, optional_params: Mapping[str, object] | None = None, - timeout: float | httpx.Timeout | None = None, primary_secret_name: str | None = None, + self, + secret_name: str, + optional_params: Mapping[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + primary_secret_name: str | None = None, ) -> Future[JsonValue]: ... @final @@ -431,11 +468,18 @@ class NativeCacheHandle: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @staticmethod def memory( - *, ttl: float = 600.0, capacity: int = 200, max_entry_bytes: int = 4194304, + *, + ttl: float = 600.0, + capacity: int = 200, + max_entry_bytes: int = 4194304, ) -> NativeCacheHandle: ... @staticmethod def redis( - url: str, *, namespace: str, ttl: float = 600.0, max_entry_bytes: int = 4194304, + url: str, + *, + namespace: str, + ttl: float = 600.0, + max_entry_bytes: int = 4194304, ) -> NativeCacheHandle: ... def get(self, key: str) -> object: ... def set(self, key: str, value: object, *, ttl: float | None = None) -> None: ... diff --git a/litellm/rust_bridge/trace/generated/models.py b/litellm/rust_bridge/trace/generated/models.py index 9987ae126c4..65aa0882962 100644 --- a/litellm/rust_bridge/trace/generated/models.py +++ b/litellm/rust_bridge/trace/generated/models.py @@ -6,277 +6,6 @@ from typing import Annotated, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field - -class ActivityAvailability(BaseModel): - model_config = ConfigDict( - frozen=True, - ) - - traces: bool = False - requests: bool = False - - -class AgentRow(BaseModel): - model_config = ConfigDict( - frozen=True, - ) - - agent_name: str - - -Count: TypeAlias = Annotated[ - int, - Field( - ..., - ge=0, - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - le=18446744073709551615, - ), -] - - -Count1: TypeAlias = Annotated[ - str, - Field( - ..., - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - pattern="^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - ), -] - - -class CountRow(BaseModel): - model_config = ConfigDict( - frozen=True, - ) - - count: int = Field(..., ge=0, le=18446744073709551615) - - -ContentSource: TypeAlias = Literal["traces", "requests"] - - -SpanCount: TypeAlias = Annotated[ - int, - Field( - ..., - ge=0, - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - le=18446744073709551615, - ), -] - - -SpanCount1: TypeAlias = Annotated[ - str, - Field( - ..., - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - pattern="^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - ), -] - - -Attribute: TypeAlias = Annotated[tuple[str, str], Field(..., max_length=2, min_length=2)] - - -Eligible: TypeAlias = Annotated[ - int, - Field( - ..., - ge=0, - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - le=18446744073709551615, - ), -] - - -Eligible1: TypeAlias = Annotated[ - str, - Field( - ..., - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - pattern="^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - ), -] - - -Selected: TypeAlias = Annotated[ - int, - Field( - ..., - ge=0, - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - le=18446744073709551615, - ), -] - - -Selected1: TypeAlias = Annotated[ - str, - Field( - ..., - json_schema_extra={ - "x-python-normalized": { - "type": "int", - "minimum": 0, - "maximum": 18446744073709551615, - } - }, - pattern="^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - ), -] - - -class ExecutionRow(BaseModel): - model_config = ConfigDict( - frozen=True, - ) - - source: ContentSource - trace_id: str - team_id: str - trace_ref: str = "" - name: str - start_time: str - span_count: int = Field(..., ge=0, le=18446744073709551615) - root_seen: int = Field(..., ge=0, le=1) - service: str = "" - attributes: tuple[Attribute, ...] = () - eligible: int = Field(..., ge=0, le=18446744073709551615) - selected: int = Field(0, ge=0, le=18446744073709551615) - selection_key: str = "" - - -class LensAccessParams(BaseModel): - model_config = ConfigDict( - extra="forbid", - frozen=True, - ) - - all_teams: Literal[0, 1] - team: str - key_hash: str - - -class LensContentParams(BaseModel): - model_config = ConfigDict( - extra="forbid", - frozen=True, - ) - - all_teams: Literal[0, 1] - team: str - key_hash: str - source: ContentSource - id: str - record_team: str - trace_ref: str - cursor: str - offset: int = Field(..., ge=0, le=4294967295) - - -class LensEvidenceParams(BaseModel): - model_config = ConfigDict( - extra="forbid", - frozen=True, - ) - - all_teams: Literal[0, 1] - team: str - key_hash: str - source: ContentSource - id: str - record_team: str - trace_ref: str - span: str - quote: str - - -ExecutionSource: TypeAlias = Literal["traces", "requests", "both"] - - -class LensSampleParams(BaseModel): - model_config = ConfigDict( - extra="forbid", - frozen=True, - ) - - all_teams: Literal[0, 1] - team: str - key_hash: str - source: ExecutionSource - start: int = Field(..., ge=0, le=18446744073709551615) - end: int = Field(..., ge=0, le=18446744073709551615) - agent_name: str - service: str - filter_keys: tuple[str, ...] - filter_values: tuple[str, ...] - selected_team: str - execution_ids: tuple[str, ...] - sample_cap: int = Field(..., ge=0, le=18446744073709551615) - sample_percent: float = Field(..., ge=0.0, le=100.0) - preview: Literal[0, 1] - after: str - limit: int = Field(..., ge=0, le=4294967295) - offset: int = Field(..., ge=0, le=18446744073709551615) - - -class PartRow(BaseModel): - model_config = ConfigDict( - frozen=True, - ) - - span_id: str - parent_span_id: str - name: str - kind: str - content: str - truncated: int = Field(..., ge=0, le=1) - - TraceTableName: TypeAlias = Literal["otel_traces", "agent_traces_by_key", "spend_logs"] @@ -419,16 +148,4 @@ class TraceQueryHelp(BaseModel): guide: str -TraceWireModels: TypeAlias = Annotated[ - ActivityAvailability - | AgentRow - | CountRow - | ExecutionRow - | LensAccessParams - | LensContentParams - | LensEvidenceParams - | LensSampleParams - | PartRow - | TraceQueryHelp, - Field(..., title="TraceWireModels"), -] +TraceWireModels: TypeAlias = Annotated[TraceQueryHelp, Field(..., title="TraceWireModels")] diff --git a/litellm/rust_bridge/trace/generated/types.py b/litellm/rust_bridge/trace/generated/types.py index b7499d38f07..ff85d0e7fad 100644 --- a/litellm/rust_bridge/trace/generated/types.py +++ b/litellm/rust_bridge/trace/generated/types.py @@ -23,7 +23,15 @@ class OwnedQueryScope(typing_extensions.TypedDict): QueryScope: TypeAlias = AllQueryScope | OwnedQueryScope -RunField: TypeAlias = Literal["name", "agent", "status", "model", "input", "trace_id"] +RunField: TypeAlias = Literal["name", "agent", "status", "model", "input", "trace_id", "service", "team"] + + +RunSortKey: TypeAlias = Literal["start_ms", "duration_ms", "span_count", "error_count", "trace_ref"] + + +class RunOrder(typing_extensions.TypedDict): + key: ReadOnly[RunSortKey] + descending: ReadOnly[bool] class RunValues(typing_extensions.TypedDict): @@ -55,6 +63,14 @@ class SpanErrorPage(typing_extensions.TypedDict): next_cursor: ReadOnly[str | None] +class SpanText(typing_extensions.TypedDict): + span_id: ReadOnly[str] + text: ReadOnly[str] + total_chars: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + version: ReadOnly[str] + contains: ReadOnly[bool] + + SpanStatus: TypeAlias = Literal["ok", "error", "unset"] @@ -89,9 +105,6 @@ class AgentRuns(typing_extensions.TypedDict): runs: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] -ReadQueryName: TypeAlias = Literal["availability", "agents", "sample", "content", "evidence"] - - class UIFields(typing_extensions.TypedDict): fields: ReadOnly[tuple[UIField, ...]] kind: ReadOnly[Literal["fields"]] @@ -190,5 +203,14 @@ class SpanDetail(typing_extensions.TypedDict): TraceWireTypes: TypeAlias = ( - QueryScope | RunField | RunValues | SpanDetail | SpanErrorPage | Trace | TraceHistogram | TracePage | ReadQueryName + QueryScope + | RunField + | RunOrder + | RunValues + | SpanDetail + | SpanErrorPage + | SpanText + | Trace + | TraceHistogram + | TracePage ) diff --git a/litellm/rust_bridge/trace/queries.py b/litellm/rust_bridge/trace/queries.py index f40e1f19944..3be7a5fa038 100644 --- a/litellm/rust_bridge/trace/queries.py +++ b/litellm/rust_bridge/trace/queries.py @@ -1,22 +1,9 @@ from collections.abc import Mapping -from dataclasses import dataclass -from typing import Final, Generic, TypeVar +from typing import Final -from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, JsonValue -from .generated.models import ( - ActivityAvailability, - AgentRow, - CountRow, - ExecutionRow, - LensAccessParams, - LensContentParams, - LensEvidenceParams, - LensSampleParams, - PartRow, - TraceQueryColumn, -) -from .generated.types import ReadQueryName +from .generated.models import TraceQueryColumn _RESPONSE_CONFIG: Final = ConfigDict(frozen=True, extra="allow") @@ -34,36 +21,3 @@ class TraceSQLResponse(BaseModel): data: tuple[Mapping[str, JsonValue], ...] rows: int | str statistics: TraceQueryStatistics - - -ParamsT: Final = TypeVar("ParamsT", bound=BaseModel) -RowT: Final = TypeVar("RowT") - - -class QueryResponse(BaseModel, Generic[RowT]): - model_config = ConfigDict(frozen=True) - data: tuple[RowT, ...] - - -@dataclass(frozen=True, slots=True) -class ReadQuery(Generic[ParamsT, RowT]): - name: ReadQueryName - parameters: type[ParamsT] - response: TypeAdapter[QueryResponse[RowT]] - - -LENS_AVAILABILITY: Final[ReadQuery[LensAccessParams, ActivityAvailability]] = ReadQuery( - "availability", LensAccessParams, TypeAdapter(QueryResponse[ActivityAvailability]) -) -LENS_AGENTS: Final[ReadQuery[LensAccessParams, AgentRow]] = ReadQuery( - "agents", LensAccessParams, TypeAdapter(QueryResponse[AgentRow]) -) -LENS_SAMPLE: Final[ReadQuery[LensSampleParams, ExecutionRow]] = ReadQuery( - "sample", LensSampleParams, TypeAdapter(QueryResponse[ExecutionRow]) -) -LENS_CONTENT: Final[ReadQuery[LensContentParams, PartRow]] = ReadQuery( - "content", LensContentParams, TypeAdapter(QueryResponse[PartRow]) -) -LENS_EVIDENCE: Final[ReadQuery[LensEvidenceParams, CountRow]] = ReadQuery( - "evidence", LensEvidenceParams, TypeAdapter(QueryResponse[CountRow]) -) diff --git a/litellm/rust_bridge/trace/storage.py b/litellm/rust_bridge/trace/storage.py index 600ce4cdea3..dfe6c377ad6 100644 --- a/litellm/rust_bridge/trace/storage.py +++ b/litellm/rust_bridge/trace/storage.py @@ -1,46 +1,28 @@ from collections.abc import Awaitable, Mapping, Sequence from dataclasses import asdict, dataclass -from typing import Final, Protocol, TypeVar, runtime_checkable +from typing import Final, Literal, Protocol, TypeAlias, TypeVar, runtime_checkable from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError from litellm.constants import AGENT_TRACING_LIST_PAGE_SIZE, OTLP_MAX_ATTRIBUTE_VALUE_BYTES from litellm.rust_bridge.loader import get_native_bridge -from litellm.rust_bridge.trace.generated.models import ( - ActivityAvailability, - AgentRow, - CountRow, - ExecutionRow, - LensAccessParams, - LensContentParams, - LensEvidenceParams, - LensSampleParams, - PartRow, -) -from litellm.rust_bridge.trace.generated.types import ReadQueryName -from litellm.rust_bridge.trace.queries import ( - LENS_AGENTS, - LENS_AVAILABILITY, - LENS_CONTENT, - LENS_EVIDENCE, - LENS_SAMPLE, - ParamsT, - ReadQuery, - RowT, -) from .generated.models import TraceQueryHelp from .generated.types import ( QueryScope, + RunOrder, RunValues, SpanDetail, SpanErrorPage, + SpanText, Trace, TraceHistogram, TracePage, ) from .queries import TraceSQLResponse +SpanPart: TypeAlias = Literal["input", "output", "error", "attributes"] + @dataclass(frozen=True, slots=True) class Tenant: @@ -54,6 +36,9 @@ class Tenant: _EMPTY_TENANT: Final = Tenant("", "") +NEWEST: Final[RunOrder] = {"key": "start_ms", "descending": True} +BY_REFERENCE: Final[RunOrder] = {"key": "trace_ref", "descending": False} + class NativeStore(Protocol): def __init__(self, config: "NativeConfig") -> None: ... @@ -65,7 +50,32 @@ class NativeStore(Protocol): def ingest(self, payload: bytes, content_type: str | None, tenant: Mapping[str, str]) -> Awaitable[int]: ... def list_traces( - self, scope: QueryScope, start_ms: int, end_ms: int, q: str, cursor: str | None, limit: int + self, + scope: QueryScope, + start_ms: int, + end_ms: int, + q: str, + cursor: str | None, + limit: int, + order: RunOrder, + trace_refs: Sequence[str], + ) -> Awaitable[JsonValue]: ... + + def count_traces( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, trace_refs: Sequence[str] + ) -> Awaitable[JsonValue]: ... + + def span_text( + self, + trace_id: str, + trace_ref: str, + span_ids: Sequence[str], + part: SpanPart, + scope: QueryScope, + offset: int, + max_chars: int | None, + tail: bool, + contains: str | None, ) -> Awaitable[JsonValue]: ... def trace_histogram( @@ -90,10 +100,6 @@ class NativeStore(Protocol): def query_help(self, scope: QueryScope, secret: str) -> Awaitable[JsonValue]: ... - def query( - self, name: ReadQueryName, parameters: Mapping[str, str | int | float | Sequence[str]] - ) -> Awaitable[str]: ... - @runtime_checkable class NativeTraces(Protocol): @@ -107,12 +113,13 @@ class NativeTraces(Protocol): ) -> list[dict[str, JsonValue]]: ... -QUERY_PARAMETERS: Final = TypeAdapter(dict[str, str | int | float | list[str]]) _SQL_RESPONSE: Final = TypeAdapter(TraceSQLResponse) _HELP_RESPONSE: Final = TypeAdapter(TraceQueryHelp) _TRACE_PAGE: Final = TypeAdapter(TracePage) _TRACE_HISTOGRAM: Final = TypeAdapter(TraceHistogram) _RUN_VALUES: Final = TypeAdapter(RunValues) +_COUNT: Final = TypeAdapter(int) +_SPAN_TEXTS: Final = TypeAdapter(tuple[SpanText, ...]) _TRACE: Final[TypeAdapter[Trace | None]] = TypeAdapter(Trace | None) _SPAN_DETAIL: Final[TypeAdapter[SpanDetail | None]] = TypeAdapter(SpanDetail | None) _SPAN_ERROR_PAGE: Final[TypeAdapter[SpanErrorPage | None]] = TypeAdapter(SpanErrorPage | None) @@ -199,10 +206,37 @@ class ClickHouseStorage: q: str = "", cursor: str | None = None, limit: int = AGENT_TRACING_LIST_PAGE_SIZE, + order: RunOrder = NEWEST, + trace_refs: Sequence[str] = (), ) -> TracePage: - result: Final = await self._native.list_traces(scope, start_ms, end_ms, q, cursor, limit) + result: Final = await self._native.list_traces( + scope, start_ms, end_ms, q, cursor, limit, order, tuple(trace_refs) + ) return _validate_query_response(_TRACE_PAGE, result) + async def count_traces( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str = "", trace_refs: Sequence[str] = () + ) -> int: + result: Final = await self._native.count_traces(scope, start_ms, end_ms, q, tuple(trace_refs)) + return _validate_query_response(_COUNT, result) + + async def span_text( + self, + trace_id: str, + trace_ref: str, + span_ids: Sequence[str], + part: SpanPart, + scope: QueryScope, + offset: int = 0, + max_chars: int | None = None, + tail: bool = False, + contains: str | None = None, + ) -> tuple[SpanText, ...]: + result: Final = await self._native.span_text( + trace_id, trace_ref, tuple(span_ids), part, scope, offset, max_chars, tail, contains + ) + return _validate_query_response(_SPAN_TEXTS, result) + async def trace_histogram( self, scope: QueryScope, start_ms: int, end_ms: int, q: str, buckets: int ) -> TraceHistogram: @@ -236,11 +270,6 @@ class ClickHouseStorage: result: Final = await self._native.get_span_error(trace_id, span_id, scope, trace_ref, cursor) return _validate_query_response(_SPAN_ERROR_PAGE, result) - async def query(self, query: ReadQuery[ParamsT, RowT], parameters: ParamsT) -> tuple[RowT, ...]: - validated: Final = query.parameters.model_validate(parameters) - result: Final = await self._native.query(query.name, QUERY_PARAMETERS.validate_python(validated.model_dump())) - return _decode_query_response(query.response, result).data - async def query_sql(self, sql: str, scope: QueryScope, secret: str) -> TraceSQLResponse: result: Final = await self._native.query_sql(sql, scope, secret) return _decode_query_response(_SQL_RESPONSE, result) @@ -248,18 +277,3 @@ class ClickHouseStorage: async def query_help(self, scope: QueryScope, secret: str) -> TraceQueryHelp: result: Final = await self._native.query_help(scope, secret) return _validate_query_response(_HELP_RESPONSE, result) - - async def lens_sample(self, parameters: LensSampleParams) -> tuple[ExecutionRow, ...]: - return await self.query(LENS_SAMPLE, parameters) - - async def lens_availability(self, parameters: LensAccessParams) -> tuple[ActivityAvailability, ...]: - return await self.query(LENS_AVAILABILITY, parameters) - - async def lens_agents(self, parameters: LensAccessParams) -> tuple[AgentRow, ...]: - return await self.query(LENS_AGENTS, parameters) - - async def lens_content(self, parameters: LensContentParams) -> tuple[PartRow, ...]: - return await self.query(LENS_CONTENT, parameters) - - async def lens_evidence(self, parameters: LensEvidenceParams) -> tuple[CountRow, ...]: - return await self.query(LENS_EVIDENCE, parameters) diff --git a/litellm/tracing/receiver.py b/litellm/tracing/receiver.py index 7e6303f0acd..1957e1d8f87 100644 --- a/litellm/tracing/receiver.py +++ b/litellm/tracing/receiver.py @@ -22,6 +22,7 @@ from litellm.constants import AGENT_TRACING_LIST_PAGE_SIZE, OTLP_MAX_BODY_BYTES, from litellm.rust_bridge.trace.generated.types import ( QueryScope, RunField, + RunOrder, RunValues, SpanDetail, SpanErrorPage, @@ -29,7 +30,7 @@ from litellm.rust_bridge.trace.generated.types import ( TraceHistogram, TracePage, ) -from litellm.rust_bridge.trace.storage import ClickHouseStorage, Tenant +from litellm.rust_bridge.trace.storage import NEWEST, ClickHouseStorage, Tenant from litellm.tracing.config import trace_storage_config from litellm.tracing.otlp_http import InvalidOTLPPayloadError, TracingPayloadTooLargeError, decompress @@ -106,9 +107,15 @@ class TraceReceiver: raise InvalidOTLPPayloadError(str(error)) from error async def list_traces( - self, scope: QueryScope, start_ms: int, end_ms: int, q: str = "", cursor: str | None = None + self, + scope: QueryScope, + start_ms: int, + end_ms: int, + q: str = "", + cursor: str | None = None, + order: RunOrder = NEWEST, ) -> TracePage: - return await self.storage.list_traces(scope, start_ms, end_ms, q, cursor, AGENT_TRACING_LIST_PAGE_SIZE) + return await self.storage.list_traces(scope, start_ms, end_ms, q, cursor, AGENT_TRACING_LIST_PAGE_SIZE, order) async def trace_histogram( self, scope: QueryScope, start_ms: int, end_ms: int, q: str, buckets: int diff --git a/scripts/generate_trace_types.py b/scripts/generate_trace_types.py index 39c93dabf49..5d751e674c1 100644 --- a/scripts/generate_trace_types.py +++ b/scripts/generate_trace_types.py @@ -169,13 +169,8 @@ def main() -> int: schema_set_matches: Final = reconcile_schemas(frozenset(path for path, _ in exported), args.check) with TemporaryDirectory(prefix="trace-codegen-") as temporary: directory: Final = Path(temporary) - types: Final = generate({**domain, "ReadQueryName": clickhouse["ReadQueryName"]}, "types", directory, config) - models: Final = generate( - {name: schema for name, schema in clickhouse.items() if name != "ReadQueryName"}, - "models", - directory, - config, - ) + types: Final = generate(domain, "types", directory, config) + models: Final = generate(clickhouse, "models", directory, config) python_results: Final = ( publish(GENERATED / "types.py", types.read_text(), args.check), publish(GENERATED / "models.py", models.read_text(), args.check), diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/ActivityAvailability.json b/scripts/trace_codegen/schemas/traces-clickhouse/ActivityAvailability.json deleted file mode 100644 index 8b7fa61f6ee..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/ActivityAvailability.json +++ /dev/null @@ -1,57 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "properties": { - "requests": { - "anyOf": [ - { - "type": "boolean" - }, - { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - { - "enum": [ - "0", - "1" - ], - "type": "string" - } - ], - "default": 0, - "x-python-normalized": { - "type": "bool" - } - }, - "traces": { - "anyOf": [ - { - "type": "boolean" - }, - { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - { - "enum": [ - "0", - "1" - ], - "type": "string" - } - ], - "default": 0, - "x-python-normalized": { - "type": "bool" - } - } - }, - "title": "ActivityAvailability", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/AgentRow.json b/scripts/trace_codegen/schemas/traces-clickhouse/AgentRow.json deleted file mode 100644 index 6e06460de09..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/AgentRow.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "properties": { - "agent_name": { - "type": "string" - } - }, - "required": [ - "agent_name" - ], - "title": "AgentRow", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/CountRow.json b/scripts/trace_codegen/schemas/traces-clickhouse/CountRow.json deleted file mode 100644 index 49737af04e2..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/CountRow.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "properties": { - "count": { - "anyOf": [ - { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - { - "pattern": "^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - "type": "string" - } - ], - "x-python-normalized": { - "maximum": 18446744073709551615, - "minimum": 0, - "type": "int" - } - } - }, - "required": [ - "count" - ], - "title": "CountRow", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/ExecutionRow.json b/scripts/trace_codegen/schemas/traces-clickhouse/ExecutionRow.json deleted file mode 100644 index 69b51eebedc..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/ExecutionRow.json +++ /dev/null @@ -1,151 +0,0 @@ -{ - "$defs": { - "ContentSource": { - "enum": [ - "traces", - "requests" - ], - "type": "string" - } - }, - "$schema": "https://json-schema.org/draft/2020-12/schema", - "properties": { - "attributes": { - "default": [], - "items": { - "maxItems": 2, - "minItems": 2, - "prefixItems": [ - { - "type": "string" - }, - { - "type": "string" - } - ], - "type": "array" - }, - "type": "array" - }, - "eligible": { - "anyOf": [ - { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - { - "pattern": "^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - "type": "string" - } - ], - "x-python-normalized": { - "maximum": 18446744073709551615, - "minimum": 0, - "type": "int" - } - }, - "name": { - "type": "string" - }, - "root_seen": { - "anyOf": [ - { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - { - "enum": [ - "0", - "1" - ], - "type": "string" - } - ], - "x-python-normalized": { - "maximum": 1, - "minimum": 0, - "type": "int" - } - }, - "selected": { - "anyOf": [ - { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - { - "pattern": "^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - "type": "string" - } - ], - "default": 0.0, - "x-python-normalized": { - "maximum": 18446744073709551615, - "minimum": 0, - "type": "int" - } - }, - "selection_key": { - "default": "", - "type": "string" - }, - "service": { - "default": "", - "type": "string" - }, - "source": { - "$ref": "#/$defs/ContentSource" - }, - "span_count": { - "anyOf": [ - { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - { - "pattern": "^(?:0|[1-9][0-9]{0,18}|1[0-7][0-9]{18}|18[0-3][0-9]{17}|184[0-3][0-9]{16}|1844[0-5][0-9]{15}|18446[0-6][0-9]{14}|184467[0-3][0-9]{13}|1844674[0-3][0-9]{12}|184467440[0-6][0-9]{10}|1844674407[0-2][0-9]{9}|18446744073[0-6][0-9]{8}|1844674407370[0-8][0-9]{6}|18446744073709[0-4][0-9]{5}|184467440737095[0-4][0-9]{4}|1844674407370955[0-0][0-9]{3}|18446744073709551[0-5][0-9]{2}|184467440737095516[0-0][0-9]{1}|1844674407370955161[0-4][0-9]{0}|18446744073709551615)$", - "type": "string" - } - ], - "x-python-normalized": { - "maximum": 18446744073709551615, - "minimum": 0, - "type": "int" - } - }, - "start_time": { - "type": "string" - }, - "team_id": { - "type": "string" - }, - "trace_id": { - "type": "string" - }, - "trace_ref": { - "default": "", - "type": "string" - } - }, - "required": [ - "source", - "trace_id", - "team_id", - "name", - "start_time", - "span_count", - "root_seen", - "eligible" - ], - "title": "ExecutionRow", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/LensAccessParams.json b/scripts/trace_codegen/schemas/traces-clickhouse/LensAccessParams.json deleted file mode 100644 index 057a306d68c..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/LensAccessParams.json +++ /dev/null @@ -1,26 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "additionalProperties": false, - "properties": { - "all_teams": { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - "key_hash": { - "type": "string" - }, - "team": { - "type": "string" - } - }, - "required": [ - "all_teams", - "team", - "key_hash" - ], - "title": "LensAccessParams", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/LensContentParams.json b/scripts/trace_codegen/schemas/traces-clickhouse/LensContentParams.json deleted file mode 100644 index 5ee5ab558ce..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/LensContentParams.json +++ /dev/null @@ -1,62 +0,0 @@ -{ - "$defs": { - "ContentSource": { - "enum": [ - "traces", - "requests" - ], - "type": "string" - } - }, - "$schema": "https://json-schema.org/draft/2020-12/schema", - "additionalProperties": false, - "properties": { - "all_teams": { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - "cursor": { - "type": "string" - }, - "id": { - "type": "string" - }, - "key_hash": { - "type": "string" - }, - "offset": { - "format": "uint32", - "maximum": 4294967295, - "minimum": 0, - "type": "integer" - }, - "record_team": { - "type": "string" - }, - "source": { - "$ref": "#/$defs/ContentSource" - }, - "team": { - "type": "string" - }, - "trace_ref": { - "type": "string" - } - }, - "required": [ - "all_teams", - "team", - "key_hash", - "source", - "id", - "record_team", - "trace_ref", - "cursor", - "offset" - ], - "title": "LensContentParams", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/LensEvidenceParams.json b/scripts/trace_codegen/schemas/traces-clickhouse/LensEvidenceParams.json deleted file mode 100644 index 07b9c216083..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/LensEvidenceParams.json +++ /dev/null @@ -1,59 +0,0 @@ -{ - "$defs": { - "ContentSource": { - "enum": [ - "traces", - "requests" - ], - "type": "string" - } - }, - "$schema": "https://json-schema.org/draft/2020-12/schema", - "additionalProperties": false, - "properties": { - "all_teams": { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - "id": { - "type": "string" - }, - "key_hash": { - "type": "string" - }, - "quote": { - "type": "string" - }, - "record_team": { - "type": "string" - }, - "source": { - "$ref": "#/$defs/ContentSource" - }, - "span": { - "type": "string" - }, - "team": { - "type": "string" - }, - "trace_ref": { - "type": "string" - } - }, - "required": [ - "all_teams", - "team", - "key_hash", - "source", - "id", - "record_team", - "trace_ref", - "span", - "quote" - ], - "title": "LensEvidenceParams", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/LensSampleParams.json b/scripts/trace_codegen/schemas/traces-clickhouse/LensSampleParams.json deleted file mode 100644 index 598f63cefb7..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/LensSampleParams.json +++ /dev/null @@ -1,127 +0,0 @@ -{ - "$defs": { - "ExecutionSource": { - "enum": [ - "traces", - "requests", - "both" - ], - "type": "string" - } - }, - "$schema": "https://json-schema.org/draft/2020-12/schema", - "additionalProperties": false, - "properties": { - "after": { - "type": "string" - }, - "agent_name": { - "type": "string" - }, - "all_teams": { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - "end": { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - "execution_ids": { - "items": { - "type": "string" - }, - "type": "array" - }, - "filter_keys": { - "items": { - "type": "string" - }, - "type": "array" - }, - "filter_values": { - "items": { - "type": "string" - }, - "type": "array" - }, - "key_hash": { - "type": "string" - }, - "limit": { - "format": "uint32", - "maximum": 4294967295, - "minimum": 0, - "type": "integer" - }, - "offset": { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - "preview": { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - "sample_cap": { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - "sample_percent": { - "format": "double", - "maximum": 100, - "minimum": 0, - "type": "number" - }, - "selected_team": { - "type": "string" - }, - "service": { - "type": "string" - }, - "source": { - "$ref": "#/$defs/ExecutionSource" - }, - "start": { - "format": "uint64", - "maximum": 18446744073709551615, - "minimum": 0, - "type": "integer" - }, - "team": { - "type": "string" - } - }, - "required": [ - "all_teams", - "team", - "key_hash", - "source", - "start", - "end", - "agent_name", - "service", - "filter_keys", - "filter_values", - "selected_team", - "execution_ids", - "sample_cap", - "sample_percent", - "preview", - "after", - "limit", - "offset" - ], - "title": "LensSampleParams", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/PartRow.json b/scripts/trace_codegen/schemas/traces-clickhouse/PartRow.json deleted file mode 100644 index 5a4d397a801..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/PartRow.json +++ /dev/null @@ -1,53 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "properties": { - "content": { - "type": "string" - }, - "kind": { - "type": "string" - }, - "name": { - "type": "string" - }, - "parent_span_id": { - "type": "string" - }, - "span_id": { - "type": "string" - }, - "truncated": { - "anyOf": [ - { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - { - "enum": [ - "0", - "1" - ], - "type": "string" - } - ], - "x-python-normalized": { - "maximum": 1, - "minimum": 0, - "type": "int" - } - } - }, - "required": [ - "span_id", - "parent_span_id", - "name", - "kind", - "content", - "truncated" - ], - "title": "PartRow", - "type": "object" -} diff --git a/scripts/trace_codegen/schemas/traces-clickhouse/ReadQueryName.json b/scripts/trace_codegen/schemas/traces-clickhouse/ReadQueryName.json deleted file mode 100644 index 179732c4b55..00000000000 --- a/scripts/trace_codegen/schemas/traces-clickhouse/ReadQueryName.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "enum": [ - "availability", - "agents", - "sample", - "content", - "evidence" - ], - "title": "ReadQueryName", - "type": "string" -} diff --git a/scripts/trace_codegen/schemas/traces/RunField.json b/scripts/trace_codegen/schemas/traces/RunField.json index 3879382fa35..ab22c4c67cd 100644 --- a/scripts/trace_codegen/schemas/traces/RunField.json +++ b/scripts/trace_codegen/schemas/traces/RunField.json @@ -6,7 +6,9 @@ "status", "model", "input", - "trace_id" + "trace_id", + "service", + "team" ], "title": "RunField", "type": "string" diff --git a/scripts/trace_codegen/schemas/traces/RunOrder.json b/scripts/trace_codegen/schemas/traces/RunOrder.json new file mode 100644 index 00000000000..211f1209072 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/RunOrder.json @@ -0,0 +1,31 @@ +{ + "$defs": { + "RunSortKey": { + "enum": [ + "start_ms", + "duration_ms", + "span_count", + "error_count", + "trace_ref" + ], + "type": "string" + } + }, + "$schema": "https://json-schema.org/draft/2020-12/schema", + "additionalProperties": false, + "description": "Runs by `key`, ties broken by `trace_ref` in the same direction.", + "properties": { + "descending": { + "type": "boolean" + }, + "key": { + "$ref": "#/$defs/RunSortKey" + } + }, + "required": [ + "key", + "descending" + ], + "title": "RunOrder", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/SpanText.json b/scripts/trace_codegen/schemas/traces/SpanText.json new file mode 100644 index 00000000000..a773ece9039 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/SpanText.json @@ -0,0 +1,33 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "properties": { + "contains": { + "type": "boolean" + }, + "span_id": { + "type": "string" + }, + "text": { + "type": "string" + }, + "total_chars": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "version": { + "description": "Uppercase hex SHA-256 of the whole part, so a reader can tell when it changed.", + "type": "string" + } + }, + "required": [ + "span_id", + "text", + "total_chars", + "version", + "contains" + ], + "title": "SpanText", + "type": "object" +} diff --git a/tests/unit/proxy/lens/test_analysis.py b/tests/unit/proxy/lens/test_analysis.py index 5231000e14b..4fd4e0c3a0d 100644 --- a/tests/unit/proxy/lens/test_analysis.py +++ b/tests/unit/proxy/lens/test_analysis.py @@ -20,6 +20,7 @@ from litellm.proxy.lens.models import ( ) from litellm.proxy.lens.state import queue_job from tests.unit.proxy.lens.test_state import NOW, issue_brief, lens, finding +from tests.unit.proxy.lens.test_sources import summary @pytest.mark.asyncio @@ -28,7 +29,7 @@ async def test_parallel_review_shares_one_model_limit_and_cleans_up(outcome: str 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) + Execution(id=str(i), trace_id=str(i), trace_ref="", summary=summary("", str(i), span_count=6)) for i in range(ANALYSIS_CONCURRENCY + 1) ) entered: Final = SimpleQueue[str]() @@ -162,9 +163,7 @@ def test_excerpt_omission_is_not_original_evidence() -> None: @pytest.mark.asyncio async def test_reviewer_sees_final_outcome_and_catalog_across_pages() -> None: - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="run", start_time="", span_count=2 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=2)) root: Final = TracePart(execution_id="run", span_id="01", name="task", kind="agent", content="Task: write a report") editor: Final = TracePart( execution_id="run", span_id="02", parent_span_id="01", name="editor", kind="agent", content="Delivered report" @@ -195,9 +194,7 @@ async def test_reviewer_sees_final_outcome_and_catalog_across_pages() -> None: async def test_reviewer_fetches_targeted_evidence_and_rejects_outside_catalog_reads() -> None: 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 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=2)) root: Final = TracePart( execution_id="run", span_id="01", name="task", kind="agent", content="Find the verified result" ) @@ -257,9 +254,7 @@ async def test_reviewer_fetches_targeted_evidence_and_rejects_outside_catalog_re async def test_reviewer_stops_repeated_read_requests() -> None: 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 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) part: Final = TracePart(execution_id="run", span_id="01", name="task", kind="agent", content="Partial export") reads: Final = SimpleQueue[int]() calls: Final = SimpleQueue[int]() @@ -294,9 +289,7 @@ def test_chunks_preserve_all_spans_and_keep_context_bounded() -> None: @pytest.mark.asyncio async def test_investigator_rejects_a_fabricated_quote() -> None: - execution: Final = Execution( - id="run1", source="traces", trace_id="t", team_id="alpha", name="search", start_time="", span_count=1 - ) + execution: Final = Execution(id="run1", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) examined: Final = Examined( execution=execution, observations=(), @@ -326,9 +319,7 @@ async def test_investigator_rejects_a_fabricated_quote() -> None: @pytest.mark.parametrize("paginated", [False, True]) @pytest.mark.parametrize("assessable", [False, True]) async def test_assessable_content_is_not_overridden_by_unknown_chunks(paginated: bool, assessable: bool) -> None: - execution: Final = Execution( - id="run1", source="traces", trace_id="t", team_id="alpha", name="review", start_time="", span_count=4 - ) + execution: Final = Execution(id="run1", trace_id="t", trace_ref="", summary=summary("", "t", span_count=4)) unknown: Final = tuple( TracePart(execution_id="run1", span_id=str(i), name="tool", kind="tool", content="x" * 8000) for i in range(3) ) @@ -360,9 +351,7 @@ async def test_assessable_content_is_not_overridden_by_unknown_chunks(paginated: @pytest.mark.asyncio async def test_investigator_keeps_final_outcome_ahead_of_repeated_model_history() -> None: - execution: Final = Execution( - id="run1", source="traces", trace_id="t", team_id="alpha", name="review", start_time="", span_count=6 - ) + execution: Final = Execution(id="run1", trace_id="t", trace_ref="", summary=summary("", "t", span_count=6)) history: Final = tuple( TracePart( execution_id="run1", span_id=str(i), name="chat", kind="llm", parent_span_id="span", content="x" * 8000 @@ -401,9 +390,7 @@ async def test_investigator_keeps_final_outcome_ahead_of_repeated_model_history( async def test_many_model_citations_are_accepted_but_quotes_are_still_verified( quote: str, check_id: str, accepted: bool ) -> None: - execution: Final = Execution( - id="run1", source="traces", trace_id="t", team_id="alpha", name="review", start_time="", span_count=1 - ) + execution: Final = Execution(id="run1", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) part: Final = TracePart(execution_id="run1", span_id="span", name="tool", kind="tool", content="timeout") attempts: Final = iter((8,)) @@ -491,9 +478,7 @@ async def test_grouping_consolidates_prior_batches_and_reports_real_progress() - @pytest.mark.asyncio @pytest.mark.parametrize("later_span", ("later", "0")) async def test_investigator_can_cite_a_later_page_or_offset(later_span: str) -> None: - execution: Final = Execution( - id="run1", source="traces", trace_id="t", team_id="alpha", name="review", start_time="", span_count=7 - ) + execution: Final = Execution(id="run1", trace_id="t", trace_ref="", summary=summary("", "t", span_count=7)) initial: Final = tuple( TracePart(execution_id="run1", span_id=str(i), name="agent", kind="agent", content="x" * 8000) for i in range(6) ) @@ -611,13 +596,7 @@ async def test_review_keeps_original_ids_in_per_run_assessments() -> None: from litellm.proxy.lens.analysis import analyze_sample execution: Final = Execution( - id="opaque-original-id", - source="requests", - trace_id="request", - team_id="", - name="call", - start_time="", - span_count=1, + id="opaque-original-id", trace_id="request", trace_ref="", summary=summary("", "request", span_count=1) ) async def read(identity: str, _cursor: str, _offset: int) -> ExecutionContent: @@ -645,15 +624,7 @@ async def test_review_keeps_original_ids_in_per_run_assessments() -> None: @pytest.mark.asyncio async def test_investigation_context_accounts_for_metadata_on_thousands_of_short_spans() -> None: executions: Final = tuple( - Execution( - id=f"run-{i}", - source="traces", - trace_id=f"trace-{i}", - team_id="", - name="Short successful task", - start_time="", - span_count=1, - ) + Execution(id=f"run-{i}", trace_id=f"trace-{i}", trace_ref="", summary=summary("", f"trace-{i}", span_count=1)) for i in range(2501) ) examined: Final = tuple( @@ -697,9 +668,7 @@ async def test_investigation_context_accounts_for_metadata_on_thousands_of_short async def test_completed_read_does_not_make_supported_review_unknown() -> None: 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 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) part: Final = TracePart(execution_id="run", span_id="s", name="task", kind="agent", content="timeout") observation: Final = Observation( check_id="retries", summary="Failed", evidence=(Evidence(execution_id="run", span_id="s", quote="timeout"),) @@ -728,9 +697,7 @@ async def test_completed_read_does_not_make_supported_review_unknown() -> None: @pytest.mark.asyncio async def test_echoed_feedback_page_does_not_skip_requested_evidence() -> None: - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) requests: Final = SimpleQueue[int]() async def read(_identity: str, _cursor: str, offset: int) -> ExecutionContent: @@ -780,9 +747,7 @@ async def test_echoed_feedback_page_does_not_skip_requested_evidence() -> None: @pytest.mark.asyncio @pytest.mark.parametrize("action", ("catalog", "observations", "feedback", "read")) async def test_empty_navigation_requires_a_final_decision(action: str) -> None: - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) examined: Final = Examined(execution=execution, observations=(), parts=(), partial=False, cannot_assess=False) calls: Final = SimpleQueue[int]() @@ -813,9 +778,7 @@ async def test_empty_navigation_requires_a_final_decision(action: str) -> None: async def test_large_feedback_history_is_accessible_without_overflowing_context(phase: str) -> None: 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 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) part: Final = TracePart(execution_id="run", span_id="span", name="task", kind="agent", content="timeout") accepted: Final = merge_finding(lens(), finding("run"), 1, NOW) prior: Final = tuple( @@ -931,9 +894,7 @@ async def test_distinct_patterns_are_consolidated_in_batches_without_losing_runs async def test_invalid_candidate_response_preserves_other_findings_and_reports_inconclusive() -> None: 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 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) part: Final = TracePart(execution_id="run", span_id="span", name="tool", kind="tool", content="timeout") item: Final = Examined(execution=execution, observations=(), parts=(part,), partial=False, cannot_assess=False) candidates: Final = tuple( @@ -968,9 +929,7 @@ async def test_invalid_candidate_response_preserves_other_findings_and_reports_i @pytest.mark.asyncio async def test_investigator_keeps_the_issue_brief() -> None: - execution: Final = Execution( - id="run1", source="traces", trace_id="t", team_id="alpha", name="search", start_time="", span_count=1 - ) + execution: Final = Execution(id="run1", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) examined: Final = Examined( execution=execution, observations=(), @@ -1044,9 +1003,7 @@ async def test_large_context_and_long_verified_quotes_do_not_silently_end_invest context: Final = "Read all recorded evidence. " * 5000 long_quote: Final = "timeout detail " * 200 - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) part: Final = TracePart(execution_id="run", span_id="span", name="tool", kind="tool", content=long_quote) reviewed: Final = Examined(execution=execution, observations=(), parts=(part,), partial=False, cannot_assess=False) expected: Final = FindingDraft.model_validate( @@ -1078,9 +1035,7 @@ async def test_large_context_and_long_verified_quotes_do_not_silently_end_invest @pytest.mark.asyncio async def test_reviewer_can_read_every_offset_of_a_long_span_before_deciding() -> None: - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) original: Final = "trace evidence! " * 16000 + "late verified failure" offsets: Final = SimpleQueue[int]() seen: Final = SimpleQueue[str]() @@ -1138,9 +1093,7 @@ async def test_reviewer_can_read_every_offset_of_a_long_span_before_deciding() - @pytest.mark.asyncio async def test_investigator_can_read_all_evidence_pages_across_successive_span_batches() -> None: - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=80 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=80)) parts: Final = tuple( TracePart( execution_id="run", diff --git a/tests/unit/proxy/lens/test_endpoints.py b/tests/unit/proxy/lens/test_endpoints.py index 54dbaac1e42..5394b52f1d9 100644 --- a/tests/unit/proxy/lens/test_endpoints.py +++ b/tests/unit/proxy/lens/test_endpoints.py @@ -9,7 +9,6 @@ import litellm from litellm import Router from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.lens.endpoints import ( - list_agents, run_settings, run_window, user_scope, @@ -19,6 +18,7 @@ from litellm.proxy.lens.endpoints import ( worker_supports_model, ) from litellm.proxy.lens.models import ActivitySelection, Lens, LensSettings, RunRequest, Scope +from tests.unit.proxy.lens.test_sources import FakeStorage, summary @pytest.fixture @@ -115,21 +115,6 @@ async def test_worker_without_active_billing_cannot_take_work(revoked: bool, key assert not await worker_supports_model(inactive, settings) -@pytest.mark.parametrize("role", (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)) -@pytest.mark.asyncio -async def test_agent_discovery_without_trace_storage_is_empty(role: LitellmUserRoles) -> None: - auth: Final = UserAPIKeyAuth(user_role=role) - assert await list_agents(auth, None) == () - - -@pytest.mark.asyncio -async def test_agent_discovery_without_trace_storage_still_requires_admin_access() -> None: - auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) - with pytest.raises(HTTPException) as error: - await list_agents(auth, None) - assert error.value.status_code == 403 - - @pytest.mark.parametrize( "role", (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.TEAM), @@ -187,7 +172,9 @@ def saved_lens() -> Lens: return Lens( id="lens", scope=Scope(all_teams=True), - settings=LensSettings(name="Support", model="analysis", context="Answer questions", agent_name="support"), + settings=LensSettings( + name="Support", model="analysis", context="Answer questions", q="agent:support status:error" + ), created_at=now, next_run_at=now, budget_month="2026-01", @@ -198,8 +185,8 @@ def test_run_now_agent_override_only_changes_the_agent_for_that_run() -> None: lens: Final = saved_lens() overridden: Final = run_settings(lens, RunRequest(agent_name="billing")) assert overridden is not None - assert overridden.agent_name == "billing" - assert overridden.model_copy(update={"agent_name": "support"}) == lens.settings + assert overridden.q == "status:error agent:billing" + assert overridden.model_copy(update={"q": lens.settings.q}) == lens.settings def test_run_now_without_overrides_keeps_the_saved_settings() -> None: @@ -285,19 +272,11 @@ def test_model_errors_reach_worker_with_status_and_redacted_provider_message(pro async def test_preview_samples_a_selection_without_investigation_settings() -> None: from litellm.proxy.lens.endpoints import Preview, preview_sample - class SelectionStorage: - async def lens_sample(self, parameters): - assert (parameters.source, parameters.agent_name, parameters.selected_team) == ("requests", "billing", "t1") - assert parameters.preview == 1 and parameters.offset == 3 - return [] - - body: Final = Preview.model_validate( - {"selection": {"source": "requests", "agent_name": "billing", "team_id": "t1"}, "offset": 3} - ) - sample: Final = await preview_sample( - body, UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), SelectionStorage() - ) - assert sample.eligible == 0 and not sample.executions + storage: Final = FakeStorage(runs=(summary("A" * 64), summary("B" * 64))) + body: Final = Preview.model_validate({"selection": {"q": "run", "sample_size": 1}}) + sample: Final = await preview_sample(body, UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), storage) + assert (sample.eligible, sample.selected, len(sample.executions)) == (2, 1, 2) + assert storage.listed[0][1] == "run" @pytest.mark.asyncio diff --git a/tests/unit/proxy/lens/test_sources.py b/tests/unit/proxy/lens/test_sources.py index 81bae7a0091..cb79bea7ed5 100644 --- a/tests/unit/proxy/lens/test_sources.py +++ b/tests/unit/proxy/lens/test_sources.py @@ -1,110 +1,345 @@ import base64 import json +from collections.abc import Sequence +from dataclasses import dataclass, field from typing import Final import pytest -from litellm.proxy.lens.models import MetadataFilter, Scope -from litellm.proxy.lens.sources import SourceReader, execution_id, parse_execution -from litellm.rust_bridge.trace.generated.models import ActivityAvailability, AgentRow, ExecutionRow -from tests.unit.proxy.lens.test_state import lens +from litellm.proxy.lens.models import ( + ActivitySelection, + Evidence, + Execution, + Job, + Scope, + execution_id, + parse_execution, +) +from litellm.proxy.lens.sources import BUDGET, OMITTED, SourceReader, lens_access, selected_count +from litellm.rust_bridge.trace.generated.types import ( + QueryScope, + RunOrder, + Span, + SpanText, + Trace, + TracePage, + TraceSummary, +) +from litellm.rust_bridge.trace.storage import BY_REFERENCE, NEWEST, SpanPart +from tests.unit.proxy.lens.test_state import NOW + +REF: Final = "A" * 64 -def test_same_trace_id_from_different_keys_is_a_distinct_execution() -> None: - assert execution_id("traces", "team", "trace", "key-one-ref") != execution_id( - "traces", "team", "trace", "key-two-ref" - ) - assert parse_execution(execution_id("traces", "team", "trace", "key-one-ref")) == ( - "traces", - "team", - "trace", - "key-one-ref", +def summary(trace_ref: str, trace_id: str = "trace", span_count: int = 1) -> TraceSummary: + return TraceSummary( + trace_id=trace_id, + trace_ref=trace_ref, + name="run", + service="svc", + input_preview="", + start_time="2026-01-01T00:00:00Z", + duration_ms=1.0, + status="ok", + span_count=span_count, + agent_count=0, + agent_invocations=0, + llm_calls=0, + tool_calls=0, + error_count=0, + input_tokens=0, + output_tokens=0, + models=(), + spend=None, ) -def test_previous_saved_findings_keep_their_execution_links() -> None: - assert parse_execution(base64.urlsafe_b64encode(json.dumps(("traces", "team", "trace")).encode()).decode()) == ( - "traces", - "team", - "trace", - "", +def span(span_id: str, parent: str | None = "root", status: str = "ok") -> Span: + return Span( + span_id=span_id, + parent_span_id=parent, + name=f"name-{span_id}", + type="tool", + agent="", + framework="", + start_offset_ms=0.0, + duration_ms=1.0, + status=status, + error=None, + error_truncated=False, + input_preview="", + model=None, + input_tokens=0, + output_tokens=0, + litellm_request_id=None, + spend=None, ) -@pytest.mark.asyncio -async def test_sample_never_returns_authentication_attributes() -> None: - class StorageResponse: - async def lens_sample(self, parameters): - assert parameters.team == "alpha" - return [ - ExecutionRow( - source="traces", - trace_id="trace", - team_id="alpha", - name="run", - start_time="", - span_count=1, - root_seen=1, - eligible=1, - attributes=( - ("litellm.api_key_hash", "opaque-oauth-bearer"), - ("environment", "production"), - ("", "invalid"), - ("oversized", "x" * 501), - ), - ) - ] +@dataclass +class FakeStorage: + """Answers the general trace reads over in-memory runs, spans and texts.""" - reader: Final = SourceReader(StorageResponse()) - sample: Final = await reader.sample(Scope(team_id="alpha"), lens().settings, 1, 2) - assert sample.executions[0].metadata == ( - MetadataFilter(key="environment", value="production"), - MetadataFilter(key="oversized", value="x" * 501), + runs: Sequence[TraceSummary] = () + spans: Sequence[Span] = () + texts: dict[tuple[str, SpanPart], str] = field(default_factory=dict) + text_reads: list[tuple[tuple[str, ...], SpanPart, int, int | None, bool]] = field(default_factory=list) + listed: list[tuple[QueryScope, str, str | None, int, RunOrder, tuple[str, ...]]] = field(default_factory=list) + + def _matching(self, q: str, trace_refs: Sequence[str]) -> list[TraceSummary]: + return [run for run in self.runs if (not trace_refs or run.get("trace_ref") in trace_refs) and q in run["name"]] + + async def list_traces( + self, + scope: QueryScope, + start_ms: int, + end_ms: int, + q: str = "", + cursor: str | None = None, + limit: int = 50, + order: RunOrder = NEWEST, + trace_refs: Sequence[str] = (), + ) -> TracePage: + self.listed.append((scope, q, cursor, limit, order, tuple(trace_refs))) + ordered: Final = sorted(self._matching(q, trace_refs), key=lambda run: run.get("trace_ref", "")) + after: Final = [run for run in ordered if cursor is None or run.get("trace_ref", "") > cursor][:limit] + return TracePage( + data=tuple(after), + next_cursor=after[-1].get("trace_ref") if len(after) == limit else None, + ) + + async def count_traces( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str = "", trace_refs: Sequence[str] = () + ) -> int: + return len(self._matching(q, trace_refs)) + + async def get_trace( + self, + trace_id: str, + scope: QueryScope, + trace_ref: str = "", + cursor: str | None = None, + page_size: int | None = None, + ) -> Trace | None: + if not self.spans: + return None + return Trace(summary=summary(trace_ref, trace_id), agents=(), spans=tuple(self.spans)) + + async def span_text( + self, + trace_id: str, + trace_ref: str, + span_ids: Sequence[str], + part: SpanPart, + scope: QueryScope, + offset: int = 0, + max_chars: int | None = None, + tail: bool = False, + contains: str | None = None, + ) -> tuple[SpanText, ...]: + self.text_reads.append((tuple(span_ids), part, offset, max_chars, tail)) + + def read(text: str) -> str: + if tail: + return text[len(text) - (max_chars or 0) :] + return text[offset : None if max_chars is None else offset + max_chars] + + return tuple( + SpanText( + span_id=span_id, + text=read(self.texts[(span_id, part)]), + total_chars=len(self.texts[(span_id, part)]), + version="0" * 64, + contains=contains is not None and contains in self.texts[(span_id, part)], + ) + for span_id in span_ids + if (span_id, part) in self.texts + ) + + +def refs(count: int) -> tuple[str, ...]: + return tuple(f"{index:064X}" for index in range(count)) + + +@pytest.mark.parametrize( + ("percent", "cap", "eligible", "selected"), + ( + (100, None, 7, 7), + (10, None, 7, 1), + (100, 3, 7, 3), + (50, 10, 7, 4), + (50, 10, 0, 0), + ), +) +def test_selection_keeps_the_rounded_up_share_within_the_cap( + percent: float, cap: int | None, eligible: int, selected: int +) -> None: + selection: Final = ActivitySelection(sample_percent=percent, sample_size=cap) + assert selected_count(selection, eligible) == selected + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("preview", "expected"), ((False, 3), (True, 7))) +async def test_sample_pages_runs_in_reference_order_up_to_its_bound(preview: bool, expected: int) -> None: + runs: Final = tuple(summary(ref) for ref in reversed(refs(7))) + storage: Final = FakeStorage(runs=runs) + reader: Final = SourceReader(storage) + selection: Final = ActivitySelection(sample_size=3) + pages: list[tuple[str, ...]] = [] # mutable-ok: collects the pages a cursor walk returns + cursor = "" # rebind-ok: follows the sample cursor until it ends + while True: + page = await reader.sample(Scope(all_teams=True), selection, 1, 2, cursor=cursor, page_size=2, preview=preview) + pages.append(tuple(e.trace_ref for e in page.executions)) + assert (page.eligible, page.selected) == (7, 3) + if page.next_cursor is None: + break + cursor = page.next_cursor + walked: Final = tuple(ref for page in pages for ref in page) # comprehension-ok: flatten pages + assert walked == refs(7)[:expected] + assert all(order == BY_REFERENCE for *_, order, _ in storage.listed) + + +@pytest.mark.asyncio +async def test_picked_executions_restrict_the_sample_to_their_runs() -> None: + storage: Final = FakeStorage(runs=tuple(summary(ref) for ref in refs(4))) + picked: Final = (execution_id(refs(4)[2], "trace"), execution_id(refs(4)[0], "trace")) + sample: Final = await SourceReader(storage).sample( + Scope(team_id="alpha"), ActivitySelection(execution_ids=picked), 1, 2 ) - assert "opaque-oauth-bearer" not in sample.model_dump_json() - assert sample.eligible == 1 + assert {e.id for e in sample.executions} == set(picked) + assert storage.listed[0][0] == {"kind": "owned", "user_id": "", "team_ids": ("alpha",)} + assert storage.listed[0][-1] == (refs(4)[2], refs(4)[0]) @pytest.mark.asyncio -async def test_agents_use_the_same_team_and_key_scope_as_samples() -> None: - class AgentStorage: - async def lens_agents(self, parameters): - assert parameters.all_teams == 0 - assert parameters.team == "alpha" - assert parameters.key_hash == "key-hash" - return (AgentRow(agent_name="research_agent"), AgentRow(agent_name="support_agent")) - - names: Final = await SourceReader(AgentStorage()).agents(Scope(team_id="alpha", api_key_hash="key-hash")) - assert names == ("research_agent", "support_agent") +async def test_a_span_within_budget_reads_whole_and_a_long_one_keeps_both_ends() -> None: + long_output: Final = "".join(f"{index:05d}" for index in range(2_000)) + storage: Final = FakeStorage( + spans=(span("root", None), span("short"), span("long", status="error")), + texts={ + ("root", "input"): "task", + ("short", "input"): "question", + ("short", "output"): "answer", + ("long", "output"): long_output, + ("long", "error"): "boom", + }, + ) + execution: Final = Execution(id=execution_id(REF, "trace"), trace_id="trace", trace_ref=REF) + content: Final = await SourceReader(storage).content(Scope(all_teams=True), execution) + by_id: Final = {part.span_id: part for part in content.parts} + assert by_id["short"].content == "Input: question\nOutput: answer\nStatus: ok " + assert not by_id["short"].truncated + long_part: Final = by_id["long"] + assert long_part.truncated + assert long_part.content == ( + "Input: \nOutput: " + + long_output[: 5_000 // 3] + + OMITTED + + long_output[-(5_000 - 5_000 // 3) :] + + "\nStatus: error boom" + ) + assert content.partial + tails: Final = [read for read in storage.text_reads if read[-1]] + assert {read[0] for read in tails} == {("long",)} @pytest.mark.asyncio -async def test_request_only_storage_is_available_for_investigation() -> None: - class RequestStorage: - async def lens_availability(self, parameters): - assert parameters.team == "alpha" - return (ActivityAvailability(traces=False, requests=True),) - - available: Final = await SourceReader(RequestStorage()).availability(Scope(team_id="alpha")) - assert available.requests - assert not available.traces +async def test_a_later_offset_reads_the_budget_window_of_the_full_text() -> None: + output: Final = "x" * BUDGET + "TAIL" + "y" * 100 + storage: Final = FakeStorage(spans=(span("root", None),), texts={("root", "output"): output}) + execution: Final = Execution(id=execution_id(REF, "trace"), trace_id="trace", trace_ref=REF) + labelled: Final = "Input: \nOutput: " + output + "\nStatus: ok " + offset: Final = len("Input: \nOutput: ") + BUDGET - 2 + content: Final = await SourceReader(storage).content(Scope(all_teams=True), execution, offset=offset) + assert content.parts[0].content == labelled[offset : offset + BUDGET] + assert content.parts[0].content.startswith("xxTAIL") + assert not content.parts[0].truncated + assert not content.partial @pytest.mark.asyncio -async def test_agent_filter_is_independent_of_service_and_metadata() -> None: - class SampleStorage: - async def lens_sample(self, parameters): - assert parameters.agent_name == "research_agent" - assert parameters.service == "shared-app" - assert parameters.filter_keys == ("enduser.id",) - assert parameters.filter_values == ("user-42",) - return [] +async def test_content_pages_forty_spans_at_a_time_in_trace_order() -> None: + spans: Final = (span("root", None), *(span(f"s{index:02d}") for index in range(45))) + storage: Final = FakeStorage(spans=spans) + execution: Final = Execution(id=execution_id(REF, "trace"), trace_id="trace", trace_ref=REF) + reader: Final = SourceReader(storage) + first: Final = await reader.content(Scope(all_teams=True), execution) + second: Final = await reader.content(Scope(all_teams=True), execution, cursor=first.next_cursor or "") + assert [p.span_id for p in (*first.parts, *second.parts)] == [s["span_id"] for s in spans] + assert second.next_cursor is None + assert [(len(ids), part, tail) for ids, part, _, _, tail in storage.text_reads] == [ + (40, "input", False), + (40, "output", False), + (40, "error", False), + (6, "input", False), + (6, "output", False), + (6, "error", False), + ] - settings: Final = lens().settings.model_copy( - update={ - "agent_name": "research_agent", - "service": "shared-app", - "filters": (MetadataFilter(key="enduser.id", value="user-42"),), + +@pytest.mark.asyncio +@pytest.mark.parametrize(("quote", "found"), (("time", True), ("boom", True), ("absent", False))) +async def test_evidence_must_appear_in_the_span_input_output_or_error(quote: str, found: bool) -> None: + storage: Final = FakeStorage(texts={("span", "output"): "timeout", ("span", "error"): "boom"}) + execution: Final = Execution(id=execution_id(REF, "trace"), trace_id="trace", trace_ref=REF) + evidence: Final = Evidence(execution_id=execution.id, span_id="span", quote=quote) + assert await SourceReader(storage).verify_evidence(Scope(all_teams=True), execution, evidence) is found + + +@pytest.mark.parametrize( + ("scope", "access"), + ( + (Scope(all_teams=True), {"kind": "all"}), + (Scope(team_id="alpha", api_key_hash="key"), {"kind": "owned", "user_id": "", "team_ids": ("alpha",)}), + (Scope(), {"kind": "owned", "user_id": "", "team_ids": ()}), + ), +) +def test_lens_scope_reads_with_the_same_access_as_traces(scope: Scope, access: QueryScope) -> None: + assert lens_access(scope) == access + + +def test_execution_ids_carry_reference_and_trace_id() -> None: + assert parse_execution(execution_id(REF, "trace:with:colons")) == (REF, "trace:with:colons") + with pytest.raises(ValueError): + parse_execution("short:trace") + + +def legacy_id(*parts: str) -> str: + return base64.urlsafe_b64encode(json.dumps(parts).encode()).decode() + + +def test_saved_filters_load_as_one_search() -> None: + selection: Final = ActivitySelection.model_validate( + { + "source": "both", + "agent_name": "research agent", + "service": "billing", + "team_id": "alpha", + "filters": [{"key": "tenant.tier", "value": "gold"}], + "execution_ids": [ + legacy_id("traces", "alpha", "trace", REF), + legacy_id("requests", "alpha", "request", ""), + legacy_id("traces", "alpha", "old"), + ], + "sample_percent": 50, } ) - assert not (await SourceReader(SampleStorage()).sample(Scope(all_teams=True), settings, 1, 2)).executions + assert selection.q == 'agent:"research agent" service:billing team:alpha attr.tenant.tier:gold' + assert selection.execution_ids == (execution_id(REF, "trace"),) + assert selection.sample_percent == 50 + + +def test_saved_job_samples_from_before_runs_load_as_absent() -> None: + job: Final = Job.model_validate( + { + "id": "job", + "created_at": NOW, + "start": NOW, + "end": NOW, + "settings": {"name": "Research", "model": "analysis", "context": "Find failures", "agent_name": "a"}, + "revision": 1, + "sample": {"executions": [{"id": "x", "source": "traces"}], "eligible": 1}, + } + ) + assert job.sample is None + assert job.settings.q == "agent:a" diff --git a/tests/unit/proxy/lens/test_worker.py b/tests/unit/proxy/lens/test_worker.py index 7983aec8af4..5729ff73206 100644 --- a/tests/unit/proxy/lens/test_worker.py +++ b/tests/unit/proxy/lens/test_worker.py @@ -19,6 +19,7 @@ from litellm.proxy.lens.models import ( from litellm.proxy.lens.state import queue_job from litellm.proxy.lens.worker import LensWorker, failure_message from tests.unit.proxy.lens.test_state import NOW, lens +from tests.unit.proxy.lens.test_sources import summary @pytest.mark.asyncio @@ -127,9 +128,7 @@ async def test_claim_without_an_identity_does_not_report_failure_for_another_inv @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(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 - ) + execution: Final = Execution(id="run", trace_id="trace", trace_ref="", summary=summary("", "trace", span_count=1)) sample: Final = Sample(executions=(execution,), eligible=1) content: Final = ExecutionContent( execution=execution, @@ -214,9 +213,7 @@ async def test_worker_saves_validation_errors_from_every_analysis_stage(purpose: import json 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 - ) + execution: Final = Execution(id="run", trace_id="trace", trace_ref="", summary=summary("", "trace", span_count=1)) sample: Final = Sample(executions=(execution,), eligible=1) content: Final = ExecutionContent( execution=execution, @@ -300,9 +297,7 @@ def test_response_validation_diagnostics_omit_input_values_and_unexpected_field_ @pytest.mark.parametrize("heartbeat_status", (401, 403, 409)) async def test_losing_the_lease_interrupts_an_in_flight_model_request(heartbeat_status: int) -> None: claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) started: Final = asyncio.Event() cancelled: Final = asyncio.Event() never: Final = asyncio.Event() @@ -358,9 +353,7 @@ async def test_losing_the_lease_interrupts_an_in_flight_model_request(heartbeat_ @pytest.mark.parametrize("failure", (429, 500, 502, 503, 504, "connection", "timeout")) async def test_transient_heartbeat_failure_recovers_without_cancelling_analysis(failure: int | str) -> None: claim: Final = Claim(lens_id="lens", job=queue_job(lens(), NOW, "job").jobs[0], findings=()) - execution: Final = Execution( - id="run", source="traces", trace_id="t", team_id="", name="task", start_time="", span_count=1 - ) + execution: Final = Execution(id="run", trace_id="t", trace_ref="", summary=summary("", "t", span_count=1)) started: Final = asyncio.Event() recovered: Final = asyncio.Event() never: Final = asyncio.Event() diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 87ba4d1f61b..e358ae52331 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -24,7 +24,7 @@ from litellm.rust_bridge.trace.errors import TraceChanged from litellm.rust_bridge.trace.generated.models import TraceQueryHelp from litellm.rust_bridge.trace.generated.types import AllQueryScope, OwnedQueryScope, QueryScope from litellm.rust_bridge.trace.queries import TraceSQLResponse -from litellm.rust_bridge.trace.storage import ClickHouseStorage, TraceStorageConfig +from litellm.rust_bridge.trace.storage import NEWEST, ClickHouseStorage, TraceStorageConfig from litellm.tracing import Tenant, TraceReceiver, TracingPayloadTooLargeError SQL_ENVELOPE: Final = { @@ -144,7 +144,9 @@ def test_trace_read_and_write_permissions( if scope is None: receiver.list_traces.assert_not_awaited() else: - receiver.list_traces.assert_awaited_once_with(scope=scope, start_ms=1, end_ms=2, q="", cursor=None) + receiver.list_traces.assert_awaited_once_with( + scope=scope, start_ms=1, end_ms=2, q="", cursor=None, order=NEWEST + ) write: Final = client.post("/v1/traces", json={}) assert write.status_code == (200 if can_write else 403), write.text @@ -257,9 +259,31 @@ def test_list_traces_passes_scope_window_query_and_cursor(client, receiver): end_ms=2, q="agent:research* -status:ok", cursor="abc", + order=NEWEST, ) +@pytest.mark.parametrize( + ("params", "order"), + [ + ({"sort_by": "duration_ms", "sort_dir": "asc"}, {"key": "duration_ms", "descending": False}), + ({"sort_by": "error_count"}, {"key": "error_count", "descending": True}), + ({"sort_dir": "asc"}, {"key": "start_ms", "descending": False}), + ], +) +def test_list_traces_orders_by_the_requested_run_key(client, receiver, params, order): + response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, **params}) + assert response.status_code == 200, response.text + assert receiver.list_traces.await_args.kwargs["order"] == order + + +@pytest.mark.parametrize("params", [{"sort_by": "trace_ref"}, {"sort_by": "spend"}, {"sort_dir": "up"}]) +def test_list_traces_rejects_orders_the_runs_table_does_not_offer(client, receiver, params): + response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, **params}) + assert response.status_code == 422 + receiver.list_traces.assert_not_awaited() + + def test_histogram_passes_scope_window_query_and_buckets(client, receiver): response = client.get("/v1/traces/histogram", params={"start_ms": 1, "end_ms": 2, "q": "x", "buckets": 12}) assert response.status_code == 200, response.text @@ -603,7 +627,7 @@ def test_lens_reads_from_the_lifespan_storage() -> None: storage: Final = MagicMock(spec=ClickHouseStorage) storage.ensure_schema = AsyncMock() - storage.lens_sample = AsyncMock(return_value=[]) + storage.count_traces = AsyncMock(return_value=0) tracing: Final = TraceReceiver(storage) @asynccontextmanager @@ -618,12 +642,13 @@ def test_lens_reads_from_the_lifespan_storage() -> None: with TestClient(app) as client: response: Final = client.post( "/lens/preview/sample", - json={"settings": {"name": "Review", "model": "analysis", "context": "Find failed executions"}}, + json={"selection": {"q": "service:checkout"}}, ) assert response.status_code == 200, response.text assert response.json()["executions"] == [] - storage.lens_sample.assert_awaited_once() - assert storage.lens_sample.await_args.args[0].all_teams == 1 + storage.count_traces.assert_awaited_once() + access, _, _, q, refs = storage.count_traces.await_args.args + assert (access, q, refs) == ({"kind": "all"}, "service:checkout", ()) def test_lens_reads_from_injected_storage_without_receiver() -> None: @@ -631,7 +656,7 @@ def test_lens_reads_from_injected_storage_without_receiver() -> None: from litellm.proxy.lens.sources import Storage storage: Final = MagicMock(spec=Storage) - storage.lens_sample = AsyncMock(return_value=[]) + storage.count_traces = AsyncMock(return_value=0) app: Final = FastAPI() app.include_router(lens_router) app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) @@ -640,12 +665,12 @@ def test_lens_reads_from_injected_storage_without_receiver() -> None: with TestClient(app) as client: response: Final = client.post( "/lens/preview/sample", - json={"settings": {"name": "Review", "model": "analysis", "context": "Find failed executions"}}, + json={"selection": {"q": "service:checkout"}}, ) assert response.status_code == 200, response.text assert response.json()["executions"] == [] - storage.lens_sample.assert_awaited_once() + storage.count_traces.assert_awaited_once() @pytest.mark.parametrize( diff --git a/tests/unit/rust_bridge/trace/test_queries.py b/tests/unit/rust_bridge/trace/test_queries.py index 1556cb5f46f..fc36f3a7ddb 100644 --- a/tests/unit/rust_bridge/trace/test_queries.py +++ b/tests/unit/rust_bridge/trace/test_queries.py @@ -2,52 +2,9 @@ from collections.abc import Mapping from typing import Final import pytest -from pydantic import JsonValue, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError -from litellm.rust_bridge.trace.generated.models import LensContentParams -from litellm.rust_bridge.trace.queries import LENS_CONTENT, LENS_EVIDENCE, TraceSQLResponse - - -@pytest.mark.parametrize("offset", (-1, 2**32)) -def test_named_query_rejects_offsets_outside_the_native_integer_range(offset: int) -> None: - with pytest.raises(ValidationError) as error: - LENS_CONTENT.parameters.model_validate( - { - "all_teams": 0, - "team": "team", - "key_hash": "", - "source": "traces", - "id": "trace", - "record_team": "team", - "trace_ref": "ref", - "cursor": "", - "offset": offset, - } - ) - assert error.value.error_count() == 1 - - -def test_named_query_rejects_parameters_for_a_different_query() -> None: - detail: Final = LensContentParams( - all_teams=0, - team="team", - key_hash="", - source="traces", - id="trace", - record_team="team", - trace_ref="ref", - cursor="", - offset=0, - ) - with pytest.raises(ValidationError) as error: - LENS_EVIDENCE.parameters.model_validate(detail) - assert error.value.error_count() == 1 - - -def test_named_query_rejects_rows_missing_required_result_fields() -> None: - with pytest.raises(ValidationError) as error: - LENS_CONTENT.response.validate_json('{"data":[{"span_id":"span","name":"name"}]}') - assert error.value.error_count() == 4 +from litellm.rust_bridge.trace.queries import TraceSQLResponse def test_sql_envelope_preserves_nested_data_large_integer_strings_and_extra_fields() -> None: @@ -62,50 +19,6 @@ def test_sql_envelope_preserves_nested_data_large_integer_strings_and_extra_fiel assert result.model_dump(mode="json", exclude_unset=True) == envelope -@pytest.mark.parametrize("count", (0, "9007199254740993", 2**64 - 1)) -def test_clickhouse_rows_normalize_numbers_and_preserve_tuples(count: int | str) -> None: - from litellm.rust_bridge.trace.queries import LENS_SAMPLE - - result: Final = LENS_SAMPLE.response.validate_json( - '{"data":[{"source":"traces","trace_id":"trace","team_id":"team","name":"run",' - '"start_time":"time","span_count":' - + (f'"{count}"' if isinstance(count, str) else str(count)) - + ',"root_seen":"1","eligible":"2","selected":2.0,"attributes":[["key","value"]]}]}' - ) - row: Final = result.data[0] - assert row.span_count == int(count) - assert row.root_seen == 1 - assert row.selected == 2 - assert row.attributes == (("key", "value"),) - assert row.service == "" - assert row.trace_ref == "" - assert row.selection_key == "" - with pytest.raises(ValidationError): - row.name = "changed" - - -def test_response_defaults_remain_normalized_when_omitted() -> None: - from litellm.rust_bridge.trace.generated.models import ActivityAvailability - from litellm.rust_bridge.trace.queries import LENS_SAMPLE - - row: Final = LENS_SAMPLE.response.validate_json( - '{"data":[{"source":"requests","trace_id":"trace","team_id":"team","name":"run",' - '"start_time":"time","span_count":"1","root_seen":1,"eligible":"2"}]}' - ).data[0] - assert row.attributes == () - assert row.selected == 0 - assert ActivityAvailability().traces is False - assert ActivityAvailability().requests is False - - -@pytest.mark.parametrize("count", (-1, "18446744073709551616", "1.5")) -def test_clickhouse_count_rejects_invalid_quoted_and_unquoted_numbers(count: int | str) -> None: - from litellm.rust_bridge.trace.queries import LENS_EVIDENCE - - with pytest.raises(ValidationError): - LENS_EVIDENCE.response.validate_python({"data": [{"count": count}]}) - - def test_dictionary_validation_keeps_required_nullable_and_optional_fields_distinct() -> None: from pydantic import TypeAdapter @@ -142,36 +55,5 @@ def test_invalid_native_response_preserves_validation_error_as_cause() -> None: from litellm.rust_bridge.trace.storage import _decode_query_response with pytest.raises(RuntimeError, match="Native trace query returned an invalid response") as error: - _decode_query_response(LENS_EVIDENCE.response, '{"data":[{"count":-1}]}') + _decode_query_response(TypeAdapter(TraceSQLResponse), '{"meta":[],"data":[],"rows":-1}') assert isinstance(error.value.__cause__, ValidationError) - - -@pytest.mark.parametrize("flag", (0, 1, "0", "1")) -def test_clickhouse_availability_normalizes_numeric_boolean_flags(flag: int | str) -> None: - from litellm.rust_bridge.trace.generated.models import ActivityAvailability - - result: Final = ActivityAvailability.model_validate({"traces": flag, "requests": flag}) - assert result.traces is (str(flag) == "1") - assert result.requests is result.traces - - -def test_response_flags_reject_values_outside_the_boolean_range() -> None: - from litellm.rust_bridge.trace.generated.models import ActivityAvailability - - with pytest.raises(ValidationError): - ActivityAvailability.model_validate({"traces": 2}) - with pytest.raises(ValidationError): - LENS_CONTENT.response.validate_python( - { - "data": [ - { - "span_id": "s", - "parent_span_id": "", - "name": "n", - "kind": "agent", - "content": "", - "truncated": "2", - } - ] - } - ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 7cdf7aa0489..716a0b26ea9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -89,7 +89,7 @@ describe("(dashboard) Layout", () => { vi.mocked(usePathname).mockReturnValue("/ui/guardrails"); }); - it("collapses the sidebar on Logs for a full-screen view and expands it again after leaving", async () => { + it("keeps the sidebar expanded on Logs instead of forcing it collapsed", async () => { const dashboard = () => ( @@ -103,10 +103,6 @@ describe("(dashboard) Layout", () => { vi.mocked(usePathname).mockReturnValue("/ui/logs"); rerender(dashboard()); - expect(screen.getByTestId("sidebar")).toHaveAttribute("data-collapsed", "true"); - - vi.mocked(usePathname).mockReturnValue("/ui/api-keys"); - rerender(dashboard()); expect(screen.getByTestId("sidebar")).toHaveAttribute("data-collapsed", "false"); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index d705089cee8..d7b265a4eab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -15,7 +15,7 @@ import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner"; import { UserBanner } from "@/components/UserBanner"; import { LiteAdminFrame } from "@/components/liteadmin/LiteAdmin"; import { UpgradeBanner } from "@/components/UpgradeBanner"; -import { routeSegmentForPathname, uiHref } from "@/utils/uiHref"; +import { uiHref } from "@/utils/uiHref"; import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext"; import { createApiClient } from "@/lib/http/client"; import { getProxyBaseUrl } from "@/components/networking"; @@ -99,17 +99,11 @@ export function AgentControlPlaneView() { ); } -const FULL_BLEED_SEGMENTS = new Set(["logs"]); - function DashboardShell({ children }: { children: React.ReactNode }) { const { accessToken } = useAuth(); const { mode } = usePluginMode(); - const routeSegment = routeSegmentForPathname(usePathname()); - const isFullBleed = FULL_BLEED_SEGMENTS.has(routeSegment); - // A manual toggle holds only for the route it was made on; full-bleed routes default to collapsed. - const [sidebarOverride, setSidebarOverride] = useState<{ segment: string; collapsed: boolean } | null>(null); - const sidebarCollapsed = sidebarOverride?.segment === routeSegment ? sidebarOverride.collapsed : isFullBleed; - const toggleSidebar = () => setSidebarOverride({ segment: routeSegment, collapsed: !sidebarCollapsed }); + const [sidebarCollapsed, setSidebarCollapsed] = useState(false); + const toggleSidebar = () => setSidebarCollapsed((collapsed) => !collapsed); const isGateway = mode === "ai-gateway"; diff --git a/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx index 16bea129f86..24b06152d9f 100644 --- a/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx +++ b/ui/litellm-dashboard/src/components/lens/LensModeSwitch.tsx @@ -30,6 +30,11 @@ function ActivityDot({ activity }: { activity: InvestigationActivity }) { ); } +export interface SetupBadge { + tab: LensTab; + label: string; +} + /** Inverted corner joining the tab's side border to the card's top border; plain CSS borders so both snap to the same pixels. */ function NotchCorner({ side, demo }: { side: "left" | "right"; demo: boolean }) { return ( @@ -54,7 +59,7 @@ export function LensModeSwitch({ activity: InvestigationActivity; demo: boolean; workers: LensList["workers"] | null; - setup?: string; + setup?: SetupBadge; }) { const connected = useWorkerConnected(workers); const workerTitle = connected ? "Worker connected" : "Connect worker"; @@ -86,9 +91,9 @@ export function LensModeSwitch({ {workerDisconnected && } {label} - {view === "investigations" && setup && ( + {setup?.tab === view && ( - {setup} + {setup.label} )} {view === "investigations" && } diff --git a/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx index c74b266714d..bf6027337cd 100644 --- a/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/LensSetup.integration.test.tsx @@ -1,4 +1,4 @@ -import { act, fireEvent, screen, within, waitFor } from "@testing-library/react"; +import { fireEvent, screen, within, waitFor } 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"; @@ -20,7 +20,7 @@ const worker = () => ({ scope: data.lenses[0].scope, }); -function serve({ enabled = false, traces = false, requests = false, connected = false } = {}) { +function serve({ enabled = false, traces = false, connected = false } = {}) { list.mockResolvedValue({ lenses: [], workers: connected ? [worker()] : [], tracing_enabled: enabled }); network.mockImplementation(async (input, init) => { const { path, method, body } = await readRequest(input, init); @@ -28,7 +28,7 @@ function serve({ enabled = false, traces = false, requests = false, connected = return enabled ? Response.json({ data: traces ? [data.runs[0].trace.summary] : [] }) : Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); - if (path === "/lens/activity/available") return Response.json({ traces, requests }); + if (path === "/lens/activity/available") return Response.json({ traces }); if (path === "/lens" && method === "POST") { const saved = { ...data.lenses[0], settings: { ...data.lenses[0].settings, ...(body as object) } }; list.mockResolvedValue({ lenses: [saved], workers: [worker()], tracing_enabled: true }); @@ -197,87 +197,6 @@ describe("Lens setup journey", () => { await connectWorkerFromSettings(user); }); - it("resumes setup from the URL and leaves only when the user chooses traces", async () => { - serve({ enabled: true, traces: true }); - const user = userEvent.setup(); - const onUrlUpdate = vi.fn(); - renderWorkspace({ searchParams: "?tab=investigations&setup=lens", onUrlUpdate }); - const intro = within(await screen.findByRole("dialog")); - expect(await intro.findByRole("button", { name: "Connect worker" })).toBeVisible(); - expect(intro.getByRole("heading", { name: "Get Lens running" })).toBeVisible(); - await user.click(intro.getByRole("button", { name: "View traces" })); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - expect(await screen.findByRole("table", { name: "Agent runs" })).toBeVisible(); - await waitFor(() => expect(setupParam(onUrlUpdate)).toBeNull()); - await user.click(screen.getByRole("tab", { name: "Investigations" })); - const guide = within(await screen.findByRole("region", { name: "Get Lens running" })); - expect(guide.getByRole("button", { name: /Connect a worker/ })).toHaveAttribute("aria-expanded", "true"); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - }); - - it("allows request-only investigations without forcing agent instrumentation", async () => { - dismissLensIntro(); - serve({ requests: true }); - const user = userEvent.setup(); - const welcome = renderWorkspace({ searchParams: "?tab=investigations" }); - expect(await screen.findByRole("region", { name: "Get Lens running" })).toBeVisible(); - expect(screen.getByRole("button", { name: "Connect worker" })).toBeEnabled(); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - welcome.unmount(); - renderWorkspace({ searchParams: "?tab=investigations&setup=lens" }); - const intro = within(await screen.findByRole("dialog")); - expect(await intro.findByRole("heading", { name: "Before you start" })).toBeVisible(); - expect(intro.getByRole("button", { name: "Connect worker" })).toBeEnabled(); - await user.click(intro.getByRole("button", { name: /Enable tracing on the gateway/ })); - expect(intro.getByRole("button", { name: "Continue with request logs" })).toBeEnabled(); - await user.click(intro.getByRole("button", { name: /Send your first trace/ })); - await user.click(intro.getByRole("button", { name: "Continue with request logs" })); - await connectWorkerFromSettings(user); - }); - - it("keeps setup recoverable when checking for a first trace fails", async () => { - serve({ enabled: true }); - const user = userEvent.setup(); - renderWorkspace({ searchParams: "?setup=lens" }); - const intro = within(await screen.findByRole("dialog")); - await intro.findByRole("button", { name: "Check for traces" }); - const normal = network.getMockImplementation()!; - network.mockImplementation(async (input, init) => { - const path = requestPath(input); - if (path === "/v1/traces") return Response.json({ detail: "Trace storage unavailable" }, { status: 503 }); - return normal(input, init); - }); - await user.click(intro.getByRole("button", { name: "Check for traces" })); - expect(await intro.findByRole("alert")).toHaveTextContent("Could not check setup"); - expect(intro.queryByRole("button", { name: "Continue to worker" })).not.toBeInTheDocument(); - serve({ enabled: true, traces: true }); - await user.click(intro.getByRole("button", { name: "Retry" })); - expect(await intro.findByRole("button", { name: "Continue to worker" })).toBeEnabled(); - expect(intro.queryByRole("alert")).not.toBeInTheDocument(); - }); - - it("keeps request-only users in investigations when an activity refresh fails", async () => { - dismissLensIntro(); - serve({ requests: true, connected: true }); - const user = userEvent.setup(); - renderWorkspace({ searchParams: "?tab=investigations" }); - expect(await screen.findByRole("button", { name: "New investigation" })).toBeEnabled(); - const normal = network.getMockImplementation()!; - network.mockImplementation((input, init) => - requestPath(input) === "/lens/activity/available" - ? Promise.resolve(Response.json({ detail: "Activity unavailable" }, { status: 503 })) - : normal(input, init), - ); - await act(() => testQueryClient.refetchQueries()); - expect(await screen.findByRole("alert")).toHaveTextContent("Could not check setup. Activity unavailable"); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - expect(screen.getByRole("button", { name: "New investigation" })).toBeDisabled(); - network.mockImplementation(normal); - await user.click(screen.getByRole("button", { name: "Retry" })); - await waitFor(() => expect(screen.getByRole("button", { name: "New investigation" })).toBeEnabled()); - expect(screen.queryByRole("alert")).not.toBeInTheDocument(); - }); - it("keeps administrator-only setup unavailable to trace viewers", async () => { serve({ enabled: true, traces: true }); renderWorkspace({ searchParams: "?setup=lens" }, "Internal User"); @@ -287,19 +206,14 @@ describe("Lens setup journey", () => { expect(network.mock.calls.some(([input]) => requestPath(input) === "/lens")).toBe(false); }); - it.each(["traces", "requests with trace errors", "requests with pending traces", "traces with activity errors"])( + it.each(["traces", "traces with activity errors"])( "finishes guided setup with %s and opens the saved investigation", async (scenario) => { - const source = scenario.startsWith("requests") ? "requests" : "traces"; - const activity = { enabled: true, traces: source === "traces", requests: source === "requests", connected: true }; - serve(activity); + serve({ enabled: true, traces: true, connected: true }); const normal = network.getMockImplementation()!; - const failingPath = scenario === "requests with trace errors" ? "/v1/traces" : "/lens/activity/available"; network.mockImplementation((input, init) => { const path = requestPath(input); - if (path === "/v1/traces" && scenario === "requests with pending traces") - return new Promise(() => {}); - if (scenario.endsWith("errors") && path === failingPath) + if (scenario.endsWith("errors") && path === "/lens/activity/available") return Promise.resolve(Response.json({ detail: "Activity unavailable" }, { status: 503 })); return normal(input, init); }); @@ -325,7 +239,7 @@ describe("Lens setup journey", () => { const requests = await Promise.all(network.mock.calls.map(([input, init]) => readRequest(input, init))); const create = requests.find((request) => request.path === "/lens" && request.method === "POST"); expect(create).toBeDefined(); - expect(create?.body).toEqual(expect.objectContaining({ name: "My first review", source })); + expect(create?.body).toEqual(expect.objectContaining({ name: "My first review", q: "" })); }, ); }); diff --git a/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx index d9a5b5826f0..c2c0952d5da 100644 --- a/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/LensWorkspace.integration.test.tsx @@ -24,7 +24,7 @@ beforeEach(() => { const path = requestPath(input); if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: false }); - return Response.json({ data: [], traces: false, requests: false }); + return Response.json({ data: [], traces: false }); }); }); @@ -177,7 +177,7 @@ describe("Lens interactive demo", () => { if (path === "/lens") return Response.json({ lenses: [saved], workers: [], tracing_enabled: true }); if (path.endsWith("/runs")) return Response.json(saved.jobs); if (path === "/v1/traces") return Response.json({ data: data.runs.map((run) => run.trace.summary) }); - return Response.json({ data: [], traces: true, requests: false }); + return Response.json({ data: [], traces: true }); }); renderWithProviders(, { searchParams: `?tab=investigations&lens=${saved.id}`, @@ -198,7 +198,7 @@ describe("Lens interactive demo", () => { if (path === "/lens") return Response.json({ lenses: [saved], workers: [], tracing_enabled: false }); if (path.endsWith("/runs")) return Response.json(saved.jobs); if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); - return Response.json({ data: [], traces: false, requests: false }); + return Response.json({ data: [], traces: false }); }); renderWithProviders(); expect(await screen.findByRole("button", { name: "Preview sample" })).toBeVisible(); @@ -223,7 +223,7 @@ describe("Lens interactive demo", () => { const path = requestPath(input); if (path === "/lens") return Response.json({ lenses: lenses(), workers: [], tracing_enabled: false }); if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); - return Response.json({ data: [], traces: false, requests: false }); + return Response.json({ data: [], traces: false }); }); renderWithProviders(); const tab = within(screen.getByRole("tablist", { name: "Lens" })).getByRole("tab", { name: "Investigations" }); @@ -243,7 +243,7 @@ describe("Lens interactive demo", () => { if (path === "/lens/agents") return Response.json([]); if (path.startsWith("/lens/preview")) return Response.json({ eligible: 0, selected: 0, executions: [] }); if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); - return Response.json({ data: [], traces: true, requests: false }); + return Response.json({ data: [], traces: true }); }); renderWithProviders(, { searchParams: "?tab=investigations&dialog=new", @@ -258,6 +258,30 @@ describe("Lens interactive demo", () => { expect(within(tab).queryByText("New")).not.toBeInTheDocument(); }); + it("marks the Traces tab while connecting another agent and clears it on back", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + network.mockImplementation(async (input) => { + const path = requestPath(input); + if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: true }); + if (path === "/v1/traces") return Response.json({ data: [createLensDemoData().runs[0].trace.summary] }); + return Response.json({ data: [], traces: true }); + }); + renderWithProviders(, { + searchParams: "?tab=traces&connect=true", + onUrlUpdate, + }); + expect(await screen.findByRole("heading", { name: "Connect another agent" })).toBeVisible(); + const tabs = within(screen.getByRole("tablist", { name: "Lens" })); + const tab = tabs.getByRole("tab", { name: /^Traces/ }); + expect(within(tab).getByText("Setup")).toBeVisible(); + expect(within(tabs.getByRole("tab", { name: /^Investigations/ })).queryByText("Setup")).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Back to traces" })); + expect(await screen.findByRole("combobox", { name: "Search runs" })).toBeVisible(); + expect(within(tab).queryByText("Setup")).not.toBeInTheDocument(); + await expectUrl(onUrlUpdate, (url) => expect(url.has("connect")).toBe(false)); + }); + it("adds a quiet Settings tab that manages the worker inline and reflects its health", async () => { const user = userEvent.setup(); const onUrlUpdate = vi.fn(); @@ -276,7 +300,7 @@ describe("Lens interactive demo", () => { if (path === "/lens") return Response.json({ lenses: [saved], workers: workers(), tracing_enabled: true }); if (path.endsWith("/runs")) return Response.json(saved.jobs); if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); - return Response.json({ data: [], traces: true, requests: false }); + return Response.json({ data: [], traces: true }); }); renderWithProviders(, { onUrlUpdate }); const tabs = within(screen.getByRole("tablist", { name: "Lens" })); @@ -311,7 +335,7 @@ describe("Lens interactive demo", () => { const path = requestPath(input); if (path === "/lens") return Response.json({ lenses: [], workers: [], tracing_enabled: true }); if (path === "/v1/traces") return Response.json({ data: [{}] }); - return Response.json({ data: [], traces: true, requests: false }); + return Response.json({ data: [], traces: true }); }); renderWithProviders(, { searchParams: "?tab=investigations", @@ -350,7 +374,7 @@ describe("Lens interactive demo", () => { if (path === "/lens/agents") return Response.json([]); if (path.startsWith("/lens/preview")) return Response.json({ eligible: 0, selected: 0, executions: [] }); if (path === "/v1/traces") return Response.json({ data: [createLensDemoData().runs[0].trace.summary] }); - return Response.json({ data: [], traces: true, requests: false }); + return Response.json({ data: [], traces: true }); }); renderWithProviders(, { searchParams: "?tab=settings", @@ -393,7 +417,7 @@ describe("Lens interactive demo", () => { network.mockImplementation(async (input) => requestPath(input) === "/lens" ? Response.json({ detail: "boom" }, { status: 500 }) - : Response.json({ data: [], traces: true, requests: false }), + : Response.json({ data: [], traces: true }), ); renderWithProviders(); const tabs = within(screen.getByRole("tablist", { name: "Lens" })); @@ -418,7 +442,7 @@ describe("Lens interactive demo", () => { const path = requestPath(input); if (path === "/lens") return Response.json({ lenses: [saved], workers: [worker], tracing_enabled: true }); if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); - return Response.json({ data: [], traces: true, requests: false }); + return Response.json({ data: [], traces: true }); }); renderWithProviders(); const tabs = within(screen.getByRole("tablist", { name: "Lens" })); @@ -449,7 +473,7 @@ describe("Lens interactive demo", () => { const path = requestPath(input); if (path === "/lens") return Response.json({ lenses: [saved], workers: workers(), tracing_enabled: true }); if (path === "/v1/traces") return Response.json({ detail: "Tracing is not enabled" }, { status: 501 }); - return Response.json({ data: [], traces: true, requests: false }); + return Response.json({ data: [], traces: true }); }); renderWithProviders(, { searchParams: "?tab=settings", diff --git a/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx b/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx index 9fa1c1df8d6..6a815e2da75 100644 --- a/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx +++ b/ui/litellm-dashboard/src/components/lens/LensWorkspace.tsx @@ -15,14 +15,14 @@ import { InvestigationsView } from "./investigations/InvestigationsView"; import { LensSettings } from "./settings/LensSettings"; import { createLensDemo } from "./data/demo/createLensDemo"; import { lensQueries } from "./data/queries"; -import { LensModeSwitch } from "./LensModeSwitch"; +import { LensModeSwitch, type SetupBadge } from "./LensModeSwitch"; import { frameCard } from "./ui/frame"; import { investigationActivity, listPollInterval } from "./model/status"; import { cn } from "@/lib/cva.config"; import { useDialogRoute, useLensRoute, type LensDialog, type LensTab } from "./route"; import { LensIntroDialog, useLensIntro } from "./onboarding/LensIntroDialog"; import { OnboardingProvider, type Onboarding } from "./onboarding/OnboardingContext"; -import { traceRefOf, useOpenTraceRouting } from "@/components/lens/traces/routing"; +import { traceRefOf, useOpenTraceRouting, useTracingSetupRoute } from "@/components/lens/traces/routing"; type WorkspaceProps = { accessToken: string; userRole: string; readOnly: boolean }; @@ -59,9 +59,15 @@ function DemoToggle({ demo, onChange }: { demo: boolean; onChange: (demo: boolea ); } -/** The inline investigation editor marks the tab so the notch says where you are, not just which tab is open. */ +/** An inline editor marks its tab so the notch says where you are, not just which tab is open. */ const SETUP_LABELS: Partial> = { new: "New", edit: "Editing", duplicate: "Duplicate" }; +function setupBadge(tab: LensTab, dialog: LensDialog | null, connecting: boolean): SetupBadge | undefined { + if (tab === "traces" && connecting) return { tab, label: "Setup" }; + const label = tab === "investigations" && dialog ? SETUP_LABELS[dialog] : undefined; + return label ? { tab, label } : undefined; +} + /** The one always-mounted `/lens` observer; every other reader is a plain cache subscriber. */ function useLensOverview(enabled: boolean, settingsOpen: boolean) { const api = useLensApi(); @@ -81,6 +87,7 @@ function LensContent({ userRole, readOnly }: Omit const { tab, lensId, demo, settingUp, setTab, setDemo, setSetup } = useLensRoute(); const { dialog, openDialog } = useDialogRoute(); const { openTrace } = useOpenTraceRouting(); + const [connecting] = useTracingSetupRoute(); const [previewTarget, setPreviewTarget] = useState(null); const canViewInvestigations = isProxyAdminTierRole(userRole); const isAdmin = isProxyAdminRole(userRole); @@ -163,7 +170,7 @@ function LensContent({ userRole, readOnly }: Omit activity={activity} demo={demo} workers={workers} - setup={activeTab === "investigations" && dialog ? SETUP_LABELS[dialog] : undefined} + setup={setupBadge(activeTab, dialog, connecting)} />

diff --git a/ui/litellm-dashboard/src/components/lens/data/LensServices.tsx b/ui/litellm-dashboard/src/components/lens/data/LensServices.tsx index a4804212012..7cedddf5272 100644 --- a/ui/litellm-dashboard/src/components/lens/data/LensServices.tsx +++ b/ui/litellm-dashboard/src/components/lens/data/LensServices.tsx @@ -18,7 +18,7 @@ export function liveLensServices(accessToken: string): LensServices { return { accessToken, lens: liveLensApi(fetchClient, apiClient, accessToken), traces: liveTracesApi(accessToken) }; } -function useLensServices(): LensServices { +export function useLensServices(): LensServices { const provided = useContext(LensServicesContext); if (!provided) throw new Error("Lens services need a LensServicesProvider above them"); return provided; diff --git a/ui/litellm-dashboard/src/components/lens/data/demo/createLensDemo.ts b/ui/litellm-dashboard/src/components/lens/data/demo/createLensDemo.ts index f8ae66c8f97..a73e72eb94a 100644 --- a/ui/litellm-dashboard/src/components/lens/data/demo/createLensDemo.ts +++ b/ui/litellm-dashboard/src/components/lens/data/demo/createLensDemo.ts @@ -20,12 +20,10 @@ function demoLensApi(data: LensDemoData): LensApi { return { scope: "demo", lenses: async () => ({ lenses: data.lenses, workers: [], tracing_enabled: true }), - activity: async () => ({ traces: true, requests: false }), + activity: async () => ({ traces: true }), runs: (lensId, offset) => found(jobs(lensId)?.slice(offset)), run: (lensId, jobId) => found(jobs(lensId)?.find((job) => job.id === jobId)), - execution: notInDemo, sample: notInDemo, - agents: notInDemo, models: async () => ({ data: [] }), modelDetails: async () => ({ data: [] }), keys: notInDemo, diff --git a/ui/litellm-dashboard/src/components/lens/data/demo/fixtures.ts b/ui/litellm-dashboard/src/components/lens/data/demo/fixtures.ts index becdff35d40..666ad56a95f 100644 --- a/ui/litellm-dashboard/src/components/lens/data/demo/fixtures.ts +++ b/ui/litellm-dashboard/src/components/lens/data/demo/fixtures.ts @@ -3,7 +3,8 @@ import type { Lens, Finding, Job, Settings } from "../../model/types"; import { withReleaseCases } from "./lensDemoLongTrace"; import { scenarios, type Scenario } from "./scenarios"; -const executionId = (traceId: string) => btoa(JSON.stringify(["traces", "", traceId])); +const demoRef = (traceId: string) => traceId.padStart(64, "0").toUpperCase(); +const executionId = (traceId: string) => `${demoRef(traceId)}:${traceId}`; const iso = (time: number) => new Date(time).toISOString(); function makeTrace(scene: Scenario, index: number, now: number) { @@ -231,13 +232,9 @@ export function createLensDemoData(now = Date.now()) { const lenses: Lens[] = definitions.map((definition) => { const settings: Settings = { name: definition.name, - agent_name: definition.agent, + q: `agent:${definition.agent}`, context: definition.context, checks: [{ id: definition.check, instruction: definition.instruction, enabled: true }], - source: "traces", - service: "", - filters: [], - team_id: "", execution_ids: [], lookback_hours: 24, sample_size: 0, @@ -252,16 +249,9 @@ export function createLensDemoData(now = Date.now()) { .filter(({ trace }) => trace.summary.name === definition.agent) .map(({ trace }) => ({ id: executionId(trace.summary.trace_id), - trace_ref: "", - metadata: [], - root_seen: true, - service: trace.summary.service, - source: "traces" as const, trace_id: trace.summary.trace_id, - team_id: "", - name: trace.summary.name, - start_time: trace.summary.start_time, - span_count: trace.summary.span_count, + trace_ref: demoRef(trace.summary.trace_id), + summary: { ...trace.summary, trace_ref: demoRef(trace.summary.trace_id) }, })); const relevant = findings.filter((f) => f.check_id === definition.check); const jobs: Job[] = [0, 1].map((day) => { diff --git a/ui/litellm-dashboard/src/components/lens/data/queries.ts b/ui/litellm-dashboard/src/components/lens/data/queries.ts index 05cb98fdb2b..359555043e2 100644 --- a/ui/litellm-dashboard/src/components/lens/data/queries.ts +++ b/ui/litellm-dashboard/src/components/lens/data/queries.ts @@ -1,6 +1,6 @@ import { infiniteQueryOptions, keepPreviousData, queryOptions } from "@tanstack/react-query"; import { hasActiveJob } from "../model/status"; -import type { Sample, Settings, ActivitySelection, Job } from "../model/types"; +import type { Sample, ActivitySelection, Job } from "../model/types"; import type { KeyPage, LensApi } from "./service"; export type { Key } from "./service"; @@ -17,17 +17,11 @@ export const lensKeys = { runs: () => [...lensKeys.all, "run"] as const, run: (scope: string, lensId: string | undefined, batchId: string) => [...lensKeys.runs(), { scope, lensId, batchId }] as const, - evidence: (scope: string, lensId: string | undefined, evidenceId: string | undefined, offset: number) => - [...lensKeys.all, "evidence", { scope, lensId, evidenceId, offset }] as const, models: (scope: string) => [...lensKeys.all, "models", { scope }] as const, modelDetails: (scope: string) => [...lensKeys.all, "model-details", { scope }] as const, activity: (scope: string) => [...lensKeys.all, "activity-available", { scope }] as const, - discoveries: () => [...lensKeys.all, "discovery"] as const, - discovery: (scope: string, source: Settings["source"], hours: number | undefined) => - [...lensKeys.discoveries(), { scope, source, hours }] as const, preview: (scope: string, selection: ActivitySelection, asOf: string) => [...lensKeys.all, "preview", { scope, selection, asOf }] as const, - agents: (scope: string) => [...lensKeys.all, "agents", { scope }] as const, analysisKeys: (scope: string, query: string) => [...lensKeys.all, "analysis-keys", { scope, query }] as const, analysisKeyInfo: (scope: string, keyId: string | undefined) => [...lensKeys.all, "analysis-key-info", { scope, keyId }] as const, @@ -48,8 +42,7 @@ export const lensQueries = { queryKey: lensKeys.activity(api.scope), queryFn: () => api.activity(), enabled: loaded, - refetchInterval: ({ state }: { state: { data?: Activity } }) => - state.data?.traces && state.data.requests ? false : 5000, + refetchInterval: ({ state }: { state: { data?: Activity } }) => (state.data?.traces ? false : 5000), }; return queryOptions(options); }, @@ -71,43 +64,12 @@ export const lensQueries = { }; return queryOptions(options); }, - evidence( - api: LensApi, - { - lensId, - evidenceId, - requestOffset, - source, - }: { - lensId: string | undefined; - evidenceId: string | undefined; - requestOffset: number; - source: string | undefined; - }, - ) { - const options = { - queryKey: lensKeys.evidence(api.scope, lensId, evidenceId, requestOffset), - enabled: !!lensId && source === "requests", - queryFn: () => api.execution(lensId as string, evidenceId ?? "", requestOffset), - }; - return queryOptions(options); - }, - discovery(api: LensApi, { value, enabled }: { value: ActivitySelection; enabled: boolean }) { - const unfiltered = { source: value.source, service: "", filters: [], lookback_hours: value.lookback_hours }; - const options = { - queryKey: lensKeys.discovery(api.scope, value.source, value.lookback_hours), - queryFn: () => api.sample(unfiltered, 0, new Date().toISOString()), - staleTime: 60000, - enabled, - }; - return queryOptions(options); - }, preview(api: LensApi, { scope, asOf, enabled }: { scope: ActivitySelection; asOf: string; enabled: boolean }) { const options = { queryKey: lensKeys.preview(api.scope, scope, asOf), - initialPageParam: 0, - queryFn: ({ pageParam }: { pageParam: number }) => api.sample(scope, pageParam, asOf), - getNextPageParam: (lastPage: Sample) => lastPage.next_offset ?? undefined, + initialPageParam: "", + queryFn: ({ pageParam }: { pageParam: string }) => api.sample(scope, pageParam, asOf), + getNextPageParam: (lastPage: Sample) => lastPage.next_cursor ?? undefined, enabled, staleTime: 30000, gcTime: 60_000, @@ -115,15 +77,6 @@ export const lensQueries = { }; return infiniteQueryOptions(options); }, - agents(api: LensApi, source: Settings["source"]) { - const options = { - queryKey: lensKeys.agents(api.scope), - queryFn: () => api.agents(), - enabled: source !== "requests", - staleTime: 60000, - }; - return queryOptions(options); - }, }; export function analysisKeysQuery(api: LensApi, query: string) { diff --git a/ui/litellm-dashboard/src/components/lens/data/service.ts b/ui/litellm-dashboard/src/components/lens/data/service.ts index 326ee398c0e..88df25c89bf 100644 --- a/ui/litellm-dashboard/src/components/lens/data/service.ts +++ b/ui/litellm-dashboard/src/components/lens/data/service.ts @@ -47,9 +47,7 @@ export interface LensApi { activity(): Promise; runs(lensId: string, offset: number): Promise; run(lensId: string, jobId: string): Promise; - execution(lensId: string, executionId: string, offset: number): Promise; - sample(selection: ActivitySelection, offset: number, asOf: string): Promise; - agents(): Promise; + sample(selection: ActivitySelection, cursor: string, asOf: string): Promise; models(): Promise<{ data: { id: string }[] }>; modelDetails(): Promise<{ data: AnalysisModelInfo[] }>; keys(alias: string, page: number, signal: AbortSignal): Promise; @@ -97,35 +95,23 @@ export function liveLensApi(client: LensClient, apiClient: ApiClient, accessToke params: { path: { lens_id: lensId, job_id: jobId } }, }), ), - execution: (lensId, executionId, offset) => - required( - client.GET("/lens/{lens_id}/executions/{execution_id}", { - headers, - params: { path: { lens_id: lensId, execution_id: executionId }, query: { offset } }, - }), - ), - sample: (selection, offset, asOf) => + sample: (selection, cursor, asOf) => required( client.POST("/lens/preview/sample", { headers, body: { - offset, + cursor, as_of: asOf, selection: { - source: selection.source, - service: selection.service ?? "", - agent_name: selection.agent_name ?? "", - filters: selection.filters ?? [], + q: selection.q ?? "", sample_size: selection.sample_size, sample_percent: selection.sample_percent ?? 100, - team_id: selection.team_id ?? "", execution_ids: [], }, lookback_hours: selection.lookback_hours ?? 24, }, }), ), - agents: () => required(client.GET("/lens/agents", { headers })), models: () => apiClient.get("/models", { accessToken }), modelDetails: () => apiClient.get("/model_group/info", { accessToken }), keys: async (alias, page, signal) => diff --git a/ui/litellm-dashboard/src/components/lens/investigations/Evidence.tsx b/ui/litellm-dashboard/src/components/lens/investigations/Evidence.tsx index fa4b716691c..90aa3c0f7b5 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/Evidence.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/Evidence.tsx @@ -1,30 +1,22 @@ "use client"; -import { useState } from "react"; -import { useQuery } from "@tanstack/react-query"; - import { Inspector } from "@/components/shared/Inspector"; -import { Button } from "@/components/ui/button"; import { RunView } from "@/components/lens/traces/detail/run/RunView"; import { useLocalRunSelection } from "@/components/lens/traces/routing"; -import { lensQueries } from "../data/queries"; -import { useLensAccessToken, useLensApi } from "../data/LensServices"; +import { useLensAccessToken } from "../data/LensServices"; import { evidenceTarget, type EvidenceTarget } from "../model/findings"; import type { EvidenceRef } from "../route"; -const SECTION = 8000; - export interface EvidenceViewProps { - readonly lensId: string; readonly evidence: EvidenceRef; /** Shown as a back link above the evidence when it is stacked over something else. */ readonly backLabel?: string; readonly onBack?: () => void; } -/** The original trace step or logged request a piece of evidence points at, filling the panel. */ -export function EvidenceView({ lensId, evidence, backLabel, onBack }: EvidenceViewProps) { +/** The original trace step a piece of evidence points at, filling the panel. */ +export function EvidenceView({ evidence, backLabel, onBack }: EvidenceViewProps) { const target = evidenceTarget(evidence.id); return (
@@ -33,29 +25,21 @@ export function EvidenceView({ lensId, evidence, backLabel, onBack }: EvidenceVi
)} - +
); } function EvidenceBody({ - lensId, target, evidence, onBack = () => {}, }: { - lensId: string; target: EvidenceTarget | null; evidence: EvidenceRef; onBack?: () => void; }) { - if (target?.source === "traces") + if (target) return ( ); - if (target?.source === "requests") return ; return

This evidence is no longer available.

; } @@ -92,40 +75,3 @@ export function TraceEvidence({ /> ); } - -function RequestEvidence({ lensId, evidenceId }: { lensId: string; evidenceId: string }) { - const api = useLensApi(); - const [offset, setOffset] = useState(0); - const request = { lensId, evidenceId, requestOffset: offset, source: "requests" }; - const evidence = useQuery(lensQueries.evidence(api, request)); - return ( -
-
-

Request evidence

-

Original logged input and output

-
-
- {evidence.isLoading &&

Loading request…

} - {evidence.error &&

{evidence.error.message}

} - {evidence.data?.parts.map((p) => ( - 
-            {p.content}
-
- ))} - {evidence.data?.parts.length === 0 &&

Request was not found or is past retention

} -
- {offset > 0 && ( - - )} - {evidence.data?.parts.some((p) => p.truncated) && ( - - )} -
-
-
- ); -} diff --git a/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.integration.test.tsx index 544a7272e14..03de704351c 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.integration.test.tsx @@ -21,7 +21,7 @@ vi.mock("@/components/lens/traces/detail/run/RunView", () => ({ const lens = { id: "lens", - settings: { name: "Support reviews", agent_name: "support_agent" }, + settings: { name: "Support reviews", q: "agent:support_agent" }, jobs: [], findings: [], } as unknown as Lens; @@ -103,7 +103,7 @@ it("keeps a feedback draft during refreshes, sends it with status changes, and r it("stacks a quote's original step over the finding and keeps the feedback draft on the way back", async () => { const user = userEvent.setup(); - const traceOf = (id: string) => btoa(JSON.stringify(["traces", "", id])); + const traceOf = (id: string) => `${"A".repeat(64)}:${id}`; const quoted: Finding = { ...finding, evidence: [ diff --git a/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.tsx b/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.tsx index 2c9ec48b13b..95c70f8351c 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/FindingDetails.tsx @@ -84,10 +84,10 @@ export function FindingDetails({ {evidenceGroups.map((group) => (
- {group.run?.name ?? evidenceTarget(group.id)?.id.slice(0, 12) ?? "Recorded run"} + {group.run?.summary?.name ?? evidenceTarget(group.id)?.id.slice(0, 12) ?? "Recorded run"} {group.quotes.length} {group.quotes.length === 1 ? "quote" : "quotes"} - {group.run ? ` · ${runTime(group.run.start_time)}` : ""} + {group.run?.summary ? ` · ${runTime(group.run.summary.start_time)}` : ""}
@@ -103,7 +103,7 @@ export function FindingDetails({ className="mt-2" onClick={() => onOpenEvidence({ id: e.execution_id, span: e.span_id })} > - {evidenceTarget(e.execution_id)?.source === "traces" ? "Open original step" : "Open request"} + Open original step
@@ -178,14 +178,7 @@ export function FindingPanelBody({ onReview={(status, reason) => onReview(owned, status, reason)} />
- {evidence && ( - setEvidence(null)} - /> - )} + {evidence && setEvidence(null)} />} ); } diff --git a/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx b/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx index 0619df164fa..da2aa5a2c67 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.integration.test.tsx @@ -9,7 +9,6 @@ import { InvestigationsView } from "./InvestigationsView"; import { LensPreviewContext } from "@/components/lens/ui/LensPreviewButton"; import { briefMarkdown } from "../model/findings"; import { findingKey } from "../model/inbox"; -import { runTime } from "../model/format"; import { type Lens, type Finding } from "../model/types"; const withPreview = (ui: React.ReactElement, open: () => void) => ( @@ -36,7 +35,7 @@ beforeEach(() => { proxy.post.mockResolvedValue({ eligible: 0, selected: 0, executions: [] }); }); -const executionId = btoa(JSON.stringify(["traces", "", "trace-42"])); +const executionId = `${"A".repeat(64)}:trace-42`; const pattern: Finding = { reason: "", suggestion: "", @@ -70,16 +69,12 @@ const lens: Lens = { scope: { all_teams: true, api_key_hash: "", team_id: "" }, settings: { context: "", - source: "traces", lookback_hours: 24, - service: "", - agent_name: "", - filters: [], + q: "", interval_minutes: 15, sample_size: 100, sample_percent: 100, concurrency: 8, - team_id: "", execution_ids: [], monthly_budget: 20, name: "Release reviews", @@ -121,16 +116,12 @@ const lens: Lens = { end: "2026-09-30T10:00:00Z", settings: { context: "", - source: "traces", lookback_hours: 24, - service: "", - agent_name: "", - filters: [], + q: "", interval_minutes: 15, sample_size: 100, sample_percent: 100, concurrency: 8, - team_id: "", execution_ids: [], monthly_budget: 20, enabled: false, @@ -145,16 +136,28 @@ const lens: Lens = { executions: [ { id: executionId, - trace_ref: "", - metadata: [], - root_seen: true, - service: "", - source: "traces", + trace_ref: "A".repeat(64), trace_id: "trace-42", - team_id: "", - name: "Release-42", - start_time: "2026-09-30 10:00:00.000", - span_count: 12, + summary: { + trace_id: "trace-42", + trace_ref: "A".repeat(64), + name: "Release-42", + service: "release", + input_preview: "", + start_time: "2026-09-30 10:00:00.000", + duration_ms: 1, + status: "ok", + span_count: 12, + agent_count: 0, + agent_invocations: 0, + llm_calls: 0, + tool_calls: 0, + error_count: 0, + input_tokens: 0, + output_tokens: 0, + models: [], + spend: null, + }, }, ], }, @@ -239,7 +242,7 @@ describe("Lens findings and runs", () => { it("shows the actual frozen run selection in the Runs tab", async () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(await screen.findByRole("tab", { name: "Agent traces" })); + await user.click(await screen.findByRole("tab", { name: "Runs" })); expect(screen.getByText("Release-42")).toBeInTheDocument(); expect(screen.getByTitle("trace-42")).toHaveTextContent("Release-42"); expect(screen.getByText(/1 selected from 1 matching runs/)).toBeInTheDocument(); @@ -289,7 +292,7 @@ it("runs with saved settings from Run now without opening setup, then accepts an ], }; if (path === "/lens/lens/runs") return lens.jobs; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; return { data: [] }; }); proxy.post.mockResolvedValue(lens); @@ -336,7 +339,7 @@ it("offers the interactive demo without starting an investigation", async () => testQueryClient.clear(); proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: false }; - return { traces: false, requests: false }; + return { traces: false }; }); const onPreview = vi.fn(); const user = userEvent.setup(); @@ -351,7 +354,6 @@ it("guides a first-time administrator into worker connection and lens setup", as testQueryClient.clear(); proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: true }; - if (path === "/lens/agents") return []; return { traces: true, requests: false, data: [] }; }); const user = userEvent.setup(); @@ -396,81 +398,6 @@ it("guides a first-time administrator into worker connection and lens setup", as expect(guide.getByRole("button", { name: "New investigation" })).toBeEnabled(); }); -it("opens the saved results of an older batch", async () => { - testQueryClient.clear(); - const older = { - ...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" }], - }; - proxy.get.mockImplementation(async (path) => { - 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(); - await screen.findByRole("option", { name: `${runTime(older.created_at)} · completed` }); - await user.selectOptions(screen.getByRole("combobox", { name: "Investigation run" }), "older"); - const investigation = within(screen.getByRole("complementary", { name: "Investigation details" })); - expect(await investigation.findByText("Earlier batch finding")).toBeVisible(); - expect(investigation.queryByText(issue.title)).not.toBeInTheDocument(); - await user.click(screen.getByRole("button", { name: "Run details" })); - expect(screen.getByText(/Took 2m 13s/)).toBeVisible(); - expect(screen.getByText("Activity window")).toBeVisible(); - await user.keyboard("{Escape}"); - await user.click(screen.getByRole("tab", { name: "History" })); - expect(within(screen.getByRole("tabpanel", { name: "History" })).getByText(/Took 2m 13s/)).toBeVisible(); -}); - -it("reads request content from the beginning after its abbreviated preview", async () => { - testQueryClient.clear(); - const requestId = btoa(JSON.stringify(["requests", "", "request-1"])); - const job = { - ...lens.jobs[0], - sample: { - eligible: 1, - executions: [{ ...lens.jobs[0].sample!.executions[0], id: requestId, source: "requests" as const }], - }, - }; - proxy.get.mockImplementation(async (path, options) => { - if (path === "/lens") return { lenses: [{ ...lens, jobs: [job] }], workers: [], tracing_enabled: true }; - if (path === "/lens/lens/runs") return [job]; - if (!path.includes("/executions/")) return { data: [] }; - const offset = Number(options.query.offset ?? 0); - return { - parts: [ - { - span_id: "request", - content: offset === 0 ? "Abbreviated preview" : `Original at ${offset}`, - truncated: true, - }, - ], - }; - }); - const user = userEvent.setup(); - renderWithProviders(); - await user.click(await screen.findByRole("tab", { name: "Agent traces" })); - const row = screen.getByRole("button", { name: /Release-42/ }); - await user.click(row); - const panel = await screen.findByRole("complementary", { name: "Run details" }); - expect(row).toHaveAttribute("aria-selected", "true"); - expect(panel).toHaveTextContent("1 / 1"); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - expect(await screen.findByText("Abbreviated preview")).toBeVisible(); - await user.click(screen.getByRole("button", { name: "Next section" })); - expect(await screen.findByText("Original at 1")).toBeVisible(); - await user.click(screen.getByRole("button", { name: "Next section" })); - expect(await screen.findByText("Original at 8001")).toBeVisible(); - await user.click(screen.getByRole("button", { name: "Previous section" })); - expect(await screen.findByText("Original at 1")).toBeVisible(); - await user.click(screen.getByRole("button", { name: "Previous section" })); - expect(await screen.findByText("Abbreviated preview")).toBeVisible(); -}); - it.each([false, true])( "directs a new user to traces when tracing_enabled=%s and there are no traces", async (enabled) => { @@ -494,7 +421,7 @@ it.each([false, true])( it("enables first-lens setup when a trace arrives without leaving Investigations", async () => { window.history.replaceState({}, "", "/lens/"); testQueryClient.clear(); - const traceCheck = vi.fn().mockResolvedValue({ traces: false, requests: false }); + const traceCheck = vi.fn().mockResolvedValue({ traces: false }); proxy.get.mockImplementation(async (path) => path === "/lens" ? { lenses: [], workers: [], tracing_enabled: true } : traceCheck(), ); @@ -504,7 +431,7 @@ it("enables first-lens setup when a trace arrives without leaving Investigations await act(async () => vi.advanceTimersByTimeAsync(50)); expect(screen.queryByText(/Your first trace is ready/)).not.toBeInTheDocument(); - traceCheck.mockResolvedValue({ traces: true, requests: false }); + traceCheck.mockResolvedValue({ traces: true }); await act(async () => vi.advanceTimersByTimeAsync(5000)); expect(screen.getByText(/Your first trace is ready/)).toBeVisible(); expect(screen.getByRole("button", { name: "Continue to worker" })).toBeEnabled(); @@ -557,41 +484,11 @@ it("shows a centered failure with a retry when investigations cannot load, then expect(screen.queryByRole("alert")).not.toBeInTheDocument(); }); -it("keeps saved investigations accessible when tracing is disabled", async () => { - testQueryClient.clear(); - proxy.get.mockImplementation(async (path) => { - if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: false }; - if (path === "/lens/lens/runs") return lens.jobs; - return { data: [] }; - }); - renderWithProviders(); - expect(await screen.findByRole("row", { name: issue.title })).toBeVisible(); - expect(screen.queryByRole("region", { name: "Get Lens running" })).not.toBeInTheDocument(); -}); - -it("allows request-only accounts to connect a worker without requiring agent traces", async () => { - window.history.replaceState({}, "", "/lens/"); - testQueryClient.clear(); - proxy.get.mockImplementation(async (path) => { - if (path === "/lens") return { lenses: [], workers: [], tracing_enabled: true }; - if (path === "/lens/activity/available") return { traces: false, requests: true }; - return { data: [] }; - }); - const user = userEvent.setup(); - renderWithProviders(); - expect(await screen.findByRole("button", { name: "Connect worker" })).toBeEnabled(); - await user.click(screen.getByRole("button", { name: /Send your first trace/ })); - const panel = within(screen.getByRole("region", { name: /Send your first trace/ })); - expect(panel.getByText(/Request logs are already available/)).toBeVisible(); - expect(panel.getByRole("button", { name: "Continue with request logs" })).toBeEnabled(); -}); - it("reopens the inline editor from a shared link and drops it from the URL on cancel", async () => { testQueryClient.clear(); proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [lens], workers: [], tracing_enabled: true }; if (path === "/lens/lens/runs") return lens.jobs; - if (path === "/lens/agents") return []; return { data: [] }; }); const onUrlUpdate = vi.fn(); @@ -654,7 +551,7 @@ it("steps across findings and investigations with J and K, skipping hidden findi }; proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [lens, twin], tracing_enabled: true, workers: [] }; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; if (path.endsWith("/runs")) return []; return { data: [] }; }); @@ -691,7 +588,7 @@ it("opens an investigation beside the list and walks from it into its findings w testQueryClient.clear(); proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [lens], tracing_enabled: true, workers: [] }; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; if (path.endsWith("/runs")) return []; return { data: [] }; }); @@ -724,7 +621,7 @@ it("lists each finding under the investigation that owns it and resolves only th const twin: Lens = { ...lens, id: "twin", settings: { ...lens.settings, name: "Twin reviews" } }; proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [lens, twin], tracing_enabled: true, workers: [] }; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; if (path.endsWith("/runs")) return []; return { data: [] }; }); @@ -747,9 +644,8 @@ it("lists investigations without edit or run controls for read-only viewers", as testQueryClient.clear(); proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [lens], tracing_enabled: true, workers: [] }; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; if (path === "/lens/lens/runs") return []; - if (path === "/lens/agents") return []; return { data: [] }; }); const user = userEvent.setup(); @@ -766,9 +662,8 @@ it("opens investigations from the keyboard without treating nested edit keys as testQueryClient.clear(); proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [lens], tracing_enabled: true, workers: [] }; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; if (path === "/lens/lens/runs") return []; - if (path === "/lens/agents") return []; return { data: [] }; }); const user = userEvent.setup(); @@ -805,10 +700,9 @@ it("opens a failed investigation's details from its row and edits only from the }; proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [{ ...lens, jobs: [job] }], workers: [], tracing_enabled: true }; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; if (path === "/lens/lens/runs") return [job]; if (path === "/lens/lens/runs/failed-run") return job; - if (path === "/lens/agents") return []; return { data: [] }; }); const user = userEvent.setup(); @@ -848,7 +742,7 @@ it("keeps a finding open to retry when its update fails", async () => { const twin: Lens = { ...lens, id: "twin", settings: { ...lens.settings, name: "Twin reviews" } }; proxy.get.mockImplementation(async (path) => { if (path === "/lens") return { lenses: [lens, twin], tracing_enabled: true, workers: [] }; - if (path === "/lens/activity/available") return { traces: true, requests: false }; + if (path === "/lens/activity/available") return { traces: true }; if (path.endsWith("/runs")) return []; return { data: [] }; }); diff --git a/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.tsx b/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.tsx index a99133d6ca2..56845958ab1 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/InvestigationsView.tsx @@ -177,7 +177,6 @@ export function InvestigationsView({ readOnly = false }: InvestigationsViewProps ready={status.ready} mode={current.mode} initial={current.mode === "new" ? undefined : current.initial} - defaultSource={!status.tracesReady && status.requestsReady ? "requests" : "traces"} onClose={closeDialog} onSave={(settings) => saveSetup(current, settings)} /> diff --git a/ui/litellm-dashboard/src/components/lens/investigations/detail/HistoryTimeline.test.ts b/ui/litellm-dashboard/src/components/lens/investigations/detail/HistoryTimeline.test.ts index a8d10b1df16..5b96828326c 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/detail/HistoryTimeline.test.ts +++ b/ui/litellm-dashboard/src/components/lens/investigations/detail/HistoryTimeline.test.ts @@ -33,17 +33,13 @@ const job: Job = { revision: 1, settings: { context: "", - source: "traces", lookback_hours: 24, - service: "", - agent_name: "", - filters: [], + q: "", enabled: false, interval_minutes: 15, sample_size: 100, sample_percent: 100, concurrency: 8, - team_id: "", execution_ids: [], monthly_budget: 20, name: "Release reviews", diff --git a/ui/litellm-dashboard/src/components/lens/investigations/detail/InvestigationDetail.tsx b/ui/litellm-dashboard/src/components/lens/investigations/detail/InvestigationDetail.tsx index 9a2bb98f29d..b9eb6730573 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/detail/InvestigationDetail.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/detail/InvestigationDetail.tsx @@ -8,7 +8,7 @@ import { InvestigationProgress } from "../InvestigationProgress"; import { StepFeed } from "../StepFeed"; import { InvestigationSummary } from "./InvestigationSummary"; import { InvestigationFailure } from "./InvestigationFailure"; -import { scopeLabel, sourceLabels } from "../../model/format"; +import { scopeLabel } from "../../model/format"; import type { OwnedFinding } from "../../model/inbox"; import { activeJob } from "../../model/status"; import { type Finding, type Lens } from "../../model/types"; @@ -51,9 +51,7 @@ export function InvestigationDetail({

{lens.settings.name}

-

- {sourceLabels[lens.settings.source ?? "traces"]} · {scopeLabel(lens.settings)} -

+

{scopeLabel(lens.settings)}

{!readOnly && }
@@ -66,7 +64,7 @@ export function InvestigationDetail({ Findings Criteria - {sourceLabels[batchSettings?.source ?? "traces"]} + Runs History {section !== "activity" && } diff --git a/ui/litellm-dashboard/src/components/lens/investigations/detail/RunNowDialog.tsx b/ui/litellm-dashboard/src/components/lens/investigations/detail/RunNowDialog.tsx index e8695c22168..1314d27f688 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/detail/RunNowDialog.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/detail/RunNowDialog.tsx @@ -14,10 +14,11 @@ import { } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; -import { lensQueries } from "../../data/queries"; -import { useLensApi } from "../../data/LensServices"; +import { useLensServices } from "../../data/LensServices"; import type { Lens, RunWindow } from "../../model/types"; -import { RUN_PRESETS, runRequest, type RunChoice, type RunPreset } from "../../model/runRequest"; +import { RUN_PRESETS, runRequest, savedAgent, type RunChoice, type RunPreset } from "../../model/runRequest"; + +const SUGGESTION_WINDOW_MS = 7 * 24 * 3_600_000; const localInput = (date: Date) => new Date(date.getTime() - date.getTimezoneOffset() * 60_000).toISOString().slice(0, 16); @@ -33,17 +34,22 @@ export function RunNowDialog({ onClose: () => void; onRun: (request: RunWindow) => Promise; }) { - const api = useLensApi(); - const agentsQuery = useQuery(lensQueries.agents(api, "traces")); - const agents = Array.isArray(agentsQuery.data) ? agentsQuery.data : []; - const now = new Date(); + const { traces } = useLensServices(); + const [now] = useState(() => new Date()); + const agentsQuery = useQuery({ + queryKey: ["lens", "run-now-agents", traces.scope], + queryFn: () => traces.values("agent", "", { startMs: now.getTime() - SUGGESTION_WINDOW_MS, endMs: now.getTime() }), + staleTime: 60_000, + }); + const agents = agentsQuery.data ?? []; + const saved = savedAgent(lens.settings.q); const [preset, setPreset] = useState(null); - const [agent, setAgent] = useState(lens.settings.agent_name ?? ""); + const [agent, setAgent] = useState(saved); const [start, setStart] = useState(localInput(new Date(now.getTime() - 3_600_000))); const [end, setEnd] = useState(localInput(now)); const [error, setError] = useState(""); const submit = async () => { - const choice: RunChoice = { preset, agent, saved: lens.settings.agent_name ?? "", start, end }; + const choice: RunChoice = { preset, agent, saved, start, end }; const request = runRequest(choice); if (typeof request === "string") { setError(request); diff --git a/ui/litellm-dashboard/src/components/lens/investigations/detail/RunsTab.tsx b/ui/litellm-dashboard/src/components/lens/investigations/detail/RunsTab.tsx index 3f02fff5c31..b0a2db0c49b 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/detail/RunsTab.tsx +++ b/ui/litellm-dashboard/src/components/lens/investigations/detail/RunsTab.tsx @@ -95,9 +95,11 @@ export function RunsTab({ lens, job }: RunsTabProps) { /> } > - {run.name} + + {run.summary?.name ?? run.trace_id} + - {runTime(run.start_time)} + {run.summary ? runTime(run.summary.start_time) : ""} {assessmentLabel(assessments.get(run.id))} @@ -141,7 +143,7 @@ export function RunsTab({ lens, job }: RunsTabProps) { )} - {(run: RunRef) => } + {(run: RunRef) => } ); diff --git a/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.test.ts b/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.test.ts index 96167b14b8f..f9fabe7aef7 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.test.ts +++ b/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.test.ts @@ -12,16 +12,12 @@ function makeLens(id: string, settings: Partial, jobs: readonl scope: { all_teams: true, api_key_hash: "", team_id: "" }, settings: { context: "", - source: "traces", lookback_hours: 24, - service: "", - agent_name: "", - filters: [], + q: "", interval_minutes: 15, sample_size: 100, sample_percent: 100, concurrency: 8, - team_id: "", execution_ids: [], monthly_budget: 20, name: id, @@ -40,9 +36,9 @@ function makeLens(id: string, settings: Partial, jobs: readonl } const lenses = [ - makeLens("Refund audit", { agent_name: "billing-agent" }, [{ status: "failed" }, { status: "completed" }]), - makeLens("Release reviews", { service: "reviewer", enabled: false }, [{ status: "completed" }]), - makeLens("Lead scoring", { filters: [{ key: "team", value: "sales" }] }), + makeLens("Refund audit", { q: "agent:billing-agent" }, [{ status: "failed" }, { status: "completed" }]), + makeLens("Release reviews", { q: 'agent:"reviewer bot"', enabled: false }, [{ status: "completed" }]), + makeLens("Lead scoring", { q: "attr.team:sales" }), ]; const names = (query: string) => filterInvestigations(lenses, query).map((lens) => lens.settings.name); @@ -52,6 +48,8 @@ describe("filterInvestigations", () => { expect(names("REFUND")).toEqual(["Refund audit"]); expect(names("sales")).toEqual(["Lead scoring"]); expect(names("reviewer")).toEqual(["Release reviews"]); + expect(names("agent:billing*")).toEqual(["Refund audit"]); + expect(names('agent:"reviewer bot"')).toEqual(["Release reviews"]); }); it("filters by the latest run's status, treating no runs as never", () => { @@ -60,17 +58,18 @@ describe("filterInvestigations", () => { expect(names("status:never")).toEqual(["Lead scoring"]); }); - it("filters by schedule and by agent, falling back to the service", () => { + it("filters by schedule and by the agent the search names", () => { expect(names("schedule:paused")).toEqual(["Release reviews"]); expect(names("-schedule:paused")).toEqual(["Refund audit", "Lead scoring"]); expect(names("agent:billing-agent")).toEqual(["Refund audit"]); - expect(names("agent:reviewer")).toEqual(["Release reviews"]); + expect(names('agent:"reviewer bot"')).toEqual(["Release reviews"]); + expect(names("agent:reviewer")).toEqual([]); }); }); describe("INVESTIGATION_INDEX values", () => { it("offers only the agents and statuses present", () => { - expect(fieldValues(INVESTIGATION_INDEX, lenses, "agent")).toEqual(["billing-agent", "reviewer"]); + expect(fieldValues(INVESTIGATION_INDEX, lenses, "agent")).toEqual(["billing-agent", "reviewer bot"]); expect(fieldValues(INVESTIGATION_INDEX, lenses, "status")).toEqual(["completed", "failed", "never"]); }); }); diff --git a/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.ts b/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.ts index 025fa73ee1a..d7abea69785 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.ts +++ b/ui/litellm-dashboard/src/components/lens/investigations/investigationQuery.ts @@ -1,3 +1,4 @@ +import { savedAgent } from "../model/runRequest"; import { Bot, CalendarClock, CircleDashed, SquareChevronRight } from "lucide-react"; import { type ClientIndex, filterItems } from "@/components/shared/search/evaluate"; @@ -22,7 +23,7 @@ export const INVESTIGATION_QUERY: QueryLanguage = { export const INVESTIGATION_INDEX: ClientIndex = { read: { name: (lens) => [lens.settings.name], - agent: (lens) => [lens.settings.agent_name, lens.settings.service].filter(Boolean), + agent: (lens) => [savedAgent(lens.settings.q)].filter(Boolean), status: (lens) => [lens.jobs[0]?.status ?? "never"], schedule: (lens) => [lens.settings.enabled ? "watching" : "paused"], }, diff --git a/ui/litellm-dashboard/src/components/lens/investigations/investigationScreen.test.ts b/ui/litellm-dashboard/src/components/lens/investigations/investigationScreen.test.ts index 6303e9e7802..bfee6a3406a 100644 --- a/ui/litellm-dashboard/src/components/lens/investigations/investigationScreen.test.ts +++ b/ui/litellm-dashboard/src/components/lens/investigations/investigationScreen.test.ts @@ -10,16 +10,12 @@ function makeLens(id: string, created_at: string): Lens { scope: { all_teams: true, api_key_hash: "", team_id: "" }, settings: { context: "", - source: "traces", lookback_hours: 24, - service: "", - agent_name: "", - filters: [], + q: "", interval_minutes: 15, sample_size: 100, sample_percent: 100, concurrency: 8, - team_id: "", execution_ids: [], monthly_budget: 20, name: `Lens ${id}`, diff --git a/ui/litellm-dashboard/src/components/lens/model/findings.ts b/ui/litellm-dashboard/src/components/lens/model/findings.ts index 2ef6929e7de..73722784e8f 100644 --- a/ui/litellm-dashboard/src/components/lens/model/findings.ts +++ b/ui/litellm-dashboard/src/components/lens/model/findings.ts @@ -9,21 +9,18 @@ export function sortedFindings(findings: Finding[]): Finding[] { } export interface EvidenceTarget { - readonly source: string; - readonly team: string; readonly id: string; - readonly traceRef?: string; + readonly traceRef: string; } -export function evidenceTarget(id: string): EvidenceTarget | null { - try { - const parsed: unknown = JSON.parse(atob(id.replace(/-/g, "+").replace(/_/g, "/"))); - if (!Array.isArray(parsed) || ![3, 4].includes(parsed.length) || !parsed.every((item) => typeof item === "string")) - return null; - return { source: parsed[0], team: parsed[1], id: parsed[2], ...(parsed[3] ? { traceRef: parsed[3] } : {}) }; - } catch { - return null; - } +const TRACE_REF_LENGTH = 64; + +/** An execution id is `:`; ids saved before that format no longer resolve. */ +export function evidenceTarget(executionId: string): EvidenceTarget | null { + const traceRef = executionId.slice(0, TRACE_REF_LENGTH); + const id = executionId.slice(TRACE_REF_LENGTH + 1); + if (traceRef.length !== TRACE_REF_LENGTH || executionId[TRACE_REF_LENGTH] !== ":" || !id) return null; + return { id, traceRef }; } export function briefMarkdown(title: string, brief: IssueBrief): string { diff --git a/ui/litellm-dashboard/src/components/lens/model/format.test.ts b/ui/litellm-dashboard/src/components/lens/model/format.test.ts index 6bea67c2bbe..d9df85dac52 100644 --- a/ui/litellm-dashboard/src/components/lens/model/format.test.ts +++ b/ui/litellm-dashboard/src/components/lens/model/format.test.ts @@ -2,10 +2,10 @@ import { describe, expect, it } from "vitest"; import { agoLabel, durationLabel, durationText, money, scopeLabel, when } from "./format"; describe("Lens labels", () => { - it("formats an empty scope and joins recorded scope constraints in their original order", () => { + it("labels a scope by its search, or all activity when there is none", () => { expect(scopeLabel({})).toBe("All activity"); - const settings = { agent_name: "research", service: "shared", filters: [{ key: "team", value: "quality" }] }; - expect(scopeLabel(settings)).toBe("research · shared · team: quality"); + expect(scopeLabel({ q: " " })).toBe("All activity"); + expect(scopeLabel({ q: " agent:research status:error " })).toBe("agent:research status:error"); }); it("formats elapsed time and treats a non-finite duration as zero", () => { diff --git a/ui/litellm-dashboard/src/components/lens/model/format.ts b/ui/litellm-dashboard/src/components/lens/model/format.ts index cd09f470dcc..de41ccb84b8 100644 --- a/ui/litellm-dashboard/src/components/lens/model/format.ts +++ b/ui/litellm-dashboard/src/components/lens/model/format.ts @@ -3,12 +3,8 @@ import type { Settings } from "./types"; export { runTime }; -export function scopeLabel(settings: Partial>): string { - return ( - [settings.agent_name, settings.service, ...(settings.filters ?? []).map((f) => `${f.key}: ${f.value}`)] - .filter(Boolean) - .join(" · ") || "All activity" - ); +export function scopeLabel(settings: Partial>): string { + return settings.q?.trim() || "All activity"; } export function durationText(seconds: number): string { @@ -40,5 +36,3 @@ export function agoLabel(thenMs: number, nowMs: number): string { export const money = (n: number) => new Intl.NumberFormat("en-US", { style: "currency", currency: "USD", maximumFractionDigits: 3 }).format(n); export const when = (value?: string | null) => (value ? runTime(value) : "Not yet"); - -export const sourceLabels = { both: "Traces and requests", requests: "LLM requests", traces: "Agent traces" }; diff --git a/ui/litellm-dashboard/src/components/lens/model/inbox.test.ts b/ui/litellm-dashboard/src/components/lens/model/inbox.test.ts index 7fd6294eaa2..9cf78ae6eed 100644 --- a/ui/litellm-dashboard/src/components/lens/model/inbox.test.ts +++ b/ui/litellm-dashboard/src/components/lens/model/inbox.test.ts @@ -35,25 +35,42 @@ const lens = ( id: string, agent: string, findings: Finding[], - { settings = {}, runs = [] }: { settings?: Partial; runs?: { id: string; service: string }[] } = {}, + { + settings = {}, + runs = [], + }: { settings?: Partial; runs?: { id: string; agents?: string[]; service: string }[] } = {}, ): Lens => ({ id, findings, - jobs: [{ sample: { executions: runs } }], + jobs: [ + { + sample: { + executions: runs.map((run) => ({ + id: run.id, + summary: { agent_names: run.agents ?? [], service: run.service }, + })), + }, + }, + ], next_run_at: "2026-10-03T12:10:00Z", - settings: { name: id, agent_name: agent, service: "", enabled: true, interval_minutes: 15, ...settings }, + settings: { name: id, q: agent ? `agent:${agent}` : "", enabled: true, interval_minutes: 15, ...settings }, }) as unknown as Lens; describe("findingAgents", () => { - it("names the agents the finding's runs were actually recorded under", () => { + it("names the agents the finding's runs were recorded under, or their service when unnamed", () => { const seen = lens("a", "", [], { runs: [ { id: "run-1", service: "support-bot" }, - { id: "run-2", service: "billing-bot" }, + { id: "run-2", agents: ["billing-bot", "refund-bot"], service: "shared" }, + { id: "run-3", agents: ["unrelated"], service: "x" }, ], }); - expect(findingAgents(seen, finding({ occurrences: ["run-2", "run-1"] }))).toEqual(["billing-bot", "support-bot"]); + expect(findingAgents(seen, finding({ occurrences: ["run-2", "run-1"] }))).toEqual([ + "billing-bot", + "refund-bot", + "support-bot", + ]); }); it("falls back to the configured agent, then to an explicit unknown, when no run says", () => { diff --git a/ui/litellm-dashboard/src/components/lens/model/inbox.ts b/ui/litellm-dashboard/src/components/lens/model/inbox.ts index f2a4d386a33..80c2a2df32a 100644 --- a/ui/litellm-dashboard/src/components/lens/model/inbox.ts +++ b/ui/litellm-dashboard/src/components/lens/model/inbox.ts @@ -1,4 +1,5 @@ -import type { Finding, Job, Lens } from "./types"; +import { savedAgent } from "./runRequest"; +import type { Execution, Finding, Job, Lens } from "./types"; export type Step = Job["steps"][number]; export type Priority = NonNullable; @@ -11,12 +12,15 @@ export function sampledExecutions(lens: Lens) { return lens.jobs.flatMap((job) => job.sample?.executions ?? []); } +const runAgents = (run: Execution): readonly string[] => { + const names = run.summary?.agent_names ?? []; + return names.length ? names : [run.summary?.service ?? ""].filter(Boolean); +}; + export function findingAgents(lens: Lens, finding: Finding): readonly string[] { - const services = new Map(sampledExecutions(lens).map((run) => [run.id, run.service] as const)); - const seen = new Set(finding.occurrences.map((id) => services.get(id)).filter((s): s is string => !!s)); - if (seen.size) return [...seen].sort(); - const configured = lens.settings.agent_name || lens.settings.service; - return [configured || UNKNOWN_AGENT]; + const agents = new Map(sampledExecutions(lens).map((run) => [run.id, runAgents(run)] as const)); + const seen = new Set(finding.occurrences.flatMap((id) => agents.get(id) ?? [])); + return seen.size ? [...seen].sort() : [savedAgent(lens.settings.q) || UNKNOWN_AGENT]; } function groupBy(items: readonly T[], key: (item: T) => string): Map { diff --git a/ui/litellm-dashboard/src/components/lens/model/progress.test.ts b/ui/litellm-dashboard/src/components/lens/model/progress.test.ts index 26ef40431b7..5e0518c33df 100644 --- a/ui/litellm-dashboard/src/components/lens/model/progress.test.ts +++ b/ui/litellm-dashboard/src/components/lens/model/progress.test.ts @@ -39,17 +39,13 @@ const job: Job = { revision: 1, settings: { context: "", - source: "traces", lookback_hours: 24, - service: "", - agent_name: "", - filters: [], + q: "", enabled: false, interval_minutes: 15, sample_size: 100, sample_percent: 100, concurrency: 8, - team_id: "", execution_ids: [], monthly_budget: 20, name: "Release reviews", diff --git a/ui/litellm-dashboard/src/components/lens/model/readiness.test.ts b/ui/litellm-dashboard/src/components/lens/model/readiness.test.ts index ce560520bfd..6b14e5ee7b2 100644 --- a/ui/litellm-dashboard/src/components/lens/model/readiness.test.ts +++ b/ui/litellm-dashboard/src/components/lens/model/readiness.test.ts @@ -21,8 +21,8 @@ describe("readiness", () => { }); it("needs a connected worker and a healthy list on top of recorded activity", () => { - const active = { ...base, activity: { traces: false, requests: true } }; - expect(readiness(active)).toMatchObject({ requestsReady: true, activityReady: true, ready: true }); + const active = { ...base, activity: { traces: true } }; + expect(readiness(active)).toMatchObject({ tracesReady: true, activityReady: true, ready: true }); expect(readiness({ ...active, connected: false }).ready).toBe(false); expect(readiness({ ...active, listError: new Error("down") }).ready).toBe(false); }); @@ -46,7 +46,7 @@ describe("initialSetupStep", () => { const offline = { ...base, connected: false }; expect(initialSetupStep(readiness(offline))).toBe(0); expect(initialSetupStep(readiness({ ...offline, tracingConfigured: true }))).toBe(1); - const recorded = { ...offline, activity: { traces: false, requests: true } }; + const recorded = { ...offline, activity: { traces: true } }; expect(initialSetupStep(readiness(recorded))).toBe(2); expect(initialSetupStep(readiness({ ...recorded, connected: true }))).toBe(3); }); diff --git a/ui/litellm-dashboard/src/components/lens/model/readiness.ts b/ui/litellm-dashboard/src/components/lens/model/readiness.ts index d62f9c8bed3..8ad94765481 100644 --- a/ui/litellm-dashboard/src/components/lens/model/readiness.ts +++ b/ui/litellm-dashboard/src/components/lens/model/readiness.ts @@ -1,7 +1,6 @@ export interface Readiness { readonly tracingEnabled: boolean; readonly tracesReady: boolean; - readonly requestsReady: boolean; readonly activityReady: boolean; readonly connected: boolean; readonly hasInvestigations: boolean; @@ -16,7 +15,7 @@ export interface ReadinessInput { readonly checked: boolean; }; readonly tracingConfigured: boolean; - readonly activity: { readonly traces: boolean; readonly requests: boolean } | undefined; + readonly activity: { readonly traces: boolean } | undefined; readonly activityError: unknown; readonly connected: boolean; readonly listError: unknown; @@ -28,12 +27,10 @@ export function readiness(input: ReadinessInput): Readiness { /** Traces confirmed straight from trace storage count as recorded activity even if the activity check fails. */ const tracesSeen = traces.recorded === true && !traces.failed; const tracesReady = tracesSeen || (activity?.traces === true && !activityError); - const requestsReady = activity?.requests === true && !activityError; - const activityReady = tracesReady || requestsReady; + const activityReady = tracesReady; return { tracingEnabled: !traces.disabled && (traces.checked || input.tracingConfigured), tracesReady, - requestsReady, activityReady, connected, hasInvestigations: input.investigations > 0, diff --git a/ui/litellm-dashboard/src/components/lens/model/runRequest.ts b/ui/litellm-dashboard/src/components/lens/model/runRequest.ts index ffb9a88f089..13fd4140373 100644 --- a/ui/litellm-dashboard/src/components/lens/model/runRequest.ts +++ b/ui/litellm-dashboard/src/components/lens/model/runRequest.ts @@ -1,5 +1,13 @@ import type { RunWindow } from "./types"; +const AGENT_TERM = /(?:^|\s)agent:(?:"([^"]*)"|(\S+))/i; + +/** The agent a saved search names, if any. */ +export function savedAgent(q: string | undefined): string { + const match = AGENT_TERM.exec(q ?? ""); + return match?.[1] ?? match?.[2] ?? ""; +} + export const RUN_PRESETS = [ { label: "Since last run", hours: null }, { label: "Last hour", hours: 1 }, diff --git a/ui/litellm-dashboard/src/components/lens/model/status.test.ts b/ui/litellm-dashboard/src/components/lens/model/status.test.ts index 2c9a2e362d1..93fc6aa742a 100644 --- a/ui/litellm-dashboard/src/components/lens/model/status.test.ts +++ b/ui/litellm-dashboard/src/components/lens/model/status.test.ts @@ -33,17 +33,13 @@ const job: Job = { revision: 1, settings: { context: "", - source: "traces", lookback_hours: 24, - service: "", - agent_name: "", - filters: [], + q: "", enabled: false, interval_minutes: 15, sample_size: 100, sample_percent: 100, concurrency: 8, - team_id: "", execution_ids: [], monthly_budget: 20, name: "Release reviews", diff --git a/ui/litellm-dashboard/src/components/lens/model/types.ts b/ui/litellm-dashboard/src/components/lens/model/types.ts index 81cb6479897..68c6575541e 100644 --- a/ui/litellm-dashboard/src/components/lens/model/types.ts +++ b/ui/litellm-dashboard/src/components/lens/model/types.ts @@ -16,20 +16,11 @@ export type Job = components["schemas"]["Job"]; export type IssueBrief = NonNullable; -export type ActivitySelection = Pick & - Partial< - Pick< - Settings, - | "service" - | "agent_name" - | "filters" - | "lookback_hours" - | "sample_percent" - | "sample_size" - | "team_id" - | "execution_ids" - > - >; +export type ActivitySelection = Partial< + Pick +>; + +export type Execution = Sample["executions"][number]; export type Worker = LensList["workers"][number]; diff --git a/ui/litellm-dashboard/src/components/lens/onboarding/OnboardingSteps.tsx b/ui/litellm-dashboard/src/components/lens/onboarding/OnboardingSteps.tsx index 165908bc77d..e424d7becf2 100644 --- a/ui/litellm-dashboard/src/components/lens/onboarding/OnboardingSteps.tsx +++ b/ui/litellm-dashboard/src/components/lens/onboarding/OnboardingSteps.tsx @@ -50,7 +50,7 @@ function StorageStep({ state, goTo }: StepProps) { function continuationLabel(state: LensReadiness) { if (state.connected) return "Continue to investigation"; - return state.tracesReady ? "Continue to worker" : "Continue with request logs"; + return "Continue to worker"; } function ActivityContinuation({ state }: { state: LensReadiness }) { @@ -63,11 +63,6 @@ function ActivityContinuation({ state }: { state: LensReadiness }) { {continuationLabel(state)}