use async_trait::async_trait;
use opentelemetry::logs::{AnyValue, LogRecord as _, Logger as _, Severity};
use opentelemetry::trace::{SpanId, TraceId};
use opentelemetry_sdk::logs::SdkLogger;
use crate::telemetry;
use super::{UsageRecord, UsageSink, UsageSinkError};
const EVENT_NAME: &str = "axond.usage";
pub struct OtlpUsageSink {
logger: SdkLogger,
}
impl OtlpUsageSink {
pub fn new() -> Result<Self, UsageSinkError> {
let logger = telemetry::usage_logger().ok_or_else(|| {
UsageSinkError::invalid(
"otlp",
"OTLP export is off; set OTEL_EXPORTER_OTLP_ENDPOINT or remove the sink",
)
})?;
Ok(Self { logger })
}
}
#[async_trait]
impl UsageSink for OtlpUsageSink {
fn name(&self) -> &'static str {
"otlp"
}
async fn record(&self, record: &UsageRecord) {
let mut log = self.logger.create_log_record();
log.set_event_name(EVENT_NAME);
log.set_severity_number(Severity::Info);
log.set_severity_text("INFO");
log.set_body(AnyValue::String(EVENT_NAME.into()));
log.add_attributes(attributes(record));
if let Some(trace_id) = record.trace_id.as_deref().and_then(parse_trace_id) {
log.set_trace_context(trace_id, SpanId::INVALID, None);
}
self.logger.emit(log);
}
}
fn attributes(record: &UsageRecord) -> Vec<(&'static str, AnyValue)> {
let mut attributes = vec![
(
"axond.schema_version",
AnyValue::Int(i64::from(record.schema_version)),
),
(
"axond.request_id",
AnyValue::String(record.request_id.clone().into()),
),
(
"axond.namespace",
AnyValue::String(record.namespace.clone().into()),
),
(
"axond.subject",
AnyValue::String(record.subject.clone().into()),
),
(
"gen_ai.request.model",
AnyValue::String(record.model.clone().into()),
),
(
"axond.target.provider",
AnyValue::String(record.target_provider.clone().into()),
),
(
"axond.target.model",
AnyValue::String(record.target_model.clone().into()),
),
(
"axond.credential_source",
AnyValue::String(record.credential_source.into()),
),
(
"axond.credential_id",
AnyValue::String(record.credential_id.clone().into()),
),
(
"axond.status",
AnyValue::String(record.status.as_str().into()),
),
(
"gen_ai.usage.input_tokens",
AnyValue::Int(clamped(record.input_tokens)),
),
(
"gen_ai.usage.cache_read_tokens",
AnyValue::Int(clamped(record.cache_read_tokens)),
),
(
"gen_ai.usage.cache_write_tokens",
AnyValue::Int(clamped(record.cache_write_tokens)),
),
(
"gen_ai.usage.output_tokens",
AnyValue::Int(clamped(record.output_tokens)),
),
(
"axond.cost_microdollars",
AnyValue::Int(clamped(record.cost_microdollars)),
),
(
"axond.catalog_version",
AnyValue::Int(clamped(record.catalog_version)),
),
(
"axond.latency_ms",
AnyValue::Int(clamped(record.latency_ms)),
),
];
if let Some(trace_id) = &record.trace_id {
attributes.push(("axond.trace_id", AnyValue::String(trace_id.clone().into())));
}
if let Some(signer_kid) = &record.signer_kid {
attributes.push((
"axond.signer_kid",
AnyValue::String(signer_kid.clone().into()),
));
}
if let Some(book) = &record.price_book {
attributes.push(("axond.price_book", AnyValue::String(book.clone().into())));
}
if let Some(checksum) = &record.price_book_checksum {
attributes.push((
"axond.price_book_checksum",
AnyValue::String(checksum.clone().into()),
));
}
if let Some(catalog) = &record.price_catalog {
attributes.push((
"axond.price_catalog",
AnyValue::String(catalog.clone().into()),
));
}
attributes
}
fn clamped(value: u64) -> i64 {
i64::try_from(value).unwrap_or(i64::MAX)
}
fn parse_trace_id(hex: &str) -> Option<TraceId> {
TraceId::from_hex(hex)
.ok()
.filter(|id| *id != TraceId::INVALID)
}
#[cfg(test)]
mod tests {
use super::super::tests::sample_record;
use super::*;
#[test]
fn attributes_carry_the_identifiers_metrics_cannot() {
let record = sample_record();
let attributes = attributes(&record);
let keys: Vec<&str> = attributes.iter().map(|(key, _)| *key).collect();
for key in [
"axond.request_id",
"axond.subject",
"axond.credential_id",
"axond.signer_kid",
"axond.trace_id",
"axond.schema_version",
] {
assert!(keys.contains(&key), "missing `{key}`");
}
assert_eq!(
attributes
.iter()
.find(|(key, _)| *key == "axond.cost_microdollars")
.map(|(_, value)| value.clone()),
Some(AnyValue::Int(640))
);
}
#[test]
fn an_static_key_omits_the_signer_kid() {
let mut record = sample_record();
record.signer_kid = None;
let keys: Vec<&str> = attributes(&record).iter().map(|(key, _)| *key).collect();
assert!(!keys.contains(&"axond.signer_kid"));
}
#[test]
fn an_untraced_record_omits_the_trace_id() {
let mut record = sample_record();
record.trace_id = None;
let keys: Vec<&str> = attributes(&record).iter().map(|(key, _)| *key).collect();
assert!(!keys.contains(&"axond.trace_id"));
}
#[test]
fn only_a_well_formed_trace_id_becomes_trace_context() {
assert!(parse_trace_id("4bf92f3577b34da6a3ce929d0e0e4736").is_some());
assert!(parse_trace_id("00000000000000000000000000000000").is_none());
assert!(parse_trace_id("not-a-trace-id").is_none());
}
#[test]
fn the_sink_refuses_to_be_built_when_export_is_off() {
if telemetry::usage_logger().is_none() {
assert!(OtlpUsageSink::new().is_err());
}
}
}