Personalizando la invocación de Tools en MCP para Incluir ProgressToken con C#

Nota: Este post se basa en la versión 0.1.0-preview.5 del paquete NuGet ModelContextProtocol. Estamos en las primeras versiones preliminares de esta librería, por lo que el código podría quedar obsoleto en cualquier momento debido a la rápida evolución del SDK 🙂

¡Hola!
Este post surge a raíz de la discusión en esta issue de GitHub, donde se planteó la necesidad de personalizar la invocación de herramientas al ser ejecutadas automáticamente por IChatClient.

El Problema

Cuando invocamos manualmente una herramienta en MCP, podemos incluir un ProgressToken sin problema, por ejemplo:

1
2
3
4
5
6
7
8
9

    var result = await client.SendRequestAsync(new JsonRpcRequest()
    {
        Method = RequestMethods.ToolsCall,
        Params = new CallToolRequestParams()
        {
            Name = progressTool.ProtocolTool.Name,
            Meta = new() { ProgressToken = new("abc123") },
        },
    }, cancellationToken);

Sin embargo, cuando usamos IChatClient y pasamos herramientas obtenidas con mcpClient.ListToolsAsync(), la invocación se realiza de manera automática por el LLM y no hay un mecanismo para incluir el ProgressToken.
Aquí vemos un ejemplo de código en el que no se puede personalizar la invocación:

1

    await foreach (var update in chatClient.GetStreamingResponseAsync(messages, new() { Tools = [.. tools] }))

Este comportamiento es problemático en escenarios del mundo real™, donde seguramente necesitaremos incluir parámetros adicionales en la invocación de una Tool. No siempre podemos confiar en que el LLM infiera correctamente todos los valores necesarios. Por ejemplo, podríamos necesitar pasar un userId, un token, o cualquier otro dato contextual que no debería ser inferido por el modelo. En este caso concreto, necesitábamos incluir el ProgressToken, que es parte de la especificación de MCP, para poder recibir notificaciones de progreso de ejecución. Este parámetro no es inferido automáticamente por el LLM, por lo que necesitamos añadirlo «a mano» en la invocación de la herramienta.

Solución

Para solventar este problema, creé una clase McpAITool que hereda de AIFunction y permite modificar la solicitud antes de ser enviada. Esto nos da la flexibilidad de incluir Meta, con un ProgressToken, en cada invocación de herramienta.

Implementación de McpAITool

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

    public class McpAITool : AIFunction
    {
        private readonly IMcpClient _client;
        private readonly Func<CallToolRequestParams,McpAITool, Task<CallToolRequestParams>>?     _modifyRequestParamsHandler;
 
        public McpAITool(IMcpClient client, Tool tool, Func<CallToolRequestParams,McpAITool,     Task<CallToolRequestParams>>? modifyRequestParamsHandler = null)
        {
            _client = client;
            _modifyRequestParamsHandler = modifyRequestParamsHandler;
            ProtocolTool = tool;
        }
 
        public McpAITool(IMcpClient client, McpClientTool mcpClientTool, Func<CallToolRequestParams,McpAITool,     Task<CallToolRequestParams>>? modifyRequestParamsHandler = null)
            : this(client, mcpClientTool.ProtocolTool, modifyRequestParamsHandler)
        {
            _client = client;
            _modifyRequestParamsHandler = modifyRequestParamsHandler;
            ProtocolTool = mcpClientTool.ProtocolTool;
        }
 
        /// <summary>Gets the protocol <see cref="Tool"/> type for this instance.</summary>
        public Tool ProtocolTool { get; }
 
        /// <inheritdoc/>
        public override string Name => ProtocolTool.Name;
 
        /// <inheritdoc/>
        public override string Description => ProtocolTool.Description ?? string.Empty;
 
        /// <inheritdoc/>
        public override JsonElement JsonSchema => ProtocolTool.InputSchema;
 
        protected override async Task<object?> InvokeCoreAsync(IEnumerable<KeyValuePair<string, object?>>     arguments, CancellationToken cancellationToken)
        {
            var @params = new CallToolRequestParams()
            {
                Name = ProtocolTool.Name,
                Arguments = arguments
                    .Where(kvp => kvp.Value is not null) 
                    .ToDictionary(
                        kvp => kvp.Key,
                        kvp => JsonSerializer.SerializeToElement(kvp.Value!)
                    )
            };
 
            if (_modifyRequestParamsHandler is not null)
            {
                @params = await _modifyRequestParamsHandler(@params,this);
            }
 
            var result = await _client.SendRequestAsync<CallToolRequestParams, CallToolResponse>(
                RequestMethods.ToolsCall,
                @params,
                cancellationToken: cancellationToken);
 
            // TODO: Review use McpJsonUtilities.JsonContext.Default.CallToolResponse
            return JsonSerializer.SerializeToElement(result);
        }
    }

Creación de McpAITool

Ahora, cuando creamos las AIFunction con ListToolsAsync(), podemos mapearlas al nuevo tipo McpAITool para que incluya el ProgressToken en cada invocación de herramienta:

1
2
3
4
5
6
7
8
9
10
11

    var tools = (await mcpClient.ListToolsAsync())
        .Select(mcpClientTool => new McpAITool(mcpClient,mcpClientTool, (@params, tool)=> Task.FromResult(new CallToolRequestParams()
        {
            Name = @params.Name,
            Arguments = @params.Arguments,
            Meta = new RequestParamsMetadata()
            {
                ProgressToken = new ProgressToken($"{tool.Name}_{Guid.NewGuid().ToString()}"))
            }
        })))
        .ToList();

Ejemplo de Program.cs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

    class Program
    {
        static async Task Main(string[] args)
        {
            using var chatClient = new AzureOpenAIClient(new Uri("https://xxxx.cognitiveservices.azure.com/"), new AzureKeyCredential("xxx"))
                .AsChatClient("gpt-4o-simple")
                .AsBuilder()
                .UseFunctionInvocation()
                .Build();
 
            var mcpClient = await McpClientFactory.CreateAsync(new McpServerConfig
            {
                Id = "everything",
                Name = "Everything",
                TransportType = TransportTypes.StdIo,
                TransportOptions = new Dictionary<string, string>
                {
                    ["command"] = "npx", ["arguments"] = "-y @modelcontextprotocol/server-everything",
                },
            });
 
            var tools = (await mcpClient.ListToolsAsync())
                .Select(mcpClientTool => new McpAITool(mcpClient,mcpClientTool, (@params, tool)=> Task.FromResult(new CallToolRequestParams()
                {
                    Name = @params.Name,
                    Arguments = @params.Arguments,
                    Meta = new RequestParamsMetadata()
                    {
                        ProgressToken = new ProgressToken($"{tool.Name}_{Guid.NewGuid().ToString()}"))
                    }
                })))
                .ToList();
 
 
            mcpClient.AddNotificationHandler(NotificationMethods.ProgressNotification, notification =>
            {
                var progressNotification = notification.Params!.Deserialize<ProgressNotification>()!;
                Console.WriteLine(progressNotification.Progress.ToString());
                return Task.CompletedTask;
            });
 
            List<ChatMessage> messages = [];
            while (true)
            {
                Console.Write("Q: ");
                var input = Console.ReadLine();
                if (input is null || input.Trim().Equals("exit", StringComparison.CurrentCultureIgnoreCase)) 
                    break;
 
                messages.Add(new ChatMessage(ChatRole.User, input));
 
                List<ChatResponseUpdate> updates = [];
                await foreach (var update in chatClient.GetStreamingResponseAsync(messages, new ChatOptions()
                               {
                                   Tools = [.. tools],
                                   ToolMode = ChatToolMode.Auto,
                               }))
                {
                    Console.Write(update);
                    updates.Add(update);
                }
                Console.WriteLine();
 
                messages.AddMessages(updates);
            }
        }
    }

Este programa básicamente hace lo siguiente:

  • Configura un cliente de MCP: Se conecta a un servidor ModelContextProtocol (MCP) para utilizar sus herramientas. Una de esas herramientas soporta notificaciones de progreso, en concreto se llama longRunningOperation.
  • Añade un ProgressToken a cada herramienta: Al listar las herramientas disponibles, se les agrega un parámetro extra llamado ProgressToken, necesario para recibir notificaciones de progreso mientras se ejecutan las herramientas.
  • Recibe notificaciones de progreso: El programa se configura para recibir y mostrar las actualizaciones de progreso mientras las herramientas de MCP se ejecutan.
  • Interactúa con el usuario: El programa pide al usuario que ingrese preguntas, luego obtiene respuestas del modelo de chat de OpenAI. Durante este proceso, las herramientas de MCP se utilizan y el progreso de su ejecución se muestra en tiempo real.

La clave del código es que el ProgressToken permite saber el estado de avance de las herramientas (que lo soporten) en ejecución, y el programa está preparado para mostrar esas actualizaciones en la consola a medida que ocurren.

Ejemplo de ejecución de Tool con notificaciones de progreso
Ejemplo de ejecución de Tool con notificaciones de progreso

Conclusión

Con esta solución, logramos personalizar la invocación de herramientas en IChatClient, permitiendo incluir ProgressToken para recibir notificaciones de progreso. Dado que estamos en las versiones preliminares del SDK de ModelContextProtocol, es probable que en futuras versiones se incluyan mecanismos para manejar esto directamente. Por ahora, esta solución nos permite seguir trabajando, dicho esto, es código de juguete y sin más pruebas que un proyecto de consola, úsalo con cautela 🙂
¡Un saludo!

Deja un comentario

Tu dirección de correo electrónico no será publicada. Los campos obligatorios están marcados con *