Skip to content

Use RawRepresentation for more faithful native multi-turn #93

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kzu merged 1 commit into from
Mar 12, 2026
Merged

Use RawRepresentation for more faithful native multi-turn #93

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 24 additions & 42 deletions src/xAI/GrokChatClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));

Expand All @@ -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>();
Expand All @@ -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;
Expand Down Expand Up @@ -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 });
Expand Down Expand Up @@ -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
{
Expand Down
85 changes: 80 additions & 5 deletions src/xAI/GrokProtocolExtensions.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand Down Expand Up @@ -214,10 +216,9 @@ 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)
Expand Down Expand Up @@ -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,
Expand Down
Loading