Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions litellm-rust/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions litellm-rust/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ repository = "https://github.com/BerriAI/litellm"
litellm-config = { path = "crates/config" }
litellm-router = { path = "crates/router" }
litellm-tracing = { path = "crates/tracing" }
litellm-traces = { path = "crates/traces" }
litellm-core = { path = "crates/core" }
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
litellm-gateway = { path = "crates/gateway" }
Expand Down
1 change: 1 addition & 0 deletions litellm-rust/crates/python-bridge/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
fancy-regex.workspace = true
litellm-tracing.workspace = true
litellm-traces.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
Expand Down
5 changes: 5 additions & 0 deletions litellm-rust/crates/python-bridge/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ mod _native {
use crate::routes::responses::{ResponsesWebSocketConnection, aresponses, responses};
#[pymodule_export]
use crate::routes::token_counter::TokenCounter;
#[pymodule_export]
use crate::routes::traces::{trace_encode_rows, trace_ensure_schema, trace_query};
#[cfg(feature = "huggingface")]
#[pymodule_export]
use crate::tokenizer::HuggingFaceEncoding;
Expand Down Expand Up @@ -107,6 +109,9 @@ mod tests {
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"trace_encode_rows",
"trace_ensure_schema",
"trace_query",
"TokenCounter",
"Tokenizer",
"gil_stats",
Expand Down
1 change: 1 addition & 0 deletions litellm-rust/crates/python-bridge/src/routes/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ pub(crate) mod messages;
pub(crate) mod ocr;
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};
Expand Down
83 changes: 83 additions & 0 deletions litellm-rust/crates/python-bridge/src/routes/traces.rs
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, &parameters).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)
}
6 changes: 6 additions & 0 deletions litellm-rust/crates/traces/AGENTS.md
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
20 changes: 20 additions & 0 deletions litellm-rust/crates/traces/Cargo.toml
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
32 changes: 32 additions & 0 deletions litellm-rust/crates/traces/config/reader.xml
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

Copy link
Copy Markdown
Contributor

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_schema creates a nondefault database, litellm_traces_reader retains grants only on default. Queries against the configured trace tables fail with access denied

Learn more

The schema creator accepts any valid database identifier in schema_statements. The reader profile grants SELECT on tables in default only. When trace_query selects another database through Connection::configured, SELECT reads tables on which that user has no grants.

Example: Creating tables in trace_test succeeds. Querying trace_test.otel_traces as litellm_traces_reader fails 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.

Devin Review


Was this helpful? React with 👍 or 👎 to provide feedback.

Comment on lines +26 to +28

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Reader grants exclude default database If this reader profile is installed with the tracing caller's default configuration from PR #43864, the tables are created in litellm but these grants cover only default. Valid trace reads then fail with access denied.

</grants>
</litellm_traces_reader>
</users>
</clickhouse>
48 changes: 48 additions & 0 deletions litellm-rust/crates/traces/migrations/0001_otel_traces.sql
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
25 changes: 25 additions & 0 deletions litellm-rust/crates/traces/migrations/0002_agent_traces.sql
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Root fields can disappear When a root span arrives in one batch and child spans arrive in a later batch, the later batch has empty root fields. Merging those batches with anyLast can overwrite the real values, leaving the trace list in PR #43864 with a blank name, input preview, or status.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Trace rollups split across calendar days

When one trace spans midnight, StartTs puts its batches into different partitions. AggregatingMergeTree cannot merge the same (TeamId, TraceId) across partitions, leaving multiple incomplete rollup rows

Learn more

The materialized view agent_traces_mv emits one aggregate row per trace and insert batch, with StartTs equal to the minimum span timestamp in that batch. AggregatingMergeTree only merges rows with the same key within one partition. A trace with batches on different dates has two distinct aggregate rows even after merges. A read that expects one row per trace receives split counts and incompatible root fields.

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.

Devin Review


Was this helpful? React with 👍 or 👎 to provide feedback.

TTL toDateTime(StartTs) + INTERVAL {trace_retention_days} DAY
21 changes: 21 additions & 0 deletions litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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, anyLastIf produces empty root values for their rollup rows. anyLast merges those rows with the root batch and can replace the root name, input, and status with empties

Learn more

The view groups spans from each incoming insert into one row per trace. For a batch with no root span, anyLastIf returns the default empty string. The destination's RootName fields use SimpleAggregateFunction(anyLast, String), which can retain the empty value during part merges rather than preserve the original root.

Example: Insert a root span named agent in batch one, then an LLM child alone in batch two. The second rollup supplies an empty RootName, and the merged trace can lose agent.

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.

Devin Review


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
43 changes: 43 additions & 0 deletions litellm-rust/crates/traces/migrations/0004_spend_logs.sql
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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, ReplacingMergeTree treats them as one row. A later merge silently discards one spend entry

Learn more

ReplacingMergeTree identifies duplicates using the ORDER BY tuple, not the version column. Converting start_time to DateTime truncates the DateTime64(3) milliseconds. Distinct request logs that reuse an ID within a second therefore have the same replacing key, and ClickHouse retains only the highest end_time on merge.

Example: Two entries for request req-1 from the same team start at 12:00:00.100 and 12:00:00.800. They share the key (team, 12:00:00, req-1) and one spend record disappears.

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.

Devin Review


Was this helpful? React with 👍 or 👎 to provide feedback.

TTL toDateTime(start_time) + INTERVAL {spend_log_retention_days} DAY
Loading
Loading