fixes

Author: reakaleek

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

@@ -16,15 +16,24 @@ public class AskAiUsecase(
 
 	public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx)
 	{
-		// Start with a placeholder model name - will be updated by transformer
-		using var activity = AskAiActivitySource.StartActivity("chat unknown", ActivityKind.Client);
+		// Start activity for the chat request - model name will be set by transformer
+		using var activity = AskAiActivitySource.StartActivity("chat", ActivityKind.Client);
 
-		// Set required GenAI semantic convention attributes
+		// Generate a correlation ID for tracking if this is a new conversation
+		// For first messages (no ThreadId), generate a temporary ID that will be updated when the provider responds
+		var correlationId = askAiRequest.ThreadId ?? $"temp-{Guid.NewGuid()}";
+
+		// Set GenAI semantic convention attributes
 		_ = (activity?.SetTag("gen_ai.operation.name", "chat"));
-		_ = (activity?.SetTag("gen_ai.request.model", "unknown")); // Will be updated by transformer
-		_ = (activity?.SetTag("gen_ai.conversation.id", askAiRequest.ThreadId ?? "pending")); // Will be updated when we receive ConversationStart
+		_ = (activity?.SetTag("gen_ai.conversation.id", correlationId)); // Will be updated when we receive ConversationStart with actual ID
 		_ = (activity?.SetTag("gen_ai.usage.input_tokens", askAiRequest.Message.Length)); // Approximate token count
 
+		// Custom attributes for tracking our abstraction layer
+		// We use custom attributes because we don't know the actual GenAI provider (OpenAI, Anthropic, etc.)
+		// or model (gpt-4, claude, etc.) - those are abstracted by AgentBuilder/LlmGateway
+		_ = (activity?.SetTag("docs.ai.gateway", streamTransformer.AgentProvider)); // agent-builder or llm-gateway
+		_ = (activity?.SetTag("docs.ai.agent_name", streamTransformer.AgentId)); // docs-agent or docs_assistant
+
 		// Add GenAI prompt event
 		_ = (activity?.AddEvent(new ActivityEvent("gen_ai.content.prompt",
 			timestamp: DateTimeOffset.UtcNow,

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

@@ -9,6 +9,16 @@ namespace Elastic.Documentation.Api.Core.AskAi;
 /// </summary>
 public interface IStreamTransformer
 {
+	/// <summary>
+	/// Get the agent/model identifier for this transformer
+	/// </summary>
+	string AgentId { get; }
+
+	/// <summary>
+	/// Get the agent provider/platform for this transformer
+	/// </summary>
+	string AgentProvider { get; }
+
 	/// <summary>
 	/// Transforms a raw SSE stream into a stream of AskAiEvent objects
 	/// </summary>

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

@@ -40,13 +40,24 @@ public abstract class StreamTransformerBase(ILogger logger) : IStreamTransformer
 	/// </summary>
 	protected abstract string GetAgentProvider();
 
+	/// <summary>
+	/// Public property to expose agent ID (implements IStreamTransformer)
+	/// </summary>
+	public string AgentId => GetAgentId();
+
+	/// <summary>
+	/// Public property to expose agent provider (implements IStreamTransformer)
+	/// </summary>
+	public string AgentProvider => GetAgentProvider();
+
 	public Task<Stream> TransformAsync(Stream rawStream, CancellationToken cancellationToken = default)
 	{
 		using var activity = StreamTransformerActivitySource.StartActivity($"chat {GetAgentId()}", ActivityKind.Client);
 		_ = (activity?.SetTag("gen_ai.operation.name", "chat"));
-		_ = (activity?.SetTag("gen_ai.request.model", GetAgentId()));
-		_ = (activity?.SetTag("gen_ai.agent.name", GetAgentId()));
-		_ = (activity?.SetTag("gen_ai.provider.name", GetAgentProvider()));
+
+		// Custom attributes for tracking our abstraction layer
+		_ = (activity?.SetTag("docs.ai.gateway", GetAgentProvider()));
+		_ = (activity?.SetTag("docs.ai.agent_name", GetAgentId()));
 
 		// Configure pipe for low-latency streaming
 		var pipeOptions = new PipeOptions(
@@ -160,8 +171,10 @@ private async Task ProcessPipeAsync(PipeReader reader, PipeWriter writer, Activi
 	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()));
+
+		// Custom attributes for tracking our abstraction layer
+		_ = (activity?.SetTag("docs.ai.gateway", GetAgentProvider()));
+		_ = (activity?.SetTag("docs.ai.agent_name", GetAgentId()));
 
 		var eventCount = 0;
 		var jsonParseErrors = 0;
@@ -201,13 +214,10 @@ protected virtual async Task ProcessStreamAsync(PipeReader reader, PipeWriter wr
 
 			if (transformedEvent != null)
 			{
-				// Update parent activity with conversation ID and model info when we receive ConversationStart events
+				// Update parent activity with conversation ID 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));
 				}
 
@@ -238,8 +248,10 @@ protected async Task WriteEventAsync(AskAiEvent? transformedEvent, PipeWriter wr
 			return;
 
 		using var activity = StreamTransformerActivitySource.StartActivity("gen_ai.agent.token");
-		_ = (activity?.SetTag("gen_ai.agent.name", GetAgentId()));
-		_ = (activity?.SetTag("gen_ai.provider.name", GetAgentProvider()));
+
+		// Custom attributes for tracking our abstraction layer
+		_ = (activity?.SetTag("docs.ai.gateway", GetAgentProvider()));
+		_ = (activity?.SetTag("docs.ai.agent_name", GetAgentId()));
 		_ = (activity?.SetTag("gen_ai.response.token_type", transformedEvent.GetType().Name));
 
 		try

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

@@ -16,18 +16,32 @@ public class StreamTransformerFactory(
 	AskAiProviderResolver providerResolver,
 	ILogger<StreamTransformerFactory> logger) : IStreamTransformer
 {
-	public async Task<Stream> TransformAsync(Stream rawStream, CancellationToken cancellationToken = default)
+	private IStreamTransformer? _resolvedTransformer;
+
+	private IStreamTransformer GetTransformer()
 	{
+		if (_resolvedTransformer != null)
+			return _resolvedTransformer;
+
 		var provider = providerResolver.ResolveProvider();
 
-		IStreamTransformer transformer = provider switch
+		_resolvedTransformer = provider switch
 		{
 			"LlmGateway" => serviceProvider.GetRequiredService<LlmGatewayStreamTransformer>(),
 			"AgentBuilder" => serviceProvider.GetRequiredService<AgentBuilderStreamTransformer>(),
 			_ => throw new InvalidOperationException($"Unknown AI provider: {provider}. Valid values are 'AgentBuilder' or 'LlmGateway'")
 		};
 
-		logger.LogDebug("Using stream transformer for provider: {Provider}", provider);
+		logger.LogDebug("Resolved stream transformer for provider: {Provider}", provider);
+		return _resolvedTransformer;
+	}
+
+	public string AgentId => GetTransformer().AgentId;
+	public string AgentProvider => GetTransformer().AgentProvider;
+
+	public async Task<Stream> TransformAsync(Stream rawStream, CancellationToken cancellationToken = default)
+	{
+		var transformer = GetTransformer();
 		return await transformer.TransformAsync(rawStream, cancellationToken);
 	}
 }

src/api/Elastic.Documentation.Api.Lambda/Program.cs

@@ -9,6 +9,7 @@
 using Elastic.Documentation.Api.Core.Search;
 using Elastic.Documentation.Api.Infrastructure;
 using Elastic.Documentation.ServiceDefaults;
+using Microsoft.Extensions.Logging;
 using OpenTelemetry;
 using OpenTelemetry.Metrics;
 using OpenTelemetry.Trace;
@@ -22,9 +23,10 @@
 	process.Refresh();
 	Console.WriteLine($"WebApplication builder created. Memory: {process.WorkingSet64 / 1024 / 1024} MB");
 
-	_ = builder.AddDocumentationServiceDefaults(ref args);
+	// Add logging configuration for Lambda
+	_ = builder.Services.AddElasticDocumentationLogging(LogLevel.Information);
 	process.Refresh();
-	Console.WriteLine($"Documentation service defaults added. Memory: {process.WorkingSet64 / 1024 / 1024} MB");
+	Console.WriteLine($"Logging configured. Memory: {process.WorkingSet64 / 1024 / 1024} MB");
 
 	_ = builder.AddElasticOpenTelemetry(edotBuilder =>
 	{