Creating an AI Agent in Java
I made a Java library to create AI assistant that can call Java functions in an easy way.
To start, install the aiagent
Maven library (or the same library from mavenCentral()
using Gradle):
<dependency>
<groupId>com.simonbrs</groupId>
<artifactId>aiagent</artifactId>
<version>1.2</version>
</dependency>
Then, create an instance of any class. For this example, I will use a simple calculator class:
public class Calculator {
private double memory = 0.0;
public double add(double a, double b) {
return a + b;
}
public double getMemory() {
return memory;
}
public void setMemory(double value) {
this.memory = value;
}
}
Register the OpenAI Agent:
import com.simonbrs.aiagent.OpenAIAgent;
String apiKey = System.getenv("OPENAI_API_KEY");
OpenAIAgent agent = new OpenAIAgent(apiKey, "gpt-4o");
Create an instance of the class and pass it to agent.registerMethods
(or, if you have an existing class, you can just pass its instance):
Calculator calculator = new Calculator();
agent.registerMethods(calculator);
And then use sendMessage
to send the agent messages. It will use the functions when appropriate. For example:
String[] questions = {
"What is 15 added to 3?",
"Store the number 42 in memory",
"Add ten to the number stored in memory",
};
for (String question : questions) {
try {
System.out.println("\nQuestion: " + question);
String response = agent.sendMessage(question).get();
System.out.println("Answer: " + response);
} catch (Exception e) {
System.err.println("Error: " + e.getMessage());
}
}
This will produce the following output (the Tool called
prints are just for clarity, the agent doesn’t print those):
Question: What is 15 added to 3?
Tool called: add(15.0, 3.0)
Answer: 15 added to 3 is 18.0.
Question: Store the number 42 in memory
Tool called: setMemory(42.0)
Answer: The number 42 has been stored in memory.
Question: Add ten to the number stored in memory
Tool called: getMemory()
Tool called: add(42.0, 10.0)
Answer: The number stored in memory is 42. Adding ten to it results in 52.0.
The agent can choose to use any number of method calls and send any number of messages to best answer the query.
A real-life example of this is when you want to allow a customer to query a complex datastore using natural language. The DTO of the datastore could be passed to the AI agent, with parameterized functions that query data in a specific way. But I imagine many more usecases are possible.
The big advantage of this library is that it makes interacting with the OpenAI API extremely simple. Normally, you’d have to create complex schemas to make such interaction possible. This library does that for you.
What are AI Agents/Assistants?
So far for the general pitch. Let’s take a step back and look at AI agents.
AI Agents, in a nutshell, are chat-based models that can modify an external environment. There are a couple of features, but the feature used here is “function calling”. When the AI decides that a function needs to be called, the assistant returns the http call with an ‘requires_action’ state. In response, we execute a piece of logic. After completion, we send a new request mentioning the existing assistant thread, which then continues the conversation. This way, the LLM can call tools on the system, which are used as input in its reasoning.
OpenAI has very good libraries for Python, but I found the Java libraries lacking.
What are the limitations of this library?
Currently, the library has the following limitations:
- Only OpenAI models are supported. I’d be happy to add more vendors when this library gets used.
- Since it relies on reflection, it wouldn’t work for obfuscated code. (to overcome this, the library could be adapted to use annotations instead of reflection)
- Right now, it doesn’t work for custom datatypes passed to functions. (overcoming this requires more elaborate schema creation and parsing in the library)
- Right now, it’s not possible to cherry-pick methods, or include an object partially. (this is by design to keep the library simple, but I’d love to hear usecases for why it could be useful)
- Right now, it purely uses the method name and its arguments to decide how to call the function. In my tests, this has been mostly sufficient, but I can imagine that it can lead to failed tool-calls in more complex scenario’s. (this is again for simplicity, though it would be a useful feature to add annotations for more complex usecases)
Why did I make this library?
I’ve had this idea of giving AI agents access to a programming language VM for a longer time. A couple weeks ago, a friend asked me to show an example of the Cursor AI Agent in action. I decided to prompt up this library. It took me a bunch of prodding to get it right, but I’m happy with how the result turned out.
What’s next?
I’m planning to create another such library for NodeJS, since I don’t really work much with Java anymore. I could use the NodeJS library in a bunch of my projects, which would add real-life usecases that can help the library become more useful.