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