use crate::completion::{AssistantContent, Message, Usage};
use crate::message::{
DocumentSourceKind, Image, MimeType, Opaque, ToolResult, ToolResultContent, UserContent,
};
use base64::Engine;
use serde::Serialize;
use std::collections::HashSet;
use std::sync::{LazyLock, Mutex};
use tracing::callsite::Identifier;
#[doc(hidden)]
pub use tracing as __tracing;
pub use tracing::field::Empty;
#[doc(hidden)]
#[macro_export]
macro_rules! __rig_canonical_completion_span {
(
target: $target:literal,
$(parent: $parent:expr,)?
name: $name:literal,
// Both blocks are spliced verbatim into `info_span!`: the header block
// must end with a trailing comma, the extras block must begin with one.
// Violating either surfaces as an `info_span!` parse error at the call
// site, not here.
{ $($header:tt)* }
{ $($extra:tt)* }
) => {
$crate::telemetry::__tracing::info_span!(
target: $target,
$(parent: $parent,)?
$name,
$($header)*
gen_ai.request.stream = $crate::telemetry::__tracing::field::Empty,
gen_ai.response.id = $crate::telemetry::__tracing::field::Empty,
gen_ai.response.model = $crate::telemetry::__tracing::field::Empty,
rig.provider_request_id = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.output_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.cache_read.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.cache_creation.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.tool_use_prompt_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.reasoning_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.input.messages = $crate::telemetry::__tracing::field::Empty,
gen_ai.output.messages = $crate::telemetry::__tracing::field::Empty
$($extra)*
)
};
}
macro_rules! new_completion_span {
($name:literal, $provider:expr, $request_model:expr, $operation:expr, $system:expr) => {
$crate::__rig_canonical_completion_span!(
target: "rig::completions",
name: $name,
{
gen_ai.operation.name = $operation,
gen_ai.provider.name = $provider,
gen_ai.request.model = $request_model,
gen_ai.system_instructions = $system,
}
{}
)
};
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GenAiOperation {
Chat,
#[deprecated(note = "use `Chat`; streaming is recorded by `SpanBuilder::streaming`")]
ChatStreaming,
GenerateContent,
#[deprecated(note = "use `Chat`; streaming is recorded by `SpanBuilder::streaming`")]
Interactions,
#[deprecated(note = "use `Chat`; streaming is recorded by `SpanBuilder::streaming`")]
InteractionsStreaming,
Embeddings,
Rerank,
Transcription,
ImageGeneration,
AudioGeneration,
}
#[allow(deprecated)]
impl GenAiOperation {
fn as_str(self) -> &'static str {
match self {
Self::Chat | Self::ChatStreaming | Self::Interactions | Self::InteractionsStreaming => {
"chat"
}
Self::GenerateContent => "generate_content",
Self::Embeddings => "embeddings",
Self::Rerank => "rerank",
Self::Transcription => "transcription",
Self::ImageGeneration => "image_generation",
Self::AudioGeneration => "audio_generation",
}
}
fn implied_streaming(self) -> Option<bool> {
matches!(self, Self::ChatStreaming | Self::InteractionsStreaming).then_some(true)
}
pub(crate) fn is_completion(self) -> bool {
matches!(
self,
Self::Chat
| Self::ChatStreaming
| Self::GenerateContent
| Self::Interactions
| Self::InteractionsStreaming
)
}
}
pub const PROVIDER_REQUEST_ID_FIELD: &str = "rig.provider_request_id";
pub const COMPLETION_PARENT_MARKER_FIELD: &str = "rig.completion_parent";
pub const COMPLETION_PARENT_REQUIRED_FIELDS: &[&str] = &[
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.request.model",
"gen_ai.system_instructions",
"gen_ai.response.id",
"gen_ai.response.model",
"gen_ai.usage.input_tokens",
"gen_ai.usage.output_tokens",
"gen_ai.usage.cache_read.input_tokens",
"gen_ai.usage.cache_creation.input_tokens",
"gen_ai.usage.tool_use_prompt_tokens",
"gen_ai.usage.reasoning_tokens",
"gen_ai.input.messages",
"gen_ai.output.messages",
];
#[macro_export]
macro_rules! completion_parent_span {
(
target: $target:literal,
parent: $parent:expr,
name: $name:literal,
operation: $operation:expr,
system_instructions: $system:expr
$(, $($extra:tt)*)?
) => {
$crate::__rig_canonical_completion_span!(
target: $target,
parent: $parent,
name: $name,
{
rig.completion_parent = true,
gen_ai.operation.name = $operation,
gen_ai.system_instructions = $system,
gen_ai.provider.name = $crate::telemetry::__tracing::field::Empty,
gen_ai.request.model = $crate::telemetry::__tracing::field::Empty,
}
{ $(, $($extra)*)? }
)
};
(
target: $target:literal,
name: $name:literal,
operation: $operation:expr,
system_instructions: $system:expr
$(, $($extra:tt)*)?
) => {
$crate::completion_parent_span!(
target: $target,
parent: $crate::telemetry::__tracing::Span::current(),
name: $name,
operation: $operation,
system_instructions: $system
$(, $($extra)*)?
)
};
}
pub use crate::completion_parent_span;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CompletionParentVerdict {
Adopt,
RejectMissingFields,
NotAParent,
}
fn missing_required_fields(metadata: &tracing::Metadata<'_>) -> Vec<&'static str> {
let fields = metadata.fields();
COMPLETION_PARENT_REQUIRED_FIELDS
.iter()
.copied()
.filter(|name| fields.field(name).is_none())
.collect()
}
fn classify_completion_parent(metadata: &tracing::Metadata<'_>) -> CompletionParentVerdict {
let fields = metadata.fields();
if fields.field(COMPLETION_PARENT_MARKER_FIELD).is_none() {
return CompletionParentVerdict::NotAParent;
}
if COMPLETION_PARENT_REQUIRED_FIELDS
.iter()
.all(|name| fields.field(name).is_some())
{
CompletionParentVerdict::Adopt
} else {
CompletionParentVerdict::RejectMissingFields
}
}
static NEAR_MISS_WARNED: LazyLock<Mutex<HashSet<Identifier>>> =
LazyLock::new(|| Mutex::new(HashSet::new()));
#[cfg(test)]
fn reset_near_miss_warnings() {
NEAR_MISS_WARNED
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clear();
}
fn warn_once_on_completion_parent_verdict(
verdict: CompletionParentVerdict,
metadata: &tracing::Metadata<'_>,
) {
match verdict {
CompletionParentVerdict::Adopt | CompletionParentVerdict::NotAParent => {}
CompletionParentVerdict::RejectMissingFields => {
let first_sighting = {
let mut warned = NEAR_MISS_WARNED
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
warned.insert(metadata.callsite())
};
if !first_sighting {
return;
}
tracing::warn!(
marker = COMPLETION_PARENT_MARKER_FIELD,
missing_fields = ?missing_required_fields(metadata),
"completion-parent span declares the marker but not every required field \
and is not adopted; provider telemetry lands on a fresh child span \
instead — declare the span with \
`rig_core::telemetry::completion_parent_span!`"
);
}
}
}
macro_rules! new_modality_span {
($name:literal, $provider:expr, $request_model:expr, $operation:expr) => {
$crate::telemetry::__tracing::info_span!(
target: "rig::modalities",
$name,
gen_ai.operation.name = $operation,
gen_ai.provider.name = $provider,
gen_ai.request.model = $request_model,
gen_ai.response.id = $crate::telemetry::__tracing::field::Empty,
gen_ai.response.model = $crate::telemetry::__tracing::field::Empty,
rig.provider_request_id = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.output_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.cache_read.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.cache_creation.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.tool_use_prompt_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.reasoning_tokens = $crate::telemetry::__tracing::field::Empty,
)
};
}
pub struct SpanBuilder<'a> {
provider: &'a str,
request_model: &'a str,
operation: GenAiOperation,
system_instructions: Option<String>,
streaming: Option<bool>,
}
impl<'a> SpanBuilder<'a> {
pub fn new(provider: &'a str, request_model: &'a str, operation: GenAiOperation) -> Self {
Self {
provider,
request_model,
operation,
system_instructions: None,
streaming: operation.implied_streaming(),
}
}
pub fn streaming(mut self, streaming: bool) -> Self {
self.streaming = Some(streaming);
self
}
pub fn system_instructions(
mut self,
system_instructions: Option<&'a str>,
record_content: bool,
) -> Self {
self.system_instructions = system_instructions_json(system_instructions, record_content);
self
}
pub fn build(self) -> tracing::Span {
if self.operation.is_completion()
&& let Some(parent) = self.adopt_completion_parent()
{
return parent;
}
let (provider, model) = (self.provider, self.request_model);
let op = self.operation.as_str();
let sys = self.system_instructions.as_deref();
#[allow(deprecated)]
let span = match self.operation {
GenAiOperation::Chat
| GenAiOperation::ChatStreaming
| GenAiOperation::Interactions
| GenAiOperation::InteractionsStreaming => {
new_completion_span!("chat", provider, model, op, sys)
}
GenAiOperation::GenerateContent => {
new_completion_span!("generate_content", provider, model, op, sys)
}
GenAiOperation::Embeddings => new_modality_span!("embeddings", provider, model, op),
GenAiOperation::Rerank => new_modality_span!("rerank", provider, model, op),
GenAiOperation::Transcription => {
new_modality_span!("transcription", provider, model, op)
}
GenAiOperation::ImageGeneration => {
new_modality_span!("image_generation", provider, model, op)
}
GenAiOperation::AudioGeneration => {
new_modality_span!("audio_generation", provider, model, op)
}
};
if self.operation.is_completion() {
self.record_streaming(&span);
}
span
}
fn record_streaming(&self, span: &tracing::Span) {
if let Some(streaming) = self.streaming {
span.record("gen_ai.request.stream", streaming);
}
}
fn adopt_completion_parent(&self) -> Option<tracing::Span> {
let current = tracing::Span::current();
let metadata = current.metadata()?;
let verdict = classify_completion_parent(metadata);
warn_once_on_completion_parent_verdict(verdict, metadata);
if verdict != CompletionParentVerdict::Adopt {
return None;
}
current.record("gen_ai.operation.name", self.operation.as_str());
current.record("gen_ai.provider.name", self.provider);
current.record("gen_ai.request.model", self.request_model);
self.record_streaming(¤t);
if let Some(system_instructions) = self.system_instructions.as_deref() {
current.record("gen_ai.system_instructions", system_instructions);
}
Some(current)
}
}
#[derive(Serialize)]
struct TelemetryChatMessage {
role: &'static str,
parts: Vec<TelemetryPart>,
}
#[derive(Serialize)]
struct TelemetryOutputMessage {
role: &'static str,
parts: Vec<TelemetryPart>,
finish_reason: &'static str,
}
#[derive(Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum TelemetryPart {
Text {
content: String,
},
ToolCall {
#[serde(skip_serializing_if = "Option::is_none")]
id: Option<String>,
name: String,
arguments: serde_json::Value,
},
ToolCallResponse {
#[serde(skip_serializing_if = "Option::is_none")]
id: Option<String>,
response: serde_json::Value,
},
Reasoning {
content: String,
},
Opaque {
kind: String,
},
Uri {
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
modality: &'static str,
uri: String,
},
File {
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
modality: &'static str,
file_id: String,
},
Blob {
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
modality: &'static str,
content: String,
},
}
fn media_part<T>(
data: &DocumentSourceKind,
media_type: Option<&T>,
modality: &'static str,
) -> Option<TelemetryPart>
where
T: MimeType,
{
let mime_type = media_type.map(|media_type| media_type.to_mime_type().to_string());
match data {
DocumentSourceKind::Url(uri) => Some(TelemetryPart::Uri {
mime_type,
modality,
uri: uri.clone(),
}),
DocumentSourceKind::FileId(file_id) => Some(TelemetryPart::File {
mime_type,
modality,
file_id: file_id.clone(),
}),
DocumentSourceKind::Base64(content) => Some(TelemetryPart::Blob {
mime_type,
modality,
content: content.clone(),
}),
DocumentSourceKind::Raw(content) => Some(TelemetryPart::Blob {
mime_type,
modality,
content: base64::engine::general_purpose::STANDARD.encode(content),
}),
DocumentSourceKind::String(content) => Some(TelemetryPart::Text {
content: content.clone(),
}),
DocumentSourceKind::Unknown => None,
}
}
fn image_part(image: &Image) -> Option<TelemetryPart> {
media_part(&image.data, image.media_type.as_ref(), "image")
}
fn opaque_part(opaque: &Opaque) -> TelemetryPart {
TelemetryPart::Opaque {
kind: opaque.kind().unwrap_or("unknown").to_owned(),
}
}
fn tool_result_response(result: &ToolResult) -> serde_json::Value {
let mut content = result
.content
.iter()
.filter_map(|content| match content {
ToolResultContent::Text(text) => Some(serde_json::Value::String(text.text.clone())),
ToolResultContent::Json { value } => Some(value.clone()),
ToolResultContent::Image(image) => {
image_part(image).and_then(|part| serde_json::to_value(part).ok())
}
})
.collect::<Vec<_>>();
if content.len() == 1 {
content.pop().unwrap_or(serde_json::Value::Null)
} else {
serde_json::Value::Array(content)
}
}
fn user_parts(content: &[UserContent]) -> Vec<TelemetryPart> {
content
.iter()
.filter_map(|content| match content {
UserContent::Text(text) => Some(TelemetryPart::Text {
content: text.text.clone(),
}),
UserContent::ToolResult(result) => Some(TelemetryPart::ToolCallResponse {
id: Some(result.call.to_string()),
response: tool_result_response(result),
}),
UserContent::Image(image) => image_part(image),
UserContent::Audio(audio) => {
media_part(&audio.data, audio.media_type.as_ref(), "audio")
}
UserContent::Video(video) => {
media_part(&video.data, video.media_type.as_ref(), "video")
}
UserContent::Document(document) => {
media_part(&document.data, document.media_type.as_ref(), "document")
}
})
.collect()
}
fn assistant_parts(content: &[AssistantContent]) -> Vec<TelemetryPart> {
content
.iter()
.flat_map(|content| match content {
AssistantContent::Text(text) => vec![TelemetryPart::Text {
content: text.text.clone(),
}],
AssistantContent::ToolCall(tool_call) => vec![TelemetryPart::ToolCall {
id: Some(tool_call.id.to_string()),
name: tool_call.function.name.clone().into(),
arguments: tool_call.function.arguments_value(),
}],
AssistantContent::Reasoning(reasoning) => vec![TelemetryPart::Reasoning {
content: reasoning.text.clone(),
}],
AssistantContent::Image(image) => image_part(image).into_iter().collect(),
AssistantContent::Opaque(opaque) => vec![opaque_part(opaque)],
})
.collect()
}
fn input_messages(messages: &[Message]) -> Vec<TelemetryChatMessage> {
messages
.iter()
.map(|message| match message {
Message::System { content } => TelemetryChatMessage {
role: "system",
parts: vec![TelemetryPart::Text {
content: content.clone(),
}],
},
Message::User { content } => TelemetryChatMessage {
role: "user",
parts: user_parts(content),
},
Message::Assistant(turn) => TelemetryChatMessage {
role: "assistant",
parts: assistant_parts(&turn.content),
},
})
.collect()
}
fn output_messages(content: &[AssistantContent]) -> Vec<TelemetryOutputMessage> {
let finish_reason = if content
.iter()
.any(|content| matches!(content, AssistantContent::ToolCall(_)))
{
"tool_call"
} else {
"unknown"
};
vec![TelemetryOutputMessage {
role: "assistant",
parts: assistant_parts(content),
finish_reason,
}]
}
pub fn system_instructions_json(instructions: Option<&str>, enabled: bool) -> Option<String> {
if !enabled {
return None;
}
instructions.and_then(|instructions| {
serde_json::to_string(&vec![TelemetryPart::Text {
content: instructions.to_string(),
}])
.ok()
})
}
pub fn record_model_input(span: &tracing::Span, messages: &[Message], enabled: bool) {
if !enabled || span.is_disabled() {
return;
}
if let Ok(messages) = serde_json::to_string(&input_messages(messages)) {
span.record("gen_ai.input.messages", messages);
}
}
pub fn record_model_output(span: &tracing::Span, content: &[AssistantContent], enabled: bool) {
if !enabled || span.is_disabled() {
return;
}
let messages = output_messages(content);
if let Ok(messages) = serde_json::to_string(&messages) {
span.record("gen_ai.output.messages", messages);
}
}
pub trait SpanCombinator {
fn record_token_usage(&self, usage: &Usage);
fn record_response(&self, response_id: Option<&str>, model: Option<&str>, usage: &Usage);
}
impl SpanCombinator for tracing::Span {
fn record_token_usage(&self, usage: &Usage) {
if self.is_disabled() {
return;
}
let fields = [
("gen_ai.usage.input_tokens", usage.input_tokens),
("gen_ai.usage.output_tokens", usage.output_tokens),
(
"gen_ai.usage.cache_read.input_tokens",
usage.cached_input_tokens,
),
(
"gen_ai.usage.cache_creation.input_tokens",
usage.cache_creation_input_tokens,
),
(
"gen_ai.usage.tool_use_prompt_tokens",
usage.tool_use_prompt_tokens,
),
("gen_ai.usage.reasoning_tokens", usage.reasoning_tokens),
];
for (field, value) in fields {
if let Some(value) = value {
self.record(field, value);
}
}
}
fn record_response(&self, response_id: Option<&str>, model: Option<&str>, usage: &Usage) {
if self.is_disabled() {
return;
}
if let Some(id) = response_id {
self.record("gen_ai.response.id", id);
}
if let Some(model) = model {
self.record("gen_ai.response.model", model);
}
self.record_token_usage(usage);
}
}
#[cfg(test)]
mod equivalence_tests;
#[cfg(test)]
mod tests;