Model MCP and code execution using MS.E.AI primitives

kzu · kzu · commit 3fc855374aa1 · 2025-12-05T12:43:20.000-03:00
Since opting for those outputs likely consumes additional output tokens, make them opt-in only.

Now both are represented with specific tools rather than our generic HostedToolCallContent. Marked its usage as experimental since it's marked such in MS.E.AI too.
diff --git a/src/Extensions.Grok/Extensions.Grok.csproj b/src/Extensions.Grok/Extensions.Grok.csproj
@@ -10,12 +10,12 @@
     <PackageLicenseFile>OSMFEULA.txt</PackageLicenseFile>
     <PackageRequireLicenseAcceptance>true</PackageRequireLicenseAcceptance>
     <EmitCompilerGeneratedFiles>true</EmitCompilerGeneratedFiles>
-    <NoWarn>MEAI001;$(NoWarn)</NoWarn>
+    <NoWarn>MEAI001;DEAI001;$(NoWarn)</NoWarn>
   </PropertyGroup>
 
   <ItemGroup>
     <PackageReference Include="Google.Protobuf" Version="3.33.1" />
-    <PackageReference Include="GrokClient" Version="1.0.0" />
+    <PackageReference Include="GrokClient" Version="1.0.2" />
     <PackageReference Include="Grpc.Net.Client" Version="2.71.0" />
     <PackageReference Include="Grpc.Tools" Version="2.76.0" PrivateAssets="all" />
     <PackageReference Include="NuGetizer" Version="1.4.5" PrivateAssets="all" />
diff --git a/src/Extensions.Grok/GrokChatClient.cs b/src/Extensions.Grok/GrokChatClient.cs
@@ -1,6 +1,5 @@
 using System.Text.Json;
 using Devlooped.Grok;
-using Google.Protobuf;
 using Grpc.Core;
 using Grpc.Net.Client;
 using Microsoft.Extensions.AI;
@@ -23,93 +22,133 @@ internal GrokChatClient(GrpcChannel channel, GrokClientOptions clientOptions, st
 
     public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
     {
-        var requestDto = MapToRequest(messages, options);
-        var protoResponse = await client.GetCompletionAsync(requestDto, cancellationToken: cancellationToken);
-        var lastOutput = protoResponse.Outputs.OrderByDescending(x => x.Index).FirstOrDefault();
+        var request = MapToRequest(messages, options);
+        var response = await client.GetCompletionAsync(request, cancellationToken: cancellationToken);
+        var lastOutput = response.Outputs.OrderByDescending(x => x.Index).FirstOrDefault();
 
         if (lastOutput == null)
         {
             return new ChatResponse()
             {
-                ResponseId = protoResponse.Id,
-                ModelId = protoResponse.Model,
-                CreatedAt = protoResponse.Created.ToDateTimeOffset(),
-                Usage = MapToUsage(protoResponse.Usage),
+                ResponseId = response.Id,
+                ModelId = response.Model,
+                CreatedAt = response.Created.ToDateTimeOffset(),
+                Usage = MapToUsage(response.Usage),
             };
         }
 
         var message = new ChatMessage(MapRole(lastOutput.Message.Role), default(string));
-        var citations = protoResponse.Citations?.Distinct().Select(MapCitation).ToList<AIAnnotation>();
+        var citations = response.Citations?.Distinct().Select(MapCitation).ToList<AIAnnotation>();
 
-        foreach (var output in protoResponse.Outputs.OrderBy(x => x.Index))
+        foreach (var output in response.Outputs.OrderBy(x => x.Index))
         {
             if (output.Message.Content is { Length: > 0 } text)
             {
-                var content = new TextContent(text)
+                // Special-case output from tools
+                if (output.Message.Role == MessageRole.RoleTool &&
+                    output.Message.ToolCalls.Count == 1 &&
+                        output.Message.ToolCalls[0] is { } toolCall)
                 {
-                    Annotations = citations
-                };
+                    if (toolCall.Type == ToolCallType.McpTool)
+                    {
+                        message.Contents.Add(new McpServerToolCallContent(toolCall.Id, toolCall.Function.Name, null)
+                        {
+                            RawRepresentation = toolCall
+                        });
+                        message.Contents.Add(new McpServerToolResultContent(toolCall.Id)
+                        {
+                            RawRepresentation = toolCall,
+                            Output = [new TextContent(text)]
+                        });
+                        continue;
+                    }
+                    else if (toolCall.Type == ToolCallType.CodeExecutionTool)
+                    {
+                        message.Contents.Add(new CodeInterpreterToolCallContent()
+                        {
+                            CallId = toolCall.Id,
+                            RawRepresentation = toolCall
+                        });
+                        message.Contents.Add(new CodeInterpreterToolResultContent()
+                        {
+                            CallId = toolCall.Id,
+                            RawRepresentation = toolCall,
+                            Outputs = [new TextContent(text)]
+                        });
+                        continue;
+                    }
+                }
+
+                var content = new TextContent(text) { Annotations = citations };
 
                 foreach (var citation in output.Message.Citations)
-                {
                     (content.Annotations ??= []).Add(MapInlineCitation(citation));
-                }
+
                 message.Contents.Add(content);
             }
 
             foreach (var toolCall in output.Message.ToolCalls)
-            {
-                if (toolCall.Type == ToolCallType.ClientSideTool)
-                {
-                    var arguments = !string.IsNullOrEmpty(toolCall.Function.Arguments)
-                        ? JsonSerializer.Deserialize<IDictionary<string, object?>>(toolCall.Function.Arguments)
-                        : null;
-
-                    var content = new FunctionCallContent(
-                            toolCall.Id,
-                            toolCall.Function.Name,
-                            arguments);
-
-                    message.Contents.Add(content);
-                }
-                else
-                {
-                    message.Contents.Add(new HostedToolCallContent(toolCall));
-                }
-            }
+                message.Contents.Add(MapToolCall(toolCall));
         }
 
         return new ChatResponse(message)
         {
-            ResponseId = protoResponse.Id,
-            ModelId = protoResponse.Model,
-            CreatedAt = protoResponse.Created.ToDateTimeOffset(),
+            ResponseId = response.Id,
+            ModelId = response.Model,
+            CreatedAt = response.Created.ToDateTimeOffset(),
             FinishReason = lastOutput != null ? MapFinishReason(lastOutput.FinishReason) : null,
-            Usage = MapToUsage(protoResponse.Usage),
+            Usage = MapToUsage(response.Usage),
         };
     }
 
+    AIContent MapToolCall(ToolCall toolCall) => toolCall.Type switch
+    {
+        ToolCallType.ClientSideTool => new FunctionCallContent(
+            toolCall.Id,
+            toolCall.Function.Name,
+            !string.IsNullOrEmpty(toolCall.Function.Arguments)
+                ? JsonSerializer.Deserialize<IDictionary<string, object?>>(toolCall.Function.Arguments)
+                : null)
+        {
+            RawRepresentation = toolCall
+        },
+        ToolCallType.McpTool => new McpServerToolCallContent(toolCall.Id, toolCall.Function.Name, null)
+        {
+            RawRepresentation = toolCall
+        },
+        ToolCallType.CodeExecutionTool => new CodeInterpreterToolCallContent()
+        {
+            CallId = toolCall.Id,
+            RawRepresentation = toolCall
+        },
+        _ => new HostedToolCallContent()
+        {
+            CallId = toolCall.Id,
+            RawRepresentation = toolCall
+        }
+    };
+
     public IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
     {
         return CompleteChatStreamingCore(messages, options, cancellationToken);
 
         async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable<ChatMessage> messages, ChatOptions? options, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
         {
-            var requestDto = MapToRequest(messages, options);
-            var call = client.GetCompletionChunk(requestDto, cancellationToken: cancellationToken);
-            
+            var request = MapToRequest(messages, options);
+            var call = client.GetCompletionChunk(request, cancellationToken: cancellationToken);
+
             await foreach (var chunk in call.ResponseStream.ReadAllAsync(cancellationToken))
             {
-                var outputChunk = chunk.Outputs[0];
-                var text = outputChunk.Delta.Content is { Length: > 0 } delta ? delta : null;
+                var output = chunk.Outputs[0];
+                var text = output.Delta.Content is { Length: > 0 } delta ? delta : null;
 
                 // Use positional arguments for ChatResponseUpdate
-                var update = new ChatResponseUpdate(MapRole(outputChunk.Delta.Role), text)
+                var update = new ChatResponseUpdate(MapRole(output.Delta.Role), text)
                 {
                     ResponseId = chunk.Id,
                     ModelId = chunk.Model,
                     CreatedAt = chunk.Created?.ToDateTimeOffset(),
-                    FinishReason = outputChunk.FinishReason != FinishReason.ReasonInvalid ? MapFinishReason(outputChunk.FinishReason) : null,
+                    FinishReason = output.FinishReason != FinishReason.ReasonInvalid ? MapFinishReason(output.FinishReason) : null,
                 };
 
                 if (chunk.Citations is { Count: > 0 } citations)
@@ -122,31 +161,11 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
                     }
 
                     foreach (var citation in citations.Distinct())
-                    {
                         (textContent.Annotations ??= []).Add(MapCitation(citation));
-                    }
                 }
 
-                foreach (var toolCall in outputChunk.Delta.ToolCalls)
-                {
-                    if (toolCall.Type == ToolCallType.ClientSideTool)
-                    {
-                        var arguments = !string.IsNullOrEmpty(toolCall.Function.Arguments)
-                            ? JsonSerializer.Deserialize<IDictionary<string, object?>>(toolCall.Function.Arguments)
-                            : null;
-
-                        var content = new FunctionCallContent(
-                                toolCall.Id,
-                                toolCall.Function.Name,
-                                arguments);
-
-                        update.Contents.Add(content);
-                    }
-                    else
-                    {
-                        update.Contents.Add(new HostedToolCallContent(toolCall));
-                    }
-                }
+                foreach (var toolCall in output.Delta.ToolCalls)
+                    update.Contents.Add(MapToolCall(toolCall));
 
                 if (update.Contents.Any())
                     yield return update;
@@ -191,6 +210,8 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
     {
         var request = new GetCompletionsRequest
         {
+            // By default always include citations in the final output if available
+            Include = { IncludeOption.InlineCitations },
             Model = options?.ModelId ?? defaultModelId,
         };
 
@@ -211,6 +232,10 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
                 {
                     gmsg.Content.Add(new Content { Text = textContent.Text });
                 }
+                else if (content.RawRepresentation is ToolCall toolCall)
+                {
+                    gmsg.ToolCalls.Add(toolCall);
+                }
                 else if (content is FunctionCallContent functionCall)
                 {
                     gmsg.ToolCalls.Add(new ToolCall
@@ -224,10 +249,6 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
                         }
                     });
                 }
-                else if (content is HostedToolCallContent serverFunction)
-                {
-                    gmsg.ToolCalls.Add(serverFunction.ToolCall);
-                }
                 else if (content is FunctionResultContent resultContent)
                 {
                     request.Messages.Add(new Message
@@ -236,19 +257,49 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
                         Content = { new Content { Text = JsonSerializer.Serialize(resultContent.Result) ?? "null" } }
                     });
                 }
+                else if (content is McpServerToolResultContent mcpResult &&
+                    mcpResult.RawRepresentation is ToolCall mcpToolCall &&
+                    mcpResult.Output is { Count: 1 } &&
+                    mcpResult.Output[0] is TextContent mcpText)
+                {
+                    request.Messages.Add(new Message
+                    {
+                        Role = MessageRole.RoleTool,
+                        ToolCalls = { mcpToolCall },
+                        Content = { new Content { Text = mcpText.Text } }
+                    });
+                }
+                else if (content is CodeInterpreterToolResultContent codeResult &&
+                    codeResult.RawRepresentation is ToolCall codeToolCall &&
+                    codeResult.Outputs is { Count: 1 } &&
+                    codeResult.Outputs[0] is TextContent codeText)
+                {
+                    request.Messages.Add(new Message
+                    {
+                        Role = MessageRole.RoleTool,
+                        ToolCalls = { codeToolCall },
+                        Content = { new Content { Text = codeText.Text } }
+                    });
+                }
             }
 
             if (gmsg.Content.Count == 0 && gmsg.ToolCalls.Count == 0)
                 continue;
 
+            // If we have only tool calls and no content, the gRPC enpoint fails, so add an empty one.
             if (gmsg.Content.Count == 0)
                 gmsg.Content.Add(new Content());
 
             request.Messages.Add(gmsg);
         }
 
+        IList<IncludeOption> includes = [IncludeOption.InlineCitations];
         if (options is GrokChatOptions grokOptions)
         {
+            // NOTE: overrides our default include for inline citations, potentially.
+            request.Include.Clear();
+            request.Include.AddRange(grokOptions.Include);
+
             if (grokOptions.Search.HasFlag(GrokSearch.X))
             {
                 (options.Tools ??= []).Insert(0, new GrokXSearchTool());
diff --git a/src/Extensions.Grok/GrokChatOptions.cs b/src/Extensions.Grok/GrokChatOptions.cs
@@ -1,4 +1,5 @@
 ﻿using System.ComponentModel;
+using Devlooped.Grok;
 using Microsoft.Extensions.AI;
 
 namespace Devlooped.Extensions.AI.Grok;
@@ -33,4 +34,8 @@ public class GrokChatOptions : ChatOptions
     /// <summary>Configures Grok's agentic search capabilities.</summary>
     /// <remarks>See https://docs.x.ai/docs/guides/tools/search-tools.</remarks>
     public GrokSearch Search { get; set; } = GrokSearch.None;
+
+    /// <summary>Additional outputs to include in responses.</summary>
+    /// <remarks>Defaults to including <see cref="IncludeOption.InlineCitations"/>.</remarks>
+    public IList<IncludeOption> Include { get; set; } = [IncludeOption.InlineCitations];
 }
diff --git a/src/Extensions.Grok/HostedToolCallContent.cs b/src/Extensions.Grok/HostedToolCallContent.cs
@@ -1,32 +1,13 @@
-﻿using System.Diagnostics;
-using System.Text.Json.Serialization;
-using Devlooped.Grok;
+﻿using System.Diagnostics.CodeAnalysis;
 using Microsoft.Extensions.AI;
 
 namespace Devlooped.Extensions.AI;
 
 /// <summary>Represents a hosted tool agentic call.</summary>
 /// <param name="toolCall">The tool call details.</param>
-[DebuggerDisplay("{DebuggerDisplay,nq}")]
-[method: JsonConstructor]
-public sealed class HostedToolCallContent(ToolCall toolCall) : AIContent
+[Experimental("DEAI001")]
+public class HostedToolCallContent : AIContent
 {
-    /// <summary>Gets the tool call details.</summary>
-    public ToolCall ToolCall => toolCall;
-
-    /// <summary>Gets a string representing this instance to display in the debugger.</summary>
-    [DebuggerBrowsable(DebuggerBrowsableState.Never)]
-    string DebuggerDisplay
-    {
-        get
-        {
-            var display = $"ToolCall = {toolCall.Id}, ";
-
-            display += toolCall.Function.Arguments is not null ?
-                $"{toolCall.Function.Name}({toolCall.Function.Arguments})" :
-                $"{toolCall.Function.Name}()";
-
-            return display;
-        }
-    }
+    /// <summary>Gets or sets the tool call ID.</summary>
+    public virtual string? CallId { get; set; }
 }
diff --git a/src/Extensions.Grok/HostedToolResultContent.cs b/src/Extensions.Grok/HostedToolResultContent.cs
@@ -0,0 +1,18 @@
+﻿using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using Microsoft.Extensions.AI;
+
+namespace Devlooped.Extensions.AI;
+
+/// <summary>Represents a hosted tool agentic call.</summary>
+/// <param name="toolCall">The tool call details.</param>
+[DebuggerDisplay("{DebuggerDisplay,nq}")]
+[Experimental("DEAI001")]
+public class HostedToolResultContent : AIContent
+{
+    /// <summary>Gets or sets the tool call ID.</summary>
+    public virtual string? CallId { get; set; }
+
+    /// <summary>Gets or sets the resulting contents from the tool.</summary>
+    public virtual IList<AIContent>? Outputs { get; set; }
+}
diff --git a/src/Tests/GrokTests.cs b/src/Tests/GrokTests.cs
diff --git a/src/Tests/Tests.csproj b/src/Tests/Tests.csproj
