Add YAML support for Agent configuration and enhance Agent class to load tools from YAML. Introduce custom prompt functionality and refactor constructor to accept additional parameters. Update SampleAgent to demonstrate YAML loading.

This commit is contained in:
amithkoujalgi 2025-10-13 14:34:03 +05:30
parent 6df57c4a23
commit cbf65fef48
No known key found for this signature in database
GPG Key ID: E29A37746AF94B70
5 changed files with 409 additions and 168 deletions

View File

@ -259,6 +259,12 @@
<artifactId>jackson-databind</artifactId> <artifactId>jackson-databind</artifactId>
<version>2.20.0</version> <version>2.20.0</version>
</dependency> </dependency>
<dependency>
<groupId>tools.jackson.dataformat</groupId>
<artifactId>jackson-dataformat-yaml</artifactId>
<version>3.0.0</version>
</dependency>
<dependency> <dependency>
<groupId>com.fasterxml.jackson.datatype</groupId> <groupId>com.fasterxml.jackson.datatype</groupId>
<artifactId>jackson-datatype-jsr310</artifactId> <artifactId>jackson-datatype-jsr310</artifactId>

View File

@ -12,10 +12,14 @@ import io.github.ollama4j.Ollama;
import io.github.ollama4j.exceptions.OllamaException; import io.github.ollama4j.exceptions.OllamaException;
import io.github.ollama4j.impl.ConsoleOutputGenerateTokenHandler; import io.github.ollama4j.impl.ConsoleOutputGenerateTokenHandler;
import io.github.ollama4j.models.chat.*; import io.github.ollama4j.models.chat.*;
import io.github.ollama4j.tools.ToolFunction;
import io.github.ollama4j.tools.Tools; import io.github.ollama4j.tools.Tools;
import java.io.InputStream;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Scanner; import java.util.Scanner;
import lombok.*;
import tools.jackson.dataformat.yaml.YAMLMapper;
public class Agent { public class Agent {
private final String name; private final String name;
@ -23,13 +27,75 @@ public class Agent {
private final Ollama ollamaClient; private final Ollama ollamaClient;
private final String model; private final String model;
private final List<OllamaChatMessage> chatHistory; private final List<OllamaChatMessage> chatHistory;
private final String customPrompt;
public Agent(String name, Ollama ollamaClient, String model, List<Tools.Tool> tools) { public Agent(
String name,
Ollama ollamaClient,
String model,
String customPrompt,
List<Tools.Tool> tools) {
this.name = name; this.name = name;
this.ollamaClient = ollamaClient; this.ollamaClient = ollamaClient;
this.chatHistory = new ArrayList<>(); this.chatHistory = new ArrayList<>();
this.tools = tools; this.tools = tools;
this.model = model; this.model = model;
this.customPrompt = customPrompt;
}
public static Agent fromYaml(String agentYaml) {
try {
YAMLMapper mapper = new YAMLMapper();
InputStream input = Agent.class.getClassLoader().getResourceAsStream(agentYaml);
if (input == null) {
java.nio.file.Path filePath = java.nio.file.Paths.get(agentYaml);
if (java.nio.file.Files.exists(filePath)) {
input = java.nio.file.Files.newInputStream(filePath);
} else {
throw new RuntimeException(
agentYaml + " not found in classpath or file system");
}
}
AgentSpec agentSpec = mapper.readValue(input, AgentSpec.class);
List<AgentToolSpec> tools = agentSpec.getTools();
for (AgentToolSpec tool : tools) {
String fqcn = tool.getToolFunctionFQCN();
if (fqcn != null && !fqcn.isEmpty()) {
try {
Class<?> clazz = Class.forName(fqcn);
Object instance = clazz.getDeclaredConstructor().newInstance();
if (instance instanceof ToolFunction) {
tool.setToolFunctionInstance((ToolFunction) instance);
} else {
throw new RuntimeException(
"Class does not implement ToolFunction: " + fqcn);
}
} catch (Exception e) {
throw new RuntimeException(
"Failed to instantiate tool function: " + fqcn, e);
}
}
}
List<Tools.Tool> agentTools = new ArrayList<>();
for (AgentToolSpec a : tools) {
Tools.Tool t = new Tools.Tool();
t.setToolFunction(a.getToolFunctionInstance());
Tools.ToolSpec ts = new Tools.ToolSpec();
ts.setName(a.getName());
ts.setDescription(a.getDescription());
ts.setParameters(a.getParameters());
t.setToolSpec(ts);
agentTools.add(t);
}
return new Agent(
agentSpec.getName(),
new Ollama(agentSpec.getHost()),
agentSpec.getModel(),
agentSpec.getCustomPrompt(),
agentTools);
} catch (Exception e) {
throw new RuntimeException("Failed to load agent from YAML", e);
}
} }
public String think(String userInput) throws OllamaException { public String think(String userInput) throws OllamaException {
@ -48,17 +114,34 @@ public class Agent {
} }
if (chatHistory.isEmpty()) { if (chatHistory.isEmpty()) {
chatHistory.add( chatHistory.add(
// new OllamaChatMessage(
// OllamaChatMessageRole.SYSTEM,
// "You are a helpful assistant named "
// + name
// + ". You only perform tasks using tools
// available for you. You"
// + " respond very precisely and you don't
// overthink or be too"
// + " creative. Do not ever reveal the tool
// specification in"
// + " terms of code or JSON or in a way that
// a software engineer"
// + " sees it. Just be careful with your
// responses and respond"
// + " like a human. Note that you only
// execute tools provided to"
// + " you. Following are the tools that you
// have access to and"
// + " you can perform right actions using
// right tools."
// + availableToolsDescription));
new OllamaChatMessage( new OllamaChatMessage(
OllamaChatMessageRole.SYSTEM, OllamaChatMessageRole.SYSTEM,
"You are a helpful assistant named " "You are a helpful assistant named "
+ name + name
+ ". You only perform tasks using tools available for you. You" + ". You only perform tasks using tools available for you. "
+ " respond very precisely and you don't overthink or be too" + customPrompt
+ " creative. Do not ever reveal the tool specification in" + ". Following are the tools that you have access to and"
+ " terms of code or JSON or in a way that a software engineer"
+ " sees it. Just be careful with your responses and respond"
+ " like a human. Note that you only execute tools provided to"
+ " you. Following are the tools that you have access to and"
+ " you can perform right actions using right tools." + " you can perform right actions using right tools."
+ availableToolsDescription)); + availableToolsDescription));
} }
@ -90,4 +173,23 @@ public class Agent {
String response = this.think(input); String response = this.think(input);
} }
} }
@Data
public static class AgentSpec {
private String name;
private String description;
private List<AgentToolSpec> tools;
private Tools.Parameters parameters;
private String host;
private String model;
private String customPrompt;
}
@Data
@Setter
@Getter
private static class AgentToolSpec extends Tools.ToolSpec {
private String toolFunctionFQCN = null;
private ToolFunction toolFunctionInstance = null;
}
} }

View File

@ -8,181 +8,218 @@
*/ */
package io.github.ollama4j.agent; package io.github.ollama4j.agent;
import io.github.ollama4j.Ollama;
import io.github.ollama4j.exceptions.OllamaException; import io.github.ollama4j.exceptions.OllamaException;
import io.github.ollama4j.tools.ToolFunction; import io.github.ollama4j.tools.ToolFunction;
import io.github.ollama4j.tools.Tools;
import java.util.ArrayList;
import java.util.List;
import java.util.Map; import java.util.Map;
/** /** Example usage of the Agent API with some dummy tool functions. */
* Example usage of the Agent API with some dummy tool functions.
*/
public class SampleAgent { public class SampleAgent {
public static void main(String[] args) throws OllamaException { public static void main(String[] args) throws OllamaException {
Ollama ollama = new Ollama("http://192.168.29.224:11434"); // Ollama ollama = new Ollama("http://192.168.29.224:11434");
ollama.setRequestTimeoutSeconds(120); // ollama.setRequestTimeoutSeconds(120);
String model = "mistral:7b"; // String model = "mistral:7b";
ollama.pullModel(model); // ollama.pullModel(model);
// List<Tools.Tool> tools = new ArrayList<>();
List<Tools.Tool> tools = new ArrayList<>(); // // Weather tool
// Weather tool // tools.add(
tools.add( // Tools.Tool.builder()
Tools.Tool.builder() // .toolSpec(
.toolSpec( // Tools.ToolSpec.builder()
Tools.ToolSpec.builder() // .name("weather-tool")
.name("weather-tool") // .description(
.description( // "Gets the current weather for a given
"Gets the current weather for a given location and" // location and"
+ " day.") // + " day.")
.parameters( // .parameters(
Tools.Parameters.of( // Tools.Parameters.of(
Map.of( // Map.of(
"location", // "location",
Tools.Property.builder() // Tools.Property.builder()
.type("string") // .type("string")
.description( // .description(
"The location for" // "The
+ " which to" // location for"
+ " get the" // + "
+ " weather.") // which to"
.required(true) // + "
.build(), // get the"
"day", // + "
Tools.Property.builder() // weather.")
.type("string") // .required(true)
.description( // .build(),
"The day of the" // "day",
+ " week for" // Tools.Property.builder()
+ " which to" // .type("string")
+ " get the" // .description(
+ " weather.") // "The day
.required(true) // of the"
.build()))) // + "
.build()) // week for"
.toolFunction(new WeatherToolFunction()) // + "
.build()); // which to"
// + "
// Calculator tool // get the"
tools.add( // + "
Tools.Tool.builder() // weather.")
.toolSpec( // .required(true)
Tools.ToolSpec.builder() // .build())))
.name("calculator-tool") // .build())
.description( // .toolFunction(new WeatherToolFunction())
"Performs a simple arithmetic operation between two" // .build());
+ " numbers.") //
.parameters( // // Calculator tool
Tools.Parameters.of( // tools.add(
Map.of( // Tools.Tool.builder()
"operation", // .toolSpec(
Tools.Property.builder() // Tools.ToolSpec.builder()
.type("string") // .name("calculator-tool")
.description( // .description(
"The arithmetic" // "Performs a simple arithmetic operation
+ " operation" // between two"
+ " to perform." // + " numbers.")
+ " One of:" // .parameters(
+ " add," // Tools.Parameters.of(
+ " subtract," // Map.of(
+ " multiply," // "operation",
+ " divide.") // Tools.Property.builder()
.required(true) // .type("string")
.build(), // .description(
"a", // "The
Tools.Property.builder() // arithmetic"
.type("number") // + "
.description( // operation"
"The first" // + " to
+ " operand.") // perform."
.required(true) // + "
.build(), // One of:"
"b", // + "
Tools.Property.builder() // add,"
.type("number") // + "
.description( // subtract,"
"The second" // + "
+ " operand.") // multiply,"
.required(true) // + "
.build()))) // divide.")
.build()) // .required(true)
.toolFunction(new CalculatorToolFunction()) // .build(),
.build()); // "a",
// Tools.Property.builder()
// Hotel Booking tool (dummy) // .type("number")
tools.add( // .description(
Tools.Tool.builder() // "The
.toolSpec( // first"
Tools.ToolSpec.builder() // + "
.name("hotel-booking-tool") // operand.")
.description( // .required(true)
"Books a hotel room in a specified city for given" // .build(),
+ " dates and number of guests.") // "b",
.parameters( // Tools.Property.builder()
Tools.Parameters.of( // .type("number")
Map.of( // .description(
"city", // "The
Tools.Property.builder() // second"
.type("string") // + "
.description( // operand.")
"The city where the" // .required(true)
+ " hotel will" // .build())))
+ " be booked.") // .build())
.required(true) // .toolFunction(new CalculatorToolFunction())
.build(), // .build());
"checkin_date", //
Tools.Property.builder() // // Hotel Booking tool (dummy)
.type("string") // tools.add(
.description( // Tools.Tool.builder()
"Check-in date" // .toolSpec(
+ " (e.g." // Tools.ToolSpec.builder()
+ " 2025-08-10).") // .name("hotel-booking-tool")
.required(true) // .description(
.build(), // "Books a hotel room in a specified city
"checkout_date", // for given"
Tools.Property.builder() // + " dates and number of guests.")
.type("string") // .parameters(
.description( // Tools.Parameters.of(
"Check-out date" // Map.of(
+ " (e.g." // "city",
+ " 2025-08-12).") // Tools.Property.builder()
.required(true) // .type("string")
.build(), // .description(
"guests", // "The city
Tools.Property.builder() // where the"
.type("number") // + "
.description( // hotel will"
"Number of guests" // + " be
+ " for the" // booked.")
+ " booking.") // .required(true)
.required(true) // .build(),
.build()))) // "checkin_date",
.build()) // Tools.Property.builder()
.toolFunction(new HotelBookingToolFunction()) // .type("string")
.build()); // .description(
// "Hotel
Agent agent = new Agent("Nimma Mirta", ollama, model, tools); // check-in date"
// + "
// (e.g."
// + "
// 2025-08-10).")
// .required(true)
// .build(),
// "checkout_date",
// Tools.Property.builder()
// .type("string")
// .description(
//
// "HotelCheck-out date"
// + "
// (e.g."
// + "
// 2025-08-12).")
// .required(true)
// .build(),
// "guests",
// Tools.Property.builder()
// .type("number")
// .description(
// "Number of
// guests"
// + "
// for the"
// + "
// booking.")
// .required(true)
// .build())))
// .build())
// .toolFunction(new HotelBookingToolFunction())
// .build());
//
// Map<String, ToolFunction> functionMap = Map.of(
// "weather-tool", new WeatherToolFunction(),
// "calculator-tool", new CalculatorToolFunction()
// );
// List<Tools.Tool> tools =
// Tools.fromYAMLFile("/Users/amithkoujalgi/Downloads/tools.yaml", functionMap);
// Agent agent = new Agent("Nimma Mirta", ollama, model, tools);
Agent agent = Agent.fromYaml("agent.yaml");
agent.runInteractive(); agent.runInteractive();
} }
} }
/** /** ToolFunction implementation that returns a dummy weekly weather forecast. */
* ToolFunction implementation that returns a dummy weekly weather forecast.
*/
class WeatherToolFunction implements ToolFunction { class WeatherToolFunction implements ToolFunction {
@Override @Override
public Object apply(Map<String, Object> arguments) { public Object apply(Map<String, Object> arguments) {
String response = String response =
"Monday: Pleasant, Tuesday: Sunny, Wednesday: Windy, Thursday: Cloudy, Friday:" "Monday: Pleasant."
+ " Rainy, Saturday: Heavy rains, Sunday: Clear"; + "Tuesday: Sunny."
+ "Wednesday: Windy."
+ "Thursday: Cloudy."
+ "Friday: Rainy."
+ "Saturday: Heavy rains."
+ "Sunday: Clear.";
return response; return response;
} }
} }
/** /** ToolFunction implementation for basic arithmetic calculations. */
* ToolFunction implementation for basic arithmetic calculations.
*/
class CalculatorToolFunction implements ToolFunction { class CalculatorToolFunction implements ToolFunction {
@Override @Override
public Object apply(Map<String, Object> arguments) { public Object apply(Map<String, Object> arguments) {
@ -213,9 +250,7 @@ class CalculatorToolFunction implements ToolFunction {
} }
} }
/** /** ToolFunction implementation simulating a hotel booking. */
* ToolFunction implementation simulating a hotel booking.
*/
class HotelBookingToolFunction implements ToolFunction { class HotelBookingToolFunction implements ToolFunction {
@Override @Override
public Object apply(Map<String, Object> arguments) { public Object apply(Map<String, Object> arguments) {

View File

@ -11,7 +11,9 @@ package io.github.ollama4j.tools;
import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode; import com.fasterxml.jackson.databind.node.ObjectNode;
import java.io.File;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
@ -19,6 +21,8 @@ import lombok.AllArgsConstructor;
import lombok.Builder; import lombok.Builder;
import lombok.Data; import lombok.Data;
import lombok.NoArgsConstructor; import lombok.NoArgsConstructor;
import tools.jackson.core.type.TypeReference;
import tools.jackson.dataformat.yaml.YAMLMapper;
public class Tools { public class Tools {
private Tools() {} private Tools() {}
@ -116,4 +120,53 @@ public class Tools {
@JsonIgnore private boolean required; @JsonIgnore private boolean required;
} }
public static List<Tool> fromJSONFile(String filePath, Map<String, ToolFunction> functionMap) {
try {
ObjectMapper mapper = new ObjectMapper();
List<Map<String, Object>> rawTools =
mapper.readValue(
new File(filePath),
new com.fasterxml.jackson.core.type.TypeReference<>() {});
List<Tool> tools = new ArrayList<>();
for (Map<String, Object> rawTool : rawTools) {
String json = mapper.writeValueAsString(rawTool);
Tool tool = mapper.readValue(json, Tool.class);
String toolName = tool.getToolSpec().getName();
for (Map.Entry<String, ToolFunction> toolFunctionEntry : functionMap.entrySet()) {
if (toolFunctionEntry.getKey().equals(toolName)) {
tool.setToolFunction(toolFunctionEntry.getValue());
}
}
tools.add(tool);
}
return tools;
} catch (Exception e) {
throw new RuntimeException("Failed to load tools from file: " + filePath, e);
}
}
public static List<Tool> fromYAMLFile(String filePath, Map<String, ToolFunction> functionMap) {
try {
YAMLMapper mapper = new YAMLMapper();
List<Map<String, Object>> rawTools =
mapper.readValue(new File(filePath), new TypeReference<>() {});
List<Tool> tools = new ArrayList<>();
for (Map<String, Object> rawTool : rawTools) {
String yaml = mapper.writeValueAsString(rawTool);
Tool tool = mapper.readValue(yaml, Tool.class);
String toolName = tool.getToolSpec().getName();
ToolFunction function = functionMap.get(toolName);
if (function != null) {
tool.setToolFunction(function);
}
tools.add(tool);
}
return tools;
} catch (Exception e) {
throw new RuntimeException("Failed to load tools from YAML file: " + filePath, e);
}
}
} }

View File

@ -0,0 +1,45 @@
name: Nimma Mitra
host: http://192.168.29.224:11434
model: mistral:7b
customPrompt: >
Only use tools and do not use your creativity.
Do not ever tell me to call the tool or how to use the tool or command myself.
You do that for me. That is why you exist.
You call tool on my behalf and give me a response from the tool.
tools:
- name: weather-tool
toolFunctionFQCN: io.github.ollama4j.agent.WeatherToolFunction
description: Gets the current weather for a given location and day.
parameters:
properties:
location:
type: string
description: The location for which to get the weather.
required: true
day:
type: string
description: The day of the week for which to get the weather.
required: true
- name: calculator-tool
toolFunctionFQCN: io.github.ollama4j.agent.CalculatorToolFunction
description: Performs a simple arithmetic operation between two numbers.
parameters:
properties:
operation:
type: string
description: The arithmetic operation to perform.
enum:
- add
- subtract
- multiply
- divide
required: true
a:
type: number
description: The first operand.
required: true
b:
type: number
description: The second operand.
required: true