AzureAICommunity.Agent.Middleware.ToolLimitMiddleware 1.0.0

<div align="center">

🛠️ AzureAICommunity - Agent - Tool Limit Middleware

Prevent runaway tool calls by enforcing global and per-tool call limits across every AI agent completion.

Getting Started · Per-Tool Limits · Inspect Usage · How It Works · Contributing

</div>

Overview

AzureAICommunity.Agent.Middleware.ToolLimitMiddleware is a lightweight guard layer for AI agent pipelines built on Microsoft.Extensions.AI. During each completion it tracks every FunctionCallContent emitted by the model and silently suppresses any calls that breach a configurable global cap or an optional per-tool cap. When calls are suppressed, a user-role message is appended to the conversation so the model is aware that limits have been reached.

✨ Features

📦 Installation

dotnet add package AzureAICommunity.Agent.Middleware.ToolLimitMiddleware

🚀 Quick Start

using System.ComponentModel;
using AzureAICommunityAgent.Middleware.ToolLimiting;
using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;
using OllamaSharp;

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

var weatherTool = AIFunctionFactory.Create(
    ([Description("The location to get the weather for.")] string location)
        => $"The weather in {location} is cloudy with a high of 15\u00b0C.",
    "GetWeather");

IChatClient client = ollamaClient
    .AsBuilder()
    .UseToolLimit(new ToolLimits { GlobalMax = 5 })
    .Build();

AIAgent agent = new ChatClientAgent(client,
    instructions: "You are a helpful assistant with access to a weather tool.",
    tools: [weatherTool]);

var response = await agent.RunAsync("What is the weather like in Amsterdam?");
Console.WriteLine(response.Text);

🔧 Per-Tool Limits

In addition to the global cap, you can restrict individual tools independently:

var weatherTool = AIFunctionFactory.Create(/* ... */, "GetWeather");
var youtubeTool = AIFunctionFactory.Create(/* ... */, "SearchVideos");

IChatClient client = ollamaClient
    .AsBuilder()
    .UseToolLimit(new ToolLimits
    {
        GlobalMax = 10,
        PerToolMax = new Dictionary<string, int>
        {
            ["GetWeather"]   = 3,
            ["SearchVideos"] = 2
        }
    })
    .Build();

AIAgent agent = new ChatClientAgent(client,
    instructions: "You are a helpful assistant.",
    tools: [weatherTool, youtubeTool]);

Any call to GetWeather beyond 3, or to SearchVideos beyond 2, is silently removed — even if the global limit has not been reached.

📊 Inspect Usage

After building the client with UseToolLimit, retrieve the tracker via GetService<IToolLimitTracker>():

IChatClient client = ollamaClient
    .AsBuilder()
    .UseToolLimit(new ToolLimits
    {
        GlobalMax = 5,
        PerToolMax = new Dictionary<string, int> { ["GetWeather"] = 3 }
    })
    .Build();

// Run the agent
AIAgent agent = new ChatClientAgent(client,
    instructions: "You are a helpful assistant with access to a weather tool.",
    tools: [weatherTool]);

var response = await agent.RunAsync("What is the weather like in Amsterdam?");
Console.WriteLine(response.Text);

// Retrieve the tracker from the pipeline
var tracker = client.GetService<IToolLimitTracker>();

ToolUsageState usage = tracker!.GetCurrentUsage();
Console.WriteLine($"Total allowed calls: {usage.TotalCalls} / {usage.GlobalLimit}");
Console.WriteLine("Per-tool usage (attempted / allowed / limit):");
foreach (var (tool, attempted) in usage.PerTool)
{
    usage.PerToolAllowed.TryGetValue(tool, out var allowed);
    usage.PerToolLimits.TryGetValue(tool, out var perMax);
    var limitText = perMax > 0 ? $" / {perMax}" : string.Empty;
    Console.WriteLine($"  {tool}: attempted={attempted}  allowed={allowed}{limitText}");
}

// Reset counters for a new session
tracker.Reset();

Example output:

Total allowed calls: 3 / 5
Per-tool usage (attempted / allowed / limit):
  GetWeather: attempted=5  allowed=3 / 3

📄 Constructor Reference

public ToolLimitMiddleware(
    IChatClient innerClient,    // Inner chat client to delegate to
    ToolLimits? limits = null   // Limits configuration (defaults: GlobalMax = 10, no per-tool limits)
)

Extension method

builder.UseToolLimit(new ToolLimits
{
    GlobalMax = 10,
    PerToolMax = new Dictionary<string, int> { ["GetWeather"] = 3, ["SearchVideos"] = 2 }
});

📋 Type Reference

ToolUsageState properties

🤝 Contributing

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

📁 Repository: https://github.com/rvinothrajendran/AgentFramework

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

👤 Author

Built and maintained by Vinoth Rajendran.

• 🐙 GitHub: github.com/rvinothrajendran — follow for more projects!
• 📺 YouTube: youtube.com/@VinothRajendran — subscribe for tutorials and demos!
• 💼 LinkedIn: linkedin.com/in/rvinothrajendran — let's connect!

📄 License

MIT