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}
29
30impl OtelObserver {
31 pub fn new(instrumentation: OtelInstrumentation) -> Self {
32 Self { turn: None, tool_definitions: Vec::new(), instrumentation }
33 }
34}
35
36impl AgentObserver for OtelObserver {
37 fn on_event(&mut self, message: &AgentEvent) {
38 match message {
39 AgentEvent::Turn(TurnEvent::Started { content }) => self.start_turn(content),
40 AgentEvent::Turn(TurnEvent::Ended { outcome }) => {
41 if let Some(turn) = self.turn.take() {
42 turn.finish(outcome);
43 }
44 }
45 AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools }) => {
46 self.tool_definitions.clone_from(tools);
47 }
48 message => {
49 if let Some(turn) = &mut self.turn {
50 turn.on_event(message, &self.instrumentation, &self.tool_definitions);
51 }
52 }
53 }
54 }
55}
56
57impl OtelObserver {
58 fn start_turn(&mut self, content: &[ContentBlock]) {
59 self.turn = None;
62
63 let mut input = self.instrumentation.capture().buffer();
64 input.set(&ContentBlock::join_text(content));
65 let mut attributes = vec![KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "invoke_agent")];
66 if let Some(text) = input.get() {
67 attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, messages_json("user", text)));
68 }
69 let builder = SpanBuilder::from_name("invoke_agent").with_kind(SpanKind::Internal).with_attributes(attributes);
70 let span_context = self.instrumentation.start_span(builder, None);
71 let span = SpanGuard::new(span_context, TURN_CANCEL_MESSAGE);
72 self.turn = Some(TurnState::new(span, input, self.instrumentation.capture()));
73 }
74}
75
76impl OtelInstrumentation {
77 fn capture(&self) -> ContentCapture {
78 ContentCapture::from_enabled(self.capture_content)
79 }
80
81 fn start_span(&self, builder: SpanBuilder, parent: Option<&Context>) -> Context {
84 let parent = parent.cloned().unwrap_or_default();
85 Context::new().with_span(self.tracer.build_with_context(builder, &parent))
86 }
87}
88
89const TURN_CANCEL_MESSAGE: &str = "turn cancelled";
90const TOOL_CANCEL_MESSAGE: &str = "turn ended before the tool completed";
91
92struct TurnState {
97 span: SpanGuard,
98 input: ContentBuffer,
99 output: ContentBuffer,
100 chat_call: Option<LlmCallState>,
101 compaction_call: Option<LlmCallState>,
102 streamed_arguments: HashMap<String, ContentBuffer>,
104 executing_tools: HashMap<String, SpanGuard>,
106}
107
108impl TurnState {
109 fn new(span: SpanGuard, input: ContentBuffer, capture: ContentCapture) -> Self {
110 Self {
111 span,
112 input,
113 output: capture.buffer(),
114 chat_call: None,
115 compaction_call: None,
116 streamed_arguments: HashMap::new(),
117 executing_tools: HashMap::new(),
118 }
119 }
120
121 fn on_event(&mut self, message: &AgentEvent, instrumentation: &OtelInstrumentation, tools: &[ToolDefinition]) {
122 match message {
123 AgentEvent::Turn(TurnEvent::LlmCallStarted { purpose, provider, model, display_name, attempt, .. }) => {
124 self.start_llm_call(
125 LlmCallStart {
126 purpose: *purpose,
127 provider: provider.as_deref(),
128 model: model.as_deref(),
129 display_name,
130 attempt: *attempt,
131 },
132 instrumentation,
133 tools,
134 );
135 }
136 AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose, outcome }) => {
137 if let Some(call) = self.llm_call_slot(*purpose).take() {
138 call.finish(outcome);
139 }
140 }
141 AgentEvent::Message(
142 MessageEvent::Text { message_id, chunk, is_complete: false }
143 | MessageEvent::Thought { message_id, chunk, is_complete: false },
144 ) => {
145 if let Some(chat) = &mut self.chat_call {
146 chat.record_response_chunk(message_id, chunk);
147 }
148 }
149 AgentEvent::Message(MessageEvent::Text { chunk, is_complete: true, .. }) => {
150 self.output.push(chunk);
151 }
152 AgentEvent::Tool(ToolEvent::Call { request, .. }) => self.on_tool_call(request, instrumentation),
153 AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id, chunk, .. }) => {
154 self.on_tool_call_update(tool_call_id, chunk);
155 }
156 AgentEvent::Tool(ToolEvent::ExecutionStarted { tool_id, tool_name }) => {
157 self.on_tool_execution_started(tool_id, tool_name, instrumentation);
158 }
159 AgentEvent::Tool(ToolEvent::Result { result, .. }) => self.on_tool_result(result, instrumentation),
160 AgentEvent::Tool(ToolEvent::Error { error, .. }) => self.on_tool_error(error),
161 _ => {}
162 }
163 }
164
165 fn finish(self, outcome: &TurnOutcome) {
166 let Self { mut span, output, chat_call, compaction_call, executing_tools, .. } = self;
167 drop(chat_call);
169 drop(compaction_call);
170 drop(executing_tools);
171 match outcome {
172 TurnOutcome::Completed => {
173 if let Some(text) = output.get() {
174 span.set_attribute(KeyValue::new(
175 semconv::GEN_AI_OUTPUT_MESSAGES,
176 messages_json("assistant", text),
177 ));
178 }
179 span.end_ok();
180 }
181 TurnOutcome::Failed { error } => span.end_error(None, error.clone()),
182 TurnOutcome::Cancelled => span.end_error(Some(ErrorKind::Cancelled), TURN_CANCEL_MESSAGE),
183 }
184 }
185
186 fn start_llm_call(
187 &mut self,
188 call: LlmCallStart<'_>,
189 instrumentation: &OtelInstrumentation,
190 tools: &[ToolDefinition],
191 ) {
192 let model_name = call.model.unwrap_or(call.display_name).to_string();
193
194 let mut metric_attributes = vec![
197 KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "chat"),
198 KeyValue::new(semconv::GEN_AI_REQUEST_MODEL, model_name.clone()),
199 ];
200
201 if let Some(provider) = call.provider {
202 metric_attributes.push(KeyValue::new(semconv::GEN_AI_PROVIDER_NAME, genai_provider_name(provider)));
203 }
204
205 if call.purpose == LlmCallPurpose::Compaction {
206 metric_attributes.push(KeyValue::new(semconv::LLM_PURPOSE, "compaction"));
207 }
208
209 let mut attributes = metric_attributes.clone();
210 attributes.push(KeyValue::new(semconv::GEN_AI_REQUEST_STREAM, true));
211 attributes.push(KeyValue::new(semconv::LLM_ATTEMPT, i64::from(call.attempt)));
212 if call.purpose == LlmCallPurpose::Chat {
215 if let Some(input) = self.input.get() {
216 attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, messages_json("user", input)));
217 }
218 if instrumentation.capture_content && !tools.is_empty() {
219 attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_DEFINITIONS, tool_definitions_json(tools)));
220 }
221 }
222
223 let name = if model_name.is_empty() { "chat".to_string() } else { format!("chat {model_name}") };
224 let builder = SpanBuilder::from_name(name).with_kind(SpanKind::Client).with_attributes(attributes);
225 let context = instrumentation.start_span(builder, Some(self.span.context()));
226 let state = LlmCallState::new(
227 context,
228 instrumentation.metrics.clone(),
229 call.purpose,
230 instrumentation.capture().buffer(),
231 metric_attributes,
232 );
233 *self.llm_call_slot(call.purpose) = Some(state);
234 }
235
236 fn on_tool_call(&mut self, request: &ToolCallRequest, instrumentation: &OtelInstrumentation) {
237 if let Some(chat) = &mut self.chat_call {
238 chat.record_tool_call_start(&request.id, &request.name);
239 }
240 let mut arguments = instrumentation.capture().buffer();
241 arguments.set(&request.arguments);
242 self.streamed_arguments.insert(request.id.clone(), arguments);
243 }
244
245 fn on_tool_call_update(&mut self, tool_call_id: &str, chunk: &str) {
246 if let Some(chat) = &mut self.chat_call {
247 chat.record_output_chunk();
248 }
249 if let Some(arguments) = self.streamed_arguments.get_mut(tool_call_id) {
250 arguments.push(chunk);
251 }
252 }
253
254 fn on_tool_execution_started(&mut self, tool_id: &str, tool_name: &str, instrumentation: &OtelInstrumentation) {
255 let mut attributes = vec![
256 KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "execute_tool"),
257 KeyValue::new(semconv::GEN_AI_TOOL_NAME, tool_name.to_string()),
258 KeyValue::new(semconv::GEN_AI_TOOL_CALL_ID, tool_id.to_string()),
259 ];
260 let arguments = self.streamed_arguments.remove(tool_id);
261
262 if let Some(text) = arguments.as_ref().and_then(ContentBuffer::get) {
263 attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_CALL_ARGUMENTS, text.to_string()));
264 }
265
266 let builder = SpanBuilder::from_name(format!("execute_tool {tool_name}"))
267 .with_kind(SpanKind::Internal)
268 .with_attributes(attributes);
269 let context = instrumentation.start_span(builder, Some(self.span.context()));
270 self.executing_tools.insert(tool_id.to_string(), SpanGuard::new(context, TOOL_CANCEL_MESSAGE));
271 }
272
273 fn on_tool_result(&mut self, result: &ToolCallResult, instrumentation: &OtelInstrumentation) {
274 self.streamed_arguments.remove(&result.id);
275 let Some(mut span) = self.executing_tools.remove(&result.id) else { return };
276
277 if let Some(content) = instrumentation.capture().content(&result.result) {
278 span.set_attribute(KeyValue::new(semconv::GEN_AI_TOOL_CALL_RESULT, content.to_string()));
279 }
280
281 span.end_ok();
282 }
283
284 fn on_tool_error(&mut self, error: &ToolCallError) {
285 self.streamed_arguments.remove(&error.id);
286 if let Some(mut span) = self.executing_tools.remove(&error.id) {
287 span.end_error(Some(ErrorKind::ToolError), error.error.clone());
288 }
289 }
290
291 fn llm_call_slot(&mut self, purpose: LlmCallPurpose) -> &mut Option<LlmCallState> {
292 match purpose {
293 LlmCallPurpose::Chat => &mut self.chat_call,
294 LlmCallPurpose::Compaction => &mut self.compaction_call,
295 }
296 }
297}
298
299fn genai_provider_name(provider: &str) -> String {
302 provider.parse::<Provider>().map_or_else(|_| provider.to_string(), |p| p.genai_provider_name().to_string())
303}
304
305#[derive(Clone, Copy)]
308struct LlmCallStart<'a> {
309 purpose: LlmCallPurpose,
310 provider: Option<&'a str>,
311 model: Option<&'a str>,
312 display_name: &'a str,
313 attempt: u32,
314}