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

.github/workflows/deploy-api-lambda.yml

@@ -50,19 +50,50 @@ jobs:
           role-to-assume: arn:aws:iam::197730964718:role/elastic-docs-v3-api-deployer-${{ inputs.environment }}
           aws-region: us-east-1
 
-      - name: Upload Lambda function
+      - name: Upload to S3 and Deploy Lambda Functions
         run: |
-          NEW_VERSION=$(aws lambda update-function-code \
-            --function-name "elastic-docs-v3-${ENVIRONMENT}-api" \
-            --zip-file "fileb://${ZIP_FILE}" \
-            --publish \
+          # Upload once to S3 with git SHA for traceability
+          S3_KEY="${ENVIRONMENT}/${GITHUB_SHA}/api-lambda.zip"
+          aws s3 cp "${ZIP_FILE}" "s3://${S3_BUCKET}/${S3_KEY}"
+
+          # Deploy to streaming-optimized Lambda
+          aws lambda update-function-code \
+            --function-name "elastic-docs-v3-${ENVIRONMENT}-lambda-stream-optimized" \
+            --s3-bucket "${S3_BUCKET}" \
+            --s3-key "${S3_KEY}"
+
+          STREAM_VERSION=$(aws lambda publish-version \
+            --function-name "elastic-docs-v3-${ENVIRONMENT}-lambda-stream-optimized" \
+            --description "Deployed from ${GITHUB_SHA}" \
             --query 'Version' \
             --output text)
-
+          
           aws lambda update-alias \
-          --function-name "elastic-docs-v3-${ENVIRONMENT}-api" \
-          --name live \
-          --function-version $NEW_VERSION
+            --function-name "elastic-docs-v3-${ENVIRONMENT}-lambda-stream-optimized" \
+            --name live \
+            --function-version $STREAM_VERSION
+
+          # Deploy to API Gateway-optimized Lambda
+          aws lambda update-function-code \
+            --function-name "elastic-docs-v3-${ENVIRONMENT}-lambda-api-gateway-optimized" \
+            --s3-bucket "${S3_BUCKET}" \
+            --s3-key "${S3_KEY}"
+
+          API_GW_VERSION=$(aws lambda publish-version \
+            --function-name "elastic-docs-v3-${ENVIRONMENT}-lambda-api-gateway-optimized" \
+            --description "Deployed from ${GITHUB_SHA}" \
+            --query 'Version' \
+            --output text)
+
+          aws lambda update-alias \
+            --function-name "elastic-docs-v3-${ENVIRONMENT}-lambda-api-gateway-optimized" \
+            --name live \
+            --function-version $API_GW_VERSION
+
+          echo "✅ Deployed to both Lambda functions from S3: ${S3_KEY}"
+          echo "   Stream Lambda: version $STREAM_VERSION (${GITHUB_SHA})"
+          echo "   API Gateway Lambda: version $API_GW_VERSION (${GITHUB_SHA})"
 
         env:
           ENVIRONMENT: ${{ inputs.environment }}
+          S3_BUCKET: elastic-docs-v3-api-lambda-artifacts

Directory.Packages.props

@@ -12,6 +12,7 @@
   </ItemGroup>
   <!-- AWS -->
   <ItemGroup>
+    <PackageVersion Include="Amazon.Lambda.APIGatewayEvents" Version="2.7.1" />
     <PackageVersion Include="Amazon.Lambda.AspNetCoreServer.Hosting" Version="1.9.0" />
     <PackageVersion Include="Amazon.Lambda.RuntimeSupport" Version="1.13.4" />
     <PackageVersion Include="Amazon.Lambda.Core" Version="2.7.1" />

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,46 @@ 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)
 	{
+		// Start activity for the chat request - DO NOT use 'using' because the stream is consumed later
+		// The activity will be passed to the transformer which will dispose it when the stream completes
+		var activity = AskAiActivitySource.StartActivity("chat", ActivityKind.Client);
+
+		// 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.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,
+			tags:
+			[
+				new KeyValuePair<string, object?>("gen_ai.prompt", askAiRequest.Message),
+				new KeyValuePair<string, object?>("gen_ai.system_instructions", 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 handle disposing the activity when streaming completes
+		var transformedStream = await streamTransformer.TransformAsync(rawStream, activity, ctx);
+
+		return transformedStream;
 	}
 }
 

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

@@ -9,11 +9,22 @@ 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>
 	/// <param name="rawStream">Raw SSE stream from gateway (Agent Builder, LLM Gateway, etc.)</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, CancellationToken cancellationToken = default);
+	Task<Stream> TransformAsync(Stream rawStream, System.Diagnostics.Activity? parentActivity, CancellationToken cancellationToken = default);
 }

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: "...", ...}]