AzureAICommunity.Agent.Middleware.TokenUsageMiddleware 1.0.0

<div align="center">

🪙 AzureAICommunity - Agent - Token Usage Middleware

Enforce per-user token quotas and capture detailed usage metrics across every AI agent completion call.

Track, throttle, and bill token consumption per user — with zero friction.

Getting Started · Quota Stores · Period Keys · Callbacks · Contributing

</div>

Overview

AzureAICommunity.Agent.Middleware.TokenUsageMiddleware is a plug-and-play quota and metering layer for AI agent pipelines built on Microsoft.Extensions.AI. Before every request it checks a user's accumulated token count against a configurable limit and throws a QuotaExceededException if exhausted. After every successful completion (streaming or non-streaming) it persists the token delta to an IQuotaStore and fires an optional onUsage callback with a TokenUsageRecord.

✨ Features

📦 Installation

dotnet add package AzureAICommunity.Agent.Middleware.TokenUsageMiddleware

🚀 Quick Start

using AzureAICommunity.Agent.Middleware.TokenUsageMiddleware;
using Microsoft.Extensions.AI;
using OllamaSharp;

IChatClient ollamaClient = new OllamaApiClient("http://localhost:11434/", "llama3.2");

var quotaStore = new InMemoryQuotaStore();

IChatClient client = ollamaClient
    .AsBuilder()
    .Use(inner => new TokenUsageMiddleware(
        inner,
        quotaStore: quotaStore,
        quotaTokens: 500,
        onUsage: async (record, ct) =>
        {
            Console.WriteLine($"[{record.UserId}] used {record.TotalTokens} tokens " +
                              $"({record.UsedTokensAfterCall}/{record.QuotaTokens} total)");
            await Task.CompletedTask;
        }))
    .Build();

var options = new ChatOptions
{
    AdditionalProperties = new() { ["user_id"] = "Vinoth" }
};

var response = await client.GetResponseAsync("What is the capital of France?", options);
Console.WriteLine(response.Message.Text);

🗄️ Quota Stores

The middleware delegates all persistence to an IQuotaStore. Two implementations are available out of the box:

InMemoryQuotaStore

Fast, zero-dependency store backed by an in-process dictionary. Suitable for development, testing, and single-process apps where quota data does not need to survive restarts.

var quotaStore = new InMemoryQuotaStore();

Custom / Persistent Store

For production or multi-process deployments, implement IQuotaStore with any backend (Redis, SQL, Azure Table Storage, etc.):

public sealed class RedisQuotaStore : IQuotaStore
{
    public long GetUsage(string userId, string periodKey) { /* ... */ }
    public void AddUsage(string userId, string periodKey, long tokens) { /* ... */ }
}

🗓️ Period Keys

The quota is scoped to a period key — a string that resets the counter. Use the built-in PeriodKeys helpers or supply any custom delegate:

// Weekly quota
var client = ollamaClient.AsBuilder()
    .Use(inner => new TokenUsageMiddleware(
        inner,
        quotaStore: quotaStore,
        quotaTokens: 1000,
        periodKeyFn: PeriodKeys.Week))
    .Build();

// Custom period (e.g. hourly)
var client = ollamaClient.AsBuilder()
    .Use(inner => new TokenUsageMiddleware(
        inner,
        quotaStore: quotaStore,
        quotaTokens: 200,
        periodKeyFn: () => DateTimeOffset.UtcNow.ToString("yyyy-MM-dd-HH")))
    .Build();

🔔 Callbacks

onUsage — Post-completion metrics

Fires after every successful completion with an immutable TokenUsageRecord snapshot:

.Use(inner => new TokenUsageMiddleware(
    inner,
    quotaStore: quotaStore,
    quotaTokens: 500,
    onUsage: async (record, ct) =>
    {
        // Forward to a billing system, database, or telemetry sink
        Console.WriteLine(
            $"User={record.UserId} Period={record.PeriodKey} Model={record.Model} " +
            $"Input={record.InputTokens} Output={record.OutputTokens} Total={record.TotalTokens} " +
            $"Used={record.UsedTokensAfterCall}/{record.QuotaTokens} Streaming={record.IsStreaming}");
        await Task.CompletedTask;
    }))

onQuotaExceeded — Pre-exception hook

Fires when the quota check fails, before QuotaExceededException is thrown, giving you a chance to log or alert:

.Use(inner => new TokenUsageMiddleware(
    inner,
    quotaStore: quotaStore,
    quotaTokens: 500,
    onQuotaExceeded: async (info, ct) =>
    {
        Console.WriteLine(
            $"[QUOTA] User={info.UserId} has used {info.UsedTokens}/{info.QuotaTokens} tokens " +
            $"in period {info.PeriodKey}. Request blocked.");
        await Task.CompletedTask;
    }))

🧑‍💻 Custom User Identification

By default the middleware reads the "user_id" key from ChatOptions.AdditionalProperties, falling back to "anonymous". Supply a custom userIdGetter delegate to integrate with your own identity system:

.Use(inner => new TokenUsageMiddleware(
    inner,
    quotaStore: quotaStore,
    quotaTokens: 500,
    userIdGetter: (messages, options) =>
    {
        // e.g. extract from a JWT claim stored in AdditionalProperties
        return options?.AdditionalProperties?.TryGetValue("sub", out var sub) == true
            ? sub?.ToString() ?? "anonymous"
            : "anonymous";
    }))

⚙️ How It Works

1. Intercept     →  middleware captures the incoming request
2. Quota check   →  GetUsage(userId, periodKey) compared against quotaTokens
3. Block         →  if used >= quota: fire onQuotaExceeded, throw QuotaExceededException
4. Delegate      →  forward to the inner IChatClient (LLM is called)
5. Record usage  →  AddUsage(userId, periodKey, totalTokens)
6. Emit record   →  fire onUsage callback with a TokenUsageRecord snapshot
7. Return        →  response (or stream) is returned to the caller

📄 Constructor Reference

public TokenUsageMiddleware(
    IChatClient inner,                                                  // Inner chat client to delegate to
    IQuotaStore quotaStore,                                             // Per-user token storage backend
    long quotaTokens,                                                   // Maximum tokens per user per period (> 0)
    Func<TokenUsageRecord, CancellationToken, Task>? onUsage = null,   // Post-completion callback
    Func<QuotaExceededInfo, CancellationToken, Task>? onQuotaExceeded = null, // Pre-exception callback
    Func<IEnumerable<ChatMessage>, ChatOptions?, string>? userIdGetter = null, // User ID extractor
    Func<string>? periodKeyFn = null                                    // Period key factory (default: PeriodKeys.Month)
)

📋 Type Reference

🤝 Contributing

Contributions are welcome! Please open an issue to discuss what you'd like to change before submitting a pull request.

1. Fork the repository
2. Create a feature branch (git checkout -b feature/my-feature)
3. Commit your changes (git commit -m 'Add my feature')
4. Push to the branch (git push origin feature/my-feature)
5. Open a Pull Request

📄 License

MIT