The following example shows how to create a MistralAITokenCounterMiddleware IMiddleware and count the token usage when chatting with MistralClientAgent.

Overview

To collect the token usage for the entire chat session, one easy solution is simply collect all the responses from agent and sum up the token usage for each response. To collect all the agent responses, we can create a middleware which simply saves all responses to a list and register it with the agent. To get the token usage information for each response, because in the example we are using MistralClientAgent, we can simply get the token usage from the response object.

Note

You can find the complete example in the Example13_OpenAIAgent_JsonMode.

• Step 1: Adding using statement

using AutoGen.Core;
using AutoGen.Mistral;

• Step 2: Create a MistralAITokenCounterMiddleware class which implements IMiddleware. This middleware will collect all the responses from the agent and sum up the token usage for each response.

public class MistralAITokenCounterMiddleware : IMiddleware
{
    private readonly List<ChatCompletionResponse> responses = new List<ChatCompletionResponse>();
    public string? Name => nameof(MistralAITokenCounterMiddleware);

    public async Task<IMessage> InvokeAsync(MiddlewareContext context, IAgent agent, CancellationToken cancellationToken = default)
    {
        var reply = await agent.GenerateReplyAsync(context.Messages, context.Options, cancellationToken);

        if (reply is IMessage<ChatCompletionResponse> message)
        {
            responses.Add(message.Content);
        }

        return reply;
    }

    public int GetCompletionTokenCount()
    {
        return responses.Sum(r => r.Usage.CompletionTokens);
    }
}

• Step 3: Create a MistralClientAgent

var apiKey = Environment.GetEnvironmentVariable("MISTRAL_API_KEY") ?? throw new Exception("Missing MISTRAL_API_KEY environment variable.");
var mistralClient = new MistralClient(apiKey);
var agent = new MistralClientAgent(
    client: mistralClient,
    name: "assistant",
    model: MistralAIModelID.OPEN_MISTRAL_7B);

• Step 4: Register the MistralAITokenCounterMiddleware with the MistralClientAgent. Note that the order of each middlewares matters. The token counter middleware needs to be registered before mistralMessageConnector because it collects response only when the responding message type is IMessage<ChatCompletionResponse> while the mistralMessageConnector will convert IMessage<ChatCompletionResponse> to one of TextMessage, ToolCallMessage or ToolCallResultMessage.

var tokenCounterMiddleware = new MistralAITokenCounterMiddleware();
var mistralMessageConnector = new MistralChatMessageConnector();
var agentWithTokenCounter = agent
    .RegisterMiddleware(tokenCounterMiddleware)
    .RegisterMiddleware(mistralMessageConnector)
    .RegisterPrintMessage();

• Step 5: Chat with the MistralClientAgent and get the token usage information from the response object.

await agentWithTokenCounter.SendAsync("write a long, tedious story");
Console.WriteLine($"Completion token count: {tokenCounterMiddleware.GetCompletionTokenCount()}");
tokenCounterMiddleware.GetCompletionTokenCount().Should().BeGreaterThan(0);

Output

When running the example, the completion token count will be printed to the console.

Completion token count: 1408 # might be different based on the response