mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
feat(traces): add Rust storage foundation (#43819)
* wip * feat(traces): establish shared Rust storage foundation * fix(traces): escape ClickHouse text parameters * test(traces): exercise response cap with bounded strings * fix(traces): remove unnecessary lint expectation * fix(traces): encode ClickHouse timestamp units in Rust * test(traces): mark exception match as a regex * refactor(traces): execute schema setup in Rust * refactor(traces): use shared logging execution wrapper * docs(traces): replace foundation README with boundary rules * fix(traces): use current bridge execution facade * fix(traces): account for protocol cast in lint budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yujonglee
GitHub
parent
264b09ac8d
commit
41df8cf4d0
25 changed files with 1181 additions and 0 deletions
19
litellm-rust/Cargo.lock
generated
19
litellm-rust/Cargo.lock
generated
|
|
@ -4075,6 +4075,7 @@ dependencies = [
|
|||
"litellm-secrets-aws",
|
||||
"litellm-secrets-types",
|
||||
"litellm-token-counter",
|
||||
"litellm-traces",
|
||||
"litellm-tracing",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
|
|
@ -4351,6 +4352,21 @@ dependencies = [
|
|||
"tiktoken-rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-traces"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-http",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"testcontainers-modules",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-tracing"
|
||||
version = "0.1.0"
|
||||
|
|
@ -5704,6 +5720,7 @@ checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029"
|
|||
dependencies = [
|
||||
"base64 0.23.1",
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2 0.4.15",
|
||||
|
|
@ -5715,6 +5732,7 @@ dependencies = [
|
|||
"hyper-util",
|
||||
"js-sys",
|
||||
"log",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"quinn",
|
||||
|
|
@ -6945,6 +6963,7 @@ dependencies = [
|
|||
"memchr",
|
||||
"parse-display",
|
||||
"pin-project-lite",
|
||||
"reqwest 0.13.5",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -107,6 +109,9 @@ mod tests {
|
|||
"aresponses",
|
||||
"ResponsesWebSocketConnection",
|
||||
"NativeDiagnosticProcessor",
|
||||
"trace_encode_rows",
|
||||
"trace_ensure_schema",
|
||||
"trace_query",
|
||||
"TokenCounter",
|
||||
"Tokenizer",
|
||||
"gil_stats",
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
83
litellm-rust/crates/python-bridge/src/routes/traces.rs
Normal file
83
litellm-rust/crates/python-bridge/src/routes/traces.rs
Normal file
|
|
@ -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)
|
||||
}
|
||||
6
litellm-rust/crates/traces/AGENTS.md
Normal file
6
litellm-rust/crates/traces/AGENTS.md
Normal file
|
|
@ -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
litellm-rust/crates/traces/Cargo.toml
Normal file
20
litellm-rust/crates/traces/Cargo.toml
Normal file
|
|
@ -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
litellm-rust/crates/traces/config/reader.xml
Normal file
32
litellm-rust/crates/traces/config/reader.xml
Normal file
|
|
@ -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>
|
||||
</grants>
|
||||
</litellm_traces_reader>
|
||||
</users>
|
||||
</clickhouse>
|
||||
48
litellm-rust/crates/traces/migrations/0001_otel_traces.sql
Normal file
48
litellm-rust/crates/traces/migrations/0001_otel_traces.sql
Normal file
|
|
@ -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
litellm-rust/crates/traces/migrations/0002_agent_traces.sql
Normal file
25
litellm-rust/crates/traces/migrations/0002_agent_traces.sql
Normal file
|
|
@ -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),
|
||||
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)
|
||||
TTL toDateTime(StartTs) + INTERVAL {trace_retention_days} DAY
|
||||
|
|
@ -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,
|
||||
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
litellm-rust/crates/traces/migrations/0004_spend_logs.sql
Normal file
43
litellm-rust/crates/traces/migrations/0004_spend_logs.sql
Normal file
|
|
@ -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)
|
||||
TTL toDateTime(start_time) + INTERVAL {spend_log_retention_days} DAY
|
||||
21
litellm-rust/crates/traces/src/error.rs
Normal file
21
litellm-rust/crates/traces/src/error.rs
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("invalid ClickHouse insert row")]
|
||||
InvalidRow,
|
||||
#[error("invalid ClickHouse HTTP URL")]
|
||||
InvalidUrl,
|
||||
#[error("database must be a nonempty SQL identifier and retention must be positive")]
|
||||
InvalidSchema,
|
||||
#[error("SQL query must not be empty")]
|
||||
EmptySql,
|
||||
#[error("ClickHouse query failed with HTTP status {0}")]
|
||||
QueryFailed(u16),
|
||||
#[error("ClickHouse schema setup failed with HTTP status {0}")]
|
||||
SchemaFailed(u16),
|
||||
#[error("ClickHouse query exceeded the response size limit")]
|
||||
ResponseTooLarge,
|
||||
#[error("ClickHouse returned an invalid or failed JSON query response")]
|
||||
InvalidResponse,
|
||||
#[error("ClickHouse query transport failed")]
|
||||
Transport,
|
||||
}
|
||||
37
litellm-rust/crates/traces/src/insert.rs
Normal file
37
litellm-rust/crates/traces/src/insert.rs
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub fn encode_rows(rows: Vec<BTreeMap<String, Value>>) -> Result<String, Error> {
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
let encoded = row
|
||||
.into_iter()
|
||||
.map(|(name, value)| insert_value(&name, value).map(|value| (name, value)))
|
||||
.collect::<Result<BTreeMap<_, _>, _>>()?;
|
||||
serde_json::to_string(&encoded).map_err(|_| Error::InvalidRow)
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map(|rows| rows.join("\n"))
|
||||
}
|
||||
|
||||
fn insert_value(name: &str, value: Value) -> Result<Value, Error> {
|
||||
let multiplier = match name {
|
||||
"Timestamp" => 1,
|
||||
"start_time" | "end_time" | "completion_start_time" => 1_000_000,
|
||||
_ => return Ok(value),
|
||||
};
|
||||
if name == "completion_start_time" && value.is_null() {
|
||||
return Ok(value);
|
||||
}
|
||||
let timestamp = value.as_i64().ok_or(Error::InvalidRow)?;
|
||||
let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * multiplier)
|
||||
.map_err(|_| Error::InvalidRow)?;
|
||||
datetime
|
||||
.format(&Rfc3339)
|
||||
.map(Value::String)
|
||||
.map_err(|_| Error::InvalidRow)
|
||||
}
|
||||
73
litellm-rust/crates/traces/src/lib.rs
Normal file
73
litellm-rust/crates/traces/src/lib.rs
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
mod error;
|
||||
mod insert;
|
||||
mod schema;
|
||||
mod sql;
|
||||
|
||||
pub use error::Error;
|
||||
pub use insert::encode_rows;
|
||||
pub use schema::{ensure_schema, schema_statements};
|
||||
pub use sql::{Parameter, execute_read};
|
||||
use url::Url;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Connection {
|
||||
url: Url,
|
||||
}
|
||||
|
||||
impl Connection {
|
||||
pub fn parse(value: &str) -> Result<Self, Error> {
|
||||
let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?;
|
||||
if !matches!(url.scheme(), "http" | "https") || url.host().is_none() {
|
||||
return Err(Error::InvalidUrl);
|
||||
}
|
||||
Ok(Self { url })
|
||||
}
|
||||
|
||||
pub fn configured(
|
||||
url: &str,
|
||||
database: &str,
|
||||
user: &str,
|
||||
password: &str,
|
||||
) -> Result<Self, Error> {
|
||||
let mut connection = Self::parse(url)?;
|
||||
connection
|
||||
.url
|
||||
.set_username(user)
|
||||
.map_err(|_| Error::InvalidUrl)?;
|
||||
connection
|
||||
.url
|
||||
.set_password(Some(password))
|
||||
.map_err(|_| Error::InvalidUrl)?;
|
||||
let pairs: Vec<_> = connection
|
||||
.url
|
||||
.query_pairs()
|
||||
.filter(|(key, _)| !matches!(key.as_ref(), "database" | "user" | "password"))
|
||||
.map(|(key, value)| (key.into_owned(), value.into_owned()))
|
||||
.collect();
|
||||
connection
|
||||
.url
|
||||
.query_pairs_mut()
|
||||
.clear()
|
||||
.extend_pairs(pairs)
|
||||
.append_pair("database", database);
|
||||
Ok(connection)
|
||||
}
|
||||
|
||||
pub fn writer(url: &str, user: &str, password: &str) -> Result<Self, Error> {
|
||||
let mut connection = Self::parse(url)?;
|
||||
connection
|
||||
.url
|
||||
.set_username(user)
|
||||
.map_err(|_| Error::InvalidUrl)?;
|
||||
connection
|
||||
.url
|
||||
.set_password(Some(password))
|
||||
.map_err(|_| Error::InvalidUrl)?;
|
||||
connection.url.set_query(None);
|
||||
Ok(connection)
|
||||
}
|
||||
|
||||
pub fn url(&self) -> &Url {
|
||||
&self.url
|
||||
}
|
||||
}
|
||||
61
litellm-rust/crates/traces/src/schema.rs
Normal file
61
litellm-rust/crates/traces/src/schema.rs
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
use litellm_http::Client;
|
||||
|
||||
use crate::Connection;
|
||||
use crate::Error;
|
||||
|
||||
const MIGRATIONS: [&str; 4] = [
|
||||
include_str!("../migrations/0001_otel_traces.sql"),
|
||||
include_str!("../migrations/0002_agent_traces.sql"),
|
||||
include_str!("../migrations/0003_agent_traces_mv.sql"),
|
||||
include_str!("../migrations/0004_spend_logs.sql"),
|
||||
];
|
||||
|
||||
pub fn schema_statements(
|
||||
database: &str,
|
||||
trace_retention_days: u32,
|
||||
spend_log_retention_days: u32,
|
||||
) -> Result<Vec<String>, Error> {
|
||||
if database.is_empty()
|
||||
|| !database
|
||||
.bytes()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
|
||||
|| trace_retention_days == 0
|
||||
|| spend_log_retention_days == 0
|
||||
{
|
||||
return Err(Error::InvalidSchema);
|
||||
}
|
||||
let database = format!("`{database}`");
|
||||
Ok(
|
||||
std::iter::once(format!("CREATE DATABASE IF NOT EXISTS {database}"))
|
||||
.chain(MIGRATIONS.iter().map(|sql| {
|
||||
sql.replace("{database}", &database)
|
||||
.replace("{trace_retention_days}", &trace_retention_days.to_string())
|
||||
.replace(
|
||||
"{spend_log_retention_days}",
|
||||
&spend_log_retention_days.to_string(),
|
||||
)
|
||||
}))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn ensure_schema(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
database: &str,
|
||||
trace_retention_days: u32,
|
||||
spend_log_retention_days: u32,
|
||||
) -> Result<(), Error> {
|
||||
for statement in schema_statements(database, trace_retention_days, spend_log_retention_days)? {
|
||||
let response = client
|
||||
.post(connection.url().clone())
|
||||
.body(statement)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::Transport)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(Error::SchemaFailed(response.status().as_u16()));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
114
litellm-rust/crates/traces/src/sql.rs
Normal file
114
litellm-rust/crates/traces/src/sql.rs
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use litellm_http::Client;
|
||||
|
||||
use crate::{Connection, Error};
|
||||
|
||||
const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum Parameter {
|
||||
Text(String),
|
||||
Integer(i64),
|
||||
Strings(Vec<String>),
|
||||
}
|
||||
|
||||
impl Parameter {
|
||||
fn encoded(&self) -> String {
|
||||
match self {
|
||||
Self::Text(value) => escaped(value),
|
||||
Self::Integer(value) => value.to_string(),
|
||||
Self::Strings(values) => format!(
|
||||
"[{}]",
|
||||
values
|
||||
.iter()
|
||||
.map(|value| format!("'{}'", escaped(value).replace('\'', "\\'")))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn escaped(value: &str) -> String {
|
||||
value
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\t', "\\t")
|
||||
.replace('\n', "\\n")
|
||||
.replace('\r', "\\r")
|
||||
.replace('\0', "\\0")
|
||||
}
|
||||
|
||||
pub async fn execute_read(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
sql: &str,
|
||||
parameters: &BTreeMap<String, Parameter>,
|
||||
) -> Result<String, Error> {
|
||||
if sql.trim().is_empty() {
|
||||
return Err(Error::EmptySql);
|
||||
}
|
||||
|
||||
let mut url = connection.url().clone();
|
||||
|
||||
let existing_pairs: Vec<(String, String)> = url
|
||||
.query_pairs()
|
||||
.filter(|(key, _)| {
|
||||
!key.starts_with("param_")
|
||||
&& !matches!(
|
||||
key.as_ref(),
|
||||
"query"
|
||||
| "readonly"
|
||||
| "default_format"
|
||||
| "max_result_rows"
|
||||
| "result_overflow_mode"
|
||||
| "max_execution_time"
|
||||
| "wait_end_of_query"
|
||||
)
|
||||
})
|
||||
.map(|(key, value)| (key.into_owned(), value.into_owned()))
|
||||
.collect();
|
||||
url.query_pairs_mut()
|
||||
.clear()
|
||||
.extend_pairs(existing_pairs)
|
||||
.append_pair("readonly", "1")
|
||||
.append_pair("max_result_rows", "1000")
|
||||
.append_pair("result_overflow_mode", "throw")
|
||||
.append_pair("max_execution_time", "10")
|
||||
.append_pair("wait_end_of_query", "1")
|
||||
.append_pair("default_format", "JSON");
|
||||
|
||||
url.query_pairs_mut().extend_pairs(
|
||||
parameters
|
||||
.iter()
|
||||
.map(|(name, value)| (format!("param_{name}"), value.encoded())),
|
||||
);
|
||||
|
||||
let request = client
|
||||
.post(url)
|
||||
.timeout(Duration::from_secs(15))
|
||||
.body(sql.to_owned());
|
||||
let mut response = request.send().await.map_err(|_| Error::Transport)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(Error::QueryFailed(response.status().as_u16()));
|
||||
}
|
||||
|
||||
let mut body = Vec::new();
|
||||
while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? {
|
||||
if body.len() + chunk.len() > MAX_RESPONSE_BYTES {
|
||||
return Err(Error::ResponseTooLarge);
|
||||
}
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
|
||||
let json: serde_json::Value =
|
||||
serde_json::from_slice(&body).map_err(|_| Error::InvalidResponse)?;
|
||||
if json.get("exception").is_some() || !json.get("data").is_some_and(serde_json::Value::is_array)
|
||||
{
|
||||
return Err(Error::InvalidResponse);
|
||||
}
|
||||
String::from_utf8(body).map_err(|_| Error::InvalidResponse)
|
||||
}
|
||||
255
litellm-rust/crates/traces/tests/admin_sql.rs
Normal file
255
litellm-rust/crates/traces/tests/admin_sql.rs
Normal file
|
|
@ -0,0 +1,255 @@
|
|||
use litellm_http::Client;
|
||||
use litellm_traces::{Connection, Error, Parameter, execute_read};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::Value;
|
||||
use std::collections::BTreeMap;
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
|
||||
const CLICKHOUSE_TAG: &str =
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
|
||||
|
||||
struct Database {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
url: String,
|
||||
admin_url: String,
|
||||
client: Client,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(CLICKHOUSE_TAG)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.with_env_var("LITELLM_TRACES_READER_PASSWORD", "test_password")
|
||||
.with_copy_to(
|
||||
"/etc/clickhouse-server/users.d/litellm-traces-reader.xml",
|
||||
include_bytes!("../config/reader.xml").to_vec(),
|
||||
)
|
||||
.start()
|
||||
.await?;
|
||||
let admin_url = format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?,
|
||||
);
|
||||
let client = Client::no_redirect_for_test();
|
||||
for sql in [
|
||||
"CREATE TABLE otel_traces (n UInt8) ENGINE = Memory",
|
||||
"INSERT INTO otel_traces VALUES (1)",
|
||||
"CREATE TABLE private_traces (n UInt8) ENGINE = Memory",
|
||||
] {
|
||||
client
|
||||
.post(&admin_url)
|
||||
.body(sql)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
}
|
||||
let url = admin_url.replacen("http://", "http://litellm_traces_reader:test_password@", 1);
|
||||
Ok(Database {
|
||||
_container: container,
|
||||
url,
|
||||
admin_url,
|
||||
client,
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn admin_sql_reads_rows_with_enforced_settings(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let connection = Connection::parse(&format!(
|
||||
"{}?readonly=0&default_format=TabSeparated&query=SELECT+2",
|
||||
database.url,
|
||||
))?;
|
||||
|
||||
let result = read(
|
||||
&database.client,
|
||||
&connection,
|
||||
"SELECT n AS answer FROM otel_traces",
|
||||
)
|
||||
.await?;
|
||||
let json: Value = serde_json::from_str(&result)?;
|
||||
assert_eq!(json["data"][0]["answer"], 1);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::table("CREATE TABLE admin_sql_test (n UInt8) ENGINE = Memory")]
|
||||
#[case::insert("INSERT INTO otel_traces VALUES (2)")]
|
||||
#[case::drop("DROP TABLE otel_traces")]
|
||||
#[case::named_collection("CREATE NAMED COLLECTION admin_sql_test AS host = 'localhost'")]
|
||||
#[case::settings("SET readonly = 0")]
|
||||
#[case::inline_settings("SELECT n FROM otel_traces SETTINGS readonly = 0")]
|
||||
#[case::time_limit("SELECT n FROM otel_traces SETTINGS max_execution_time = 0")]
|
||||
#[case::row_limit("SELECT n FROM otel_traces SETTINGS max_result_rows = 0")]
|
||||
#[case::byte_limit("SELECT n FROM otel_traces SETTINGS max_result_bytes = 0")]
|
||||
#[case::memory_limit("SELECT n FROM otel_traces SETTINGS max_memory_usage = 0")]
|
||||
#[case::other_table("SELECT * FROM private_traces")]
|
||||
#[tokio::test]
|
||||
async fn reader_rejects_writes_and_privilege_escalation(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
#[case] sql: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let connection = Connection::parse(&format!("{}?readonly=0", database.url))?;
|
||||
|
||||
let result = read(&database.client, &connection, sql).await;
|
||||
|
||||
assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}");
|
||||
let rows = read(&database.client, &connection, "SELECT n FROM otel_traces").await?;
|
||||
let json: Value = serde_json::from_str(&rows)?;
|
||||
assert_eq!(json["data"], serde_json::json!([{ "n": 1 }]));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn admin_sql_rejects_errors_after_output_starts(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let connection = Connection::parse(&format!(
|
||||
"{}?max_block_size=1&buffer_size=1&http_write_exception_in_output_format=1\
|
||||
&send_progress_in_http_headers=1&http_headers_progress_interval_ms=0",
|
||||
database.admin_url,
|
||||
))?;
|
||||
|
||||
let result = read(
|
||||
&database.client,
|
||||
&connection,
|
||||
"SELECT sleepEachRow(0.2), throwIf(number = 2) FROM numbers(5)",
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(Error::InvalidResponse)),
|
||||
"expected an error embedded in a successful HTTP response: {result:?}"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn admin_sql_enforces_result_row_limit(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let connection = Connection::parse(&format!(
|
||||
"{}?max_result_rows=0&result_overflow_mode=throw&wait_end_of_query=1",
|
||||
database.url,
|
||||
))?;
|
||||
|
||||
let result = read(
|
||||
&database.client,
|
||||
&connection,
|
||||
"SELECT number FROM numbers(1001)",
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn admin_sql_enforces_response_byte_limit(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let connection = Connection::parse(&database.admin_url)?;
|
||||
|
||||
let result = read(
|
||||
&database.client,
|
||||
&connection,
|
||||
"SELECT repeat('x', 512 * 1024) AS payload FROM numbers(9)",
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(Error::ResponseTooLarge)), "{result:?}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::plain("test_password", "test_password")]
|
||||
#[case::encoded("p@ss/word%", "p%40ss%2Fword%25")]
|
||||
#[tokio::test]
|
||||
async fn admin_sql_authenticates_url_credentials(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
#[case] password: &str,
|
||||
#[case] encoded_password: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
database
|
||||
.client
|
||||
.post(&database.admin_url)
|
||||
.body(format!(
|
||||
"CREATE USER sql_reader IDENTIFIED WITH plaintext_password BY '{password}'"
|
||||
))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
let connection = Connection::parse(&database.admin_url.replacen(
|
||||
"http://",
|
||||
&format!("http://sql_reader:{encoded_password}@"),
|
||||
1,
|
||||
))?;
|
||||
|
||||
let result = read(
|
||||
&database.client,
|
||||
&connection,
|
||||
"SELECT currentUser() AS username",
|
||||
)
|
||||
.await?;
|
||||
let json: Value = serde_json::from_str(&result)?;
|
||||
|
||||
assert_eq!(json["data"][0]["username"], "sql_reader");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read(client: &Client, connection: &Connection, sql: &str) -> Result<String, Error> {
|
||||
execute_read(client, connection, sql, &BTreeMap::new()).await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::sql("'; DROP TABLE otel_traces; --")]
|
||||
#[case::escapes("back\\slash\ttab\nline\0null")]
|
||||
#[tokio::test]
|
||||
async fn query_parameters_preserve_values_and_replace_url_parameters(
|
||||
#[case] value: &str,
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let connection = Connection::parse(&format!("{}?param_value=wrong", database.url))?;
|
||||
let values = vec![
|
||||
"a'b".to_owned(),
|
||||
"back\\slash".to_owned(),
|
||||
"line\nbreak".to_owned(),
|
||||
"雪".to_owned(),
|
||||
];
|
||||
let parameters = BTreeMap::from([
|
||||
("value".to_owned(), Parameter::Text(value.into())),
|
||||
("teams".to_owned(), Parameter::Strings(values.clone())),
|
||||
("number".to_owned(), Parameter::Integer(-42)),
|
||||
]);
|
||||
let body = execute_read(&database.client, &connection,
|
||||
"SELECT {value:String} AS value, {teams:Array(String)} AS teams, toInt32({number:Int64}) AS number",
|
||||
¶meters).await?;
|
||||
let json: Value = serde_json::from_str(&body)?;
|
||||
assert_eq!(json["data"][0]["value"], value);
|
||||
assert_eq!(json["data"][0]["teams"], serde_json::json!(values));
|
||||
assert_eq!(json["data"][0]["number"], -42);
|
||||
assert!(
|
||||
read(&database.client, &connection, "SELECT n FROM otel_traces")
|
||||
.await
|
||||
.is_ok()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
40
litellm-rust/crates/traces/tests/insert.rs
Normal file
40
litellm-rust/crates/traces/tests/insert.rs
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_traces::encode_rows;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
#[rstest]
|
||||
#[case::span("Timestamp", json!(1_234_567_890), json!("1970-01-01T00:00:01.23456789Z"))]
|
||||
#[case::start("start_time", json!(1_234), json!("1970-01-01T00:00:01.234Z"))]
|
||||
#[case::end("end_time", json!(2_345), json!("1970-01-01T00:00:02.345Z"))]
|
||||
#[case::completion("completion_start_time", json!(1_345), json!("1970-01-01T00:00:01.345Z"))]
|
||||
#[case::absent_completion("completion_start_time", Value::Null, Value::Null)]
|
||||
#[case::before_epoch("Timestamp", json!(-1), json!("1969-12-31T23:59:59.999999999Z"))]
|
||||
fn insert_encoding_preserves_timestamp_precision_and_other_fields(
|
||||
#[case] field: &str,
|
||||
#[case] value: Value,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
let rows = vec![BTreeMap::from([
|
||||
(field.to_owned(), value),
|
||||
("SpanAttributes".into(), json!({"message": "a\nb\\c\"雪"})),
|
||||
("InputTokens".into(), json!(42)),
|
||||
])];
|
||||
let encoded = encode_rows(rows).expect("valid row");
|
||||
let actual: Value = serde_json::from_str(&encoded).expect("JSONEachRow record");
|
||||
assert_eq!(
|
||||
actual,
|
||||
json!({
|
||||
field: expected, "SpanAttributes": {"message": "a\nb\\c\"雪"}, "InputTokens": 42
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::fractional(json!(1.25))]
|
||||
#[case::out_of_range(json!(u64::MAX))]
|
||||
#[case::null(Value::Null)]
|
||||
fn insert_encoding_rejects_invalid_span_timestamps(#[case] timestamp: Value) {
|
||||
assert!(encode_rows(vec![BTreeMap::from([("Timestamp".into(), timestamp)])]).is_err());
|
||||
}
|
||||
101
litellm-rust/crates/traces/tests/migrations.rs
Normal file
101
litellm-rust/crates/traces/tests/migrations.rs
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_traces::{Connection, encode_rows, ensure_schema, execute_read, schema_statements};
|
||||
use rstest::rstest;
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
|
||||
const CLICKHOUSE_TAG: &str =
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn schema_supports_span_rollups_and_spend_joins() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(CLICKHOUSE_TAG)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.start()
|
||||
.await?;
|
||||
let url = format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?
|
||||
);
|
||||
let client = Client::no_redirect_for_test();
|
||||
let writer = Connection::writer(&url, "default", "")?;
|
||||
ensure_schema(&client, &writer, "trace_test", 7, 14).await?;
|
||||
ensure_schema(&client, &writer, "trace_test", 7, 14).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let span = serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1", "ParentSpanId": "",
|
||||
"ServiceName": "proxy", "SpanName": "request", "Input": "hello world",
|
||||
"ResourceAttributes": {"litellm.team_id": "team-1", "litellm.api_key_hash": "hash-1"},
|
||||
"SpanAttributes": {"gen_ai.response.id": "response-1", "gen_ai.usage.input_tokens": "12"}
|
||||
}))?;
|
||||
let spend = serde_json::from_value(serde_json::json!({
|
||||
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1", "spend": 0.125,
|
||||
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + 100,
|
||||
"completion_start_time": null
|
||||
}))?;
|
||||
for (table, row) in [("otel_traces", span), ("spend_logs", spend)] {
|
||||
client
|
||||
.post(&url)
|
||||
.query(&[
|
||||
(
|
||||
"query",
|
||||
format!("INSERT INTO trace_test.{table} FORMAT JSONEachRow"),
|
||||
),
|
||||
("date_time_input_format", "best_effort".into()),
|
||||
])
|
||||
.body(encode_rows(vec![row])?)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
}
|
||||
let connection = Connection::configured(&url, "trace_test", "default", "")?;
|
||||
let body = execute_read(&client, &connection,
|
||||
"SELECT o.TeamId, o.ApiKeyHash, o.ObservationType, o.InputPreview, s.spend, \
|
||||
toString(toUnixTimestamp64Nano(o.Timestamp)) AS timestamp_ns, \
|
||||
toString(toUnixTimestamp64Milli(s.start_time)) AS start_ms \
|
||||
FROM otel_traces o JOIN spend_logs s ON o.LiteLLMRequestId = s.response_id AND o.TeamId = s.team_id",
|
||||
&BTreeMap::new()).await?;
|
||||
let response: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_eq!(
|
||||
response["data"],
|
||||
serde_json::json!([{
|
||||
"TeamId": "team-1", "ApiKeyHash": "hash-1", "ObservationType": "agent",
|
||||
"InputPreview": "hello world", "spend": 0.125,
|
||||
"timestamp_ns": timestamp.to_string(), "start_ms": (timestamp / 1_000_000).to_string()
|
||||
}])
|
||||
);
|
||||
let body = execute_read(
|
||||
&client,
|
||||
&connection,
|
||||
"SELECT toUInt32(sum(SpanCount)) AS spans, toUInt32(sum(InputTokens)) AS tokens \
|
||||
FROM agent_traces WHERE TeamId = 'team-1' AND TraceId = 'trace-1'",
|
||||
&BTreeMap::new(),
|
||||
)
|
||||
.await?;
|
||||
let response: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_eq!(
|
||||
response["data"],
|
||||
serde_json::json!([{"spans": 1, "tokens": 12}])
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty("", 7, 14)]
|
||||
#[case::sql("db; DROP DATABASE default", 7, 14)]
|
||||
#[case::trace_retention("traces", 0, 14)]
|
||||
#[case::spend_retention("traces", 7, 0)]
|
||||
fn schema_rejects_invalid_configuration(
|
||||
#[case] database: &str,
|
||||
#[case] traces: u32,
|
||||
#[case] spend: u32,
|
||||
) {
|
||||
assert!(schema_statements(database, traces, spend).is_err());
|
||||
}
|
||||
11
litellm-rust/crates/traces/tests/queries.rs
Normal file
11
litellm-rust/crates/traces/tests/queries.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
use litellm_traces::Connection;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::http("http://localhost:8123", true)]
|
||||
#[case::https("https://localhost:8443", true)]
|
||||
#[case::tcp("tcp://localhost:9000", false)]
|
||||
#[case::missing_host("http://", false)]
|
||||
fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool) {
|
||||
assert_eq!(Connection::parse(value).is_ok(), expected);
|
||||
}
|
||||
|
|
@ -20,6 +20,19 @@ class RustUpstreamError(Exception): ...
|
|||
class ForkedAfterNativeRuntimeStarted(RuntimeError): ...
|
||||
class ProcessReservedForForking(RuntimeError): ...
|
||||
|
||||
def trace_encode_rows(rows: Sequence[Mapping[str, JsonValue]]) -> str: ...
|
||||
def trace_ensure_schema(
|
||||
url: str, database: str, user: str, password: str, trace_retention_days: int, spend_log_retention_days: int
|
||||
) -> Future[None]: ...
|
||||
def trace_query(
|
||||
url: str,
|
||||
database: str,
|
||||
user: str,
|
||||
password: str,
|
||||
sql: str,
|
||||
parameters: Mapping[str, str | int | Sequence[str]],
|
||||
) -> Future[str]: ...
|
||||
|
||||
@final
|
||||
class NativeDiagnosticProcessor:
|
||||
def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ...
|
||||
|
|
@ -338,6 +351,9 @@ __all__ = [
|
|||
"process_state_started",
|
||||
"reserve_process_for_forking",
|
||||
"responses",
|
||||
"trace_encode_rows",
|
||||
"trace_ensure_schema",
|
||||
"trace_query",
|
||||
"transcription",
|
||||
]
|
||||
|
||||
|
|
|
|||
69
litellm/rust_bridge/traces.py
Normal file
69
litellm/rust_bridge/traces.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from typing import Final, Protocol, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
from litellm.rust_bridge.loader import get_native_bridge
|
||||
|
||||
|
||||
class NativeTraces(Protocol):
|
||||
def trace_encode_rows(self, rows: Sequence[Mapping[str, JsonValue]]) -> str: ...
|
||||
|
||||
def trace_ensure_schema(
|
||||
self,
|
||||
url: str,
|
||||
database: str,
|
||||
user: str,
|
||||
password: str,
|
||||
trace_retention_days: int,
|
||||
spend_log_retention_days: int,
|
||||
) -> Awaitable[None]: ...
|
||||
|
||||
def trace_query(
|
||||
self,
|
||||
url: str,
|
||||
database: str,
|
||||
user: str,
|
||||
password: str,
|
||||
sql: str,
|
||||
parameters: Mapping[str, str | int | Sequence[str]],
|
||||
) -> Awaitable[str]: ...
|
||||
|
||||
|
||||
class QueryResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
data: list[dict[str, JsonValue]]
|
||||
|
||||
|
||||
def _native() -> NativeTraces:
|
||||
native: Final = get_native_bridge()
|
||||
if native is None:
|
||||
raise RuntimeError("Agent tracing requires the Rust extension")
|
||||
return cast(NativeTraces, native) # cast-ok: the native extension is validated against this protocol at call sites
|
||||
|
||||
|
||||
async def ensure_schema(
|
||||
url: str,
|
||||
database: str,
|
||||
user: str,
|
||||
password: str,
|
||||
trace_retention_days: int,
|
||||
spend_log_retention_days: int,
|
||||
) -> None:
|
||||
await _native().trace_ensure_schema(url, database, user, password, trace_retention_days, spend_log_retention_days)
|
||||
|
||||
|
||||
async def query(
|
||||
url: str,
|
||||
database: str,
|
||||
user: str,
|
||||
password: str,
|
||||
sql: str,
|
||||
parameters: Mapping[str, str | int | Sequence[str]],
|
||||
) -> list[dict[str, JsonValue]]:
|
||||
result: Final = await _native().trace_query(url, database, user, password, sql, parameters)
|
||||
return QueryResponse.model_validate_json(result).data
|
||||
|
||||
|
||||
def encode_rows(rows: Sequence[Mapping[str, JsonValue]]) -> bytes:
|
||||
return _native().trace_encode_rows(rows).encode("utf-8")
|
||||
78
tests/test_litellm_rust/test_traces.py
Normal file
78
tests/test_litellm_rust/test_traces.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
import base64
|
||||
import json
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.rust_bridge.traces import encode_rows, ensure_schema, query
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_reader_projects_connection_and_parameters(recording_server: RecordingServer) -> None:
|
||||
recording_server.enqueue(ResponseSpec(body={"data": [{"trace_id": "trace-1"}]}))
|
||||
rows: Final = await query(
|
||||
recording_server.base_url + "?database=wrong&user=wrong&password=wrong",
|
||||
"trace_test",
|
||||
"reader",
|
||||
"p@ss/word%",
|
||||
"SELECT {trace_id:String} AS trace_id",
|
||||
{"trace_id": "trace-1"},
|
||||
)
|
||||
request: Final = recording_server.requests[0]
|
||||
parameters: Final = parse_qs(urlsplit(request.path).query)
|
||||
assert rows == [{"trace_id": "trace-1"}]
|
||||
assert request.raw_body == b"SELECT {trace_id:String} AS trace_id"
|
||||
assert parameters["database"] == ["trace_test"]
|
||||
assert parameters["param_trace_id"] == ["trace-1"]
|
||||
assert parameters["readonly"] == ["1"]
|
||||
assert "user" not in parameters
|
||||
assert "password" not in parameters
|
||||
assert request.headers["authorization"] == "Basic " + base64.b64encode(b"reader:p@ss/word%").decode()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_reader_rejects_success_status_with_embedded_error(recording_server: RecordingServer) -> None:
|
||||
recording_server.enqueue(ResponseSpec(body={"data": [], "exception": "query failed"}))
|
||||
with pytest.raises(RuntimeError, match="invalid or failed JSON"):
|
||||
await query(recording_server.base_url, "trace_test", "reader", "password", "SELECT 1", {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("database,retention", [("db; DROP DATABASE default", 7), ("traces", 0)])
|
||||
async def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None:
|
||||
with pytest.raises(ValueError, match=r"database.*retention"):
|
||||
await ensure_schema("http://localhost:8123", database, "writer", "password", retention, 14)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schema_setup_uses_writer_credentials_and_rejects_failed_statement(recording_server: RecordingServer) -> None:
|
||||
recording_server.expected_requests = 2
|
||||
recording_server.enqueue(ResponseSpec(body=""))
|
||||
recording_server.enqueue(ResponseSpec(status=403, body="denied"))
|
||||
with pytest.raises(RuntimeError, match="schema setup failed with HTTP status 403"):
|
||||
await ensure_schema(
|
||||
recording_server.base_url + "?database=wrong&readonly=1",
|
||||
"trace_test",
|
||||
"writer",
|
||||
"p@ss/word%",
|
||||
7,
|
||||
14,
|
||||
)
|
||||
assert len(recording_server.requests) == 2
|
||||
assert recording_server.requests[0].raw_body.startswith(b"CREATE DATABASE IF NOT EXISTS")
|
||||
assert recording_server.requests[1].raw_body.startswith(b"CREATE TABLE IF NOT EXISTS")
|
||||
assert "readonly" not in parse_qs(urlsplit(recording_server.requests[0].path).query)
|
||||
assert recording_server.requests[0].headers["authorization"] == "Basic " + base64.b64encode(
|
||||
b"writer:p@ss/word%"
|
||||
).decode()
|
||||
|
||||
|
||||
def test_insert_encoding_preserves_nanoseconds_through_bridge() -> None:
|
||||
assert json.loads(encode_rows([{"Timestamp": 1_234_567_890, "Input": "hello"}])) == {
|
||||
"Input": "hello",
|
||||
"Timestamp": "1970-01-01T00:00:01.23456789Z",
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue