Fix the ConversationID handling of LLMGateway

Author: reakaleek

docs-builder.sln

@@ -2,7 +2,7 @@
 import { AiProviderSelector } from './AiProviderSelector'
 import { AskAiSuggestions } from './AskAiSuggestions'
 import { ChatMessageList } from './ChatMessageList'
-import { useChatActions, useChatMessages } from './chat.store'
+import {useChatActions, useChatMessages, useConversationId} from './chat.store'
 import {
     useEuiOverflowScroll,
     EuiButtonEmpty,
@@ -68,6 +68,7 @@ export const Chat = () => {
     const scrollRef = useRef<HTMLDivElement>(null)
     const lastMessageStatusRef = useRef<string | null>(null)
     const [inputValue, setInputValue] = useState('')
+    const conversationId = useConversationId()
 
     const dynamicScrollableStyles = css`
         ${scrollableStyles}
@@ -124,6 +125,7 @@ export const Chat = () => {
             css={containerStyles}
         >
             <EuiSpacer size="m" />
+            ConversationId: {conversationId}
 
             {messages.length > 0 && (
                 <NewConversationHeader onClick={clearChat} />

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/Chat.tsx

@@ -129,7 +129,7 @@ private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, string
 	/// 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, string? conversationId, Activity? parentActivity, CancellationToken cancellationToken)
+	protected virtual async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, string? conversationId, Activity? parentActivity, CancellationToken cancellationToken)
 	{
 		using var activity = StreamTransformerActivitySource.StartActivity(nameof(ProcessStreamAsync));
 
@@ -255,7 +255,7 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, stri
 	/// <summary>
 	/// Write a transformed event to the output stream
 	/// </summary>
-	private async Task WriteEventAsync(AskAiEvent? transformedEvent, PipeWriter writer, CancellationToken cancellationToken)
+	protected async Task WriteEventAsync(AskAiEvent? transformedEvent, PipeWriter writer, CancellationToken cancellationToken)
 	{
 		if (transformedEvent == null)
 			return;

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

@@ -68,7 +68,7 @@ public static LlmGatewayRequest CreateFromRequest(AskAiRequest request) =>
 				new ChatInput("user", AskAiRequest.SystemPrompt),
 				new ChatInput("user", request.Message)
 			],
-			ThreadId: request.ConversationId ?? "elastic-docs-" + Guid.NewGuid()
+			ThreadId: request.ConversationId ?? Guid.NewGuid().ToString()
 		);
 }
 

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

@@ -2,6 +2,8 @@
 // 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 System.IO.Pipelines;
 using System.Text.Json;
 using Elastic.Documentation.Api.Core.AskAi;
 using Microsoft.Extensions.Logging;
@@ -15,6 +17,36 @@ public class LlmGatewayStreamTransformer(ILogger<LlmGatewayStreamTransformer> lo
 {
 	protected override string GetAgentId() => LlmGatewayAskAiGateway.ModelName;
 	protected override string GetAgentProvider() => LlmGatewayAskAiGateway.ProviderName;
+
+	/// <summary>
+	/// Override to emit ConversationStart event when conversationId is null (new conversation)
+	/// </summary>
+	protected override async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, string? conversationId, Activity? parentActivity, CancellationToken cancellationToken)
+	{
+		// If conversationId is null, generate a new one and emit ConversationStart event
+		// This matches the ThreadId format used in LlmGatewayAskAiGateway
+		var actualConversationId = conversationId;
+		if (conversationId == null)
+		{
+			actualConversationId = Guid.NewGuid().ToString();
+			var timestamp = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
+			var conversationStartEvent = new AskAiEvent.ConversationStart(
+				Id: Guid.NewGuid().ToString(),
+				Timestamp: timestamp,
+				ConversationId: actualConversationId
+			);
+
+			// Set activity tags for the new conversation
+			_ = parentActivity?.SetTag("gen_ai.conversation.id", actualConversationId);
+			Logger.LogDebug("LLM Gateway conversation started: {ConversationId}", actualConversationId);
+
+			// Write the ConversationStart event to the stream
+			await WriteEventAsync(conversationStartEvent, writer, cancellationToken);
+		}
+
+		// Continue with normal stream processing using the actual conversation ID
+		await base.ProcessStreamAsync(reader, writer, actualConversationId, parentActivity, cancellationToken);
+	}
 	protected override AskAiEvent? TransformJsonEvent(string? conversationId, string? eventType, JsonElement json)
 	{
 		// LLM Gateway format: ["custom", {type: "...", ...}]