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
This commit is contained in:
Aaron Powell 2025-03-27 15:25:25 +11:00 • committed by GitHub
commit 3c72c348d6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 260 additions and 47 deletions

View file

@ -1,7 +1,9 @@
using Microsoft.Extensions.AI;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Utils;
using ModelContextProtocol.Utils.Json;
using System.Runtime.InteropServices;
using System.Text.Json;
namespace ModelContextProtocol;
@ -101,4 +103,28 @@ public static class AIContentExtensions
Convert.ToBase64String(dataContent.Data.ToArray());
#endif
}
internal static Content ToContent(this AIContent content) =>
content switch
{
TextContent textContent => new()
{
Text = textContent.Text,
Type = "text",
},
DataContent dataContent => new()
{
Data = dataContent.GetBase64Data(),
MimeType = dataContent.MediaType,
Type =
dataContent.HasTopLevelMediaType("image") ? "image" :
dataContent.HasTopLevelMediaType("audio") ? "audio" :
"resource",
},
_ => new()
{
Text = JsonSerializer.Serialize(content, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object))),
Type = "text",
}
};
}

View file

@ -22,7 +22,7 @@ internal sealed class AIFunctionMcpServerTool : McpServerTool
public static new AIFunctionMcpServerTool Create(
Delegate method,
string? name,
string? description,
string? description,
IServiceProvider? services)
{
Throw.IfNull(method);
@ -34,7 +34,7 @@ internal sealed class AIFunctionMcpServerTool : McpServerTool
/// Creates an <see cref="McpServerTool"/> instance for a method, specified via a <see cref="Delegate"/> instance.
/// </summary>
public static new AIFunctionMcpServerTool Create(
MethodInfo method,
MethodInfo method,
object? target,
string? name,
string? description,
@ -195,57 +195,49 @@ internal sealed class AIFunctionMcpServerTool : McpServerTool
};
}
switch (result)
return result switch
{
case null:
return new()
{
Content = []
};
case string text:
return new()
{
Content = [new() { Text = text, Type = "text" }]
};
case TextContent textContent:
return new()
{
Content = [new() { Text = textContent.Text, Type = "text" }]
};
case DataContent dataContent:
return new()
{
Content = [new()
{
Data = dataContent.GetBase64Data(),
MimeType = dataContent.MediaType,
Type = dataContent.HasTopLevelMediaType("image") ? "image" : "resource",
}]
};
case string[] texts:
return new()
{
Content = texts
.Select(x => new Content() { Type = "text", Text = x ?? string.Empty })
.ToList()
};
AIContent aiContent => new()
{
Content = [aiContent.ToContent()]
},
null => new()
{
Content = []
},
string text => new()
{
Content = [new() { Text = text, Type = "text" }]
},
Content content => new()
{
Content = [content]
},
IEnumerable<string> texts => new()
{
Content = [.. texts.Select(x => new Content() { Type = "text", Text = x ?? string.Empty })]
},
IEnumerable<AIContent> contentItems => new()
{
Content = [.. contentItems.Select(static item => item.ToContent())]
},
IEnumerable<Content> contents => new()
{
Content = [.. contents]
},
CallToolResponse callToolResponse => callToolResponse,
// TODO https://github.com/modelcontextprotocol/csharp-sdk/issues/69:
// Add specialization for annotations.
default:
return new()
{
Content = [new()
_ => new()
{
Content = [new()
{
Text = JsonSerializer.Serialize(result, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object))),
Type = "text"
}]
};
}
},
};
}
}

View file

@ -0,0 +1,194 @@
using Microsoft.Extensions.AI;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using Moq;
namespace ModelContextProtocol.Tests.Server;
public class McpServerToolReturnTests
{
[Fact]
public async Task CanReturnCollectionOfAIContent()
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return new List<AIContent>() {
new TextContent("text"),
new DataContent("data:image/png;base64,1234"),
new DataContent("data:audio/wav;base64,1234")
};
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Equal(3, result.Content.Count);
Assert.Equal("text", result.Content[0].Text);
Assert.Equal("text", result.Content[0].Type);
Assert.Equal("1234", result.Content[1].Data);
Assert.Equal("image/png", result.Content[1].MimeType);
Assert.Equal("image", result.Content[1].Type);
Assert.Equal("1234", result.Content[2].Data);
Assert.Equal("audio/wav", result.Content[2].MimeType);
Assert.Equal("audio", result.Content[2].Type);
}
[Theory]
[InlineData("text", "text")]
[InlineData("data:image/png;base64,1234", "image")]
[InlineData("data:audio/wav;base64,1234", "audio")]
public async Task CanReturnSingleAIContent(string data, string type)
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return type switch
{
"text" => (AIContent)new TextContent(data),
"image" => new DataContent(data),
"audio" => new DataContent(data),
_ => throw new ArgumentException("Invalid type")
};
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Single(result.Content);
Assert.Equal(type, result.Content[0].Type);
if (type != "text")
{
Assert.NotNull(result.Content[0].MimeType);
Assert.Equal(data.Split(',').Last(), result.Content[0].Data);
}
else
{
Assert.Null(result.Content[0].MimeType);
Assert.Equal(data, result.Content[0].Text);
}
}
[Fact]
public async Task CanReturnNullAIContent()
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return (string?)null;
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Empty(result.Content);
}
[Fact]
public async Task CanReturnString()
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return "42";
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Single(result.Content);
Assert.Equal("42", result.Content[0].Text);
Assert.Equal("text", result.Content[0].Type);
}
[Fact]
public async Task CanReturnCollectionOfStrings()
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return new List<string>() { "42", "43" };
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Equal(2, result.Content.Count);
Assert.Equal("42", result.Content[0].Text);
Assert.Equal("text", result.Content[0].Type);
Assert.Equal("43", result.Content[1].Text);
Assert.Equal("text", result.Content[1].Type);
}
[Fact]
public async Task CanReturnMcpContent()
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return new Content { Text = "42", Type = "text" };
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Single(result.Content);
Assert.Equal("42", result.Content[0].Text);
Assert.Equal("text", result.Content[0].Type);
}
[Fact]
public async Task CanReturnCollectionOfMcpContent()
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return new List<Content>() { new() { Text = "42", Type = "text" }, new() { Data = "1234", Type = "image", MimeType = "image/png" } };
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Equal(2, result.Content.Count);
Assert.Equal("42", result.Content[0].Text);
Assert.Equal("text", result.Content[0].Type);
Assert.Equal("1234", result.Content[1].Data);
Assert.Equal("image", result.Content[1].Type);
Assert.Equal("image/png", result.Content[1].MimeType);
Assert.Null(result.Content[1].Text);
}
[Fact]
public async Task CanReturnCallToolResponse()
{
CallToolResponse response = new()
{
Content = [new() { Text = "text", Type = "text" }, new() { Data = "1234", Type = "image" }]
};
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return response;
});
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Same(response, result);
Assert.Equal(2, result.Content.Count);
Assert.Equal("text", result.Content[0].Text);
Assert.Equal("text", result.Content[0].Type);
Assert.Equal("1234", result.Content[1].Data);
Assert.Equal("image", result.Content[1].Type);
}
}

View file

@ -1,4 +1,5 @@
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using Moq;