Add support for hybrid lambda architecture for both buffered and streaming responses

src/api/Elastic.Documentation.Api.Core/AskAi/AskAiUsecase.cs

@@ -2,6 +2,7 @@
 // Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
 // See the LICENSE file in the project root for more information
 
+using System.Diagnostics;
 using Microsoft.Extensions.Logging;
 
 namespace Elastic.Documentation.Api.Core.AskAi;
@@ -11,11 +12,35 @@ public class AskAiUsecase(
 	IStreamTransformer streamTransformer,
 	ILogger<AskAiUsecase> logger)
 {
+	private static readonly ActivitySource AskAiActivitySource = new("Elastic.Documentation.Api.AskAi");
+
 	public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx)
 	{
+		using var activity = AskAiActivitySource.StartActivity("gen_ai.agent");
+
+		// We'll determine the actual agent name after we know which provider is being used
+		_ = (activity?.SetTag("gen_ai.request.input", askAiRequest.Message));
+		_ = (activity?.SetTag("gen_ai.request.conversation_id", askAiRequest.ThreadId ?? "new-conversation"));
+
+		// Add GenAI inference operation details event
+		_ = (activity?.AddEvent(new ActivityEvent("gen_ai.client.inference.operation.details",
+			timestamp: DateTimeOffset.UtcNow,
+			tags:
+			[
+				new KeyValuePair<string, object?>("gen_ai.operation.name", "chat"),
+				new KeyValuePair<string, object?>("gen_ai.conversation.id", askAiRequest.ThreadId ?? "pending"), // Will be updated when we receive ConversationStart
+				new KeyValuePair<string, object?>("gen_ai.input.messages", $"[{{\"role\":\"user\",\"content\":\"{askAiRequest.Message}\"}}]"),
+				new KeyValuePair<string, object?>("gen_ai.system_instructions", $"[{{\"type\":\"text\",\"content\":\"{AskAiRequest.SystemPrompt}\"}}]")
+			])));
+
 		logger.LogDebug("Processing AskAiRequest: {Request}", askAiRequest);
+
 		var rawStream = await askAiGateway.AskAi(askAiRequest, ctx);
-		return await streamTransformer.TransformAsync(rawStream, ctx);
+
+		// The stream transformer will set the correct agent name, model name and provider
+		var transformedStream = await streamTransformer.TransformAsync(rawStream, ctx);
+
+		return transformedStream;
 	}
 }
 

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

@@ -15,6 +15,15 @@ namespace Elastic.Documentation.Api.Infrastructure.Adapters.AskAi;
 
 public class AgentBuilderAskAiGateway(HttpClient httpClient, IParameterProvider parameterProvider, ILogger<AgentBuilderAskAiGateway> logger) : IAskAiGateway<Stream>
 {
+	/// <summary>
+	/// Model name used by Agent Builder (from AgentId)
+	/// </summary>
+	public const string ModelName = "docs-agent";
+
+	/// <summary>
+	/// Provider name for tracing
+	/// </summary>
+	public const string ProviderName = "agent-builder";
 	public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx = default)
 	{
 		// Only include conversation_id if threadId is provided (subsequent requests)

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

@@ -16,6 +16,8 @@ namespace Elastic.Documentation.Api.Infrastructure.Adapters.AskAi;
 /// </summary>
 public class AgentBuilderStreamTransformer(ILogger<AgentBuilderStreamTransformer> logger) : StreamTransformerBase(logger)
 {
+	protected override string GetAgentId() => AgentBuilderAskAiGateway.ModelName;
+	protected override string GetAgentProvider() => AgentBuilderAskAiGateway.ProviderName;
 	protected override AskAiEvent? TransformJsonEvent(string? eventType, JsonElement json)
 	{
 		var type = eventType ?? "message";

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

@@ -12,6 +12,15 @@ namespace Elastic.Documentation.Api.Infrastructure.Adapters.AskAi;
 
 public class LlmGatewayAskAiGateway(HttpClient httpClient, GcpIdTokenProvider tokenProvider, LlmGatewayOptions options) : IAskAiGateway<Stream>
 {
+	/// <summary>
+	/// Model name used by LLM Gateway (from PlatformContext.UseCase)
+	/// </summary>
+	public const string ModelName = "docs_assistant";
+
+	/// <summary>
+	/// Provider name for tracing
+	/// </summary>
+	public const string ProviderName = "llm-gateway";
 	public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx = default)
 	{
 		var llmGatewayRequest = LlmGatewayRequest.CreateFromRequest(askAiRequest);

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

@@ -16,6 +16,8 @@ namespace Elastic.Documentation.Api.Infrastructure.Adapters.AskAi;
 /// </summary>
 public class LlmGatewayStreamTransformer(ILogger<LlmGatewayStreamTransformer> logger) : StreamTransformerBase(logger)
 {
+	protected override string GetAgentId() => LlmGatewayAskAiGateway.ModelName;
+	protected override string GetAgentProvider() => LlmGatewayAskAiGateway.ProviderName;
 	protected override AskAiEvent? TransformJsonEvent(string? eventType, JsonElement json)
 	{
 		// LLM Gateway format: ["custom", {type: "...", ...}]

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

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information
 
 using System.Buffers;
+using System.Diagnostics;
 using System.IO.Pipelines;
 using System.Runtime.CompilerServices;
 using System.Text;
@@ -26,15 +27,42 @@ public abstract class StreamTransformerBase(ILogger logger) : IStreamTransformer
 {
 	protected ILogger Logger { get; } = logger;
 
+	// ActivitySource for tracing streaming operations
+	private static readonly ActivitySource StreamTransformerActivitySource = new("Elastic.Documentation.Api.StreamTransformer");
+
+	/// <summary>
+	/// Get the agent ID for this transformer
+	/// </summary>
+	protected abstract string GetAgentId();
+
+	/// <summary>
+	/// Get the agent provider/platform for this transformer
+	/// </summary>
+	protected abstract string GetAgentProvider();
+
 	public Task<Stream> TransformAsync(Stream rawStream, CancellationToken cancellationToken = default)
 	{
-		var pipe = new Pipe();
+		using var activity = StreamTransformerActivitySource.StartActivity("gen_ai.agent");
+		_ = (activity?.SetTag("gen_ai.agent.name", GetAgentId()));
+		_ = (activity?.SetTag("gen_ai.provider.name", GetAgentProvider()));
+
+		// Configure pipe for low-latency streaming
+		var pipeOptions = new PipeOptions(
+			minimumSegmentSize: 1024, // Smaller segments for faster processing
+			pauseWriterThreshold: 64 * 1024, // 64KB high water mark
+			resumeWriterThreshold: 32 * 1024, // 32KB low water mark
+			readerScheduler: PipeScheduler.Inline,
+			writerScheduler: PipeScheduler.Inline,
+			useSynchronizationContext: false
+		);
+
+		var pipe = new Pipe(pipeOptions);
 		var reader = PipeReader.Create(rawStream);
 
 		// Start processing task to transform and write events to pipe
 		// Note: We intentionally don't await this task as we need to return the stream immediately
 		// The pipe handles synchronization and backpressure between producer and consumer
-		_ = ProcessPipeAsync(reader, pipe.Writer, cancellationToken);
+		_ = ProcessPipeAsync(reader, pipe.Writer, activity, cancellationToken);
 
 		// Return the read side of the pipe as a stream
 		return Task.FromResult(pipe.Reader.AsStream());
@@ -44,16 +72,21 @@ public Task<Stream> TransformAsync(Stream rawStream, CancellationToken cancellat
 	/// Process the pipe reader and write transformed events to the pipe writer.
 	/// This runs concurrently with the consumer reading from the output stream.
 	/// </summary>
-	private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, CancellationToken cancellationToken)
+	private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, Activity? parentActivity, CancellationToken cancellationToken)
 	{
+		using var activity = StreamTransformerActivitySource.StartActivity("StreamTransformer.ProcessPipeAsync");
+		_ = (activity?.SetTag("transformer.type", GetType().Name));
+
 		try
 		{
-			await ProcessStreamAsync(reader, writer, cancellationToken);
+			await ProcessStreamAsync(reader, writer, parentActivity, cancellationToken);
 		}
 		catch (OperationCanceledException ex)
 		{
 			// Cancellation is expected and not an error - log as debug
 			Logger.LogDebug("Stream processing was cancelled.");
+			_ = (activity?.SetTag("gen_ai.response.error", true));
+			_ = (activity?.SetTag("gen_ai.response.error_type", "OperationCanceledException"));
 			try
 			{
 				await writer.CompleteAsync(ex);
@@ -68,6 +101,9 @@ private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, Cancel
 		catch (Exception ex)
 		{
 			Logger.LogError(ex, "Error transforming stream. Stream processing will be terminated.");
+			_ = (activity?.SetTag("gen_ai.response.error", true));
+			_ = (activity?.SetTag("gen_ai.response.error_type", ex.GetType().Name));
+			_ = (activity?.SetTag("gen_ai.response.error_message", ex.Message));
 			try
 			{
 				await writer.CompleteAsync(ex);
@@ -96,10 +132,18 @@ private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, Cancel
 	/// Process the raw stream and write transformed events to the pipe writer.
 	/// Default implementation parses SSE events and JSON, then calls TransformJsonEvent.
 	/// </summary>
-	protected virtual async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, CancellationToken cancellationToken)
+	protected virtual async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Activity? parentActivity, CancellationToken cancellationToken)
 	{
+		using var activity = StreamTransformerActivitySource.StartActivity("gen_ai.agent.stream");
+		_ = (activity?.SetTag("gen_ai.agent.name", GetAgentId()));
+		_ = (activity?.SetTag("gen_ai.provider.name", GetAgentProvider()));
+
+		var eventCount = 0;
+		var jsonParseErrors = 0;
+
 		await foreach (var sseEvent in ParseSseEventsAsync(reader, cancellationToken))
 		{
+			eventCount++;
 			AskAiEvent? transformedEvent = null;
 
 			try
@@ -113,14 +157,29 @@ protected virtual async Task ProcessStreamAsync(PipeReader reader, PipeWriter wr
 			}
 			catch (JsonException ex)
 			{
+				jsonParseErrors++;
 				Logger.LogError(ex, "Failed to parse JSON from SSE event: {Data}", sseEvent.Data);
 			}
 
 			if (transformedEvent != null)
 			{
+				// Update parent activity with conversation ID and model info when we receive ConversationStart events
+				if (transformedEvent is AskAiEvent.ConversationStart conversationStart)
+				{
+					_ = (parentActivity?.SetTag("gen_ai.conversation.id", conversationStart.ConversationId));
+					_ = (parentActivity?.SetTag("gen_ai.request.model", GetAgentId()));
+					_ = (parentActivity?.SetTag("gen_ai.agent.name", GetAgentId()));
+					_ = (parentActivity?.SetTag("gen_ai.provider.name", GetAgentProvider()));
+					_ = (activity?.SetTag("gen_ai.conversation.id", conversationStart.ConversationId));
+				}
+
 				await WriteEventAsync(transformedEvent, writer, cancellationToken);
 			}
 		}
+
+		// Set metrics on the activity using GenAI conventions
+		_ = (activity?.SetTag("gen_ai.response.token_count", eventCount));
+		_ = (activity?.SetTag("gen_ai.response.error_count", jsonParseErrors));
 	}
 
 	/// <summary>
@@ -140,11 +199,28 @@ protected async Task WriteEventAsync(AskAiEvent? transformedEvent, PipeWriter wr
 		if (transformedEvent == null)
 			return;
 
+		using var activity = StreamTransformerActivitySource.StartActivity("gen_ai.agent.token");
+		_ = (activity?.SetTag("gen_ai.agent.name", GetAgentId()));
+		_ = (activity?.SetTag("gen_ai.provider.name", GetAgentProvider()));
+		_ = (activity?.SetTag("gen_ai.response.token_type", transformedEvent.GetType().Name));
+
+		// Add GenAI response event for each token/chunk
+		_ = (activity?.AddEvent(new ActivityEvent("gen_ai.client.inference.operation.details",
+			timestamp: DateTimeOffset.UtcNow,
+			tags:
+			[
+				new KeyValuePair<string, object?>("gen_ai.operation.name", "chat"),
+				new KeyValuePair<string, object?>("gen_ai.response.model", GetAgentId()),
+				new KeyValuePair<string, object?>("gen_ai.output.messages", JsonSerializer.Serialize(transformedEvent, AskAiEventJsonContext.Default.AskAiEvent))
+			])));
+
 		// Serialize as base AskAiEvent type to include the type discriminator
 		var json = JsonSerializer.Serialize<AskAiEvent>(transformedEvent, AskAiEventJsonContext.Default.AskAiEvent);
 		var sseData = $"data: {json}\n\n";
 		var bytes = Encoding.UTF8.GetBytes(sseData);
 
+		_ = (activity?.SetTag("gen_ai.response.token_size", bytes.Length));
+
 		// Write to pipe and flush immediately for real-time streaming
 		_ = await writer.WriteAsync(bytes, cancellationToken);
 		_ = await writer.FlushAsync(cancellationToken);
@@ -158,13 +234,22 @@ protected 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))
@@ -188,6 +273,7 @@ protected async IAsyncEnumerable<SseEvent> ParseSseEventsAsync(
 				{
 					if (dataBuilder.Length > 0)
 					{
+						eventsParsed++;
 						yield return new SseEvent(currentEvent, dataBuilder.ToString());
 						currentEvent = null;
 						_ = dataBuilder.Clear();
@@ -204,11 +290,17 @@ protected async IAsyncEnumerable<SseEvent> ParseSseEventsAsync(
 				// 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;
 			}
 		}
+
+		// Set metrics on the activity using GenAI conventions
+		_ = (activity?.SetTag("gen_ai.response.token_count", eventsParsed));
+		_ = (activity?.SetTag("gen_ai.request.input_size", totalBytesRead));
+		_ = (activity?.SetTag("gen_ai.model.operation_count", readOperations));
 	}
 
 	/// <summary>

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

@@ -27,7 +27,16 @@
 		_ = askAiGroup.MapPost("/stream", async (AskAiRequest askAiRequest, AskAiUsecase askAiUsecase, Cancel ctx) =>
 		{
 			var stream = await askAiUsecase.AskAi(askAiRequest, ctx);
-			return Results.Stream(stream, "text/event-stream");
+
+			// Configure response headers for optimal streaming
+			var response = Results.Stream(stream, "text/event-stream");
+
+			// Add headers to prevent buffering
+			response.Response.Headers.Add("Cache-Control", "no-cache");
+			response.Response.Headers.Add("Connection", "keep-alive");
+			response.Response.Headers.Add("X-Accel-Buffering", "no"); // Disable nginx buffering
+
+			return response;
 		});
 	}
 

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

@@ -58,7 +58,14 @@
 		{
 			options.SerializerOptions.TypeInfoResolverChain.Insert(0, ApiJsonContext.Default);
 		});
-		_ = services.AddHttpClient();
+
+		// Configure HttpClient for streaming optimization
+		_ = services.AddHttpClient(client =>
+		{
+			// Disable response buffering for streaming
+			client.DefaultRequestHeaders.Connection.Add("keep-alive");
+			client.Timeout = TimeSpan.FromMinutes(10); // Longer timeout for streaming
+		});
 		// Register AppEnvironment as a singleton for dependency injection
 		_ = services.AddSingleton(new AppEnvironment { Current = appEnv });
 		AddParameterProvider(services, appEnv);