Skip to content

Commit 37c0b3d

Browse files
kzukzu
committed
Add basic MCP tool calling support
See https://docs.x.ai/docs/guides/tools/remote-mcp-tools
1 parent 246a46f commit 37c0b3d

File tree

5 files changed

+145
-47
lines changed

5 files changed

+145
-47
lines changed

‎src/Extensions.Grok/Extensions.Grok.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
<PackageLicenseFile>OSMFEULA.txt</PackageLicenseFile>
1111
<PackageRequireLicenseAcceptance>true</PackageRequireLicenseAcceptance>
1212
<EmitCompilerGeneratedFiles>true</EmitCompilerGeneratedFiles>
13+
<NoWarn>MEAI001;$(NoWarn)</NoWarn>
1314
</PropertyGroup>
1415

1516
<ItemGroup>

‎src/Extensions.Grok/GrokChatClient.cs‎

Lines changed: 94 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,62 @@ internal GrokChatClient(GrpcChannel channel, GrokClientOptions clientOptions, st
2323
public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
2424
{
2525
var requestDto = MapToRequest(messages, options);
26-
2726
var protoResponse = await client.GetCompletionAsync(requestDto, cancellationToken: cancellationToken);
27+
var lastOutput = protoResponse.Outputs.OrderByDescending(x => x.Index).FirstOrDefault();
2828

29-
var chatMessages = protoResponse.Outputs
30-
.Select(x => MapToChatMessage(x.Message, protoResponse.Citations))
31-
.Where(x => x.Contents.Count > 0)
32-
.ToList();
29+
if (lastOutput == null)
30+
{
31+
return new ChatResponse()
32+
{
33+
ResponseId = protoResponse.Id,
34+
ModelId = protoResponse.Model,
35+
CreatedAt = protoResponse.Created.ToDateTimeOffset(),
36+
Usage = MapToUsage(protoResponse.Usage),
37+
};
38+
}
3339

34-
var lastOutput = protoResponse.Outputs.LastOrDefault();
40+
var message = new ChatMessage(MapRole(lastOutput.Message.Role), default(string));
41+
var citations = protoResponse.Citations?.Distinct().Select(MapCitation).ToList<AIAnnotation>();
3542

36-
return new ChatResponse(chatMessages)
43+
foreach (var output in protoResponse.Outputs.OrderBy(x => x.Index))
44+
{
45+
if (output.Message.Content is { Length: > 0 } text)
46+
{
47+
var content = new TextContent(text)
48+
{
49+
Annotations = citations
50+
};
51+
52+
foreach (var citation in output.Message.Citations)
53+
{
54+
(content.Annotations ??= []).Add(MapInlineCitation(citation));
55+
}
56+
message.Contents.Add(content);
57+
}
58+
59+
foreach (var toolCall in output.Message.ToolCalls)
60+
{
61+
if (toolCall.Type == ToolCallType.ClientSideTool)
62+
{
63+
var arguments = !string.IsNullOrEmpty(toolCall.Function.Arguments)
64+
? JsonSerializer.Deserialize<IDictionary<string, object?>>(toolCall.Function.Arguments)
65+
: null;
66+
67+
var content = new FunctionCallContent(
68+
toolCall.Id,
69+
toolCall.Function.Name,
70+
arguments);
71+
72+
message.Contents.Add(content);
73+
}
74+
else
75+
{
76+
message.Contents.Add(new HostedToolCallContent(toolCall));
77+
}
78+
}
79+
}
80+
81+
return new ChatResponse(message)
3782
{
3883
ResponseId = protoResponse.Id,
3984
ModelId = protoResponse.Model,
@@ -80,7 +125,7 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
80125

81126
foreach (var citation in citations.Distinct())
82127
{
83-
(textContent.Annotations ??= []).Add(new CitationAnnotation { Url = new(citation) });
128+
(textContent.Annotations ??= []).Add(MapCitation(citation));
84129
}
85130
}
86131

@@ -89,46 +134,37 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
89134
}
90135
}
91136

92-
static ChatMessage MapToChatMessage(CompletionMessage message, IList<string>? citations = null)
137+
static CitationAnnotation MapInlineCitation(InlineCitation citation) => citation.CitationCase switch
93138
{
94-
var chatMessage = new ChatMessage() { Role = MapRole(message.Role) };
139+
InlineCitation.CitationOneofCase.WebCitation => new CitationAnnotation { Url = new(citation.WebCitation.Url) },
140+
InlineCitation.CitationOneofCase.XCitation => new CitationAnnotation { Url = new(citation.XCitation.Url) },
141+
InlineCitation.CitationOneofCase.CollectionsCitation => new CitationAnnotation
142+
{
143+
FileId = citation.CollectionsCitation.FileId,
144+
Snippet = citation.CollectionsCitation.ChunkContent,
145+
ToolName = "file_search",
146+
},
147+
_ => new CitationAnnotation()
148+
};
95149

96-
if (message.Content is { Length: > 0 } text)
150+
static CitationAnnotation MapCitation(string citation)
151+
{
152+
var url = new Uri(citation);
153+
if (url.Scheme != "collections")
154+
return new CitationAnnotation { Url = url };
155+
156+
// Special-case collection citations so we get better metadata
157+
var collection = url.Host;
158+
var file = url.AbsolutePath[7..];
159+
return new CitationAnnotation
97160
{
98-
var textContent = new TextContent(text);
99-
if (citations is { Count: > 0 })
100-
{
101-
foreach (var citation in citations.Distinct())
161+
ToolName = "collections_search",
162+
FileId = file,
163+
AdditionalProperties = new AdditionalPropertiesDictionary
102164
{
103-
(textContent.Annotations ??= []).Add(new CitationAnnotation { Url = new(citation) });
165+
{ "collection_id", collection }
104166
}
105-
}
106-
chatMessage.Contents.Add(textContent);
107-
}
108-
109-
foreach (var toolCall in message.ToolCalls)
110-
{
111-
// Only include client-side tools in the response messages
112-
if (toolCall.Type == ToolCallType.ClientSideTool)
113-
{
114-
var arguments = !string.IsNullOrEmpty(toolCall.Function.Arguments)
115-
? JsonSerializer.Deserialize<IDictionary<string, object?>>(toolCall.Function.Arguments)
116-
: null;
117-
118-
var content = new FunctionCallContent(
119-
toolCall.Id,
120-
toolCall.Function.Name,
121-
arguments);
122-
123-
chatMessage.Contents.Add(content);
124-
}
125-
else
126-
{
127-
chatMessage.Contents.Add(new HostedToolCallContent(toolCall));
128-
}
129-
}
130-
131-
return chatMessage;
167+
};
132168
}
133169

134170
GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOptions? options)
@@ -258,12 +294,25 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
258294

259295
if (fileSearch.Inputs?.OfType<HostedVectorStoreContent>() is { } vectorStores)
260296
toolProto.CollectionIds.AddRange(vectorStores.Select(x => x.VectorStoreId).Distinct());
261-
297+
262298
if (fileSearch.MaximumResultCount is { } maxResults)
263299
toolProto.Limit = maxResults;
264-
300+
265301
request.Tools.Add(new Tool { CollectionsSearch = toolProto });
266302
}
303+
else if (tool is HostedMcpServerTool mcpTool)
304+
{
305+
request.Tools.Add(new Tool
306+
{
307+
Mcp = new MCP
308+
{
309+
Authorization = mcpTool.AuthorizationToken,
310+
ServerLabel = mcpTool.ServerName,
311+
ServerUrl = mcpTool.ServerAddress,
312+
AllowedToolNames = { mcpTool.AllowedTools ?? Array.Empty<string>() }
313+
}
314+
});
315+
}
267316
}
268317
}
269318

‎src/Tests/DotEnv.cs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Runtime.CompilerServices;
4+
using System.Runtime.InteropServices;
5+
using System.Text;
6+
7+
class DotEnv
8+
{
9+
[ModuleInitializer]
10+
public static void Init()
11+
{
12+
// Load environment variables from .env files in current dir and above.
13+
DotNetEnv.Env.TraversePath().Load();
14+
15+
// Load environment variables from user profile directory.
16+
var userEnv = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".env");
17+
if (File.Exists(userEnv))
18+
DotNetEnv.Env.Load(userEnv);
19+
}
20+
}

‎src/Tests/GrokTests.cs‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,8 @@ public async Task GrokInvokesToolAndSearch()
5656

5757
var grok = new GrokClient(Configuration["XAI_API_KEY"]!).AsIChatClient("grok-4")
5858
.AsBuilder()
59-
.UseLogging(output.AsLoggerFactory())
6059
.UseFunctionInvocation()
60+
.UseLogging(output.AsLoggerFactory())
6161
.Build();
6262

6363
var getDateCalls = 0;
@@ -292,4 +292,31 @@ public async Task GrokInvokesHostedCollectionSearch()
292292
response.Messages.SelectMany(x => x.Contents).OfType<HostedToolCallContent>(),
293293
x => x.ToolCall.Type == Devlooped.Grok.ToolCallType.CollectionsSearchTool);
294294
}
295+
296+
[SecretsFact("XAI_API_KEY", "GITHUB_TOKEN")]
297+
public async Task GrokInvokesHostedMcp()
298+
{
299+
var messages = new Chat()
300+
{
301+
{ "user", "When was GrokClient v1.0.0 released on the devlooped/GrokClient repo? Respond with just the date, in YYYY-MM-DD format." },
302+
};
303+
304+
var grok = new GrokClient(Configuration["XAI_API_KEY"]!).AsIChatClient("grok-4-fast");
305+
306+
var options = new ChatOptions
307+
{
308+
Tools = [new HostedMcpServerTool("GitHub", "https://api.githubcopilot.com/mcp/") {
309+
AuthorizationToken = Configuration["GITHUB_TOKEN"]!,
310+
AllowedTools = ["list_releases"],
311+
}]
312+
};
313+
314+
var response = await grok.GetResponseAsync(messages, options);
315+
var text = response.Text;
316+
317+
Assert.Equal("2025-11-29", text);
318+
Assert.Contains(
319+
response.Messages.SelectMany(x => x.Contents).OfType<HostedToolCallContent>(),
320+
x => x.ToolCall.Type == Devlooped.Grok.ToolCallType.McpTool);
321+
}
295322
}

‎src/Tests/Tests.csproj‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
<PropertyGroup>
44
<TargetFramework>net10.0</TargetFramework>
5-
<NoWarn>OPENAI001;$(NoWarn)</NoWarn>
5+
<NoWarn>OPENAI001;MEAI001;$(NoWarn)</NoWarn>
66
<LangVersion>Preview</LangVersion>
77
<EmitCompilerGeneratedFiles>true</EmitCompilerGeneratedFiles>
88
<RootNamespace>Devlooped</RootNamespace>
@@ -28,6 +28,7 @@
2828
<PackageReference Include="Microsoft.Extensions.Logging.Console" Version="$(NetVersion)" />
2929

3030
<PackageReference Include="Tomlyn.Extensions.Configuration" Version="1.0.6" />
31+
<PackageReference Include="DotNetEnv" Version="3.1.1" />
3132
</ItemGroup>
3233

3334
<ItemGroup>

0 commit comments

Comments
 (0)