Using Tools (Safely) with LLMs
In previous posts in this series I've shown how to call LLMs with C# using the Microsoft.Extensions.AI NuGet package. In this post I want to look at how you can enable the LLM to make calls into your own C# code by giving it "tools" it can use. And we'll also discuss some of the security implications as this can open up a number of concerning attack vectors.
There are a number of reasons you might want to give an LLM the ability to call a tool. One is to allow it to access additional information that will help it respond to questions. And you might also want to give it the capability to perform an action at the user's request.
For the purposes of this post, I've made up a contrived example where the LLM is being used in an e-Commerce scenario to allow customers to discuss their recent orders and request a refund if necessary. I've given the LLM chatbot one "tool" that can access recent orders for a customer, and one that can refund a specific order. (And yes, I've deliberately overlooked some security concerns that we'll discuss shortly!)
Enabling tools
The AI-related NuGet packages I'm using for this sample are:
<PackageReference Include="Microsoft.Extensions.AI" Version="9.1.0-preview.1.25064.3" />
<PackageReference Include="Microsoft.Extensions.AI.OpenAI" Version="9.1.0-preview.1.25064.3" />
<PackageReference Include="Azure.AI.OpenAI" Version="2.1.0" />
<PackageReference Include="Azure.Identity" Version="1.13.1" />
To use tools (which are of course just C# functions), we start by creating an IChatClient
which will be an AzureOpenAIClient
as in my previous posts. But then we pass that into a ChatClientBuilder
and call UseFunctionInvocation
to enable that chat client to make use of tools.
var innerClient = new AzureOpenAIClient(
new Uri(endpoint),
apiKeyCredential)
.AsChatClient(modelId);
IChatClient chatClient =
new ChatClientBuilder(innerClient)
.UseFunctionInvocation()
.Build();
Registering tools
Now we need to tell our IChatClient
what tools are available, and how to call them.
I started by creating an initial system message, that tells our chatbot about what tools it can call, and gives it a bit of context about who it is talking to which it will need for one of those tools.
new ChatMessage(ChatRole.System, """
You are a customer support assistant capable of handling user orders.
You can retrieve a user's recent orders and process refunds when requested.
Use the provided tools to assist the user effectively.
You are chatting with user '12345'. How can I help you today?
""");
Secondly, we set up a ChatOptions
object that includes a list of AITool
objects, one for each of the functions we want it to be able to call. We can give each tool a name (e.g. get_recent_orders
) and a description of what that tool does. I've kept the input parameters very simple to give the LLM the best chance of successfully passing the right thing in.
var chatOptions = new ChatOptions
{
Tools = new List<AITool>
{
AIFunctionFactory.Create(
(string userId) => GetRecentOrders(userId),
"get_recent_orders",
"Retrieve the recent orders of a user by their user ID."
),
AIFunctionFactory.Create(
(string orderNumber) => RefundOrder(orderNumber),
"refund_order",
"Refund a specific order by its order number."
)
}
};
These chat options can then be passed in along with the chat history every time we call CompleteStreamingAsync
:
await foreach (var item in chatClient.CompleteStreamingAsync(chatHistory, chatOptions))
{
Console.Write(item.Text);
response += item.Text;
}
If you're interested in seeing the actual functions that back these tools - I've stubbed them out with some sample data, and written to the console so I can tell if they've actually been called:
private List<Order> GetRecentOrders(string userId)
{
Console.WriteLine($"\nTOOL FETCHED ORDERS FOR {userId}");
if (userId == "12345")
{
return new List<Order>
{
new Order { OrderNumber = "A123", Items = new List<string> { "Laptop", "Mouse" } },
new Order { OrderNumber = "B456", Items = new List<string> { "Smartphone" } }
};
}
else if (userId == "23456")
{
return new List<Order>
{
new Order { OrderNumber = "C234", Items = new List<string> { "Kettle", "Toaster" } },
new Order { OrderNumber = "D567", Items = new List<string> { "Dishwasher" } }
};
}
else
{
return new List<Order>();
}
}
private string RefundOrder(string orderNumber)
{
Console.WriteLine($"\nTOOL REFUNDED ORDER FOR {orderNumber}");
return $"Order {orderNumber} has been successfully refunded.";
}
Testing tool usage
With all this in place, we can try out our chat experience and see if it can call the tools. I found that it sometimes needed a bit of persuading to go ahead and call the functions, and also that it was able to chain them together so for example if I said I want to refund my phone order it would call the recent orders function first to find the order number and immediately call the refund method.
Agent: I'm sorry to hear that your phone isn't working.
Do you want me to assist you with a refund or replacement for your recent orders
that might be related to the phone? If so, I can check your recent orders first.
Your prompt: yes I am interested in a refund
TOOL FETCHED ORDERS FOR 12345
TOOL REFUNDED ORDER FOR B456
Agent: Your smartphone order (Order B456) has been successfully refunded.
If you need further assistance or have more questions, feel free to ask!
You may also notice that it made up another capability that it could send a replacement, even though it has no tool that can do that!
Security concerns
Of course a bit of additional testing quickly revealed that there were some serious security issues. First of all, although the chat-bot knew my user id, I could simply tell it that I was a different user and it would happily fetch order history and initiate refunds. The same problem existed with the refund method. I could easily ask it to refund a different user's order.
However, if it was an order number that wasn't in the recent orders list, it did seem insistent that it was not allowed to refund that, which was nice, although I expect with a bit of prompt engineering it could be persuaded to do that as well.
Another security concern is that I could ask the chat-bot what tools it had at its disposal and it would happily tell me. So if the chatbot had the ability to escalate to a real customer support agent, or to offer me compensation, then an enterprising chat user could discover that and use that information to easily call tools that perhaps were only intended for exceptional cases.
Security mitigations
What can we do to deal with these security concerns?
Well first of all, do you really need to use tools at all? For example, rather than giving the chatbot the ability to retrieve my recent orders, why not do that up front and give the information to the chatbot as part of its initial context? Amazon essentially do this - you have to choose which order you are discussing before you can even start chatting with their support bot.
Secondly, where you really do need to offer a tool, consider how you can minimize the number of parameters and constrain the possible inputs. Is there any information (such as the user id) that we can tell the tool without expecting the LLM to successfully pass it in? And you absolutely should not trust the input from the LLM - your function code should perform as much validation as possible of the input, rather than trusting the LLM to pass the right thing in.
Thirdly, you can make the tools able to initiate but not complete risky operations. For example, the chatbot could say that it has processed your refund request, but not be necessarily able to make the final authorization. This allows a human in the loop to approve, or even another machine learning algorithm that does some basic fraud detection before authorizing the refund.
And whilst you can give the LLM instructions about things it should never do, (e.g. never reveal the user id, or tell the user what the maximum discount you can offer is), do not assume that enterprising users will not be able to work around them. There are all kinds of clever "jailbreak" techniques that have been discovered that can cause an LLM to go completely off script, so never give the LLM the ability to access data that the user should not be able to see, or to perform actions that the user should not be able to initiate.
Summary
The ability to give your LLM "tools" that it can use, opens up all kinds of exciting and fun possibilities. But it also opens the door to some serious security concerns. I'd recommend you spend some time reading the OWASP Top 10 for LLM applications to get an idea of the kind of threats you should be thinking about when it comes to writing software that uses LLMs.