Update OllamaAPIIntegrationTest.java

This commit is contained in:
Amith Koujalgi 2025-03-25 20:40:58 +05:30
parent 1155a9be9f
commit 252ea50717

View File

@ -58,6 +58,7 @@ public class OllamaAPIIntegrationTest {
boolean useExternalOllamaHost = Boolean.parseBoolean(System.getenv("USE_EXTERNAL_OLLAMA_HOST")); boolean useExternalOllamaHost = Boolean.parseBoolean(System.getenv("USE_EXTERNAL_OLLAMA_HOST"));
String ollamaHost = System.getenv("OLLAMA_HOST"); String ollamaHost = System.getenv("OLLAMA_HOST");
if (useExternalOllamaHost) { if (useExternalOllamaHost) {
LOG.info("Using external Ollama host...");
api = new OllamaAPI(ollamaHost); api = new OllamaAPI(ollamaHost);
} else { } else {
throw new RuntimeException( throw new RuntimeException(
@ -73,6 +74,7 @@ public class OllamaAPIIntegrationTest {
portBindings.add(mappedPort + ":" + internalPort); portBindings.add(mappedPort + ":" + internalPort);
ollama.setPortBindings(portBindings); ollama.setPortBindings(portBindings);
ollama.start(); ollama.start();
LOG.info("Using Testcontainer Ollama host...");
api = new OllamaAPI("http://" + ollama.getHost() + ":" + ollama.getMappedPort(internalPort)); api = new OllamaAPI("http://" + ollama.getHost() + ":" + ollama.getMappedPort(internalPort));
} }
api.setRequestTimeoutSeconds(120); api.setRequestTimeoutSeconds(120);
@ -152,12 +154,12 @@ public class OllamaAPIIntegrationTest {
@Order(6) @Order(6)
void testAskModelWithStructuredOutput() void testAskModelWithStructuredOutput()
throws OllamaBaseException, IOException, InterruptedException, URISyntaxException { throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
api.pullModel(CHAT_MODEL_QWEN_SMALL); api.pullModel(CHAT_MODEL_LLAMA3);
int timeHour = 6; int timeHour = 6;
boolean isNightTime = false; boolean isNightTime = false;
String prompt = "The Sun is shining, and its " + timeHour + " in the morning right now. So, its daytime."; String prompt = "The Sun is shining, and its " + timeHour + ". Its daytime.";
Map<String, Object> format = new HashMap<>(); Map<String, Object> format = new HashMap<>();
format.put("type", "object"); format.put("type", "object");
@ -177,22 +179,21 @@ public class OllamaAPIIntegrationTest {
}); });
format.put("required", Arrays.asList("timeHour", "isNightTime")); format.put("required", Arrays.asList("timeHour", "isNightTime"));
OllamaResult result = api.generate(CHAT_MODEL_QWEN_SMALL, prompt, format); OllamaResult result = api.generate(CHAT_MODEL_LLAMA3, prompt, format);
assertNotNull(result); assertNotNull(result);
assertNotNull(result.getResponse()); assertNotNull(result.getResponse());
assertFalse(result.getResponse().isEmpty()); assertFalse(result.getResponse().isEmpty());
assertEquals(result.getStructuredResponse().get("timeHour").toString(), assertEquals(timeHour,
result.getStructuredResponse().get("timeHour").toString()); result.getStructuredResponse().get("timeHour"));
assertEquals(result.getStructuredResponse().get("isNightTime").toString(), assertEquals(isNightTime,
result.getStructuredResponse().get("isNightTime").toString()); result.getStructuredResponse().get("isNightTime"));
System.out.println(result.getResponse());
TimeOfDay timeOfDay = result.as(TimeOfDay.class); TimeOfDay timeOfDay = result.as(TimeOfDay.class);
assertEquals(timeOfDay.getTimeHour(), timeHour); assertEquals(timeHour, timeOfDay.getTimeHour());
assertEquals(timeOfDay.isNightTime(), isNightTime); assertEquals(isNightTime, timeOfDay.isNightTime());
} }
@Test @Test