Merge remote-tracking branch 'origin/main'

README.md

@@ -91,7 +91,13 @@ For simplest way to get started, I prefer to use the Ollama docker setup.
 Start the Ollama docker container:
 
 ```shell
-docker run -v ~/ollama:/root/.ollama -p 11434:11434 ollama/ollama
+docker run -it -v ~/ollama:/root/.ollama -p 11434:11434 ollama/ollama
+```
+
+With GPUs
+
+```shell
+docker run -it --gpus=all -v ~/ollama:/root/.ollama -p 11434:11434 ollama/ollama
 ```
 
 Instantiate `OllamaAPI`
@@ -347,6 +353,7 @@ Find the full `Javadoc` (API specifications) [here](https://amithkoujalgi.github
       conversational memory
     - `stream`: Add support for streaming responses from the model
 - [x] Setup logging
+- [ ] Use lombok
 - [ ] Add test cases
 - [ ] Handle exceptions better (maybe throw more appropriate exceptions)
 

src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaAPI.java

@@ -3,9 +3,6 @@ package io.github.amithkoujalgi.ollama4j.core;
 import io.github.amithkoujalgi.ollama4j.core.exceptions.OllamaBaseException;
 import io.github.amithkoujalgi.ollama4j.core.models.*;
 import io.github.amithkoujalgi.ollama4j.core.utils.Utils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStream;
@@ -17,10 +14,10 @@ import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-/**
+/** The base Ollama API class. */
- * The base Ollama API class.
- */
 @SuppressWarnings("DuplicatedCode")
 public class OllamaAPI {
 
@@ -43,6 +40,7 @@ public class OllamaAPI {
 
   /**
    * Set/unset logging of responses
+   *
    * @param verbose true/false
    */
   public void setVerbose(boolean verbose) {
@@ -54,38 +52,59 @@ public class OllamaAPI {
    *
    * @return the list
    */
-    public List<Model> listModels() throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
+  public List<Model> listModels()
+      throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
     String url = this.host + "/api/tags";
     HttpClient httpClient = HttpClient.newHttpClient();
-        HttpRequest httpRequest = HttpRequest.newBuilder().uri(new URI(url)).header("Accept", "application/json").header("Content-type", "application/json").GET().build();
+    HttpRequest httpRequest =
-        HttpResponse<String> response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
+        HttpRequest.newBuilder()
+            .uri(new URI(url))
+            .header("Accept", "application/json")
+            .header("Content-type", "application/json")
+            .GET()
+            .build();
+    HttpResponse<String> response =
+        httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
     int statusCode = response.statusCode();
     String responseString = response.body();
     if (statusCode == 200) {
-            return Utils.getObjectMapper().readValue(responseString, ListModelsResponse.class).getModels();
+      return Utils.getObjectMapper()
+          .readValue(responseString, ListModelsResponse.class)
+          .getModels();
     } else {
       throw new OllamaBaseException(statusCode + " - " + responseString);
     }
   }
 
   /**
-     * Pull a model on the Ollama server from the list of <a href="https://ollama.ai/library">available models</a>.
+   * Pull a model on the Ollama server from the list of <a
+   * href="https://ollama.ai/library">available models</a>.
    *
    * @param model the name of the model
    */
-    public void pullModel(String model) throws OllamaBaseException, IOException, URISyntaxException, InterruptedException {
+  public void pullModel(String model)
+      throws OllamaBaseException, IOException, URISyntaxException, InterruptedException {
     String url = this.host + "/api/pull";
     String jsonData = String.format("{\"name\": \"%s\"}", model);
-        HttpRequest request = HttpRequest.newBuilder().uri(new URI(url)).POST(HttpRequest.BodyPublishers.ofString(jsonData)).header("Accept", "application/json").header("Content-type", "application/json").build();
+    HttpRequest request =
+        HttpRequest.newBuilder()
+            .uri(new URI(url))
+            .POST(HttpRequest.BodyPublishers.ofString(jsonData))
+            .header("Accept", "application/json")
+            .header("Content-type", "application/json")
+            .build();
     HttpClient client = HttpClient.newHttpClient();
-        HttpResponse<InputStream> response = client.send(request, HttpResponse.BodyHandlers.ofInputStream());
+    HttpResponse<InputStream> response =
+        client.send(request, HttpResponse.BodyHandlers.ofInputStream());
     int statusCode = response.statusCode();
     InputStream responseBodyStream = response.body();
     String responseString = "";
-        try (BufferedReader reader = new BufferedReader(new InputStreamReader(responseBodyStream, StandardCharsets.UTF_8))) {
+    try (BufferedReader reader =
+        new BufferedReader(new InputStreamReader(responseBodyStream, StandardCharsets.UTF_8))) {
       String line;
       while ((line = reader.readLine()) != null) {
-                ModelPullResponse modelPullResponse = Utils.getObjectMapper().readValue(line, ModelPullResponse.class);
+        ModelPullResponse modelPullResponse =
+            Utils.getObjectMapper().readValue(line, ModelPullResponse.class);
         if (verbose) {
           logger.info(modelPullResponse.getStatus());
         }
@@ -102,10 +121,17 @@ public class OllamaAPI {
    * @param modelName the model
    * @return the model details
    */
-    public ModelDetail getModelDetails(String modelName) throws IOException, OllamaBaseException, InterruptedException {
+  public ModelDetail getModelDetails(String modelName)
+      throws IOException, OllamaBaseException, InterruptedException {
     String url = this.host + "/api/show";
     String jsonData = String.format("{\"name\": \"%s\"}", modelName);
-        HttpRequest request = HttpRequest.newBuilder().uri(URI.create(url)).header("Accept", "application/json").header("Content-type", "application/json").POST(HttpRequest.BodyPublishers.ofString(jsonData)).build();
+    HttpRequest request =
+        HttpRequest.newBuilder()
+            .uri(URI.create(url))
+            .header("Accept", "application/json")
+            .header("Content-type", "application/json")
+            .POST(HttpRequest.BodyPublishers.ofString(jsonData))
+            .build();
     HttpClient client = HttpClient.newHttpClient();
     HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
     int statusCode = response.statusCode();
@@ -118,16 +144,24 @@ public class OllamaAPI {
   }
 
   /**
-     * Create a custom model from a model file.
+   * Create a custom model from a model file. Read more about custom model file creation <a
-     * Read more about custom model file creation <a href="https://github.com/jmorganca/ollama/blob/main/docs/modelfile.md">here</a>.
+   * href="https://github.com/jmorganca/ollama/blob/main/docs/modelfile.md">here</a>.
    *
    * @param modelName the name of the custom model to be created.
    * @param modelFilePath the path to model file that exists on the Ollama server.
    */
-    public void createModel(String modelName, String modelFilePath) throws IOException, InterruptedException, OllamaBaseException {
+  public void createModel(String modelName, String modelFilePath)
+      throws IOException, InterruptedException, OllamaBaseException {
     String url = this.host + "/api/create";
-        String jsonData = String.format("{\"name\": \"%s\", \"path\": \"%s\"}", modelName, modelFilePath);
+    String jsonData =
-        HttpRequest request = HttpRequest.newBuilder().uri(URI.create(url)).header("Accept", "application/json").header("Content-Type", "application/json").POST(HttpRequest.BodyPublishers.ofString(jsonData, StandardCharsets.UTF_8)).build();
+        String.format("{\"name\": \"%s\", \"path\": \"%s\"}", modelName, modelFilePath);
+    HttpRequest request =
+        HttpRequest.newBuilder()
+            .uri(URI.create(url))
+            .header("Accept", "application/json")
+            .header("Content-Type", "application/json")
+            .POST(HttpRequest.BodyPublishers.ofString(jsonData, StandardCharsets.UTF_8))
+            .build();
     HttpClient client = HttpClient.newHttpClient();
     HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
     int statusCode = response.statusCode();
@@ -135,7 +169,8 @@ public class OllamaAPI {
     if (statusCode != 200) {
       throw new OllamaBaseException(statusCode + " - " + responseString);
     }
-        // FIXME: Ollama API returns HTTP status code 200 for model creation failure cases. Correct this if the issue is fixed in the Ollama API server.
+    // FIXME: Ollama API returns HTTP status code 200 for model creation failure cases. Correct this
+    // if the issue is fixed in the Ollama API server.
     if (responseString.contains("error")) {
       throw new OllamaBaseException(responseString);
     }
@@ -148,12 +183,20 @@ public class OllamaAPI {
    * Delete a model from Ollama server.
    *
    * @param name the name of the model to be deleted.
-     * @param ignoreIfNotPresent - ignore errors if the specified model is not present on Ollama server.
+   * @param ignoreIfNotPresent - ignore errors if the specified model is not present on Ollama
+   *     server.
    */
-    public void deleteModel(String name, boolean ignoreIfNotPresent) throws IOException, InterruptedException, OllamaBaseException {
+  public void deleteModel(String name, boolean ignoreIfNotPresent)
+      throws IOException, InterruptedException, OllamaBaseException {
     String url = this.host + "/api/delete";
     String jsonData = String.format("{\"name\": \"%s\"}", name);
-        HttpRequest request = HttpRequest.newBuilder().uri(URI.create(url)).method("DELETE", HttpRequest.BodyPublishers.ofString(jsonData, StandardCharsets.UTF_8)).header("Accept", "application/json").header("Content-type", "application/json").build();
+    HttpRequest request =
+        HttpRequest.newBuilder()
+            .uri(URI.create(url))
+            .method("DELETE", HttpRequest.BodyPublishers.ofString(jsonData, StandardCharsets.UTF_8))
+            .header("Accept", "application/json")
+            .header("Content-type", "application/json")
+            .build();
     HttpClient client = HttpClient.newHttpClient();
     HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
     int statusCode = response.statusCode();
@@ -173,16 +216,24 @@ public class OllamaAPI {
    * @param prompt text to generate embeddings for
    * @return embeddings
    */
-    public List<Double> generateEmbeddings(String model, String prompt) throws IOException, InterruptedException, OllamaBaseException {
+  public List<Double> generateEmbeddings(String model, String prompt)
+      throws IOException, InterruptedException, OllamaBaseException {
     String url = this.host + "/api/embeddings";
     String jsonData = String.format("{\"model\": \"%s\", \"prompt\": \"%s\"}", model, prompt);
     HttpClient httpClient = HttpClient.newHttpClient();
-        HttpRequest request = HttpRequest.newBuilder().uri(URI.create(url)).header("Accept", "application/json").header("Content-type", "application/json").POST(HttpRequest.BodyPublishers.ofString(jsonData)).build();
+    HttpRequest request =
+        HttpRequest.newBuilder()
+            .uri(URI.create(url))
+            .header("Accept", "application/json")
+            .header("Content-type", "application/json")
+            .POST(HttpRequest.BodyPublishers.ofString(jsonData))
+            .build();
     HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
     int statusCode = response.statusCode();
     String responseBody = response.body();
     if (statusCode == 200) {
-            EmbeddingResponse embeddingResponse = Utils.getObjectMapper().readValue(responseBody, EmbeddingResponse.class);
+      EmbeddingResponse embeddingResponse =
+          Utils.getObjectMapper().readValue(responseBody, EmbeddingResponse.class);
       return embeddingResponse.getEmbedding();
     } else {
       throw new OllamaBaseException(statusCode + " - " + responseBody);
@@ -196,36 +247,53 @@ public class OllamaAPI {
    * @param promptText the prompt/question text
    * @return OllamaResult - that includes response text and time taken for response
    */
-    public OllamaResult ask(String ollamaModelType, String promptText) throws OllamaBaseException, IOException, InterruptedException {
+  public OllamaResult ask(String ollamaModelType, String promptText)
+      throws OllamaBaseException, IOException, InterruptedException {
     OllamaRequestModel ollamaRequestModel = new OllamaRequestModel(ollamaModelType, promptText);
     long startTime = System.currentTimeMillis();
     HttpClient httpClient = HttpClient.newHttpClient();
     URI uri = URI.create(this.host + "/api/generate");
-        HttpRequest request = HttpRequest.newBuilder(uri).POST(HttpRequest.BodyPublishers.ofString(Utils.getObjectMapper().writeValueAsString(ollamaRequestModel))).header("Content-Type", "application/json").build();
+    HttpRequest request =
-        HttpResponse<InputStream> response = httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream());
+        HttpRequest.newBuilder(uri)
+            .POST(
+                HttpRequest.BodyPublishers.ofString(
+                    Utils.getObjectMapper().writeValueAsString(ollamaRequestModel)))
+            .header("Content-Type", "application/json")
+            .build();
+    HttpResponse<InputStream> response =
+        httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream());
     int statusCode = response.statusCode();
     InputStream responseBodyStream = response.body();
     StringBuilder responseBuffer = new StringBuilder();
-        try (BufferedReader reader = new BufferedReader(new InputStreamReader(responseBodyStream, StandardCharsets.UTF_8))) {
+    try (BufferedReader reader =
+        new BufferedReader(new InputStreamReader(responseBodyStream, StandardCharsets.UTF_8))) {
       String line;
       while ((line = reader.readLine()) != null) {
-                OllamaResponseModel ollamaResponseModel = Utils.getObjectMapper().readValue(line, OllamaResponseModel.class);
+        if (statusCode == 404) {
+          OllamaErrorResponseModel ollamaResponseModel =
+              Utils.getObjectMapper().readValue(line, OllamaErrorResponseModel.class);
+          responseBuffer.append(ollamaResponseModel.getError());
+        } else {
+          OllamaResponseModel ollamaResponseModel =
+              Utils.getObjectMapper().readValue(line, OllamaResponseModel.class);
           if (!ollamaResponseModel.getDone()) {
             responseBuffer.append(ollamaResponseModel.getResponse());
           }
         }
       }
+    }
     if (statusCode != 200) {
-            throw new OllamaBaseException(statusCode + " - " + responseBuffer);
+      throw new OllamaBaseException(responseBuffer.toString());
     } else {
       long endTime = System.currentTimeMillis();
-            return new OllamaResult(responseBuffer.toString().trim(), endTime - startTime);
+      return new OllamaResult(responseBuffer.toString().trim(), endTime - startTime, statusCode);
     }
   }
 
   /**
-     * Ask a question to a model running on Ollama server and get a callback handle that can be used to check for status and get the response from the model later.
+   * Ask a question to a model running on Ollama server and get a callback handle that can be used
-     * This would be an async/non-blocking call.
+   * to check for status and get the response from the model later. This would be an
+   * async/non-blocking call.
    *
    * @param ollamaModelType the ollama model to ask the question to
    * @param promptText the prompt/question text
@@ -235,7 +303,8 @@ public class OllamaAPI {
     OllamaRequestModel ollamaRequestModel = new OllamaRequestModel(ollamaModelType, promptText);
     HttpClient httpClient = HttpClient.newHttpClient();
     URI uri = URI.create(this.host + "/api/generate");
-        OllamaAsyncResultCallback ollamaAsyncResultCallback = new OllamaAsyncResultCallback(httpClient, uri, ollamaRequestModel);
+    OllamaAsyncResultCallback ollamaAsyncResultCallback =
+        new OllamaAsyncResultCallback(httpClient, uri, ollamaRequestModel);
     ollamaAsyncResultCallback.start();
     return ollamaAsyncResultCallback;
   }

src/main/java/io/github/amithkoujalgi/ollama4j/core/models/OllamaAsyncResultCallback.java

@@ -2,7 +2,6 @@ package io.github.amithkoujalgi.ollama4j.core.models;
 
 import io.github.amithkoujalgi.ollama4j.core.exceptions.OllamaBaseException;
 import io.github.amithkoujalgi.ollama4j.core.utils.Utils;
-
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStream;
@@ -23,9 +22,13 @@ public class OllamaAsyncResultCallback extends Thread {
   private final Queue<String> queue = new LinkedList<>();
   private String result;
   private boolean isDone;
+  private boolean succeeded;
+
+  private int httpStatusCode;
   private long responseTime = 0;
 
-    public OllamaAsyncResultCallback(HttpClient client, URI uri, OllamaRequestModel ollamaRequestModel) {
+  public OllamaAsyncResultCallback(
+      HttpClient client, URI uri, OllamaRequestModel ollamaRequestModel) {
     this.client = client;
     this.ollamaRequestModel = ollamaRequestModel;
     this.uri = uri;
@@ -38,44 +41,89 @@ public class OllamaAsyncResultCallback extends Thread {
   public void run() {
     try {
       long startTime = System.currentTimeMillis();
-            HttpRequest request = HttpRequest.newBuilder(uri).POST(HttpRequest.BodyPublishers.ofString(Utils.getObjectMapper().writeValueAsString(ollamaRequestModel))).header("Content-Type", "application/json").build();
+      HttpRequest request =
-            HttpResponse<InputStream> response = client.send(request, HttpResponse.BodyHandlers.ofInputStream());
+          HttpRequest.newBuilder(uri)
+              .POST(
+                  HttpRequest.BodyPublishers.ofString(
+                      Utils.getObjectMapper().writeValueAsString(ollamaRequestModel)))
+              .header("Content-Type", "application/json")
+              .build();
+      HttpResponse<InputStream> response =
+          client.send(request, HttpResponse.BodyHandlers.ofInputStream());
       int statusCode = response.statusCode();
+      this.httpStatusCode = statusCode;
 
       InputStream responseBodyStream = response.body();
-            String responseString = "";
+      try (BufferedReader reader =
-            try (BufferedReader reader = new BufferedReader(new InputStreamReader(responseBodyStream, StandardCharsets.UTF_8))) {
+          new BufferedReader(new InputStreamReader(responseBodyStream, StandardCharsets.UTF_8))) {
         String line;
         StringBuilder responseBuffer = new StringBuilder();
         while ((line = reader.readLine()) != null) {
-                    OllamaResponseModel ollamaResponseModel = Utils.getObjectMapper().readValue(line, OllamaResponseModel.class);
+          if (statusCode == 404) {
+            OllamaErrorResponseModel ollamaResponseModel =
+                Utils.getObjectMapper().readValue(line, OllamaErrorResponseModel.class);
+            queue.add(ollamaResponseModel.getError());
+            responseBuffer.append(ollamaResponseModel.getError());
+          } else {
+            OllamaResponseModel ollamaResponseModel =
+                Utils.getObjectMapper().readValue(line, OllamaResponseModel.class);
             queue.add(ollamaResponseModel.getResponse());
             if (!ollamaResponseModel.getDone()) {
               responseBuffer.append(ollamaResponseModel.getResponse());
             }
           }
+        }
         reader.close();
+
         this.isDone = true;
+        this.succeeded = true;
         this.result = responseBuffer.toString();
         long endTime = System.currentTimeMillis();
         responseTime = endTime - startTime;
       }
       if (statusCode != 200) {
-                throw new OllamaBaseException(statusCode + " - " + responseString);
+        throw new OllamaBaseException(this.result);
       }
     } catch (IOException | InterruptedException | OllamaBaseException e) {
       this.isDone = true;
-            this.result = "FAILED! " + e.getMessage();
+      this.succeeded = false;
+      this.result = "[FAILED] " + e.getMessage();
     }
   }
 
+  /**
+   * Returns the status of the thread. This does not indicate that the request was successful or a
+   * failure, rather it is just a status flag to indicate if the thread is active or ended.
+   *
+   * @return boolean - status
+   */
   public boolean isComplete() {
     return isDone;
   }
 
+  /**
+   * Returns the HTTP response status code for the request that was made to Ollama server.
+   *
+   * @return int - the status code for the request
+   */
+  public int getHttpStatusCode() {
+    return httpStatusCode;
+  }
+
+  /**
+   * Returns the status of the request. Indicates if the request was successful or a failure. If the
+   * request was a failure, the `getResponse()` method will return the error message.
+   *
+   * @return boolean - status
+   */
+  public boolean isSucceeded() {
+    return succeeded;
+  }
+
   /**
    * Returns the final response when the execution completes. Does not return intermediate results.
-     * @return response text
+   *
+   * @return String - response text
    */
   public String getResponse() {
     return result;
@@ -86,8 +134,9 @@ public class OllamaAsyncResultCallback extends Thread {
   }
 
   /**
-     * Returns the response time in seconds.
+   * Returns the response time in milliseconds.
-     * @return response time in seconds
+   *
+   * @return long - response time in milliseconds.
    */
   public long getResponseTime() {
     return responseTime;

src/main/java/io/github/amithkoujalgi/ollama4j/core/models/OllamaErrorResponseModel.java

@@ -0,0 +1,18 @@
+package io.github.amithkoujalgi.ollama4j.core.models;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.List;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class OllamaErrorResponseModel {
+  private String error;
+
+  public String getError() {
+    return error;
+  }
+
+  public void setError(String error) {
+    this.error = error;
+  }
+}

src/main/java/io/github/amithkoujalgi/ollama4j/core/models/OllamaResult.java

@@ -1,20 +1,57 @@
 package io.github.amithkoujalgi.ollama4j.core.models;
 
+import static io.github.amithkoujalgi.ollama4j.core.utils.Utils.getObjectMapper;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+
+/** The type Ollama result. */
 @SuppressWarnings("unused")
 public class OllamaResult {
-    private String response;
+  private final String response;
+
+  private int httpStatusCode;
+
   private long responseTime = 0;
 
-    public OllamaResult(String response, long responseTime) {
+  public OllamaResult(String response, long responseTime, int httpStatusCode) {
     this.response = response;
     this.responseTime = responseTime;
+    this.httpStatusCode = httpStatusCode;
   }
 
+  /**
+   * Get the response text
+   *
+   * @return String - response text
+   */
   public String getResponse() {
     return response;
   }
 
+  /**
+   * Get the response time in milliseconds.
+   *
+   * @return long - response time in milliseconds
+   */
   public long getResponseTime() {
     return responseTime;
   }
+
+  /**
+   * Get the response status code.
+   *
+   * @return int - response status code
+   */
+  public int getHttpStatusCode() {
+    return httpStatusCode;
+  }
+
+  @Override
+  public String toString() {
+    try {
+      return getObjectMapper().writerWithDefaultPrettyPrinter().writeValueAsString(this);
+    } catch (JsonProcessingException e) {
+      throw new RuntimeException(e);
+    }
+  }
 }

src/test/java/io/github/amithkoujalgi/ollama4j/unittests/TestMockedAPIs.java

@@ -1,19 +1,18 @@
 package io.github.amithkoujalgi.ollama4j.unittests;
 
+import static org.mockito.Mockito.*;
+
 import io.github.amithkoujalgi.ollama4j.core.OllamaAPI;
 import io.github.amithkoujalgi.ollama4j.core.exceptions.OllamaBaseException;
 import io.github.amithkoujalgi.ollama4j.core.models.ModelDetail;
 import io.github.amithkoujalgi.ollama4j.core.models.OllamaAsyncResultCallback;
 import io.github.amithkoujalgi.ollama4j.core.models.OllamaResult;
 import io.github.amithkoujalgi.ollama4j.core.types.OllamaModelType;
-import org.junit.jupiter.api.Test;
-import org.mockito.Mockito;
-
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
-
+import org.junit.jupiter.api.Test;
-import static org.mockito.Mockito.*;
+import org.mockito.Mockito;
 
 public class TestMockedAPIs {
   @Test
@@ -101,7 +100,7 @@ public class TestMockedAPIs {
     String model = OllamaModelType.LLAMA2;
     String prompt = "some prompt text";
     try {
-            when(ollamaAPI.ask(model, prompt)).thenReturn(new OllamaResult("", 0));
+      when(ollamaAPI.ask(model, prompt)).thenReturn(new OllamaResult("", 0, 200));
       ollamaAPI.ask(model, prompt);
       verify(ollamaAPI, times(1)).ask(model, prompt);
     } catch (IOException | OllamaBaseException | InterruptedException e) {
@@ -114,7 +113,8 @@ public class TestMockedAPIs {
     OllamaAPI ollamaAPI = Mockito.mock(OllamaAPI.class);
     String model = OllamaModelType.LLAMA2;
     String prompt = "some prompt text";
-        when(ollamaAPI.askAsync(model, prompt)).thenReturn(new OllamaAsyncResultCallback(null, null, null));
+    when(ollamaAPI.askAsync(model, prompt))
+        .thenReturn(new OllamaAsyncResultCallback(null, null, null));
     ollamaAPI.askAsync(model, prompt);
     verify(ollamaAPI, times(1)).askAsync(model, prompt);
   }