Fix setting conversationId

Author: reakaleek

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

@@ -4,7 +4,6 @@
 
 using System.Diagnostics;
 using System.Text.Json;
-using Elastic.Documentation.Api.Core;
 using Microsoft.Extensions.Logging;
 
 namespace Elastic.Documentation.Api.Core.AskAi;
@@ -23,6 +22,8 @@ public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx)
 		_ = activity?.SetTag("gen_ai.operation.name", "chat");
 		_ = activity?.SetTag("gen_ai.provider.name", streamTransformer.AgentProvider); // agent-builder or llm-gateway
 		_ = activity?.SetTag("gen_ai.agent.id", streamTransformer.AgentId); // docs-agent or docs_assistant
+		if (askAiRequest.ThreadId is not null)
+			_ = activity?.SetTag("gen_ai.conversation.id", askAiRequest.ThreadId);
 		var inputMessages = new[]
 		{
 			new InputMessage("user", [new MessagePart("text", askAiRequest.Message)])
@@ -33,7 +34,7 @@ public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx)
 		logger.LogInformation("Streaming AskAI response");
 		var rawStream = await askAiGateway.AskAi(askAiRequest, ctx);
 		// The stream transformer will handle disposing the activity when streaming completes
-		var transformedStream = await streamTransformer.TransformAsync(rawStream, activity, ctx);
+		var transformedStream = await streamTransformer.TransformAsync(rawStream, askAiRequest.ThreadId, activity, ctx);
 		return transformedStream;
 	}
 }

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

@@ -23,8 +23,9 @@ public interface IStreamTransformer
 	/// Transforms a raw SSE stream into a stream of AskAiEvent objects
 	/// </summary>
 	/// <param name="rawStream">Raw SSE stream from gateway (Agent Builder, LLM Gateway, etc.)</param>
+	/// <param name="threadId">Thread/conversation ID (if known)</param>
 	/// <param name="parentActivity">Parent activity to track the streaming operation (will be disposed when stream completes)</param>
 	/// <param name="cancellationToken">Cancellation token</param>
 	/// <returns>Stream containing SSE-formatted AskAiEvent objects</returns>
-	Task<Stream> TransformAsync(Stream rawStream, System.Diagnostics.Activity? parentActivity, CancellationToken cancellationToken = default);
+	Task<Stream> TransformAsync(Stream rawStream, string? threadId, System.Diagnostics.Activity? parentActivity, CancellationToken cancellationToken = default);
 }

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

@@ -18,7 +18,7 @@ public class AgentBuilderStreamTransformer(ILogger<AgentBuilderStreamTransformer
 {
 	protected override string GetAgentId() => AgentBuilderAskAiGateway.ModelName;
 	protected override string GetAgentProvider() => AgentBuilderAskAiGateway.ProviderName;
-	protected override AskAiEvent? TransformJsonEvent(string? eventType, JsonElement json)
+	protected override AskAiEvent? TransformJsonEvent(string? threadId, string? eventType, JsonElement json)
 	{
 		var type = eventType ?? "message";
 		var timestamp = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();

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

@@ -18,7 +18,7 @@ public class LlmGatewayStreamTransformer(ILogger<LlmGatewayStreamTransformer> lo
 {
 	protected override string GetAgentId() => LlmGatewayAskAiGateway.ModelName;
 	protected override string GetAgentProvider() => LlmGatewayAskAiGateway.ProviderName;
-	protected override AskAiEvent? TransformJsonEvent(string? eventType, JsonElement json)
+	protected override AskAiEvent? TransformJsonEvent(string? threadId, string? eventType, JsonElement json)
 	{
 		// LLM Gateway format: ["custom", {type: "...", ...}]
 		if (json.ValueKind != JsonValueKind.Array || json.GetArrayLength() < 2)
@@ -34,12 +34,13 @@ public class LlmGatewayStreamTransformer(ILogger<LlmGatewayStreamTransformer> lo
 		var id = message.GetProperty("id").GetString()!;
 		var messageData = message.GetProperty("data");
 
+		// LLM gateway does not emit conversation start events with thread IDs
+		// so we create a synthetic conversation start event here
+		if (threadId is null)
+			return new AskAiEvent.ConversationStart(id, timestamp, Guid.NewGuid().ToString());
+
 		return type switch
 		{
-			"agent_start" =>
-				// LLM Gateway doesn't provide conversation ID, so generate one
-				new AskAiEvent.ConversationStart(id, timestamp, Guid.NewGuid().ToString()),
-
 			"ai_message_chunk" when messageData.TryGetProperty("content", out var content) =>
 				new AskAiEvent.MessageChunk(id, timestamp, content.GetString()!),
 

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

@@ -51,7 +51,7 @@ public abstract class StreamTransformerBase(ILogger logger) : IStreamTransformer
 	/// </summary>
 	public string AgentProvider => GetAgentProvider();
 
-	public Task<Stream> TransformAsync(Stream rawStream, Activity? parentActivity, Cancel cancellationToken = default)
+	public Task<Stream> TransformAsync(Stream rawStream, string? threadId, Activity? parentActivity, Cancel cancellationToken = default)
 	{
 		// Configure pipe for low-latency streaming
 		var pipeOptions = new PipeOptions(
@@ -70,7 +70,7 @@ public Task<Stream> TransformAsync(Stream rawStream, Activity? parentActivity, C
 		// 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
 		// Pass parent activity - it will be disposed when streaming completes
-		_ = ProcessPipeAsync(reader, pipe.Writer, parentActivity, cancellationToken);
+		_ = ProcessPipeAsync(reader, pipe.Writer, threadId, parentActivity, cancellationToken);
 
 		// Return the read side of the pipe as a stream
 		return Task.FromResult(pipe.Reader.AsStream());
@@ -80,13 +80,13 @@ public Task<Stream> TransformAsync(Stream rawStream, Activity? parentActivity, C
 	/// 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, Activity? parentActivity, CancellationToken cancellationToken)
+	private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, string? threadId, Activity? parentActivity, CancellationToken cancellationToken)
 	{
 		try
 		{
 			try
 			{
-				await ProcessStreamAsync(reader, writer, parentActivity, cancellationToken);
+				await ProcessStreamAsync(reader, writer, threadId, parentActivity, cancellationToken);
 			}
 			catch (OperationCanceledException ex)
 			{
@@ -131,7 +131,7 @@ private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, Activi
 	/// Default implementation parses SSE events and JSON, then calls TransformJsonEvent.
 	/// </summary>
 	/// <returns>Stream processing result with metrics and captured output</returns>
-	private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Activity? parentActivity, CancellationToken cancellationToken)
+	private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, string? threadId, Activity? parentActivity, CancellationToken cancellationToken)
 	{
 		using var activity = StreamTransformerActivitySource.StartActivity(nameof(ProcessStreamAsync));
 
@@ -141,7 +141,6 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 		List<MessagePart> outputMessageParts = [];
 		await foreach (var sseEvent in ParseSseEventsAsync(reader, cancellationToken))
 		{
-			using var parseActivity = StreamTransformerActivitySource.StartActivity("AskAI Event");
 			AskAiEvent? transformedEvent;
 			try
 			{
@@ -150,7 +149,7 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 				var root = doc.RootElement;
 
 				// Subclass transforms JsonElement to AskAiEvent
-				transformedEvent = TransformJsonEvent(sseEvent.EventType, root);
+				transformedEvent = TransformJsonEvent(threadId, sseEvent.EventType, root);
 			}
 			catch (JsonException ex)
 			{
@@ -166,9 +165,11 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 				continue;
 			}
 
+			using var parseActivity = StreamTransformerActivitySource.StartActivity("AskAI Event");
+
 			// Set event type tag on parse_event activity
 			_ = parseActivity?.SetTag("ask_ai.event.type", transformedEvent.GetType().Name);
-			_ = parseActivity?.SetTag("gen_ai.respone.id", transformedEvent.Id);
+			_ = parseActivity?.SetTag("gen_ai.response.id", transformedEvent.Id);
 
 			switch (transformedEvent)
 			{
@@ -247,10 +248,11 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 	/// Transform a parsed JSON event into an AskAiEvent.
 	/// Subclasses implement provider-specific transformation logic.
 	/// </summary>
+	/// <param name="threadId">The conversation/thread ID, if available</param>
 	/// <param name="eventType">The SSE event type (from "event:" field), or null if not present</param>
 	/// <param name="json">The parsed JSON data from the "data:" field</param>
 	/// <returns>The transformed AskAiEvent, or null to skip this event</returns>
-	protected abstract AskAiEvent? TransformJsonEvent(string? eventType, JsonElement json);
+	protected abstract AskAiEvent? TransformJsonEvent(string? threadId, string? eventType, JsonElement json);
 
 	/// <summary>
 	/// Write a transformed event to the output stream