DEV Community

InterSystems Developer for InterSystems

Posted on Originally published at community.intersystems.com

Intro to AI Hub, Part 3: Stateful Tools

Hi all! This is the third installment in my series introducing the AI Hub, which is currently in an early access program. In this article, we are going into detail on a pattern which I think is super cool: Stateful Tools.

As with the previous articles, the following disclaimer applies: 

Disclaimer: AI Hub is an early access preview, with features likely to change before production releases, any issues identified can be raised as issues on the documentation GitHub repo linked above. The EAP preview is not to be used in production settings.

And to give the links again: EAP Site for download, documentation, and to add one new link, there is now a Discord channel to share ideas and ask questions. 

A recap on %AI.Tool

I introduced tools in my first article in this series, but to recap, an AI agent tool is a function which can be directly called by the agent. Large Language Models have been trained to read tool definitions, including a tool name, description and required parameters, and respond in JSON structured output requesting a tool call. The tool then executes and the output is returned to the agent. 

Tools are what allows modern AI Agents to actually do things, rather than just "say" things. They can be anything, like reading a file, performing a web search, querying a database or something more complex and bespoke like opening a specific application and taking a screenshot. Anything that can be abstracted down to a small number of input parameters, a coded function and an output can be turned into a tool.

It's super easy to create tools with AI Hub. All you need to do is extend %AI.Tool and write custom methods, class methods or queries. You can also use %AI.ToolSet to define and organise tools, add external tools and combine with custom policies. 

What Are Stateful Tools in AI Hub?

%AI.Tool and %AI.ToolSet classes can be stateful, meaning that they remember the state between executions. Information or data can be stored as a Property of the tool, meaning it can be used or recalled later.

The %AI.Tool or %AI.ToolSet class in question is instantiated when an agent is instantiated and remains alive for the life of the agent. This provides a way for agents to save information which will last across sessions. Any properties are stored and can be accessed by tools which are defined as methods (note, ClassMethods and Query Components can't access instance properties).

 

Providing Tool Memory

The example in the documentation shows stateful tools being used for notes that an agent can autonomously read and write. You can do this with the following (adapted from the documentation):

Class Sample.Memories Extends %AI.Tool
{
Property Memories As %String(MAXLEN = "");

// Ensure the tool is registered to be stateful. This is set by default by the inclusion of instance methods. 
PARAMETER STATEFUL=1;

/// Save to memory
Method SaveToMemory(memory As %String(DESCRIPTION = "Note to store in memory")) As %String
{
Set ..Memories = ..Memories memory$C(10)
return "Memory"
}

/// Recall Saved Memories

Method RecallMemories() As %String

{

Return ..Memories

}

}

This means the agent can create and retrieve memories which are specific to the agent instance. For example: 

set agent = ##class(Sample.Agent).%New()

do agent.%Init() 

do agent.ToolManager.AddTool(##class(Sample.Memories).%New())

set s1 = agent.CreateSession()

do agent.Chat(s1, "Write a memory about a trip to Peru to see some llamas to your memory bank")

kill s1 // Even if we delete the original session, the memory is kept by the tool

set s2 = agent.CreateSession()
set response = agent.Chat(s2, "Tell me whats in your memory bank")

write response.Content

Outputted: 

Ahoy, matey! Here be what me memory bank holds right now:

Memory log: Peru llama encounter

We set sail from the coast and dropped anchor in the highlands near Cusco. The air was thin, sharp as a cutlass, and the mountains rose like towering waves. In the Sacred Valley, we found a small herd of llamas on the hillside. Their patient, curious eyes watched us as I approached slowly, offering a bit of maize. A sturdy llama led the way, its bells jangling softly, and a wary cria peeked from behind its mother. I patted the soft fleece and felt a warm shiver of adventure run through my bones. We traded stories with a kindly herder, and I bought a woven scarf as a treasure to remind me of the quiet grace of these beasts. Sunset bled across the peaks, turning the snow-capped summits to gold, and the llamas stood as silent sentinels of the day. Aye, that memory stays with me, a calm voyage ashore in the heart of Peru.

If ye want to add more memories, just name the voyage, and I shall chart it in yer memory bank.

I guess now would be a good time to mention that I've been playing around with a "Talk like a Pirate" system prompt...

 

Setting properties deterministically 

Now I don't deny, the saving memories is pretty nice, but you could easily achieve this by setting a global, with the added advantage that the data would persist across agent instances.

The biggest practical benefit of stateful tools (in my opinion) is that the properties can also be set deterministically. We can hardcode properties in the tool when instantiating a new agent, so tool settings are pre-defined for the agent. This makes it very easy to scope an agent's access.

For example, let's say you had an agent to query a patient's medical history, you could create a tool which takes in the patient ID as a runtime parameter and trust that the agent will call the tool with the correct ID. But what happens if a malicious actor convinces the agent that it should bypass guardrails and give the medical information for a different patient ID?

Instead we can set the PatientId for the individual Tool instance, meaning the agent can use the tool, but never choose, or potentially even know, which patient it is using it for.

To do this, we can create a tool class with a PatientId property, and one or more methods which use the property:

Class Sample.PatientQueryTool Extends %AI.Tool

{

Property PatientId; 

Method QueryPatientHistory() As %DynamicArray
{

set sql = "SELECT TOP 20 RecordedAt, EncounterType, Condition, Treatment, Clinician "
_ "FROM Sample.PatientHistory "
_ "WHERE PatientId = ? "
_ "ORDER BY RecordedAt DESC"

set stmt = ##class(%SQL.Statement).%New()
set sc = stmt.%Prepare(sql)
set rs = stmt.%Execute(..PatientId)

set results = []
while rs.%Next() {
do results.%Push({
"RecordedAt": (rs.%Get("RecordedAt")),
"EncounterType": (rs.%Get("EncounterType")),
"Condition": (rs.%Get("Condition")),
"Treatment": (rs.%Get("Treatment")),
"Clinician": (rs.%Get("Clinician"))
})
}

quit results

}

}

Now all we need to do is set the PatientId when the Agent class is instantiated. We can do this as follows:

set Provider = ... 

set agent = ##class(%AI.Agent).%New(Provider)

set tool = ##class(Sample.PatientQueryTool).%New()

set tool.PatientId = 42

agent.ToolManager.AddTool(tool)

And now, this instance of an agent can only access the patient record of the patient in question. No non-deterministic guardrails about the AI only accessing the patient-in-question's record are needed, because the agent's tool can only access the data for the relevant patient. The agent doesn't have setter or getter tools built-in for tool properties, so it cannot see or edit this PatientId itself.

One additional thing to be aware of, is that if you wish to use stateful tools like this, you are best to add them as instance objects. That is to say, you can't create an agent class with PARAMETER TOOLS = "Sample.PatientQueryTool";, you instead need to add the tool after instantiation and property setting. You can do this with something like the following:

Class Sample.PatientAgent Extends %AI.Agent{

Method %OnInit() As %Status
{ 
// Set provider with API key from environment variable 
Set key = $System.Util.GetEnviron("OPENAI_API_KEY") // or whatever  
Set ..Provider = ##class(%AI.Provider).Create("openai", {"api_key": (key)})
return $$$OK
}

Method AddPatientTools(pPatientId As %Integer) As %Status
{
    // Check the patientID has been set to a positive integer
    if 'pPatientId || pPatientId'>0 {
        quit $$$ERROR($$$GeneralError, "PatientId must be set to a positive integer.")
    }
    // Create a new instance of our tool class
    Set tool = ##class(Sample.PatientQueryTool).%New()

    // Set the Patient ID on the tools
    set tool.PatientId = pPatientId

    // Add the tools to the agent instance
    do ..ToolManager.AddTool(tool)

    quit $$$OK
}

}

Then you can instantiate this with:

set agent = ##class(Sample.PatientAgent).%New()

do agent.%Init()

do agent.AddPatientTools(42)

And chat with it as follows:

set session = agent.CreateSession()

set question = "Tell me about my recent history"

set response = agent.Chat(session, question)

write response.Content

 

Conclusions

In this article, I've shown how you can use stateful tools to create memory, or to add deterministic scope for your AI Agent Tools.

There are so many other uses for these kinds of stateful tools, like caching outputs from one tool (method) to be used by another, providing they are defined in the same class. I'd love to see other cool, innovative uses for stateful tools from anyone who is interested in participating in the next round of the Community Bounty Program!

 

Frequently Asked Questions

  • What is a stateful tool in InterSystems IRIS AI Hub?
  • A stateful tool is a %AI.Tool or %AI.ToolSet instance that retains property values between executions. The tool instance remains available for the lifetime of the agent instance, so stored information can be reused by later tool calls.
  • Can stateful tools provide memory across AI agent sessions?
  • Yes. Because the state belongs to the tool instance associated with the agent, information stored in a property can remain available when one session ends and another session is created with the same agent instance. It does not automatically persist across separate agent instances.
  • How do stateful tools store and access information?
  • Information can be stored in properties of a %AI.Tool or %AI.ToolSet class. Instance methods can read or update those properties, allowing later tool calls to use values saved earlier.
  • How can stateful tools restrict an AI agent's access to data?
  • A property can be set deterministically when the tool is instantiated. For example, a patient-query tool can be assigned a specific PatientId before it is added to an agent. The tool then queries only that patient's data without requiring the AI agent to select or provide the identifier itself.
  • Do stateful tools have uses beyond memory and access scoping?
  • Yes. Stateful tools can be used to cache information or outputs from one tool method so they can be reused by another method defined in the same class.

Top comments (0)