1use crate::content_capture::{ContentBuffer, ContentCapture};
2use crate::content_json::{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, 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, messages_json("user", 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 { purpose, provider, model, display_name, attempt, .. }) => {
125 self.start_llm_call(
126 LlmCallStart {
127 purpose: *purpose,
128 provider: provider.as_deref(),
129 model: model.as_deref(),
130 display_name,
131 attempt: *attempt,
132 },
133 instrumentation,
134 tools,
135 );
136 }
137 AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose, outcome }) => {
138 if let Some(call) = self.llm_call_slot(*purpose).take() {
139 call.finish(outcome);
140 }
141 }
142 AgentEvent::Message(
143 MessageEvent::Text { message_id, chunk, is_complete: false }
144 | MessageEvent::Thought { message_id, chunk, is_complete: false },
145 ) => {
146 if let Some(chat) = &mut self.chat_call {
147 chat.record_response_chunk(message_id, chunk);
148 }
149 }
150 AgentEvent::Message(MessageEvent::Text { chunk, is_complete: true, .. }) => {
151 self.output.push(chunk);
152 }
153 AgentEvent::Tool(ToolEvent::Call { request, .. }) => self.on_tool_call(request, instrumentation),
154 AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id, chunk, .. }) => {
155 self.on_tool_call_update(tool_call_id, chunk);
156 }
157 AgentEvent::Tool(ToolEvent::ExecutionStarted { tool_id, tool_name }) => {
158 self.on_tool_execution_started(tool_id, tool_name, instrumentation);
159 }
160 AgentEvent::Tool(ToolEvent::Result { result, .. }) => self.on_tool_result(result, instrumentation),
161 AgentEvent::Tool(ToolEvent::Error { error, .. }) => self.on_tool_error(error),
162 _ => {}
163 }
164 }
165
166 fn finish(self, outcome: &TurnOutcome) {
167 let Self { mut span, output, chat_call, compaction_call, executing_tools, .. } = self;
168 drop(chat_call);
170 drop(compaction_call);
171 drop(executing_tools);
172 match outcome {
173 TurnOutcome::Completed => {
174 if let Some(text) = output.get() {
175 span.set_attribute(KeyValue::new(
176 semconv::GEN_AI_OUTPUT_MESSAGES,
177 messages_json("assistant", text),
178 ));
179 }
180 span.end_ok();
181 }
182 TurnOutcome::Failed { error } => span.end_error(None, error.clone()),
183 TurnOutcome::Cancelled => span.end_error(Some(ErrorKind::Cancelled), TURN_CANCEL_MESSAGE),
184 }
185 }
186
187 fn start_llm_call(
188 &mut self,
189 call: LlmCallStart<'_>,
190 instrumentation: &OtelInstrumentation,
191 tools: &[ToolDefinition],
192 ) {
193 let model_name = call.model.unwrap_or(call.display_name).to_string();
194
195 let mut metric_attributes = vec![
198 KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "chat"),
199 KeyValue::new(semconv::GEN_AI_REQUEST_MODEL, model_name.clone()),
200 ];
201
202 if let Some(provider) = call.provider {
203 metric_attributes.push(KeyValue::new(semconv::GEN_AI_PROVIDER_NAME, genai_provider_name(provider)));
204 }
205
206 if call.purpose == LlmCallPurpose::Compaction {
207 metric_attributes.push(KeyValue::new(semconv::LLM_PURPOSE, "compaction"));
208 }
209
210 let mut attributes = metric_attributes.clone();
211 attributes.push(KeyValue::new(semconv::GEN_AI_REQUEST_STREAM, true));
212 attributes.push(KeyValue::new(semconv::LLM_ATTEMPT, i64::from(call.attempt)));
213 if call.purpose == LlmCallPurpose::Chat {
216 if let Some(input) = self.input.get() {
217 attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, messages_json("user", input)));
218 }
219 if instrumentation.capture_content && !tools.is_empty() {
220 attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_DEFINITIONS, tool_definitions_json(tools)));
221 }
222 }
223
224 let name = if model_name.is_empty() { "chat".to_string() } else { format!("chat {model_name}") };
225 let builder = SpanBuilder::from_name(name).with_kind(SpanKind::Client).with_attributes(attributes);
226 let context = instrumentation.start_span(builder, Some(self.span.context()));
227 let state = LlmCallState::new(
228 context,
229 instrumentation.metrics.clone(),
230 call.purpose,
231 instrumentation.capture().buffer(),
232 metric_attributes,
233 );
234 *self.llm_call_slot(call.purpose) = Some(state);
235 }
236
237 fn on_tool_call(&mut self, request: &ToolCallRequest, instrumentation: &OtelInstrumentation) {
238 if let Some(chat) = &mut self.chat_call {
239 chat.record_tool_call_start(&request.id, &request.name);
240 }
241 let mut arguments = instrumentation.capture().buffer();
242 arguments.set(&request.arguments);
243 self.streamed_arguments.insert(request.id.clone(), arguments);
244 }
245
246 fn on_tool_call_update(&mut self, tool_call_id: &str, chunk: &str) {
247 if let Some(chat) = &mut self.chat_call {
248 chat.record_output_chunk();
249 }
250 if let Some(arguments) = self.streamed_arguments.get_mut(tool_call_id) {
251 arguments.push(chunk);
252 }
253 }
254
255 fn on_tool_execution_started(&mut self, tool_id: &str, tool_name: &str, instrumentation: &OtelInstrumentation) {
256 let mut attributes = vec![
257 KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "execute_tool"),
258 KeyValue::new(semconv::GEN_AI_TOOL_NAME, tool_name.to_string()),
259 KeyValue::new(semconv::GEN_AI_TOOL_CALL_ID, tool_id.to_string()),
260 ];
261 let arguments = self.streamed_arguments.remove(tool_id);
262
263 if let Some(text) = arguments.as_ref().and_then(ContentBuffer::get) {
264 attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_CALL_ARGUMENTS, text.to_string()));
265 }
266
267 let builder = SpanBuilder::from_name(format!("execute_tool {tool_name}"))
268 .with_kind(SpanKind::Internal)
269 .with_attributes(attributes);
270 let context = instrumentation.start_span(builder, Some(self.span.context()));
271 self.executing_tools.insert(tool_id.to_string(), SpanGuard::new(context, TOOL_CANCEL_MESSAGE));
272 }
273
274 fn on_tool_result(&mut self, result: &ToolCallResult, instrumentation: &OtelInstrumentation) {
275 self.streamed_arguments.remove(&result.id);
276 let Some(mut span) = self.executing_tools.remove(&result.id) else { return };
277
278 if let Some(content) = instrumentation.capture().content(&result.result) {
279 span.set_attribute(KeyValue::new(semconv::GEN_AI_TOOL_CALL_RESULT, content.to_string()));
280 }
281
282 span.end_ok();
283 }
284
285 fn on_tool_error(&mut self, error: &ToolCallError) {
286 self.streamed_arguments.remove(&error.id);
287 if let Some(mut span) = self.executing_tools.remove(&error.id) {
288 span.end_error(Some(ErrorKind::ToolError), error.error.clone());
289 }
290 }
291
292 fn llm_call_slot(&mut self, purpose: LlmCallPurpose) -> &mut Option<LlmCallState> {
293 match purpose {
294 LlmCallPurpose::Chat => &mut self.chat_call,
295 LlmCallPurpose::Compaction => &mut self.compaction_call,
296 }
297 }
298}
299
300fn genai_provider_name(provider: &str) -> String {
303 provider.parse::<Provider>().map_or_else(|_| provider.to_string(), |p| p.genai_provider_name().to_string())
304}
305
306#[derive(Clone, Copy)]
309struct LlmCallStart<'a> {
310 purpose: LlmCallPurpose,
311 provider: Option<&'a str>,
312 model: Option<&'a str>,
313 display_name: &'a str,
314 attempt: u32,
315}