Use activities instead of ActivityEvent

Author: reakaleek

src/api/Elastic.Documentation.Api.Infrastructure/Adapters/AskAi/AgentBuilderStreamTransformer.cs

@@ -26,9 +26,7 @@ public class AgentBuilderStreamTransformer(ILogger<AgentBuilderStreamTransformer
 
 		// Special handling for error events - they may have a different structure
 		if (type == "error")
-		{
 			return ParseErrorEventFromRoot(id, timestamp, json);
-		}
 
 		// Most Agent Builder events have data nested in a "data" property
 		if (!json.TryGetProperty("data", out var innerData))
@@ -76,7 +74,7 @@ public class AgentBuilderStreamTransformer(ILogger<AgentBuilderStreamTransformer
 		return null;
 	}
 
-	private AskAiEvent.Reasoning ParseReasoningEvent(string id, long timestamp, JsonElement innerData)
+	private static AskAiEvent.Reasoning ParseReasoningEvent(string id, long timestamp, JsonElement innerData)
 	{
 		// Agent Builder sends: {"data":{"reasoning":"..."}}
 		var message = innerData.TryGetProperty("reasoning", out var reasoningProp)
@@ -86,7 +84,7 @@ private AskAiEvent.Reasoning ParseReasoningEvent(string id, long timestamp, Json
 		return new AskAiEvent.Reasoning(id, timestamp, message ?? "Thinking...");
 	}
 
-	private AskAiEvent.ToolResult ParseToolResultEvent(string id, long timestamp, JsonElement innerData)
+	private static AskAiEvent.ToolResult ParseToolResultEvent(string id, long timestamp, JsonElement innerData)
 	{
 		// Extract tool_call_id and results
 		var toolCallId = innerData.TryGetProperty("tool_call_id", out var tcId) ? tcId.GetString() : id;
@@ -99,7 +97,7 @@ private AskAiEvent.ToolResult ParseToolResultEvent(string id, long timestamp, Js
 		return new AskAiEvent.ToolResult(id, timestamp, toolCallId ?? id, result);
 	}
 
-	private AskAiEvent ParseToolCallEvent(string id, long timestamp, JsonElement innerData)
+	private static AskAiEvent ParseToolCallEvent(string id, long timestamp, JsonElement innerData)
 	{
 		// Extract fields from Agent Builder's tool_call structure
 		var toolCallId = innerData.TryGetProperty("tool_call_id", out var tcId) ? tcId.GetString() : id;
@@ -128,16 +126,13 @@ private AskAiEvent ParseToolCallEvent(string id, long timestamp, JsonElement inn
 		return new AskAiEvent.ToolCall(id, timestamp, toolCallId ?? id, toolId ?? "unknown", args);
 	}
 
-	private AskAiEvent.ErrorEvent ParseErrorEventFromRoot(string id, long timestamp, JsonElement root)
+	private static AskAiEvent.ErrorEvent ParseErrorEventFromRoot(string id, long timestamp, JsonElement root)
 	{
 		// Agent Builder sends: {"error":{"code":"...","message":"...","meta":{...}}}
 		var errorMessage = root.TryGetProperty("error", out var errorProp) &&
 						   errorProp.TryGetProperty("message", out var msgProp)
 			? msgProp.GetString()
 			: null;
-
-		Logger.LogError("Error event received from Agent Builder: {ErrorMessage}", errorMessage ?? "Unknown error");
-
 		return new AskAiEvent.ErrorEvent(id, timestamp, errorMessage ?? "Unknown error occurred");
 	}
 }

src/api/Elastic.Documentation.Api.Infrastructure/Adapters/AskAi/StreamTransformerBase.cs

@@ -21,16 +21,6 @@ namespace Elastic.Documentation.Api.Infrastructure.Adapters.AskAi;
 /// <param name="Data">The accumulated data from all "data:" fields</param>
 public record SseEvent(string? EventType, string Data);
 
-/// <summary>
-/// Holds the results of processing a stream, including metrics for OpenTelemetry
-/// </summary>
-public class StreamProcessingResult
-{
-	public bool HasErrorEvents { get; set; }
-	public List<string> FinishReasons { get; } = [];
-	public List<object> OutputMessages { get; } = [];
-}
-
 /// <summary>
 /// Base class for stream transformers that handles common streaming logic
 /// </summary>
@@ -144,11 +134,17 @@ private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, Activi
 	private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Activity? parentActivity, CancellationToken cancellationToken)
 	{
 		using var activity = StreamTransformerActivitySource.StartActivity("transform_stream");
-		var result = new StreamProcessingResult();
+
+		if (parentActivity?.Id != null)
+			_ = activity?.SetParentId(parentActivity.Id);
+
 		List<MessagePart> outputMessageParts = [];
 		await foreach (var sseEvent in ParseSseEventsAsync(reader, cancellationToken))
 		{
-			AskAiEvent? transformedEvent = null;
+			using var parseActivity = StreamTransformerActivitySource.StartActivity("parse_event");
+			// parseActivity automatically inherits from Activity.Current (transform_stream)
+
+			AskAiEvent? transformedEvent;
 			try
 			{
 				// Parse JSON once in base class
@@ -162,91 +158,79 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 			{
 				Logger.LogError(ex, "Failed to parse JSON from SSE event for transformer {TransformerType}. EventType: {EventType}, Data: {Data}",
 					GetType().Name, sseEvent.EventType, sseEvent.Data);
-				var commonTags = new ActivityTagsCollection([
-					new KeyValuePair<string, object?>("sse.type", sseEvent.EventType),
-					new KeyValuePair<string, object?>("sse.data", sseEvent.Data)
-				]);
-
-				// Add error event to activity for JSON parsing failures
-				_ = activity?.AddEvent(
-					new ActivityEvent(
-						name: "Failed to parse JSON",
-						timestamp: DateTimeOffset.UtcNow,
-						commonTags
-						)
-					);
-				_ = activity?.SetStatus(ActivityStatusCode.Error, "Failed to parse JSON from SSE event");
-				_ = activity?.AddException(ex);
-				throw; // Re-throw to be handled by caller
+				throw;
 			}
 
 			if (transformedEvent == null)
 				continue;
 
+			// Set event type tag on parse_event activity
+			_ = parseActivity?.SetTag("ask_ai.event", transformedEvent.GetType().Name);
+
 			switch (transformedEvent)
 			{
 				case AskAiEvent.ConversationStart conversationStart:
 					{
 						_ = parentActivity?.SetTag("gen_ai.conversation.id", conversationStart.ConversationId);
 						_ = activity?.SetTag("gen_ai.conversation.id", conversationStart.ConversationId);
-						_ = activity?.AddEvent(new ActivityEvent("conversation_start"));
 						break;
 					}
 				case AskAiEvent.Reasoning reasoning:
 					{
-						_ = activity?.AddEvent(new ActivityEvent("conversation_start"));
 						outputMessageParts.Add(new MessagePart("reasoning", reasoning.Message ?? string.Empty));
 						break;
 					}
 				case AskAiEvent.Chunk:
 					{
-						_ = activity?.AddEvent(new ActivityEvent("message_chunk"));
+						// Event type already tagged above
 						break;
 					}
 
 				case AskAiEvent.ErrorEvent errorEvent:
 					{
-						_ = activity?.AddEvent(new ActivityEvent("error"));
 						_ = activity?.SetStatus(ActivityStatusCode.Error, "AI provider error event");
 						_ = activity?.SetTag("error.type", "AIProviderError");
 						_ = activity?.SetTag("error.message", errorEvent.Message);
+						_ = parseActivity?.SetStatus(ActivityStatusCode.Error, errorEvent.Message);
 						break;
 					}
 				case AskAiEvent.ToolCall:
 					{
-						_ = activity?.AddEvent(new ActivityEvent("tool_call"));
+						// Event type already tagged above
 						break;
 					}
 				case AskAiEvent.SearchToolCall searchToolCall:
 					{
-						_ = activity?.AddEvent(new ActivityEvent("search_tool_call", tags: [
-							new KeyValuePair<string, object?>("search.query", searchToolCall.SearchQuery)
-						]));
+						_ = parseActivity?.SetTag("search.query", searchToolCall.SearchQuery);
 						break;
 					}
 				case AskAiEvent.ToolResult toolResult:
 					{
-						_ = activity?.AddEvent(new ActivityEvent("tool_result", tags: [
-							new KeyValuePair<string, object?>("tool.result_summary", toolResult.Result)
-						]));
+						_ = parseActivity?.SetTag("tool.result_summary", toolResult.Result);
 						break;
 					}
 				case AskAiEvent.ChunkComplete chunkComplete:
 					{
-						var parts = outputMessageParts.Prepend(new MessagePart("text", chunkComplete.FullContent)).ToArray();
-						var outputMessages = new OutputMessage("assistant", parts, "stop");
-						_ = activity?.SetTag("gen_ai.output.messages", JsonSerializer.Serialize(outputMessages, ApiJsonContext.Default.OutputMessage));
-						_ = activity?.AddEvent(new ActivityEvent("message_chunk_complete"));
+						outputMessageParts.Add(new MessagePart("text", chunkComplete.FullContent));
 						break;
 					}
 				case AskAiEvent.ConversationEnd:
 					{
-						_ = activity?.AddEvent(new ActivityEvent("conversation_end"));
+						// Event type already tagged above
 						break;
 					}
 			}
 			await WriteEventAsync(transformedEvent, writer, cancellationToken);
 		}
+
+		// Set output messages tag once after all events are processed
+		if (outputMessageParts.Count > 0)
+		{
+			var outputMessages = new OutputMessage("assistant", outputMessageParts.ToArray(), "stop");
+			var outputMessagesJson = JsonSerializer.Serialize(outputMessages, ApiJsonContext.Default.OutputMessage);
+			_ = parentActivity?.SetTag("gen_ai.output.messages", outputMessagesJson);
+			_ = activity?.SetTag("gen_ai.output.messages", outputMessagesJson);
+		}
 	}
 
 	/// <summary>
@@ -261,14 +245,10 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 	/// <summary>
 	/// Write a transformed event to the output stream
 	/// </summary>
-	protected async Task WriteEventAsync(AskAiEvent? transformedEvent, PipeWriter writer, CancellationToken cancellationToken)
+	private async Task WriteEventAsync(AskAiEvent? transformedEvent, PipeWriter writer, CancellationToken cancellationToken)
 	{
 		if (transformedEvent == null)
 			return;
-
-		// Don't create spans for each token - too noisy and expensive
-		// Just track metrics at the stream level
-
 		try
 		{
 			// Serialize as base AskAiEvent type to include the type discriminator
@@ -292,26 +272,17 @@ protected async Task WriteEventAsync(AskAiEvent? transformedEvent, PipeWriter wr
 	/// Parse Server-Sent Events (SSE) from a PipeReader following the W3C SSE specification.
 	/// This method handles the standard SSE format with event:, data:, and comment lines.
 	/// </summary>
-	protected async IAsyncEnumerable<SseEvent> ParseSseEventsAsync(
+	private static async IAsyncEnumerable<SseEvent> ParseSseEventsAsync(
 		PipeReader reader,
 		[EnumeratorCancellation] CancellationToken cancellationToken)
 	{
-		using var activity = StreamTransformerActivitySource.StartActivity("gen_ai.agent.parse");
-		_ = (activity?.SetTag("gen_ai.agent.name", GetAgentId()));
-		_ = (activity?.SetTag("gen_ai.provider.name", GetAgentProvider()));
-
 		string? currentEvent = null;
 		var dataBuilder = new StringBuilder();
-		var eventsParsed = 0;
-		var readOperations = 0;
-		var totalBytesRead = 0L;
 
 		while (!cancellationToken.IsCancellationRequested)
 		{
-			readOperations++;
 			var result = await reader.ReadAsync(cancellationToken);
 			var buffer = result.Buffer;
-			totalBytesRead += buffer.Length;
 
 			// Process all complete lines in the buffer
 			while (TryReadLine(ref buffer, out var line))
@@ -322,47 +293,33 @@ protected async IAsyncEnumerable<SseEvent> ParseSseEventsAsync(
 
 				// Event type line
 				if (line.StartsWith("event:", StringComparison.Ordinal))
-				{
-					currentEvent = line.Substring(6).Trim();
-				}
+					currentEvent = line[6..].Trim();
 				// Data line
 				else if (line.StartsWith("data:", StringComparison.Ordinal))
-				{
-					_ = dataBuilder.Append(line.Substring(5).Trim());
-				}
+					_ = dataBuilder.Append(line[5..].Trim());
 				// Empty line - marks end of event
 				else if (string.IsNullOrEmpty(line))
 				{
-					if (dataBuilder.Length > 0)
-					{
-						eventsParsed++;
-						yield return new SseEvent(currentEvent, dataBuilder.ToString());
-						currentEvent = null;
-						_ = dataBuilder.Clear();
-					}
+					if (dataBuilder.Length <= 0)
+						continue;
+					yield return new SseEvent(currentEvent, dataBuilder.ToString());
+					currentEvent = null;
+					_ = dataBuilder.Clear();
 				}
 			}
 
 			// Tell the PipeReader how much of the buffer we consumed
 			reader.AdvanceTo(buffer.Start, buffer.End);
 
 			// Stop reading if there's no more data coming
-			if (result.IsCompleted)
-			{
-				// Yield any remaining event that hasn't been terminated with an empty line
-				if (dataBuilder.Length > 0)
-				{
-					eventsParsed++;
-					yield return new SseEvent(currentEvent, dataBuilder.ToString());
-				}
-				break;
-			}
-		}
+			if (!result.IsCompleted)
+				continue;
 
-		// Set metrics on the activity using GenAI conventions
-		_ = (activity?.SetTag("gen_ai.response.sse_event_count", eventsParsed));
-		_ = (activity?.SetTag("gen_ai.response.bytes_received", totalBytesRead));
-		_ = (activity?.SetTag("gen_ai.response.read_operations", readOperations));
+			// Yield any remaining event that hasn't been terminated with an empty line
+			if (dataBuilder.Length > 0)
+				yield return new SseEvent(currentEvent, dataBuilder.ToString());
+			break;
+		}
 	}
 
 	/// <summary>

src/api/Elastic.Documentation.Api.Infrastructure/OpenTelemetryExtensions.cs

@@ -5,6 +5,7 @@
 using Elastic.OpenTelemetry;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
 using OpenTelemetry;
 using OpenTelemetry.Instrumentation.AspNetCore;
 using OpenTelemetry.Instrumentation.Http;
@@ -13,7 +14,7 @@
 
 namespace Elastic.Documentation.Api.Infrastructure;
 
-public static class OpenTelemetryExtensions
+public static partial class OpenTelemetryExtensions
 {
 	/// <summary>
 	/// Configures Elastic OpenTelemetry (EDOT) for the Docs API.