Use RawRepresentation for more faithful native multi-turn

Rather than converting back and forth when the same messages we returned are passed back to continue a conversation, we leverage the message's RawRepresentation to bring back directly the original gRPC message.

This also requires doing a more 1:1 conversion of completion messages into individual response messages, rather than having a single response message with multiple contents. This would align better with the official xAI docs on response too.

Author: kzu

src/xAI/GrokChatClient.cs

@@ -45,11 +45,11 @@ public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messag
             ResponseId = response.Id,
             ModelId = response.Model,
             CreatedAt = response.Created?.ToDateTimeOffset(),
-            FinishReason = lastOutput != null ? MapFinishReason(lastOutput.FinishReason) : null,
-            Usage = MapToUsage(response.Usage),
+            FinishReason = lastOutput != null ? lastOutput.FinishReason.Convert() : null,
+            Usage = response.Usage.Convert(),
         };
 
-        var citations = response.Citations?.Distinct().Select(MapCitation).ToList<AIAnnotation>();
+        var citations = response.Citations?.Distinct().Select(x => x.FromCitationUrl()).ToList<AIAnnotation>();
 
         ((List<ChatMessage>)result.Messages).AddRange(response.Outputs.AsChatMessages(citations));
 
@@ -73,12 +73,12 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
                 // Use positional arguments for ChatResponseUpdate
                 var update = new ChatResponseUpdate
                 {
-                    Role = MapRole(output.Delta.Role),
+                    Role = output.Delta.Role.Convert(),
                     ResponseId = chunk.Id,
                     ModelId = chunk.Model,
                     CreatedAt = chunk.Created?.ToDateTimeOffset(),
                     RawRepresentation = chunk,
-                    FinishReason = output.FinishReason != FinishReason.ReasonInvalid ? MapFinishReason(output.FinishReason) : null,
+                    FinishReason = output.FinishReason != FinishReason.ReasonInvalid ? output.FinishReason.Convert() : null,
                 };
 
                 var citations = chunk.Citations?.Distinct().Select(MapCitation).ToList<AIAnnotation>();
@@ -101,7 +101,7 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
                     text is not null)
                     update.Contents.Add(new TextContent(text));
 
-                if (MapToUsage(chunk.Usage) is { } usage)
+                if (chunk.Usage.Convert() is { } usage)
                     update.Contents.Add(new UsageContent(usage) { RawRepresentation = chunk.Usage });
 
                 yield return update;
@@ -149,10 +149,27 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
 
         foreach (var message in messages)
         {
-            var gmsg = new Message { Role = MapRole(message.Role) };
+            if (message.RawRepresentation is Message input)
+            {
+                request.Messages.Add(input);
+                continue;
+            }
+            else if (message.RawRepresentation is CompletionMessage completion)
+            {
+                request.Messages.Add(completion.AsMessage());
+                continue;
+            }
+
+            var gmsg = new Message { Role = message.Role.Convert() };
 
             foreach (var content in message.Contents)
             {
+                if (content.RawRepresentation is CompletionMessage completion)
+                {
+                    request.Messages.Add(completion.AsMessage());
+                    continue;
+                }
+
                 if (content is TextContent textContent && !string.IsNullOrEmpty(textContent.Text))
                 {
                     gmsg.Content.Add(new Content { Text = textContent.Text });
@@ -271,41 +288,6 @@ codeResult.RawRepresentation is ToolCall codeToolCall &&
         return request;
     }
 
-    static MessageRole MapRole(ChatRole role) => role switch
-    {
-        _ when role == ChatRole.System => MessageRole.RoleSystem,
-        _ when role == ChatRole.User => MessageRole.RoleUser,
-        _ when role == ChatRole.Assistant => MessageRole.RoleAssistant,
-        _ when role == ChatRole.Tool => MessageRole.RoleTool,
-        _ => MessageRole.RoleUser
-    };
-
-    static ChatRole MapRole(MessageRole role) => role switch
-    {
-        MessageRole.RoleSystem => ChatRole.System,
-        MessageRole.RoleUser => ChatRole.User,
-        MessageRole.RoleAssistant => ChatRole.Assistant,
-        MessageRole.RoleTool => ChatRole.Tool,
-        _ => ChatRole.Assistant
-    };
-
-    static ChatFinishReason? MapFinishReason(FinishReason finishReason) => finishReason switch
-    {
-        FinishReason.ReasonStop => ChatFinishReason.Stop,
-        FinishReason.ReasonMaxLen => ChatFinishReason.Length,
-        FinishReason.ReasonToolCalls => ChatFinishReason.ToolCalls,
-        FinishReason.ReasonMaxContext => ChatFinishReason.Length,
-        FinishReason.ReasonTimeLimit => ChatFinishReason.Length,
-        _ => null
-    };
-
-    static UsageDetails? MapToUsage(SamplingUsage usage) => usage == null ? null : new()
-    {
-        InputTokenCount = usage.PromptTokens,
-        OutputTokenCount = usage.CompletionTokens,
-        TotalTokenCount = usage.TotalTokens
-    };
-
     /// <inheritdoc />
     public object? GetService(Type serviceType, object? serviceKey = null) => serviceType switch
     {

src/xAI/GrokProtocolExtensions.cs

@@ -161,11 +161,13 @@ grokSearch.City is not null ||
 
     static IEnumerable<ChatMessage> ToChatMessages(IEnumerable<CompletionMessage> messages, List<AIAnnotation>? citations = default)
     {
-        ChatMessage? message = null;
-
         foreach (var completion in messages)
         {
-            message ??= new(ChatRole.Assistant, (string?)null);
+            ChatMessage message = new(ChatRole.Assistant, (string?)null)
+            {
+                RawRepresentation = completion
+            };
+
             var annotations = citations;
             if (completion.Citations.Count > 0)
             {
@@ -214,11 +216,10 @@ static IEnumerable<ChatMessage> ToChatMessages(IEnumerable<CompletionMessage> me
             //        RawRepresentation = completion
             //    });
             //}
-        }
 
-        if (message is not null)
             yield return message;
     }
+    }
 
     internal static IEnumerable<AIContent> AsContents(this IEnumerable<ToolCall> toolCalls, string? content = default, List<AIAnnotation>? annotations = default)
     {
@@ -356,6 +357,80 @@ static IEnumerable<CitationAnnotation> AsCitations(CollectionSearchItem item)
         _ => [new CitationAnnotation { RawRepresentation = citation }]
     };
 
+    internal static Message AsMessage(this CompletionMessage completion)
+    {
+        var message = new Message
+        {
+            Role = completion.Role,
+            EncryptedContent = completion.EncryptedContent,
+            ReasoningContent = completion.ReasoningContent,
+        };
+
+        if (!string.IsNullOrEmpty(completion.Content))
+            message.Content.Add(new Content { Text = completion.Content });
+
+        message.ToolCalls.AddRange(completion.ToolCalls);
+
+        return message;
+    }
+
+    internal static MessageRole Convert(this ChatRole role) => role switch
+    {
+        _ when role == ChatRole.System => MessageRole.RoleSystem,
+        _ when role == ChatRole.User => MessageRole.RoleUser,
+        _ when role == ChatRole.Assistant => MessageRole.RoleAssistant,
+        _ when role == ChatRole.Tool => MessageRole.RoleTool,
+        _ => MessageRole.RoleUser
+    };
+
+    internal static ChatRole Convert(this MessageRole role) => role switch
+    {
+        MessageRole.RoleSystem => ChatRole.System,
+        MessageRole.RoleUser => ChatRole.User,
+        MessageRole.RoleAssistant => ChatRole.Assistant,
+        MessageRole.RoleTool => ChatRole.Tool,
+        _ => ChatRole.Assistant
+    };
+
+    internal static ChatFinishReason? Convert(this FinishReason finishReason) => finishReason switch
+    {
+        FinishReason.ReasonStop => ChatFinishReason.Stop,
+        FinishReason.ReasonMaxLen => ChatFinishReason.Length,
+        FinishReason.ReasonToolCalls => ChatFinishReason.ToolCalls,
+        FinishReason.ReasonMaxContext => ChatFinishReason.Length,
+        FinishReason.ReasonTimeLimit => ChatFinishReason.Length,
+        _ => null
+    };
+
+    internal static UsageDetails? Convert(this SamplingUsage usage) => usage == null ? null : new()
+    {
+        InputTokenCount = usage.PromptTokens,
+        OutputTokenCount = usage.CompletionTokens,
+        TotalTokenCount = usage.TotalTokens
+    };
+
+    internal static CitationAnnotation FromCitationUrl(this string citationUrl)
+    {
+        var url = new Uri(citationUrl);
+        if (url.Scheme != "collections")
+            return new CitationAnnotation { Url = url };
+
+        // Special-case collection citations so we get better metadata
+        var collection = url.Host;
+        var file = url.AbsolutePath[7..];
+
+        return new CitationAnnotation
+        {
+            AdditionalProperties = new AdditionalPropertiesDictionary
+            {
+                { "collection_id", collection }
+            },
+            FileId = file,
+            ToolName = "collections_search",
+            Url = new Uri($"collections://{collection}/files/{file}"),
+        };
+    }
+
     [JsonSourceGenerationOptions(JsonSerializerDefaults.Web,
         UseStringEnumConverter = true,
         UnmappedMemberHandling = JsonUnmappedMemberHandling.Skip,