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