Merge pull request #75 from ollama4j/model-listing

`listModelsFromLibrary` API

Makefile

@@ -11,9 +11,9 @@ doxygen:
 	doxygen Doxyfile
 
 list-releases:
-	curl 'https://central.sonatype.com/api/internal/browse/component/versions?sortField=normalizedVersion&sortDirection=asc&page=0&size=12&filter=namespace%3Aio.github.amithkoujalgi%2Cname%3Aollama4j' \
+	curl 'https://central.sonatype.com/api/internal/browse/component/versions?sortField=normalizedVersion&sortDirection=desc&page=0&size=20&filter=namespace%3Aio.github.ollama4j%2Cname%3Aollama4j' \
       --compressed \
-      --silent | jq '.components[].version'
+      --silent | jq -r '.components[].version'
 
 build-docs:
 	npm i --prefix docs && npm run build --prefix docs

README.md

@@ -279,7 +279,9 @@ If you like or are using this project to build your own, please give us a star.
 - `ollama-translator`: Minecraft 1.20.6 spigot plugin allows to easily break language barriers by using ollama on the
   server to translate all messages into a specfic target language.
     - https://github.com/liebki/ollama-translator
-- `Another Minecraft Mod`: https://www.reddit.com/r/fabricmc/comments/1e65x5s/comment/ldr2vcf/
+- `AI Player`: A minecraft mod which aims to add a "second player" into the game which will actually be intelligent.
+    - https://github.com/shasankp000/AI-Player
+    - https://www.reddit.com/r/fabricmc/comments/1e65x5s/comment/ldr2vcf/
 - `Ollama4j Web UI`: A web UI for Ollama written in Java using Spring Boot and Vaadin framework and
   Ollama4j.
     - https://github.com/ollama4j/ollama4j-web-ui
@@ -289,6 +291,8 @@ If you like or are using this project to build your own, please give us a star.
     - https://github.com/wyona/katie-backend
 - `TeleLlama3 Bot`: A Question-Answering Telegram Bot.
     - https://git.hiast.edu.sy/mohamadbashar.disoki/telellama3-bot
+- `moqui-wechat`: A wechat plugin
+    - https://github.com/heguangyong/moqui-wechat
 
 ## Traction
 

docs/docs/apis-generate/custom-roles.md

@@ -0,0 +1,65 @@
+---
+sidebar_position: 8
+---
+
+# Custom Roles
+
+Allows to manage custom roles (apart from the base roles) for chat interactions with the models.
+
+_Particularly helpful when you would need to use different roles that the newer models support other than the base
+roles._
+
+_Base roles are `SYSTEM`, `USER`, `ASSISTANT`, `TOOL`._
+
+### Usage
+
+#### Add new role
+
+```java
+import io.github.ollama4j.OllamaAPI;
+import io.github.ollama4j.models.chat.OllamaChatMessageRole;
+
+public class Main {
+
+    public static void main(String[] args) {
+        String host = "http://localhost:11434/";
+        OllamaAPI ollamaAPI = new OllamaAPI(host);
+
+        OllamaChatMessageRole customRole = ollamaAPI.addCustomRole("custom-role");
+    }
+}
+```
+
+#### List roles
+
+```java
+import io.github.ollama4j.OllamaAPI;
+import io.github.ollama4j.models.chat.OllamaChatMessageRole;
+
+public class Main {
+
+    public static void main(String[] args) {
+        String host = "http://localhost:11434/";
+        OllamaAPI ollamaAPI = new OllamaAPI(host);
+
+        List<OllamaChatMessageRole> roles = ollamaAPI.listRoles();
+    }
+}
+```
+
+#### Get role
+
+```java
+import io.github.ollama4j.OllamaAPI;
+import io.github.ollama4j.models.chat.OllamaChatMessageRole;
+
+public class Main {
+
+    public static void main(String[] args) {
+        String host = "http://localhost:11434/";
+        OllamaAPI ollamaAPI = new OllamaAPI(host);
+
+        List<OllamaChatMessageRole> roles = ollamaAPI.getRole("custom-role");
+    }
+}
+```

docs/docs/apis-generate/generate-embeddings.md

@@ -35,7 +35,7 @@ public class Main {
 }
 ```
 
-Or, using the `OllamaEmbedResponseModel`:
+Or, using the `OllamaEmbedRequestModel`:
 
 ```java
 import io.github.ollama4j.OllamaAPI;

pom.xml

@@ -63,6 +63,10 @@
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-javadoc-plugin</artifactId>
                 <version>3.5.0</version>
+                <configuration>
+                    <!-- to disable the "missing" warnings. Remove the doclint to enable warnings-->
+                    <doclint>all,-missing</doclint>
+                </configuration>
                 <executions>
                     <execution>
                         <id>attach-javadocs</id>
@@ -136,6 +140,11 @@
             <version>${lombok.version}</version>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.jsoup</groupId>
+            <artifactId>jsoup</artifactId>
+            <version>1.18.1</version>
+        </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-databind</artifactId>

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

@@ -1,12 +1,10 @@
 package io.github.ollama4j;
 
 import io.github.ollama4j.exceptions.OllamaBaseException;
+import io.github.ollama4j.exceptions.RoleNotFoundException;
 import io.github.ollama4j.exceptions.ToolInvocationException;
 import io.github.ollama4j.exceptions.ToolNotFoundException;
-import io.github.ollama4j.models.chat.OllamaChatMessage;
-import io.github.ollama4j.models.chat.OllamaChatRequest;
-import io.github.ollama4j.models.chat.OllamaChatRequestBuilder;
-import io.github.ollama4j.models.chat.OllamaChatResult;
+import io.github.ollama4j.models.chat.*;
 import io.github.ollama4j.models.embeddings.OllamaEmbedRequestModel;
 import io.github.ollama4j.models.embeddings.OllamaEmbeddingResponseModel;
 import io.github.ollama4j.models.embeddings.OllamaEmbeddingsRequestModel;
@@ -20,8 +18,6 @@ import io.github.ollama4j.tools.*;
 import io.github.ollama4j.utils.Options;
 import io.github.ollama4j.utils.Utils;
 import lombok.Setter;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import java.io.*;
 import java.net.URI;
@@ -34,11 +30,19 @@ import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.time.Duration;
 import java.util.*;
+import java.util.stream.Collectors;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.jsoup.Jsoup;
+import org.jsoup.nodes.Document;
+import org.jsoup.nodes.Element;
+import org.jsoup.select.Elements;
 
 /**
  * The base Ollama API class.
  */
-@SuppressWarnings("DuplicatedCode")
+@SuppressWarnings({"DuplicatedCode", "resource"})
 public class OllamaAPI {
 
     private static final Logger logger = LoggerFactory.getLogger(OllamaAPI.class);
@@ -99,12 +103,7 @@ public class OllamaAPI {
         HttpClient httpClient = HttpClient.newHttpClient();
         HttpRequest httpRequest = null;
         try {
-            httpRequest =
-                    getRequestBuilderDefault(new URI(url))
-                            .header("Accept", "application/json")
-                            .header("Content-type", "application/json")
-                            .GET()
-                            .build();
+            httpRequest = getRequestBuilderDefault(new URI(url)).header("Accept", "application/json").header("Content-type", "application/json").GET().build();
         } catch (URISyntaxException e) {
             throw new RuntimeException(e);
         }
@@ -123,19 +122,17 @@ public class OllamaAPI {
     /**
      * Provides a list of running models and details about each model currently loaded into memory.
      *
-     * @return ModelsProcessResponse
+     * @return ModelsProcessResponse containing details about the running models
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws OllamaBaseException  if the response indicates an error status
      */
     public ModelsProcessResponse ps() throws IOException, InterruptedException, OllamaBaseException {
         String url = this.host + "/api/ps";
         HttpClient httpClient = HttpClient.newHttpClient();
         HttpRequest httpRequest = null;
         try {
-            httpRequest =
-                    getRequestBuilderDefault(new URI(url))
-                            .header("Accept", "application/json")
-                            .header("Content-type", "application/json")
-                            .GET()
-                            .build();
+            httpRequest = getRequestBuilderDefault(new URI(url)).header("Accept", "application/json").header("Content-type", "application/json").GET().build();
         } catch (URISyntaxException e) {
             throw new RuntimeException(e);
         }
@@ -144,36 +141,72 @@ public class OllamaAPI {
         int statusCode = response.statusCode();
         String responseString = response.body();
         if (statusCode == 200) {
-            return Utils.getObjectMapper()
-                    .readValue(responseString, ModelsProcessResponse.class);
+            return Utils.getObjectMapper().readValue(responseString, ModelsProcessResponse.class);
         } else {
             throw new OllamaBaseException(statusCode + " - " + responseString);
         }
     }
 
     /**
-     * List available models from Ollama server.
+     * Lists available models from the Ollama server.
      *
-     * @return the list
+     * @return a list of models available on the server
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    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 =
-                getRequestBuilderDefault(new URI(url))
-                        .header("Accept", "application/json")
-                        .header("Content-type", "application/json")
-                        .GET()
-                        .build();
-        HttpResponse<String> response =
-                httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
+        HttpRequest httpRequest = getRequestBuilderDefault(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);
+        }
+    }
+
+    /**
+     * Retrieves a list of models from the Ollama library. This method fetches the available models directly from Ollama
+     * library page, including model details such as the name, pull count, popular tags, tag count, and the time when model was updated.
+     *
+     * @return A list of {@link LibraryModel} objects representing the models available in the Ollama library.
+     * @throws OllamaBaseException  If the HTTP request fails or the response is not successful (non-200 status code).
+     * @throws IOException          If an I/O error occurs during the HTTP request or response processing.
+     * @throws InterruptedException If the thread executing the request is interrupted.
+     * @throws URISyntaxException   If there is an error creating the URI for the HTTP request.
+     */
+    public List<LibraryModel> listModelsFromLibrary() throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
+        String url = "https://ollama.com/library";
+        HttpClient httpClient = HttpClient.newHttpClient();
+        HttpRequest httpRequest = getRequestBuilderDefault(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();
+        List<LibraryModel> models = new ArrayList<>();
+        if (statusCode == 200) {
+            Document doc = Jsoup.parse(responseString);
+            Elements modelSections = doc.selectXpath("//*[@id='repo']/ul/li/a");
+            for (Element e : modelSections) {
+                LibraryModel model = new LibraryModel();
+                Elements names = e.select("div > h2 > span");
+                Elements pullCounts = e.select("div:nth-of-type(2) > p > span:first-of-type > span:first-of-type");
+                Elements popularTags = e.select("div > div > span");
+                Elements tagCount = e.select("div:nth-of-type(2) > p > span:nth-of-type(2) > span:first-of-type");
+                Elements updatedAt = e.select("div:nth-of-type(2) > p > span:nth-of-type(3) > span:nth-of-type(2)");
+
+                model.setName(names.first().text());
+                model.setPullCount(pullCounts.first().text());
+                model.setPopularTags(popularTags.stream().map(Element::text).collect(Collectors.toList()));
+                model.setNumTags(Integer.parseInt(tagCount.first().text()));
+                model.setUpdatedAt(updatedAt.first().text());
+                models.add(model);
+            }
+            return models;
         } else {
             throw new OllamaBaseException(statusCode + " - " + responseString);
         }
@@ -184,29 +217,24 @@ public class OllamaAPI {
      * href="https://ollama.ai/library">available models</a>.
      *
      * @param modelName the name of the model
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    public void pullModel(String modelName)
-            throws OllamaBaseException, IOException, URISyntaxException, InterruptedException {
+    public void pullModel(String modelName) throws OllamaBaseException, IOException, URISyntaxException, InterruptedException {
         String url = this.host + "/api/pull";
         String jsonData = new ModelRequest(modelName).toString();
-        HttpRequest request =
-                getRequestBuilderDefault(new URI(url))
-                        .POST(HttpRequest.BodyPublishers.ofString(jsonData))
-                        .header("Accept", "application/json")
-                        .header("Content-type", "application/json")
-                        .build();
+        HttpRequest request = getRequestBuilderDefault(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());
                 }
@@ -222,17 +250,15 @@ public class OllamaAPI {
      *
      * @param modelName the model
      * @return the model details
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    public ModelDetail getModelDetails(String modelName)
-            throws IOException, OllamaBaseException, InterruptedException, URISyntaxException {
+    public ModelDetail getModelDetails(String modelName) throws IOException, OllamaBaseException, InterruptedException, URISyntaxException {
         String url = this.host + "/api/show";
         String jsonData = new ModelRequest(modelName).toString();
-        HttpRequest request =
-                getRequestBuilderDefault(new URI(url))
-                        .header("Accept", "application/json")
-                        .header("Content-type", "application/json")
-                        .POST(HttpRequest.BodyPublishers.ofString(jsonData))
-                        .build();
+        HttpRequest request = getRequestBuilderDefault(new URI(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();
@@ -250,17 +276,15 @@ public class OllamaAPI {
      *
      * @param modelName     the name of the custom model to be created.
      * @param modelFilePath the path to model file that exists on the Ollama server.
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    public void createModelWithFilePath(String modelName, String modelFilePath)
-            throws IOException, InterruptedException, OllamaBaseException, URISyntaxException {
+    public void createModelWithFilePath(String modelName, String modelFilePath) throws IOException, InterruptedException, OllamaBaseException, URISyntaxException {
         String url = this.host + "/api/create";
         String jsonData = new CustomModelFilePathRequest(modelName, modelFilePath).toString();
-        HttpRequest request =
-                getRequestBuilderDefault(new URI(url))
-                        .header("Accept", "application/json")
-                        .header("Content-Type", "application/json")
-                        .POST(HttpRequest.BodyPublishers.ofString(jsonData, StandardCharsets.UTF_8))
-                        .build();
+        HttpRequest request = getRequestBuilderDefault(new URI(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();
@@ -284,17 +308,15 @@ public class OllamaAPI {
      *
      * @param modelName         the name of the custom model to be created.
      * @param modelFileContents the path to model file that exists on the Ollama server.
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    public void createModelWithModelFileContents(String modelName, String modelFileContents)
-            throws IOException, InterruptedException, OllamaBaseException, URISyntaxException {
+    public void createModelWithModelFileContents(String modelName, String modelFileContents) throws IOException, InterruptedException, OllamaBaseException, URISyntaxException {
         String url = this.host + "/api/create";
         String jsonData = new CustomModelFileContentsRequest(modelName, modelFileContents).toString();
-        HttpRequest request =
-                getRequestBuilderDefault(new URI(url))
-                        .header("Accept", "application/json")
-                        .header("Content-Type", "application/json")
-                        .POST(HttpRequest.BodyPublishers.ofString(jsonData, StandardCharsets.UTF_8))
-                        .build();
+        HttpRequest request = getRequestBuilderDefault(new URI(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();
@@ -315,17 +337,15 @@ public class OllamaAPI {
      *
      * @param modelName          the name of the model to be deleted.
      * @param ignoreIfNotPresent ignore errors if the specified model is not present on Ollama server.
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    public void deleteModel(String modelName, boolean ignoreIfNotPresent)
-            throws IOException, InterruptedException, OllamaBaseException, URISyntaxException {
+    public void deleteModel(String modelName, boolean ignoreIfNotPresent) throws IOException, InterruptedException, OllamaBaseException, URISyntaxException {
         String url = this.host + "/api/delete";
         String jsonData = new ModelRequest(modelName).toString();
-        HttpRequest request =
-                getRequestBuilderDefault(new URI(url))
-                        .method("DELETE", HttpRequest.BodyPublishers.ofString(jsonData, StandardCharsets.UTF_8))
-                        .header("Accept", "application/json")
-                        .header("Content-type", "application/json")
-                        .build();
+        HttpRequest request = getRequestBuilderDefault(new URI(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();
@@ -344,11 +364,13 @@ public class OllamaAPI {
      * @param model  name of model to generate embeddings from
      * @param prompt text to generate embeddings for
      * @return embeddings
-     * @deprecated Use {@link #embed(String, List<String>)} instead.
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @deprecated Use {@link #embed(String, List)} instead.
      */
     @Deprecated
-    public List<Double> generateEmbeddings(String model, String prompt)
-            throws IOException, InterruptedException, OllamaBaseException {
+    public List<Double> generateEmbeddings(String model, String prompt) throws IOException, InterruptedException, OllamaBaseException {
         return generateEmbeddings(new OllamaEmbeddingsRequestModel(model, prompt));
     }
 
@@ -357,6 +379,9 @@ public class OllamaAPI {
      *
      * @param modelRequest request for '/api/embeddings' endpoint
      * @return embeddings
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      * @deprecated Use {@link #embed(OllamaEmbedRequestModel)} instead.
      */
     @Deprecated
@@ -364,17 +389,13 @@ public class OllamaAPI {
         URI uri = URI.create(this.host + "/api/embeddings");
         String jsonData = modelRequest.toString();
         HttpClient httpClient = HttpClient.newHttpClient();
-        HttpRequest.Builder requestBuilder =
-                getRequestBuilderDefault(uri)
-                        .header("Accept", "application/json")
-                        .POST(HttpRequest.BodyPublishers.ofString(jsonData));
+        HttpRequest.Builder requestBuilder = getRequestBuilderDefault(uri).header("Accept", "application/json").POST(HttpRequest.BodyPublishers.ofString(jsonData));
         HttpRequest request = requestBuilder.build();
         HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
         int statusCode = response.statusCode();
         String responseBody = response.body();
         if (statusCode == 200) {
-            OllamaEmbeddingResponseModel embeddingResponse =
-                    Utils.getObjectMapper().readValue(responseBody, OllamaEmbeddingResponseModel.class);
+            OllamaEmbeddingResponseModel embeddingResponse = Utils.getObjectMapper().readValue(responseBody, OllamaEmbeddingResponseModel.class);
             return embeddingResponse.getEmbedding();
         } else {
             throw new OllamaBaseException(statusCode + " - " + responseBody);
@@ -387,9 +408,11 @@ public class OllamaAPI {
      * @param model  name of model to generate embeddings from
      * @param inputs text/s to generate embeddings for
      * @return embeddings
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
-    public OllamaEmbedResponseModel embed(String model, List<String> inputs)
-            throws IOException, InterruptedException, OllamaBaseException {
+    public OllamaEmbedResponseModel embed(String model, List<String> inputs) throws IOException, InterruptedException, OllamaBaseException {
         return embed(new OllamaEmbedRequestModel(model, inputs));
     }
 
@@ -398,26 +421,23 @@ public class OllamaAPI {
      *
      * @param modelRequest request for '/api/embed' endpoint
      * @return embeddings
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
-    public OllamaEmbedResponseModel embed(OllamaEmbedRequestModel modelRequest)
-            throws IOException, InterruptedException, OllamaBaseException {
+    public OllamaEmbedResponseModel embed(OllamaEmbedRequestModel modelRequest) throws IOException, InterruptedException, OllamaBaseException {
         URI uri = URI.create(this.host + "/api/embed");
         String jsonData = Utils.getObjectMapper().writeValueAsString(modelRequest);
         HttpClient httpClient = HttpClient.newHttpClient();
 
-        HttpRequest request = HttpRequest.newBuilder(uri)
-                .header("Accept", "application/json")
-                .POST(HttpRequest.BodyPublishers.ofString(jsonData))
-                .build();
+        HttpRequest request = HttpRequest.newBuilder(uri).header("Accept", "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) {
-            OllamaEmbedResponseModel embeddingResponse =
-                    Utils.getObjectMapper().readValue(responseBody, OllamaEmbedResponseModel.class);
-            return embeddingResponse;
+            return Utils.getObjectMapper().readValue(responseBody, OllamaEmbedResponseModel.class);
         } else {
             throw new OllamaBaseException(statusCode + " - " + responseBody);
         }
@@ -434,9 +454,11 @@ public class OllamaAPI {
      *                      details on the options</a>
      * @param streamHandler optional callback consumer that will be applied every time a streamed response is received. If not set, the stream parameter of the request is set to false.
      * @return OllamaResult that includes response text and time taken for response
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
-    public OllamaResult generate(String model, String prompt, boolean raw, Options options, OllamaStreamHandler streamHandler)
-            throws OllamaBaseException, IOException, InterruptedException {
+    public OllamaResult generate(String model, String prompt, boolean raw, Options options, OllamaStreamHandler streamHandler) throws OllamaBaseException, IOException, InterruptedException {
         OllamaGenerateRequest ollamaRequestModel = new OllamaGenerateRequest(model, prompt);
         ollamaRequestModel.setRaw(raw);
         ollamaRequestModel.setOptions(options.getOptionsMap());
@@ -453,13 +475,14 @@ public class OllamaAPI {
      * @param raw     In some cases, you may wish to bypass the templating system and provide a full prompt. In this case, you can use the raw parameter to disable templating. Also note that raw mode will not return a context.
      * @param options Additional options or configurations to use when generating the response.
      * @return {@link OllamaResult}
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
-    public OllamaResult generate(String model, String prompt, boolean raw, Options options)
-            throws OllamaBaseException, IOException, InterruptedException {
+    public OllamaResult generate(String model, String prompt, boolean raw, Options options) throws OllamaBaseException, IOException, InterruptedException {
         return generate(model, prompt, raw, options, null);
     }
 
-
     /**
      * Generates response using the specified AI model and prompt (in blocking mode), and then invokes a set of tools
      * on the generated response.
@@ -468,13 +491,11 @@ public class OllamaAPI {
      * @param prompt  The input text or prompt to provide to the AI model.
      * @param options Additional options or configurations to use when generating the response.
      * @return {@link OllamaToolsResult} An OllamaToolsResult object containing the response from the AI model and the results of invoking the tools on that output.
-     * @throws OllamaBaseException  If there is an error related to the Ollama API or service.
-     * @throws IOException          If there is an error related to input/output operations.
-     * @throws InterruptedException If the method is interrupted while waiting for the AI model
-     *                              to generate the response or for the tools to be invoked.
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
-    public OllamaToolsResult generateWithTools(String model, String prompt, Options options)
-            throws OllamaBaseException, IOException, InterruptedException, ToolInvocationException {
+    public OllamaToolsResult generateWithTools(String model, String prompt, Options options) throws OllamaBaseException, IOException, InterruptedException, ToolInvocationException {
         boolean raw = true;
         OllamaToolsResult toolResult = new OllamaToolsResult();
         Map<ToolFunctionCallSpec, Object> toolResults = new HashMap<>();
@@ -495,7 +516,6 @@ public class OllamaAPI {
         return toolResult;
     }
 
-
     /**
      * Generate response for 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. This would be
@@ -509,9 +529,7 @@ public class OllamaAPI {
         OllamaGenerateRequest ollamaRequestModel = new OllamaGenerateRequest(model, prompt);
         ollamaRequestModel.setRaw(raw);
         URI uri = URI.create(this.host + "/api/generate");
-        OllamaAsyncResultStreamer ollamaAsyncResultStreamer =
-                new OllamaAsyncResultStreamer(
-                        getRequestBuilderDefault(uri), ollamaRequestModel, requestTimeoutSeconds);
+        OllamaAsyncResultStreamer ollamaAsyncResultStreamer = new OllamaAsyncResultStreamer(getRequestBuilderDefault(uri), ollamaRequestModel, requestTimeoutSeconds);
         ollamaAsyncResultStreamer.start();
         return ollamaAsyncResultStreamer;
     }
@@ -528,10 +546,11 @@ public class OllamaAPI {
      *                      details on the options</a>
      * @param streamHandler optional callback consumer that will be applied every time a streamed response is received. If not set, the stream parameter of the request is set to false.
      * @return OllamaResult that includes response text and time taken for response
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
-    public OllamaResult generateWithImageFiles(
-            String model, String prompt, List<File> imageFiles, Options options, OllamaStreamHandler streamHandler)
-            throws OllamaBaseException, IOException, InterruptedException {
+    public OllamaResult generateWithImageFiles(String model, String prompt, List<File> imageFiles, Options options, OllamaStreamHandler streamHandler) throws OllamaBaseException, IOException, InterruptedException {
         List<String> images = new ArrayList<>();
         for (File imageFile : imageFiles) {
             images.add(encodeFileToBase64(imageFile));
@@ -545,10 +564,12 @@ public class OllamaAPI {
      * Convenience method to call Ollama API without streaming responses.
      * <p>
      * Uses {@link #generateWithImageFiles(String, String, List, Options, OllamaStreamHandler)}
+     *
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
-    public OllamaResult generateWithImageFiles(
-            String model, String prompt, List<File> imageFiles, Options options)
-            throws OllamaBaseException, IOException, InterruptedException {
+    public OllamaResult generateWithImageFiles(String model, String prompt, List<File> imageFiles, Options options) throws OllamaBaseException, IOException, InterruptedException {
         return generateWithImageFiles(model, prompt, imageFiles, options, null);
     }
 
@@ -564,10 +585,12 @@ public class OllamaAPI {
      *                      details on the options</a>
      * @param streamHandler optional callback consumer that will be applied every time a streamed response is received. If not set, the stream parameter of the request is set to false.
      * @return OllamaResult that includes response text and time taken for response
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    public OllamaResult generateWithImageURLs(
-            String model, String prompt, List<String> imageURLs, Options options, OllamaStreamHandler streamHandler)
-            throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
+    public OllamaResult generateWithImageURLs(String model, String prompt, List<String> imageURLs, Options options, OllamaStreamHandler streamHandler) throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
         List<String> images = new ArrayList<>();
         for (String imageURL : imageURLs) {
             images.add(encodeByteArrayToBase64(Utils.loadImageBytesFromUrl(imageURL)));
@@ -581,14 +604,16 @@ public class OllamaAPI {
      * Convenience method to call Ollama API without streaming responses.
      * <p>
      * Uses {@link #generateWithImageURLs(String, String, List, Options, OllamaStreamHandler)}
+     *
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
+     * @throws URISyntaxException   if the URI for the request is malformed
      */
-    public OllamaResult generateWithImageURLs(String model, String prompt, List<String> imageURLs,
-                                              Options options)
-            throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
+    public OllamaResult generateWithImageURLs(String model, String prompt, List<String> imageURLs, Options options) throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
         return generateWithImageURLs(model, prompt, imageURLs, options, null);
     }
 
-
     /**
      * Ask a question to a model based on a given message stack (i.e. a chat history). Creates a synchronous call to the api
      * 'api/chat'.
@@ -599,6 +624,9 @@ public class OllamaAPI {
      * @throws OllamaBaseException  any response code than 200 has been returned
      * @throws IOException          in case the responseStream can not be read
      * @throws InterruptedException in case the server is not reachable or network issues happen
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
     public OllamaChatResult chat(String model, List<OllamaChatMessage> messages) throws OllamaBaseException, IOException, InterruptedException {
         OllamaChatRequestBuilder builder = OllamaChatRequestBuilder.getInstance(model);
@@ -615,6 +643,9 @@ public class OllamaAPI {
      * @throws OllamaBaseException  any response code than 200 has been returned
      * @throws IOException          in case the responseStream can not be read
      * @throws InterruptedException in case the server is not reachable or network issues happen
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
     public OllamaChatResult chat(OllamaChatRequest request) throws OllamaBaseException, IOException, InterruptedException {
         return chat(request, null);
@@ -631,6 +662,9 @@ public class OllamaAPI {
      * @throws OllamaBaseException  any response code than 200 has been returned
      * @throws IOException          in case the responseStream can not be read
      * @throws InterruptedException in case the server is not reachable or network issues happen
+     * @throws OllamaBaseException  if the response indicates an error status
+     * @throws IOException          if an I/O error occurs during the HTTP request
+     * @throws InterruptedException if the operation is interrupted
      */
     public OllamaChatResult chat(OllamaChatRequest request, OllamaStreamHandler streamHandler) throws OllamaBaseException, IOException, InterruptedException {
         OllamaChatEndpointCaller requestCaller = new OllamaChatEndpointCaller(host, basicAuth, requestTimeoutSeconds, verbose);
@@ -648,6 +682,37 @@ public class OllamaAPI {
         toolRegistry.addFunction(toolSpecification.getFunctionName(), toolSpecification.getToolDefinition());
     }
 
+    /**
+     * Adds a custom role.
+     *
+     * @param roleName the name of the custom role to be added
+     * @return the newly created OllamaChatMessageRole
+     */
+    public OllamaChatMessageRole addCustomRole(String roleName) {
+        return OllamaChatMessageRole.newCustomRole(roleName);
+    }
+
+    /**
+     * Lists all available roles.
+     *
+     * @return a list of available OllamaChatMessageRole objects
+     */
+    public List<OllamaChatMessageRole> listRoles() {
+        return OllamaChatMessageRole.getRoles();
+    }
+
+    /**
+     * Retrieves a specific role by name.
+     *
+     * @param roleName the name of the role to retrieve
+     * @return the OllamaChatMessageRole associated with the given name
+     * @throws RoleNotFoundException if the role with the specified name does not exist
+     */
+    public OllamaChatMessageRole getRole(String roleName) throws RoleNotFoundException {
+        return OllamaChatMessageRole.getRole(roleName);
+    }
+
+
     // technical private methods //
 
     private static String encodeFileToBase64(File file) throws IOException {
@@ -658,11 +723,8 @@ public class OllamaAPI {
         return Base64.getEncoder().encodeToString(bytes);
     }
 
-    private OllamaResult generateSyncForOllamaRequestModel(
-            OllamaGenerateRequest ollamaRequestModel, OllamaStreamHandler streamHandler)
-            throws OllamaBaseException, IOException, InterruptedException {
-        OllamaGenerateEndpointCaller requestCaller =
-                new OllamaGenerateEndpointCaller(host, basicAuth, requestTimeoutSeconds, verbose);
+    private OllamaResult generateSyncForOllamaRequestModel(OllamaGenerateRequest ollamaRequestModel, OllamaStreamHandler streamHandler) throws OllamaBaseException, IOException, InterruptedException {
+        OllamaGenerateEndpointCaller requestCaller = new OllamaGenerateEndpointCaller(host, basicAuth, requestTimeoutSeconds, verbose);
         OllamaResult result;
         if (streamHandler != null) {
             ollamaRequestModel.setStream(true);
@@ -680,10 +742,7 @@ public class OllamaAPI {
      * @return HttpRequest.Builder
      */
     private HttpRequest.Builder getRequestBuilderDefault(URI uri) {
-        HttpRequest.Builder requestBuilder =
-                HttpRequest.newBuilder(uri)
-                        .header("Content-Type", "application/json")
-                        .timeout(Duration.ofSeconds(requestTimeoutSeconds));
+        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder(uri).header("Content-Type", "application/json").timeout(Duration.ofSeconds(requestTimeoutSeconds));
         if (isBasicAuthCredentialsSet()) {
             requestBuilder.header("Authorization", getBasicAuthHeaderValue());
         }
@@ -709,7 +768,6 @@ public class OllamaAPI {
         return basicAuth != null;
     }
 
-
     private Object invokeTool(ToolFunctionCallSpec toolFunctionCallSpec) throws ToolInvocationException {
         try {
             String methodName = toolFunctionCallSpec.getName();

src/main/java/io/github/ollama4j/exceptions/RoleNotFoundException.java

@@ -0,0 +1,8 @@
+package io.github.ollama4j.exceptions;
+
+public class RoleNotFoundException extends Exception {
+
+    public RoleNotFoundException(String s) {
+        super(s);
+    }
+}

src/main/java/io/github/ollama4j/models/chat/OllamaChatMessage.java

@@ -8,6 +8,7 @@ import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import io.github.ollama4j.utils.FileToBase64Serializer;
 
 import java.util.List;
+
 import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.NoArgsConstructor;
@@ -33,13 +34,13 @@ public class OllamaChatMessage {
 
     @JsonSerialize(using = FileToBase64Serializer.class)
     private List<byte[]> images;
-    
-      @Override
-  public String toString() {
-    try {
-      return getObjectMapper().writerWithDefaultPrettyPrinter().writeValueAsString(this);
-    } catch (JsonProcessingException e) {
-      throw new RuntimeException(e);
+
+    @Override
+    public String toString() {
+        try {
+            return getObjectMapper().writerWithDefaultPrettyPrinter().writeValueAsString(this);
+        } catch (JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
     }
-  }
 }

src/main/java/io/github/ollama4j/models/chat/OllamaChatMessageRole.java

@@ -1,20 +1,53 @@
 package io.github.ollama4j.models.chat;
 
 import com.fasterxml.jackson.annotation.JsonValue;
+import io.github.ollama4j.exceptions.RoleNotFoundException;
+import lombok.Getter;
+
+import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Defines the possible Chat Message roles.
  */
-public enum OllamaChatMessageRole {
-    SYSTEM("system"),
-    USER("user"),
-    ASSISTANT("assistant"),
-    TOOL("tool");
+@Getter
+public class OllamaChatMessageRole {
+    private static final List<OllamaChatMessageRole> roles = new ArrayList<>();
+
+    public static final OllamaChatMessageRole SYSTEM = new OllamaChatMessageRole("system");
+    public static final OllamaChatMessageRole USER = new OllamaChatMessageRole("user");
+    public static final OllamaChatMessageRole ASSISTANT = new OllamaChatMessageRole("assistant");
+    public static final OllamaChatMessageRole TOOL = new OllamaChatMessageRole("tool");
 
     @JsonValue
-    private String roleName;
+    private final String roleName;
 
-    private OllamaChatMessageRole(String roleName){
+    private OllamaChatMessageRole(String roleName) {
         this.roleName = roleName;
+        roles.add(this);
+    }
+
+    public static OllamaChatMessageRole newCustomRole(String roleName) {
+        OllamaChatMessageRole customRole = new OllamaChatMessageRole(roleName);
+        roles.add(customRole);
+        return customRole;
+    }
+
+    public static List<OllamaChatMessageRole> getRoles() {
+        return new ArrayList<>(roles);
+    }
+
+    public static OllamaChatMessageRole getRole(String roleName) throws RoleNotFoundException {
+        for (OllamaChatMessageRole role : roles) {
+            if (role.roleName.equals(roleName)) {
+                return role;
+            }
+        }
+        throw new RoleNotFoundException("Invalid role name: " + roleName);
+    }
+
+    @Override
+    public String toString() {
+        return roleName;
     }
 }

src/main/java/io/github/ollama4j/models/chat/OllamaChatResult.java

@@ -8,12 +8,11 @@ import io.github.ollama4j.models.response.OllamaResult;
  * Specific chat-API result that contains the chat history sent to the model and appends the answer as {@link OllamaChatResult} given by the
  * {@link OllamaChatMessageRole#ASSISTANT} role.
  */
-public class OllamaChatResult extends OllamaResult{
+public class OllamaChatResult extends OllamaResult {
 
     private List<OllamaChatMessage> chatHistory;
 
-    public OllamaChatResult(String response, long responseTime, int httpStatusCode,
-            List<OllamaChatMessage> chatHistory) {
+    public OllamaChatResult(String response, long responseTime, int httpStatusCode, List<OllamaChatMessage> chatHistory) {
         super(response, responseTime, httpStatusCode);
         this.chatHistory = chatHistory;
         appendAnswerToChatHistory(response);
@@ -21,12 +20,10 @@ public class OllamaChatResult extends OllamaResult{
 
     public List<OllamaChatMessage> getChatHistory() {
         return chatHistory;
-    } 
+    }
 
-    private void appendAnswerToChatHistory(String answer){
+    private void appendAnswerToChatHistory(String answer) {
         OllamaChatMessage assistantMessage = new OllamaChatMessage(OllamaChatMessageRole.ASSISTANT, answer);
         this.chatHistory.add(assistantMessage);
     }
-    
-    
 }

src/main/java/io/github/ollama4j/models/response/LibraryModel.java

@@ -0,0 +1,15 @@
+package io.github.ollama4j.models.response;
+import java.util.ArrayList;
+import java.util.List;
+
+import lombok.Data;
+
+@Data
+public class LibraryModel {
+
+    private String name;
+    private String pullCount;
+    private int numTags;
+    private List<String> popularTags = new ArrayList<>();
+    private String updatedAt;
+}

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

@@ -2,6 +2,10 @@ package io.github.ollama4j.unittests;
 
 import io.github.ollama4j.OllamaAPI;
 import io.github.ollama4j.exceptions.OllamaBaseException;
+import io.github.ollama4j.exceptions.RoleNotFoundException;
+import io.github.ollama4j.models.chat.OllamaChatMessageRole;
+import io.github.ollama4j.models.embeddings.OllamaEmbedRequestModel;
+import io.github.ollama4j.models.embeddings.OllamaEmbedResponseModel;
 import io.github.ollama4j.models.response.ModelDetail;
 import io.github.ollama4j.models.response.OllamaAsyncResultStreamer;
 import io.github.ollama4j.models.response.OllamaResult;
@@ -14,7 +18,9 @@ import java.io.IOException;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 
+import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.*;
 
 class TestMockedAPIs {
@@ -97,6 +103,34 @@ class TestMockedAPIs {
         }
     }
 
+    @Test
+    void testEmbed() {
+        OllamaAPI ollamaAPI = Mockito.mock(OllamaAPI.class);
+        String model = OllamaModelType.LLAMA2;
+        List<String> inputs = List.of("some prompt text");
+        try {
+            when(ollamaAPI.embed(model, inputs)).thenReturn(new OllamaEmbedResponseModel());
+            ollamaAPI.embed(model, inputs);
+            verify(ollamaAPI, times(1)).embed(model, inputs);
+        } catch (IOException | OllamaBaseException | InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Test
+    void testEmbedWithEmbedRequestModel() {
+        OllamaAPI ollamaAPI = Mockito.mock(OllamaAPI.class);
+        String model = OllamaModelType.LLAMA2;
+        List<String> inputs = List.of("some prompt text");
+        try {
+            when(ollamaAPI.embed(new OllamaEmbedRequestModel(model, inputs))).thenReturn(new OllamaEmbedResponseModel());
+            ollamaAPI.embed(new OllamaEmbedRequestModel(model, inputs));
+            verify(ollamaAPI, times(1)).embed(new OllamaEmbedRequestModel(model, inputs));
+        } catch (IOException | OllamaBaseException | InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     @Test
     void testAsk() {
         OllamaAPI ollamaAPI = Mockito.mock(OllamaAPI.class);
@@ -161,4 +195,68 @@ class TestMockedAPIs {
         ollamaAPI.generateAsync(model, prompt, false);
         verify(ollamaAPI, times(1)).generateAsync(model, prompt, false);
     }
+
+    @Test
+    void testAddCustomRole() {
+        OllamaAPI ollamaAPI = mock(OllamaAPI.class);
+        String roleName = "custom-role";
+        OllamaChatMessageRole expectedRole = OllamaChatMessageRole.newCustomRole(roleName);
+        when(ollamaAPI.addCustomRole(roleName)).thenReturn(expectedRole);
+        OllamaChatMessageRole customRole = ollamaAPI.addCustomRole(roleName);
+        assertEquals(expectedRole, customRole);
+        verify(ollamaAPI, times(1)).addCustomRole(roleName);
+    }
+
+    @Test
+    void testListRoles() {
+        OllamaAPI ollamaAPI = Mockito.mock(OllamaAPI.class);
+        OllamaChatMessageRole role1 = OllamaChatMessageRole.newCustomRole("role1");
+        OllamaChatMessageRole role2 = OllamaChatMessageRole.newCustomRole("role2");
+        List<OllamaChatMessageRole> expectedRoles = List.of(role1, role2);
+        when(ollamaAPI.listRoles()).thenReturn(expectedRoles);
+        List<OllamaChatMessageRole> actualRoles = ollamaAPI.listRoles();
+        assertEquals(expectedRoles, actualRoles);
+        verify(ollamaAPI, times(1)).listRoles();
+    }
+
+    @Test
+    void testGetRoleNotFound() {
+        OllamaAPI ollamaAPI = mock(OllamaAPI.class);
+        String roleName = "non-existing-role";
+        try {
+            when(ollamaAPI.getRole(roleName)).thenThrow(new RoleNotFoundException("Role not found"));
+        } catch (RoleNotFoundException exception) {
+            throw new RuntimeException("Failed to run test: testGetRoleNotFound");
+        }
+        try {
+            ollamaAPI.getRole(roleName);
+            fail("Expected RoleNotFoundException not thrown");
+        } catch (RoleNotFoundException exception) {
+            assertEquals("Role not found", exception.getMessage());
+        }
+        try {
+            verify(ollamaAPI, times(1)).getRole(roleName);
+        } catch (RoleNotFoundException exception) {
+            throw new RuntimeException("Failed to run test: testGetRoleNotFound");
+        }
+    }
+
+    @Test
+    void testGetRoleFound() {
+        OllamaAPI ollamaAPI = mock(OllamaAPI.class);
+        String roleName = "existing-role";
+        OllamaChatMessageRole expectedRole = OllamaChatMessageRole.newCustomRole(roleName);
+        try {
+            when(ollamaAPI.getRole(roleName)).thenReturn(expectedRole);
+        } catch (RoleNotFoundException exception) {
+            throw new RuntimeException("Failed to run test: testGetRoleFound");
+        }
+        try {
+            OllamaChatMessageRole actualRole = ollamaAPI.getRole(roleName);
+            assertEquals(expectedRole, actualRole);
+            verify(ollamaAPI, times(1)).getRole(roleName);
+        } catch (RoleNotFoundException exception) {
+            throw new RuntimeException("Failed to run test: testGetRoleFound");
+        }
+    }
 }