Update OllamaAPIIntegrationTest.java

This commit is contained in:
amithkoujalgi 2025-03-25 00:20:25 +05:30
parent b9b18271a1
commit be08f11027
No known key found for this signature in database
GPG Key ID: E29A37746AF94B70

View File

@ -1,5 +1,6 @@
package io.github.ollama4j.integrationtests;
import com.fasterxml.jackson.annotation.JsonProperty;
import io.github.ollama4j.OllamaAPI;
import io.github.ollama4j.exceptions.OllamaBaseException;
import io.github.ollama4j.models.chat.*;
@ -153,29 +154,28 @@ public class OllamaAPIIntegrationTest {
throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
api.pullModel(CHAT_MODEL_QWEN_SMALL);
int age = 28;
boolean available = false;
int timeHour = 6;
boolean isNightTime = false;
String prompt = "Batman is " + age + " years old and is " + (available ? "available" : "not available")
+ " because he is busy saving Gotham City. Respond using JSON";
String prompt = "The Sun is shining, and its " + timeHour + " in the morning right now. So, its daytime.";
Map<String, Object> format = new HashMap<>();
format.put("type", "object");
format.put("properties", new HashMap<String, Object>() {
{
put("age", new HashMap<String, Object>() {
put("timeHour", new HashMap<String, Object>() {
{
put("type", "integer");
}
});
put("available", new HashMap<String, Object>() {
put("isNightTime", new HashMap<String, Object>() {
{
put("type", "boolean");
}
});
}
});
format.put("required", Arrays.asList("age", "available"));
format.put("required", Arrays.asList("timeHour", "isNightTime"));
OllamaResult result = api.generate(CHAT_MODEL_QWEN_SMALL, prompt, format);
@ -183,14 +183,16 @@ public class OllamaAPIIntegrationTest {
assertNotNull(result.getResponse());
assertFalse(result.getResponse().isEmpty());
assertEquals(result.getStructuredResponse().get("age").toString(),
result.getStructuredResponse().get("age").toString());
assertEquals(result.getStructuredResponse().get("available").toString(),
result.getStructuredResponse().get("available").toString());
assertEquals(result.getStructuredResponse().get("timeHour").toString(),
result.getStructuredResponse().get("timeHour").toString());
assertEquals(result.getStructuredResponse().get("isNightTime").toString(),
result.getStructuredResponse().get("isNightTime").toString());
Person person = result.as(Person.class);
assertEquals(person.getAge(), age);
assertEquals(person.isAvailable(), available);
System.out.println(result.getResponse());
TimeOfDay timeOfDay = result.as(TimeOfDay.class);
assertEquals(timeOfDay.getTimeHour(), timeHour);
assertEquals(timeOfDay.isNightTime(), isNightTime);
}
@Test
@ -669,7 +671,9 @@ public class OllamaAPIIntegrationTest {
@Data
@AllArgsConstructor
@NoArgsConstructor
class Person {
private int age;
private boolean available;
class TimeOfDay {
@JsonProperty("timeHour")
private int timeHour;
@JsonProperty("isNightTime")
private boolean nightTime;
}