Repository navigation
feat(traces): add Rust storage foundation #43819
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
e4b893a
28e2407
50dc095
835118f
c120adb
350d06a
1da1437
829e9a7
1f1c985
c3f3b14
1f3eb17
b0518f2
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,83 @@ | ||
| use std::collections::BTreeMap; | ||
|
|
||
| use litellm_http::ClientVariant; | ||
| use litellm_traces::{Connection, Error, Parameter}; | ||
| use pyo3::{ | ||
| exceptions::{PyRuntimeError, PyValueError}, | ||
| prelude::*, | ||
| }; | ||
|
|
||
| fn map_error(error: Error) -> PyErr { | ||
| match error { | ||
| Error::InvalidRow | Error::InvalidSchema | Error::EmptySql => { | ||
| PyValueError::new_err(error.to_string()) | ||
| } | ||
| Error::InvalidUrl | ||
| | Error::QueryFailed(_) | ||
| | Error::SchemaFailed(_) | ||
| | Error::ResponseTooLarge | ||
| | Error::InvalidResponse | ||
| | Error::Transport => PyRuntimeError::new_err(error.to_string()), | ||
| } | ||
| } | ||
|
|
||
| #[pyfunction] | ||
| pub fn trace_ensure_schema<'py>( | ||
| py: Python<'py>, | ||
| url: &str, | ||
| database: String, | ||
| user: &str, | ||
| password: &str, | ||
| trace_retention_days: u32, | ||
| spend_log_retention_days: u32, | ||
| ) -> PyResult<Bound<'py, PyAny>> { | ||
| let connection = Connection::writer(url, user, password).map_err(map_error)?; | ||
| let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; | ||
| crate::execution::run_async( | ||
| py, | ||
| async move { | ||
| litellm_traces::ensure_schema( | ||
| &client, | ||
| &connection, | ||
| &database, | ||
| trace_retention_days, | ||
| spend_log_retention_days, | ||
| ) | ||
| .await | ||
| }, | ||
| map_error, | ||
| ) | ||
| } | ||
|
|
||
| #[pyfunction] | ||
| pub fn trace_query<'py>( | ||
| py: Python<'py>, | ||
| url: &str, | ||
| database: &str, | ||
| user: &str, | ||
| password: &str, | ||
| sql: String, | ||
| #[pyo3(from_py_with = litellm_host_python::from_py_argument)] parameters: BTreeMap< | ||
| String, | ||
| Parameter, | ||
| >, | ||
| ) -> PyResult<Bound<'py, PyAny>> { | ||
| let connection = Connection::configured(url, database, user, password).map_err(map_error)?; | ||
| let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; | ||
| crate::execution::run_async( | ||
| py, | ||
| async move { litellm_traces::execute_read(&client, &connection, &sql, ¶meters).await }, | ||
| map_error, | ||
| ) | ||
| } | ||
|
|
||
| #[pyfunction] | ||
| pub fn trace_encode_rows( | ||
| py: Python<'_>, | ||
| #[pyo3(from_py_with = litellm_host_python::from_py_argument)] rows: Vec< | ||
| BTreeMap<String, serde_json::Value>, | ||
| >, | ||
| ) -> PyResult<String> { | ||
| py.detach(|| litellm_traces::encode_rows(rows)) | ||
| .map_err(map_error) | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| - Rust owns OTLP wire decoding, ClickHouse schema, row encoding, named reads, connection validation and transport | ||
| - Keep this crate independent of Python; PyO3 conversion and public Python exceptions belong in `python-bridge` | ||
| - Keep the SQL migrations here as the only ClickHouse schema definition | ||
| - Use typed query parameters and a dedicated SELECT-only reader with server-side limits | ||
| - Bound insert time and encoded bytes; make retry deduplication behavior explicit for supported ClickHouse versions | ||
| - Test storage behavior through the crate's public API against ClickHouse |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| [package] | ||
| name = "litellm-traces" | ||
| version = "0.1.0" | ||
| edition.workspace = true | ||
| license.workspace = true | ||
| repository.workspace = true | ||
|
|
||
| [dependencies] | ||
| time = { workspace = true, features = ["formatting"] } | ||
| litellm-http.workspace = true | ||
| serde.workspace = true | ||
| serde_json.workspace = true | ||
| thiserror.workspace = true | ||
| url.workspace = true | ||
|
|
||
| [dev-dependencies] | ||
| litellm-http = { workspace = true, features = ["test-support"] } | ||
| rstest.workspace = true | ||
| testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] } | ||
| tokio.workspace = true |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,32 @@ | ||
| <clickhouse> | ||
| <profiles> | ||
| <litellm_traces_reader> | ||
| <readonly>1</readonly> | ||
| <max_execution_time>10</max_execution_time> | ||
| <max_result_rows>1000</max_result_rows> | ||
| <max_result_bytes>4194304</max_result_bytes> | ||
| <result_overflow_mode>throw</result_overflow_mode> | ||
| <max_memory_usage>268435456</max_memory_usage> | ||
| <constraints> | ||
| <readonly><readonly/></readonly> | ||
| <max_execution_time><readonly/></max_execution_time> | ||
| <max_result_rows><readonly/></max_result_rows> | ||
| <max_result_bytes><readonly/></max_result_bytes> | ||
| <result_overflow_mode><readonly/></result_overflow_mode> | ||
| <max_memory_usage><readonly/></max_memory_usage> | ||
| </constraints> | ||
| </litellm_traces_reader> | ||
| </profiles> | ||
| <users> | ||
| <litellm_traces_reader> | ||
| <password from_env="LITELLM_TRACES_READER_PASSWORD"/> | ||
| <networks><ip>::/0</ip></networks> | ||
| <profile>litellm_traces_reader</profile> | ||
| <grants> | ||
| <query>GRANT SELECT ON default.otel_traces</query> | ||
| <query>GRANT SELECT ON default.agent_traces</query> | ||
| <query>GRANT SELECT ON default.spend_logs</query> | ||
|
Comment on lines
+26
to
+28
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| </grants> | ||
| </litellm_traces_reader> | ||
| </users> | ||
| </clickhouse> | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| CREATE TABLE IF NOT EXISTS {database}.otel_traces | ||
| ( | ||
| Timestamp DateTime64(9) CODEC(Delta, ZSTD(1)), | ||
| TraceId String CODEC(ZSTD(1)), | ||
| SpanId String CODEC(ZSTD(1)), | ||
| ParentSpanId String CODEC(ZSTD(1)), | ||
| TraceState String CODEC(ZSTD(1)), | ||
| SpanName LowCardinality(String) CODEC(ZSTD(1)), | ||
| SpanKind LowCardinality(String) CODEC(ZSTD(1)), | ||
| ServiceName LowCardinality(String) CODEC(ZSTD(1)), | ||
| ResourceAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), | ||
| ScopeName String CODEC(ZSTD(1)), | ||
| ScopeVersion String CODEC(ZSTD(1)), | ||
| SpanAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), | ||
| Duration UInt64 CODEC(ZSTD(1)), | ||
| StatusCode LowCardinality(String) CODEC(ZSTD(1)), | ||
| StatusMessage String CODEC(ZSTD(1)), | ||
| `Events.Timestamp` Array(DateTime64(9)) CODEC(ZSTD(1)), | ||
| `Events.Name` Array(LowCardinality(String)) CODEC(ZSTD(1)), | ||
| `Events.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), | ||
| `Links.TraceId` Array(String) CODEC(ZSTD(1)), | ||
| `Links.SpanId` Array(String) CODEC(ZSTD(1)), | ||
| `Links.TraceState` Array(String) CODEC(ZSTD(1)), | ||
| `Links.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), | ||
| TeamId LowCardinality(String) DEFAULT ResourceAttributes['litellm.team_id'], | ||
| ApiKeyHash String DEFAULT ResourceAttributes['litellm.api_key_hash'], | ||
| ObservationType LowCardinality(String) DEFAULT multiIf( | ||
| ParentSpanId = '', 'agent', | ||
| SpanAttributes['gen_ai.operation.name'] = 'invoke_agent', 'agent', | ||
| SpanAttributes['gen_ai.operation.name'] IN ('chat', 'text_completion', 'generate_content'), 'llm', | ||
| SpanAttributes['gen_ai.operation.name'] = 'execute_tool', 'tool', | ||
| 'chain'), | ||
| AgentName LowCardinality(String) DEFAULT SpanAttributes['gen_ai.agent.name'], | ||
| LiteLLMRequestId String DEFAULT SpanAttributes['gen_ai.response.id'], | ||
| Model LowCardinality(String) DEFAULT SpanAttributes['gen_ai.request.model'], | ||
| InputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.input_tokens']), | ||
| OutputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.output_tokens']), | ||
| Input String CODEC(ZSTD(3)), | ||
| Output String CODEC(ZSTD(3)), | ||
| InputPreview String DEFAULT substring(Input, 1, 240), | ||
| INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1, | ||
| INDEX idx_req_id LiteLLMRequestId TYPE bloom_filter(0.01) GRANULARITY 1 | ||
| ) | ||
| ENGINE = MergeTree | ||
| PARTITION BY toDate(Timestamp) | ||
| ORDER BY (TeamId, ServiceName, toDateTime(Timestamp), TraceId) | ||
| TTL toDateTime(Timestamp) + INTERVAL {trace_retention_days} DAY | ||
| SETTINGS ttl_only_drop_parts = 1 |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| CREATE TABLE IF NOT EXISTS {database}.agent_traces | ||
| ( | ||
| TeamId LowCardinality(String), | ||
| TraceId String, | ||
| StartTs SimpleAggregateFunction(min, DateTime64(9)), | ||
| EndTs SimpleAggregateFunction(max, DateTime64(9)), | ||
| ServiceName SimpleAggregateFunction(any, LowCardinality(String)), | ||
| RootName SimpleAggregateFunction(anyLast, String), | ||
| RootInput SimpleAggregateFunction(anyLast, String), | ||
| RootStatus SimpleAggregateFunction(anyLast, String), | ||
|
Comment on lines
+8
to
+10
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| SpanCount SimpleAggregateFunction(sum, UInt64), | ||
| AgentCount SimpleAggregateFunction(sum, UInt64), | ||
| LlmCount SimpleAggregateFunction(sum, UInt64), | ||
| ToolCount SimpleAggregateFunction(sum, UInt64), | ||
| ErrorCount SimpleAggregateFunction(sum, UInt64), | ||
| InputTokens SimpleAggregateFunction(sum, UInt64), | ||
| OutputTokens SimpleAggregateFunction(sum, UInt64), | ||
| Models SimpleAggregateFunction(groupUniqArrayArray, Array(String)), | ||
| AgentNames SimpleAggregateFunction(groupUniqArrayArray, Array(String)), | ||
| RequestIds SimpleAggregateFunction(groupArrayArray, Array(String)) | ||
| ) | ||
| ENGINE = AggregatingMergeTree | ||
| PARTITION BY toDate(StartTs) | ||
| ORDER BY (TeamId, TraceId) | ||
|
Comment on lines
+23
to
+24
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🔴 Trace rollups split across calendar days When one trace spans midnight, Learn moreThe materialized view agent_traces_mv emits one aggregate row per trace and insert batch, with Example: Team A inserts the root at 23:59 and an LLM span at 00:01 the next day with the same TraceId. The agent and LLM counts remain in separate rows instead of forming one trace. Recommended fix: Choose a partition expression stable for all batches of a trace, or explicitly aggregate across partitions on every read and ensure retention matches the intended whole-trace lifetime. Was this helpful? React with 👍 or 👎 to provide feedback. |
||
| TTL toDateTime(StartTs) + INTERVAL {trace_retention_days} DAY | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| CREATE MATERIALIZED VIEW IF NOT EXISTS {database}.agent_traces_mv TO {database}.agent_traces AS | ||
| SELECT | ||
| TeamId, TraceId, | ||
| min(Timestamp) AS StartTs, | ||
| max(Timestamp + toIntervalNanosecond(Duration)) AS EndTs, | ||
| any(ServiceName) AS ServiceName, | ||
| anyLastIf(SpanName, ParentSpanId = '') AS RootName, | ||
| anyLastIf(InputPreview, ParentSpanId = '') AS RootInput, | ||
| anyLastIf(StatusCode, ParentSpanId = '') AS RootStatus, | ||
|
Comment on lines
+7
to
+9
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🟡 Missing-root batches erase trace root fields When later batches contain no root span, Learn moreThe view groups spans from each incoming insert into one row per trace. For a batch with no root span, Example: Insert a root span named Recommended fix: Store a nullable or explicitly conditional root aggregate that ignores absent root spans when combining batches, and test separate root and child inserts followed by merges. Was this helpful? React with 👍 or 👎 to provide feedback. |
||
| count() AS SpanCount, | ||
| countIf(ObservationType = 'agent') AS AgentCount, | ||
| countIf(ObservationType = 'llm') AS LlmCount, | ||
| countIf(ObservationType = 'tool') AS ToolCount, | ||
| countIf(StatusCode = 'STATUS_CODE_ERROR') AS ErrorCount, | ||
| sum(InputTokens) AS InputTokens, | ||
| sum(OutputTokens) AS OutputTokens, | ||
| groupUniqArrayIf(toString(Model), Model != '') AS Models, | ||
| groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames, | ||
| groupArrayIf(LiteLLMRequestId, LiteLLMRequestId != '') AS RequestIds | ||
| FROM {database}.otel_traces | ||
| GROUP BY TeamId, TraceId | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,43 @@ | ||
| CREATE TABLE IF NOT EXISTS {database}.spend_logs | ||
| ( | ||
| request_id String, | ||
| response_id String, | ||
| call_type LowCardinality(String), | ||
| api_key String, | ||
| key_alias String, | ||
| team_id LowCardinality(String), | ||
| team_alias String, | ||
| organization_id String, | ||
| user String, | ||
| end_user String, | ||
| model LowCardinality(String), | ||
| model_group LowCardinality(String), | ||
| model_id String, | ||
| custom_llm_provider LowCardinality(String), | ||
| api_base String, | ||
| spend Float64, | ||
| prompt_tokens UInt32, | ||
| completion_tokens UInt32, | ||
| total_tokens UInt32, | ||
| cache_read_tokens UInt32, | ||
| cache_write_tokens UInt32, | ||
| start_time DateTime64(3), | ||
| end_time DateTime64(3), | ||
| completion_start_time Nullable(DateTime64(3)), | ||
| status LowCardinality(String), | ||
| error_str String, | ||
| cache_hit Bool, | ||
| session_id String, | ||
| trace_id String, | ||
| span_id String, | ||
| request_tags Array(String), | ||
| metadata String CODEC(ZSTD(3)), | ||
| messages String CODEC(ZSTD(3)), | ||
| response String CODEC(ZSTD(3)), | ||
| INDEX idx_response_id response_id TYPE bloom_filter(0.001) GRANULARITY 1, | ||
| INDEX idx_trace_id trace_id TYPE bloom_filter(0.001) GRANULARITY 1 | ||
| ) | ||
| ENGINE = ReplacingMergeTree(end_time) | ||
| PARTITION BY toYYYYMM(start_time) | ||
| ORDER BY (team_id, toDateTime(start_time), request_id) | ||
|
Comment on lines
+40
to
+42
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🔴 Same-second spend logs collapse into one When distinct spend logs share a team, second, and request ID, Learn moreReplacingMergeTree identifies duplicates using the ORDER BY tuple, not the version column. Converting Example: Two entries for request Recommended fix: Preserve millisecond precision in the ordering key and determine whether request ID is unique across all spend-log events. Use a key that identifies actual retry duplicates rather than unrelated events. Was this helpful? React with 👍 or 👎 to provide feedback. |
||
| TTL toDateTime(start_time) + INTERVAL {spend_log_retention_days} DAY | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🔴 Reader cannot query custom trace databases
When
ensure_schemacreates a nondefault database,litellm_traces_readerretains grants only ondefault. Queries against the configured trace tables fail with access deniedLearn more
The schema creator accepts any valid database identifier in schema_statements. The reader profile grants SELECT on tables in
defaultonly. Whentrace_queryselects another database through Connection::configured, SELECT reads tables on which that user has no grants.Example: Creating tables in
trace_testsucceeds. Queryingtrace_test.otel_tracesaslitellm_traces_readerfails with access denied.Recommended fix: Match reader grants to the configured schema database during deployment or provisioning, and test a nondefault database with the restricted reader.
Was this helpful? React with 👍 or 👎 to provide feedback.