Fix AOT errors

Author: reakaleek

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

@@ -3,6 +3,8 @@
 // See the LICENSE file in the project root for more information
 
 using System.Diagnostics;
+using System.Text.Json;
+using Elastic.Documentation.Api.Core;
 using Microsoft.Extensions.Logging;
 
 namespace Elastic.Documentation.Api.Core.AskAi;
@@ -21,18 +23,11 @@ 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
-		var inputMessages = new object[]
+		var inputMessages = new[]
 		{
-			new
-			{
-				role = "user",
-				parts = new object[]
-				{
-					new { type = "text", content = askAiRequest.Message }
-				}
-			}
+			new InputMessage("user", [new MessagePart("text", askAiRequest.Message)])
 		};
-		var inputMessagesJson = System.Text.Json.JsonSerializer.Serialize(inputMessages);
+		var inputMessagesJson = JsonSerializer.Serialize(inputMessages, ApiJsonContext.Default.InputMessageArray);
 		_ = activity?.SetTag("gen_ai.input.messages", inputMessagesJson);
 		logger.LogInformation("Streaming AskAI response");
 		var rawStream = await askAiGateway.AskAi(askAiRequest, ctx);

src/api/Elastic.Documentation.Api.Core/SerializationContext.cs

@@ -8,9 +8,21 @@
 
 namespace Elastic.Documentation.Api.Core;
 
+/// <summary>
+/// Types for OpenTelemetry telemetry serialization (AOT-compatible)
+/// </summary>
+public record MessagePart(string Type, string Content);
+
+public record InputMessage(string Role, MessagePart[] Parts);
+
+public record OutputMessage(string Role, MessagePart[] Parts, string FinishReason);
 
 [JsonSerializable(typeof(AskAiRequest))]
 [JsonSerializable(typeof(SearchRequest))]
 [JsonSerializable(typeof(SearchResponse))]
+[JsonSerializable(typeof(InputMessage))]
+[JsonSerializable(typeof(OutputMessage))]
+[JsonSerializable(typeof(MessagePart))]
+[JsonSerializable(typeof(InputMessage[]))]
 [JsonSourceGenerationOptions(PropertyNamingPolicy = JsonKnownNamingPolicy.CamelCase)]
 public partial class ApiJsonContext : JsonSerializerContext;

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

@@ -8,6 +8,7 @@
 using System.Runtime.CompilerServices;
 using System.Text;
 using System.Text.Json;
+using Elastic.Documentation.Api.Core;
 using Elastic.Documentation.Api.Core.AskAi;
 using Microsoft.Extensions.Logging;
 
@@ -144,6 +145,7 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 	{
 		using var activity = StreamTransformerActivitySource.StartActivity("transform_stream");
 		var result = new StreamProcessingResult();
+		List<MessagePart> outputMessageParts = [];
 		await foreach (var sseEvent in ParseSseEventsAsync(reader, cancellationToken))
 		{
 			AskAiEvent? transformedEvent = null;
@@ -181,8 +183,6 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 			if (transformedEvent == null)
 				continue;
 
-			List<object> outputMessageParts = [];
-
 			switch (transformedEvent)
 			{
 				case AskAiEvent.ConversationStart conversationStart:
@@ -195,8 +195,7 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 				case AskAiEvent.Reasoning reasoning:
 					{
 						_ = activity?.AddEvent(new ActivityEvent("conversation_start"));
-						outputMessageParts.Add((type: "reasoning", content: reasoning.Message));
-
+						outputMessageParts.Add(new MessagePart("reasoning", reasoning.Message ?? string.Empty));
 						break;
 					}
 				case AskAiEvent.Chunk:
@@ -234,14 +233,10 @@ private async Task ProcessStreamAsync(PipeReader reader, PipeWriter writer, Acti
 					}
 				case AskAiEvent.ChunkComplete chunkComplete:
 					{
-						var parts = outputMessageParts.Prepend((type: "text", content: chunkComplete.FullContent)).ToArray();
-						var outputMessages = new
-						{
-							role = "assistant",
-							parts,
-							finish_reason = "stop"
-						};
-						_ = activity?.SetTag("gen_ai.output.messages", JsonSerializer.Serialize(outputMessages));
+						var parts = outputMessageParts.Prepend(new MessagePart("text", chunkComplete.FullContent)).ToArray();
+						var outputMessages = new OutputMessage("assistant", parts, "stop");
+						_ = activity?.SetTag("gen_ai.output.messages", JsonSerializer.Serialize(outputMessages, ApiJsonContext.Default.OutputMessage));
+						_ = activity?.AddEvent(new ActivityEvent("message_chunk_complete"));
 						break;
 					}
 				case AskAiEvent.ConversationEnd: