rig/telemetry/
mod.rs

1//! This module primarily concerns being able to orchestrate telemetry across a given pipeline or workflow.
2//! This includes tracing, being able to send traces to an OpenTelemetry collector, setting up your
3//! agents with the correct tracing style so you can emit the right traces for platforms like Langfuse,
4//! and more.
5
6use crate::completion::GetTokenUsage;
7use serde::Serialize;
8
9pub trait ProviderRequestExt {
10    type InputMessage: Serialize;
11
12    fn get_input_messages(&self) -> Vec<Self::InputMessage>;
13    fn get_system_prompt(&self) -> Option<String>;
14    fn get_model_name(&self) -> String;
15    fn get_prompt(&self) -> Option<String>;
16}
17
18pub trait ProviderResponseExt {
19    type OutputMessage: Serialize;
20    type Usage: Serialize;
21
22    fn get_response_id(&self) -> Option<String>;
23
24    fn get_response_model_name(&self) -> Option<String>;
25
26    fn get_output_messages(&self) -> Vec<Self::OutputMessage>;
27
28    fn get_text_response(&self) -> Option<String>;
29
30    fn get_usage(&self) -> Option<Self::Usage>;
31}
32
33/// A trait designed specifically to be used with Spans for the purpose of recording telemetry.
34/// Nearly all methods
35pub trait SpanCombinator {
36    fn record_token_usage<U>(&self, usage: &U)
37    where
38        U: GetTokenUsage;
39
40    fn record_response_metadata<R>(&self, response: &R)
41    where
42        R: ProviderResponseExt;
43
44    fn record_model_input<T>(&self, messages: &T)
45    where
46        T: Serialize;
47
48    fn record_model_output<T>(&self, messages: &T)
49    where
50        T: Serialize;
51}
52
53impl SpanCombinator for tracing::Span {
54    fn record_token_usage<U>(&self, usage: &U)
55    where
56        U: GetTokenUsage,
57    {
58        if let Some(usage) = usage.token_usage() {
59            self.record("gen_ai.usage.input_tokens", usage.input_tokens);
60            self.record("gen_ai.usage.output_tokens", usage.output_tokens);
61        }
62    }
63
64    fn record_response_metadata<R>(&self, response: &R)
65    where
66        R: ProviderResponseExt,
67    {
68        if let Some(id) = response.get_response_id() {
69            self.record("gen_ai.response.id", id);
70        }
71
72        if let Some(model_name) = response.get_response_model_name() {
73            self.record("gen_ai.response.model_name", model_name);
74        }
75    }
76
77    fn record_model_input<T>(&self, input: &T)
78    where
79        T: Serialize,
80    {
81        let input_as_json_string =
82            serde_json::to_string(input).expect("Serializing a Rust type to JSON should not break");
83
84        self.record("gen_ai.input.messages", input_as_json_string);
85    }
86
87    fn record_model_output<T>(&self, input: &T)
88    where
89        T: Serialize,
90    {
91        let input_as_json_string =
92            serde_json::to_string(input).expect("Serializing a Rust type to JSON should not break");
93
94        self.record("gen_ai.input.messages", input_as_json_string);
95    }
96}
rig/telemetry/mod.rs

rig/telemetry/
mod.rs
