Skip to content

Commit 3c72c34

Browse files
aaronpowellaaronpowell
authored
Adding support for returning collections from tools (#116)
* Adding support for returning collections from tools This will mean we can return multiple AIContent objects from a tool, such as a mixed text/image set Contributes to #68 * Code review feedback * Adding tests and some more cases * More cases in switch and tests
1 parent 5171ccc commit 3c72c34

File tree

4 files changed

+260
-47
lines changed

4 files changed

+260
-47
lines changed

‎src/ModelContextProtocol/AIContentExtensions.cs

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
using Microsoft.Extensions.AI;
22
using ModelContextProtocol.Protocol.Types;
33
using ModelContextProtocol.Utils;
4+
using ModelContextProtocol.Utils.Json;
45
using System.Runtime.InteropServices;
6+
using System.Text.Json;
57

68
namespace ModelContextProtocol;
79

@@ -101,4 +103,28 @@ internal static string GetBase64Data(this DataContent dataContent)
101103
Convert.ToBase64String(dataContent.Data.ToArray());
102104
#endif
103105
}
106+
107+
internal static Content ToContent(this AIContent content) =>
108+
content switch
109+
{
110+
TextContent textContent => new()
111+
{
112+
Text = textContent.Text,
113+
Type = "text",
114+
},
115+
DataContent dataContent => new()
116+
{
117+
Data = dataContent.GetBase64Data(),
118+
MimeType = dataContent.MediaType,
119+
Type =
120+
dataContent.HasTopLevelMediaType("image") ? "image" :
121+
dataContent.HasTopLevelMediaType("audio") ? "audio" :
122+
"resource",
123+
},
124+
_ => new()
125+
{
126+
Text = JsonSerializer.Serialize(content, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object))),
127+
Type = "text",
128+
}
129+
};
104130
}

‎src/ModelContextProtocol/Server/AIFunctionMcpServerTool.cs

Lines changed: 38 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ internal sealed class AIFunctionMcpServerTool : McpServerTool
2222
public static new AIFunctionMcpServerTool Create(
2323
Delegate method,
2424
string? name,
25-
string? description,
25+
string? description,
2626
IServiceProvider? services)
2727
{
2828
Throw.IfNull(method);
@@ -34,7 +34,7 @@ internal sealed class AIFunctionMcpServerTool : McpServerTool
3434
/// Creates an <see cref="McpServerTool"/> instance for a method, specified via a <see cref="Delegate"/> instance.
3535
/// </summary>
3636
public static new AIFunctionMcpServerTool Create(
37-
MethodInfo method,
37+
MethodInfo method,
3838
object? target,
3939
string? name,
4040
string? description,
@@ -195,57 +195,49 @@ public override async Task<CallToolResponse> InvokeAsync(
195195
};
196196
}
197197

198-
switch (result)
198+
return result switch
199199
{
200-
case null:
201-
return new()
202-
{
203-
Content = []
204-
};
205-
206-
case string text:
207-
return new()
208-
{
209-
Content = [new() { Text = text, Type = "text" }]
210-
};
211-
212-
case TextContent textContent:
213-
return new()
214-
{
215-
Content = [new() { Text = textContent.Text, Type = "text" }]
216-
};
217-
218-
case DataContent dataContent:
219-
return new()
220-
{
221-
Content = [new()
222-
{
223-
Data = dataContent.GetBase64Data(),
224-
MimeType = dataContent.MediaType,
225-
Type = dataContent.HasTopLevelMediaType("image") ? "image" : "resource",
226-
}]
227-
};
228-
229-
case string[] texts:
230-
return new()
231-
{
232-
Content = texts
233-
.Select(x => new Content() { Type = "text", Text = x ?? string.Empty })
234-
.ToList()
235-
};
200+
AIContent aiContent => new()
201+
{
202+
Content = [aiContent.ToContent()]
203+
},
204+
null => new()
205+
{
206+
Content = []
207+
},
208+
string text => new()
209+
{
210+
Content = [new() { Text = text, Type = "text" }]
211+
},
212+
Content content => new()
213+
{
214+
Content = [content]
215+
},
216+
IEnumerable<string> texts => new()
217+
{
218+
Content = [.. texts.Select(x => new Content() { Type = "text", Text = x ?? string.Empty })]
219+
},
220+
IEnumerable<AIContent> contentItems => new()
221+
{
222+
Content = [.. contentItems.Select(static item => item.ToContent())]
223+
},
224+
IEnumerable<Content> contents => new()
225+
{
226+
Content = [.. contents]
227+
},
228+
CallToolResponse callToolResponse => callToolResponse,
236229

237230
// TODO https://github.com/modelcontextprotocol/csharp-sdk/issues/69:
238231
// Add specialization for annotations.
239-
240-
default:
241-
return new()
242-
{
243-
Content = [new()
232+
_ => new()
233+
{
234+
Content = [new()
244235
{
245236
Text = JsonSerializer.Serialize(result, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object))),
246237
Type = "text"
247238
}]
248-
};
249-
}
239+
},
240+
};
250241
}
242+
251243
}

‎tests/ModelContextProtocol.Tests/Server/McpServerToolReturnTests.cs

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
using Microsoft.Extensions.AI;
2+
using ModelContextProtocol.Protocol.Types;
3+
using ModelContextProtocol.Server;
4+
using Moq;
5+
6+
namespace ModelContextProtocol.Tests.Server;
7+
public class McpServerToolReturnTests
8+
{
9+
[Fact]
10+
public async Task CanReturnCollectionOfAIContent()
11+
{
12+
Mock<IMcpServer> mockServer = new();
13+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
14+
{
15+
Assert.Same(mockServer.Object, server);
16+
return new List<AIContent>() {
17+
new TextContent("text"),
18+
new DataContent("data:image/png;base64,1234"),
19+
new DataContent("data:audio/wav;base64,1234")
20+
};
21+
});
22+
23+
var result = await tool.InvokeAsync(
24+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
25+
TestContext.Current.CancellationToken);
26+
27+
Assert.Equal(3, result.Content.Count);
28+
29+
Assert.Equal("text", result.Content[0].Text);
30+
Assert.Equal("text", result.Content[0].Type);
31+
32+
Assert.Equal("1234", result.Content[1].Data);
33+
Assert.Equal("image/png", result.Content[1].MimeType);
34+
Assert.Equal("image", result.Content[1].Type);
35+
36+
Assert.Equal("1234", result.Content[2].Data);
37+
Assert.Equal("audio/wav", result.Content[2].MimeType);
38+
Assert.Equal("audio", result.Content[2].Type);
39+
}
40+
41+
[Theory]
42+
[InlineData("text", "text")]
43+
[InlineData("data:image/png;base64,1234", "image")]
44+
[InlineData("data:audio/wav;base64,1234", "audio")]
45+
public async Task CanReturnSingleAIContent(string data, string type)
46+
{
47+
Mock<IMcpServer> mockServer = new();
48+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
49+
{
50+
Assert.Same(mockServer.Object, server);
51+
return type switch
52+
{
53+
"text" => (AIContent)new TextContent(data),
54+
"image" => new DataContent(data),
55+
"audio" => new DataContent(data),
56+
_ => throw new ArgumentException("Invalid type")
57+
};
58+
});
59+
60+
var result = await tool.InvokeAsync(
61+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
62+
TestContext.Current.CancellationToken);
63+
64+
Assert.Single(result.Content);
65+
Assert.Equal(type, result.Content[0].Type);
66+
67+
if (type != "text")
68+
{
69+
Assert.NotNull(result.Content[0].MimeType);
70+
Assert.Equal(data.Split(',').Last(), result.Content[0].Data);
71+
}
72+
else
73+
{
74+
Assert.Null(result.Content[0].MimeType);
75+
Assert.Equal(data, result.Content[0].Text);
76+
}
77+
}
78+
79+
[Fact]
80+
public async Task CanReturnNullAIContent()
81+
{
82+
Mock<IMcpServer> mockServer = new();
83+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
84+
{
85+
Assert.Same(mockServer.Object, server);
86+
return (string?)null;
87+
});
88+
var result = await tool.InvokeAsync(
89+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
90+
TestContext.Current.CancellationToken);
91+
Assert.Empty(result.Content);
92+
}
93+
94+
[Fact]
95+
public async Task CanReturnString()
96+
{
97+
Mock<IMcpServer> mockServer = new();
98+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
99+
{
100+
Assert.Same(mockServer.Object, server);
101+
return "42";
102+
});
103+
var result = await tool.InvokeAsync(
104+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
105+
TestContext.Current.CancellationToken);
106+
Assert.Single(result.Content);
107+
Assert.Equal("42", result.Content[0].Text);
108+
Assert.Equal("text", result.Content[0].Type);
109+
}
110+
111+
[Fact]
112+
public async Task CanReturnCollectionOfStrings()
113+
{
114+
Mock<IMcpServer> mockServer = new();
115+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
116+
{
117+
Assert.Same(mockServer.Object, server);
118+
return new List<string>() { "42", "43" };
119+
});
120+
var result = await tool.InvokeAsync(
121+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
122+
TestContext.Current.CancellationToken);
123+
Assert.Equal(2, result.Content.Count);
124+
Assert.Equal("42", result.Content[0].Text);
125+
Assert.Equal("text", result.Content[0].Type);
126+
Assert.Equal("43", result.Content[1].Text);
127+
Assert.Equal("text", result.Content[1].Type);
128+
}
129+
130+
[Fact]
131+
public async Task CanReturnMcpContent()
132+
{
133+
Mock<IMcpServer> mockServer = new();
134+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
135+
{
136+
Assert.Same(mockServer.Object, server);
137+
return new Content { Text = "42", Type = "text" };
138+
});
139+
var result = await tool.InvokeAsync(
140+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
141+
TestContext.Current.CancellationToken);
142+
Assert.Single(result.Content);
143+
Assert.Equal("42", result.Content[0].Text);
144+
Assert.Equal("text", result.Content[0].Type);
145+
}
146+
147+
[Fact]
148+
public async Task CanReturnCollectionOfMcpContent()
149+
{
150+
Mock<IMcpServer> mockServer = new();
151+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
152+
{
153+
Assert.Same(mockServer.Object, server);
154+
return new List<Content>() { new() { Text = "42", Type = "text" }, new() { Data = "1234", Type = "image", MimeType = "image/png" } };
155+
});
156+
var result = await tool.InvokeAsync(
157+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
158+
TestContext.Current.CancellationToken);
159+
Assert.Equal(2, result.Content.Count);
160+
Assert.Equal("42", result.Content[0].Text);
161+
Assert.Equal("text", result.Content[0].Type);
162+
Assert.Equal("1234", result.Content[1].Data);
163+
Assert.Equal("image", result.Content[1].Type);
164+
Assert.Equal("image/png", result.Content[1].MimeType);
165+
Assert.Null(result.Content[1].Text);
166+
}
167+
168+
[Fact]
169+
public async Task CanReturnCallToolResponse()
170+
{
171+
CallToolResponse response = new()
172+
{
173+
Content = [new() { Text = "text", Type = "text" }, new() { Data = "1234", Type = "image" }]
174+
};
175+
176+
Mock<IMcpServer> mockServer = new();
177+
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
178+
{
179+
Assert.Same(mockServer.Object, server);
180+
return response;
181+
});
182+
var result = await tool.InvokeAsync(
183+
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
184+
TestContext.Current.CancellationToken);
185+
186+
Assert.Same(response, result);
187+
188+
Assert.Equal(2, result.Content.Count);
189+
Assert.Equal("text", result.Content[0].Text);
190+
Assert.Equal("text", result.Content[0].Type);
191+
Assert.Equal("1234", result.Content[1].Data);
192+
Assert.Equal("image", result.Content[1].Type);
193+
}
194+
}

‎tests/ModelContextProtocol.Tests/Server/McpServerToolTests.cs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
using Microsoft.Extensions.DependencyInjection;
1+
using Microsoft.Extensions.AI;
2+
using Microsoft.Extensions.DependencyInjection;
23
using ModelContextProtocol.Protocol.Types;
34
using ModelContextProtocol.Server;
45
using Moq;

0 commit comments

Comments
 (0)