1use crate::content_capture::{ContentBuffer, ContentCapture};
2use crate::content_json::{input_messages_json, output_messages_json, tool_definitions_json};
3use crate::gen_ai_metrics::GenAiMetrics;
4use crate::genai_constants as semconv;
5use crate::llm_call_state::LlmCallState;
6use crate::span_guard::{ErrorKind, SpanGuard};
7use aether_core::events::{AgentEvent, AgentObserver, LlmCallPurpose, MessageEvent, ToolEvent, TurnEvent, TurnOutcome};
8use llm::catalog::Provider;
9use llm::{ContentBlock, ModelPricing, ToolCallError, ToolCallRequest, ToolCallResult, ToolDefinition};
10use opentelemetry::trace::{SpanBuilder, SpanKind, TraceContextExt, Tracer as _};
11use opentelemetry::{Context, KeyValue};
12use opentelemetry_sdk::trace::SdkTracer;
13use std::collections::HashMap;
14
15pub struct OtelObserver {
18 turn: Option<TurnState>,
19 tool_definitions: Vec<ToolDefinition>,
20 instrumentation: OtelInstrumentation,
21}
22
23#[derive(Clone)]
24pub struct OtelInstrumentation {
25 pub tracer: SdkTracer,
26 pub metrics: GenAiMetrics,
27 pub capture_content: bool,
28 pub root_parent: Option<Context>,
29}
30
31impl OtelObserver {
32 pub fn new(instrumentation: OtelInstrumentation) -> Self {
33 Self { turn: None, tool_definitions: Vec::new(), instrumentation }
34 }
35}
36
37impl AgentObserver for OtelObserver {
38 fn on_event(&mut self, message: &AgentEvent) {
39 match message {
40 AgentEvent::Turn(TurnEvent::Started { content }) => self.start_turn(content),
41 AgentEvent::Turn(TurnEvent::Ended { outcome }) => {
42 if let Some(turn) = self.turn.take() {
43 turn.finish(outcome);
44 }
45 }
46 AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools }) => {
47 self.tool_definitions.clone_from(tools);
48 }
49 message => {
50 if let Some(turn) = &mut self.turn {
51 turn.on_event(message, &self.instrumentation, &self.tool_definitions);
52 }
53 }
54 }
55 }
56}
57
58impl OtelObserver {
59 fn start_turn(&mut self, content: &[ContentBlock]) {
60 self.turn = None;
63
64 let mut input = self.instrumentation.capture().buffer();
65 input.set(&ContentBlock::join_text(content));
66 let mut attributes = vec![KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "invoke_agent")];
67 if let Some(text) = input.get() {
68 attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, input_messages_json(text)));
69 }
70 let builder = SpanBuilder::from_name("invoke_agent").with_kind(SpanKind::Internal).with_attributes(attributes);
71 let span_context = self.instrumentation.start_span(builder, self.instrumentation.root_parent.as_ref());
72 let span = SpanGuard::new(span_context, TURN_CANCEL_MESSAGE);
73 self.turn = Some(TurnState::new(span, input, self.instrumentation.capture()));
74 }
75}
76
77impl OtelInstrumentation {
78 fn capture(&self) -> ContentCapture {
79 ContentCapture::from_enabled(self.capture_content)
80 }
81
82 fn start_span(&self, builder: SpanBuilder, parent: Option<&Context>) -> Context {
85 let parent = parent.cloned().unwrap_or_default();
86 Context::new().with_span(self.tracer.build_with_context(builder, &parent))
87 }
88}
89
90const TURN_CANCEL_MESSAGE: &str = "turn cancelled";
91const TOOL_CANCEL_MESSAGE: &str = "turn ended before the tool completed";
92
93struct TurnState {
98 span: SpanGuard,
99 input: ContentBuffer,
100 output: ContentBuffer,
101 chat_call: Option<LlmCallState>,
102 compaction_call: Option<LlmCallState>,
103 streamed_arguments: HashMap<String, ContentBuffer>,
105 executing_tools: HashMap<String, SpanGuard>,
107}
108
109impl TurnState {
110 fn new(span: SpanGuard, input: ContentBuffer, capture: ContentCapture) -> Self {
111 Self {
112 span,
113 input,
114 output: capture.buffer(),
115 chat_call: None,
116 compaction_call: None,
117 streamed_arguments: HashMap::new(),
118 executing_tools: HashMap::new(),
119 }
120 }
121
122 fn on_event(&mut self, message: &AgentEvent, instrumentation: &OtelInstrumentation, tools: &[ToolDefinition]) {
123 match message {
124 AgentEvent::Turn(TurnEvent::LlmCallStarted {
125 purpose,
126 provider,
127 model,
128 display_name,
129 pricing,
130 attempt,
131 ..
132 }) => {
133 self.start_llm_call(
134 LlmCallStart {
135 purpose: *purpose,
136 provider: provider.as_deref(),
137 model: model.as_deref(),
138 display_name,
139 pricing: *pricing,
140 attempt: *attempt,
141 },
142 instrumentation,
143 tools,
144 );
145 }
146 AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose, outcome }) => {
147 if let Some(call) = self.llm_call_slot(*purpose).take() {
148 call.finish(outcome);
149 }
150 }
151 AgentEvent::Message(
152 MessageEvent::Text { message_id, chunk, is_complete: false }
153 | MessageEvent::Thought { message_id, chunk, is_complete: false },
154 ) => {
155 if let Some(chat) = &mut self.chat_call {
156 chat.record_response_chunk(message_id, chunk);
157 }
158 }
159 AgentEvent::Message(MessageEvent::Text { chunk, is_complete: true, .. }) => {
160 self.output.push(chunk);
161 }
162 AgentEvent::Tool(ToolEvent::Call { request, .. }) => self.on_tool_call(request, instrumentation),
163 AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id, chunk, .. }) => {
164 self.on_tool_call_update(tool_call_id, chunk);
165 }
166 AgentEvent::Tool(ToolEvent::ExecutionStarted { tool_id, tool_name }) => {
167 self.on_tool_execution_started(tool_id, tool_name, instrumentation);
168 }
169 AgentEvent::Tool(ToolEvent::Result { result, .. }) => self.on_tool_result(result, instrumentation),
170 AgentEvent::Tool(ToolEvent::Error { error, .. }) => self.on_tool_error(error),
171 _ => {}
172 }
173 }
174
175 fn finish(self, outcome: &TurnOutcome) {
176 let Self { mut span, output, chat_call, compaction_call, executing_tools, .. } = self;
177 drop(chat_call);
179 drop(compaction_call);
180 drop(executing_tools);
181 match outcome {
182 TurnOutcome::Completed => {
183 if let Some(text) = output_messages_json(output.get(), &[], None) {
184 span.set_attribute(KeyValue::new(semconv::GEN_AI_OUTPUT_MESSAGES, text));
185 }
186 span.end_ok();
187 }
188 TurnOutcome::Failed { error } => span.end_error(None, error.clone()),
189 TurnOutcome::Cancelled => span.end_error(Some(ErrorKind::Cancelled), TURN_CANCEL_MESSAGE),
190 }
191 }
192
193 fn start_llm_call(
194 &mut self,
195 call: LlmCallStart<'_>,
196 instrumentation: &OtelInstrumentation,
197 tools: &[ToolDefinition],
198 ) {
199 let model_name = call.model.unwrap_or(call.display_name).to_string();
200
201 let mut metric_attributes = vec![
204 KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "chat"),
205 KeyValue::new(semconv::GEN_AI_REQUEST_MODEL, model_name.clone()),
206 ];
207
208 if let Some(provider) = call.provider {
209 metric_attributes.push(KeyValue::new(semconv::GEN_AI_PROVIDER_NAME, genai_provider_name(provider)));
210 }
211
212 if call.purpose == LlmCallPurpose::Compaction {
213 metric_attributes.push(KeyValue::new(semconv::LLM_PURPOSE, "compaction"));
214 }
215
216 let mut attributes = metric_attributes.clone();
217 attributes.push(KeyValue::new(semconv::GEN_AI_REQUEST_STREAM, true));
218 attributes.push(KeyValue::new(semconv::LLM_ATTEMPT, i64::from(call.attempt)));
219 if let Some(pricing) = call.pricing {
220 attributes.extend(pricing_attributes(pricing));
221 }
222 if call.purpose == LlmCallPurpose::Chat {
225 if let Some(input) = self.input.get() {
226 attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, input_messages_json(input)));
227 }
228 if instrumentation.capture_content && !tools.is_empty() {
229 attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_DEFINITIONS, tool_definitions_json(tools)));
230 }
231 }
232
233 let name = if model_name.is_empty() { "chat".to_string() } else { format!("chat {model_name}") };
234 let builder = SpanBuilder::from_name(name).with_kind(SpanKind::Client).with_attributes(attributes);
235 let context = instrumentation.start_span(builder, Some(self.span.context()));
236 let state = LlmCallState::new(
237 context,
238 instrumentation.metrics.clone(),
239 call.purpose,
240 instrumentation.capture(),
241 metric_attributes,
242 );
243 *self.llm_call_slot(call.purpose) = Some(state);
244 }
245
246 fn on_tool_call(&mut self, request: &ToolCallRequest, instrumentation: &OtelInstrumentation) {
247 if let Some(chat) = &mut self.chat_call {
248 chat.record_tool_call_start(request);
249 }
250 let mut arguments = instrumentation.capture().buffer();
251 arguments.set(&request.arguments);
252 self.streamed_arguments.insert(request.id.clone(), arguments);
253 }
254
255 fn on_tool_call_update(&mut self, tool_call_id: &str, chunk: &str) {
256 if let Some(chat) = &mut self.chat_call {
257 chat.record_tool_call_update(tool_call_id, chunk);
258 }
259 if let Some(arguments) = self.streamed_arguments.get_mut(tool_call_id) {
260 arguments.push(chunk);
261 }
262 }
263
264 fn on_tool_execution_started(&mut self, tool_id: &str, tool_name: &str, instrumentation: &OtelInstrumentation) {
265 let mut attributes = vec![
266 KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "execute_tool"),
267 KeyValue::new(semconv::GEN_AI_TOOL_NAME, tool_name.to_string()),
268 KeyValue::new(semconv::GEN_AI_TOOL_CALL_ID, tool_id.to_string()),
269 ];
270 let arguments = self.streamed_arguments.remove(tool_id);
271
272 if let Some(text) = arguments.as_ref().and_then(ContentBuffer::get) {
273 attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_CALL_ARGUMENTS, text.to_string()));
274 }
275
276 let builder = SpanBuilder::from_name(format!("execute_tool {tool_name}"))
277 .with_kind(SpanKind::Internal)
278 .with_attributes(attributes);
279 let context = instrumentation.start_span(builder, Some(self.span.context()));
280 self.executing_tools.insert(tool_id.to_string(), SpanGuard::new(context, TOOL_CANCEL_MESSAGE));
281 }
282
283 fn on_tool_result(&mut self, result: &ToolCallResult, instrumentation: &OtelInstrumentation) {
284 self.streamed_arguments.remove(&result.id);
285 let Some(mut span) = self.executing_tools.remove(&result.id) else { return };
286
287 if let Some(content) = instrumentation.capture().content(&result.result) {
288 span.set_attribute(KeyValue::new(semconv::GEN_AI_TOOL_CALL_RESULT, content.to_string()));
289 }
290
291 span.end_ok();
292 }
293
294 fn on_tool_error(&mut self, error: &ToolCallError) {
295 self.streamed_arguments.remove(&error.id);
296 if let Some(mut span) = self.executing_tools.remove(&error.id) {
297 span.end_error(Some(ErrorKind::ToolError), error.error.clone());
298 }
299 }
300
301 fn llm_call_slot(&mut self, purpose: LlmCallPurpose) -> &mut Option<LlmCallState> {
302 match purpose {
303 LlmCallPurpose::Chat => &mut self.chat_call,
304 LlmCallPurpose::Compaction => &mut self.compaction_call,
305 }
306 }
307}
308
309fn genai_provider_name(provider: &str) -> String {
312 provider.parse::<Provider>().map_or_else(|_| provider.to_string(), |p| p.genai_provider_name().to_string())
313}
314
315#[derive(Clone, Copy)]
318struct LlmCallStart<'a> {
319 purpose: LlmCallPurpose,
320 provider: Option<&'a str>,
321 model: Option<&'a str>,
322 display_name: &'a str,
323 pricing: Option<ModelPricing>,
324 attempt: u32,
325}
326
327fn pricing_attributes(pricing: ModelPricing) -> Vec<KeyValue> {
328 let mut attributes = vec![
329 KeyValue::new(semconv::AI_INPUT_TOKEN_PRICE, pricing.input_per_million / TOKENS_PER_MILLION),
330 KeyValue::new(semconv::AI_OUTPUT_TOKEN_PRICE, pricing.output_per_million / TOKENS_PER_MILLION),
331 ];
332
333 if let Some(price) = pricing.cache_read_per_million {
334 attributes.push(KeyValue::new(semconv::AI_CACHE_READ_TOKEN_PRICE, price / TOKENS_PER_MILLION));
335 }
336
337 if let Some(price) = pricing.cache_write_per_million {
338 attributes.push(KeyValue::new(semconv::AI_CACHE_WRITE_TOKEN_PRICE, price / TOKENS_PER_MILLION));
339 }
340
341 attributes
342}
343
344const TOKENS_PER_MILLION: f64 = 1_000_000.0;