mirror of
https://github.com/amithkoujalgi/ollama4j.git
synced 2025-05-15 03:47:13 +02:00
Update OllamaAPIIntegrationTest.java
This commit is contained in:
parent
1155a9be9f
commit
252ea50717
@ -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
|
||||||
|
Loading…
x
Reference in New Issue
Block a user