OllamaAPI.java

Go to the documentation of this file.

package io.github.ollama4j;
 
import com.fasterxml.jackson.core.JsonParseException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
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.*;
import io.github.ollama4j.models.embeddings.OllamaEmbedRequestModel;
import io.github.ollama4j.models.embeddings.OllamaEmbedResponseModel;
import io.github.ollama4j.models.embeddings.OllamaEmbeddingResponseModel;
import io.github.ollama4j.models.embeddings.OllamaEmbeddingsRequestModel;
import io.github.ollama4j.models.generate.OllamaGenerateRequest;
import io.github.ollama4j.models.generate.OllamaStreamHandler;
import io.github.ollama4j.models.generate.OllamaTokenHandler;
import io.github.ollama4j.models.ps.ModelsProcessResponse;
import io.github.ollama4j.models.request.*;
import io.github.ollama4j.models.response.*;
import io.github.ollama4j.tools.*;
import io.github.ollama4j.tools.annotations.OllamaToolService;
import io.github.ollama4j.tools.annotations.ToolProperty;
import io.github.ollama4j.tools.annotations.ToolSpec;
import io.github.ollama4j.utils.Options;
import io.github.ollama4j.utils.Utils;
import lombok.Setter;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
 
import java.io.*;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.http.HttpClient;
import java.net.http.HttpConnectTimeoutException;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.time.Duration;
import java.util.*;
import java.util.stream.Collectors;

@SuppressWarnings({ "DuplicatedCode", "resource" })
public class OllamaAPI {
 
    private static final Logger logger = LoggerFactory.getLogger(OllamaAPI.class);
    private final String host;
    @Setter
    private long requestTimeoutSeconds = 10;
    @Setter
    private boolean verbose = true;
 
    @Setter
    private int maxChatToolCallRetries = 3;
 
    private Auth auth;
 
    private int numberOfRetriesForModelPull = 0;
 
    public void setNumberOfRetriesForModelPull(int numberOfRetriesForModelPull) {
        this.numberOfRetriesForModelPull = numberOfRetriesForModelPull;
    }
 
    private final ToolRegistry toolRegistry = new ToolRegistry();

    public OllamaAPI() {
        this.host = "http://localhost:11434";
    }

    public OllamaAPI(String host) {
        if (host.endsWith("/")) {
            this.host = host.substring(0, host.length() - 1);
        } else {
            this.host = host;
        }
        if (this.verbose) {
            logger.info("Ollama API initialized with host: " + this.host);
        }
    }

    public void setBasicAuth(String username, String password) {
        this.auth = new BasicAuth(username, password);
    }

    public void setBearerAuth(String bearerToken) {
        this.auth = new BearerAuth(bearerToken);
    }

    public boolean ping() {
        String url = this.host + "/api/tags";
        HttpClient httpClient = HttpClient.newHttpClient();
        HttpRequest httpRequest = null;
        try {
            httpRequest = getRequestBuilderDefault(new URI(url)).header("Accept", "application/json")
                    .header("Content-type", "application/json").GET().build();
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
        HttpResponse<String> response = null;
        try {
            response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
        } catch (HttpConnectTimeoutException e) {
            return false;
        } catch (IOException | InterruptedException e) {
            throw new RuntimeException(e);
        }
        int statusCode = response.statusCode();
        return statusCode == 200;
    }

    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();
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
        HttpResponse<String> response = null;
        response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
        int statusCode = response.statusCode();
        String responseString = response.body();
        if (statusCode == 200) {
            return Utils.getObjectMapper().readValue(responseString, ModelsProcessResponse.class);
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
    }

    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());
        int statusCode = response.statusCode();
        String responseString = response.body();
        if (statusCode == 200) {
            return Utils.getObjectMapper().readValue(responseString, ListModelsResponse.class).getModels();
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
    }

    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 > div > span");
                Elements desc = e.select("div > p");
                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 totalTags = e.select("div:nth-of-type(2) > p > span:nth-of-type(2) > span:first-of-type");
                Elements lastUpdatedTime = e
                        .select("div:nth-of-type(2) > p > span:nth-of-type(3) > span:nth-of-type(2)");
 
                if (names.first() == null || names.isEmpty()) {
                    // if name cannot be extracted, skip.
                    continue;
                }
                Optional.ofNullable(names.first()).map(Element::text).ifPresent(model::setName);
                model.setDescription(Optional.ofNullable(desc.first()).map(Element::text).orElse(""));
                model.setPopularTags(Optional.of(popularTags)
                        .map(tags -> tags.stream().map(Element::text).collect(Collectors.toList()))
                        .orElse(new ArrayList<>()));
                model.setPullCount(Optional.ofNullable(pullCounts.first()).map(Element::text).orElse(""));
                model.setTotalTags(
                        Optional.ofNullable(totalTags.first()).map(Element::text).map(Integer::parseInt).orElse(0));
                model.setLastUpdated(Optional.ofNullable(lastUpdatedTime.first()).map(Element::text).orElse(""));
 
                models.add(model);
            }
            return models;
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
    }

    public LibraryModelDetail getLibraryModelDetails(LibraryModel libraryModel)
            throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
        String url = String.format("https://ollama.com/library/%s/tags", libraryModel.getName());
        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<LibraryModelTag> libraryModelTags = new ArrayList<>();
        if (statusCode == 200) {
            Document doc = Jsoup.parse(responseString);
            Elements tagSections = doc
                    .select("html > body > main > div > section > div > div > div:nth-child(n+2) > div");
            for (Element e : tagSections) {
                Elements tags = e.select("div > a > div");
                Elements tagsMetas = e.select("div > span");
 
                LibraryModelTag libraryModelTag = new LibraryModelTag();
 
                if (tags.first() == null || tags.isEmpty()) {
                    // if tag cannot be extracted, skip.
                    continue;
                }
                libraryModelTag.setName(libraryModel.getName());
                Optional.ofNullable(tags.first()).map(Element::text).ifPresent(libraryModelTag::setTag);
                libraryModelTag.setSize(Optional.ofNullable(tagsMetas.first()).map(element -> element.text().split("•"))
                        .filter(parts -> parts.length > 1).map(parts -> parts[1].trim()).orElse(""));
                libraryModelTag
                        .setLastUpdated(Optional.ofNullable(tagsMetas.first()).map(element -> element.text().split("•"))
                                .filter(parts -> parts.length > 1).map(parts -> parts[2].trim()).orElse(""));
                libraryModelTags.add(libraryModelTag);
            }
            LibraryModelDetail libraryModelDetail = new LibraryModelDetail();
            libraryModelDetail.setModel(libraryModel);
            libraryModelDetail.setTags(libraryModelTags);
            return libraryModelDetail;
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
    }

    public LibraryModelTag findModelTagFromLibrary(String modelName, String tag)
            throws OllamaBaseException, IOException, URISyntaxException, InterruptedException {
        List<LibraryModel> libraryModels = this.listModelsFromLibrary();
        LibraryModel libraryModel = libraryModels.stream().filter(model -> model.getName().equals(modelName))
                .findFirst().orElseThrow(
                        () -> new NoSuchElementException(String.format("Model by name '%s' not found", modelName)));
        LibraryModelDetail libraryModelDetail = this.getLibraryModelDetails(libraryModel);
        LibraryModelTag libraryModelTag = libraryModelDetail.getTags().stream()
                .filter(tagName -> tagName.getTag().equals(tag)).findFirst()
                .orElseThrow(() -> new NoSuchElementException(
                        String.format("Tag '%s' for model '%s' not found", tag, modelName)));
        return libraryModelTag;
    }

    public void pullModel(String modelName)
            throws OllamaBaseException, IOException, URISyntaxException, InterruptedException {
        if (numberOfRetriesForModelPull == 0) {
            this.doPullModel(modelName);
        } else {
            int numberOfRetries = 0;
            while (numberOfRetries < numberOfRetriesForModelPull) {
                try {
                    this.doPullModel(modelName);
                    return;
                } catch (OllamaBaseException e) {
                    logger.error("Failed to pull model " + modelName + ", retrying...");
                    numberOfRetries++;
                }
            }
            throw new OllamaBaseException(
                    "Failed to pull model " + modelName + " after " + numberOfRetriesForModelPull + " retries");
        }
    }
 
    private void doPullModel(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();
        HttpClient client = HttpClient.newHttpClient();
        HttpResponse<InputStream> response = client.send(request, HttpResponse.BodyHandlers.ofInputStream());
        int statusCode = response.statusCode();
        InputStream responseBodyStream = response.body();
        String responseString = "";
        boolean success = false; // Flag to check the pull success.
        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);
                if (modelPullResponse != null && modelPullResponse.getStatus() != null) {
                    if (verbose) {
                        logger.info(modelName + ": " + modelPullResponse.getStatus());
                    }
                    // Check if status is "success" and set success flag to true.
                    if ("success".equalsIgnoreCase(modelPullResponse.getStatus())) {
                        success = true;
                    }
                } else {
                    logger.error("Received null or invalid status for model pull.");
                }
            }
        }
        if (!success) {
            logger.error("Model pull failed or returned invalid status.");
            throw new OllamaBaseException("Model pull failed or returned invalid status.");
        }
        if (statusCode != 200) {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
    }
 
    public String getVersion() throws URISyntaxException, IOException, InterruptedException, OllamaBaseException {
        String url = this.host + "/api/version";
        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();
        if (statusCode == 200) {
            return Utils.getObjectMapper().readValue(responseString, OllamaVersion.class).getVersion();
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
    }

    public void pullModel(LibraryModelTag libraryModelTag)
            throws OllamaBaseException, IOException, URISyntaxException, InterruptedException {
        String tagToPull = String.format("%s:%s", libraryModelTag.getName(), libraryModelTag.getTag());
        pullModel(tagToPull);
    }

    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();
        HttpClient client = HttpClient.newHttpClient();
        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        int statusCode = response.statusCode();
        String responseBody = response.body();
        if (statusCode == 200) {
            return Utils.getObjectMapper().readValue(responseBody, ModelDetail.class);
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseBody);
        }
    }

    @Deprecated
    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();
        HttpClient client = HttpClient.newHttpClient();
        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        int statusCode = response.statusCode();
        String responseString = response.body();
        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.
        if (responseString.contains("error")) {
            throw new OllamaBaseException(responseString);
        }
        if (verbose) {
            logger.info(responseString);
        }
    }

    @Deprecated
    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();
        HttpClient client = HttpClient.newHttpClient();
        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        int statusCode = response.statusCode();
        String responseString = response.body();
        if (statusCode != 200) {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
        if (responseString.contains("error")) {
            throw new OllamaBaseException(responseString);
        }
        if (verbose) {
            logger.info(responseString);
        }
    }

    public void createModel(CustomModelRequest customModelRequest)
            throws IOException, InterruptedException, OllamaBaseException, URISyntaxException {
        String url = this.host + "/api/create";
        String jsonData = customModelRequest.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();
        HttpClient client = HttpClient.newHttpClient();
        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        int statusCode = response.statusCode();
        String responseString = response.body();
        if (statusCode != 200) {
            throw new OllamaBaseException(statusCode + " - " + responseString);
        }
        if (responseString.contains("error")) {
            throw new OllamaBaseException(responseString);
        }
        if (verbose) {
            logger.info(responseString);
        }
    }

    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();
        HttpClient client = HttpClient.newHttpClient();
        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        int statusCode = response.statusCode();
        String responseBody = response.body();
        if (statusCode == 404 && responseBody.contains("model") && responseBody.contains("not found")) {
            return;
        }
        if (statusCode != 200) {
            throw new OllamaBaseException(statusCode + " - " + responseBody);
        }
    }

    @Deprecated
    public List<Double> generateEmbeddings(String model, String prompt)
            throws IOException, InterruptedException, OllamaBaseException {
        return generateEmbeddings(new OllamaEmbeddingsRequestModel(model, prompt));
    }

    @Deprecated
    public List<Double> generateEmbeddings(OllamaEmbeddingsRequestModel modelRequest)
            throws IOException, InterruptedException, OllamaBaseException {
        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 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);
            return embeddingResponse.getEmbedding();
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseBody);
        }
    }

    public OllamaEmbedResponseModel embed(String model, List<String> inputs)
            throws IOException, InterruptedException, OllamaBaseException {
        return embed(new OllamaEmbedRequestModel(model, inputs));
    }

    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();
 
        HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
        int statusCode = response.statusCode();
        String responseBody = response.body();
 
        if (statusCode == 200) {
            return Utils.getObjectMapper().readValue(responseBody, OllamaEmbedResponseModel.class);
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseBody);
        }
    }

    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());
        return generateSyncForOllamaRequestModel(ollamaRequestModel, streamHandler);
    }

    public OllamaResult generate(String model, String prompt, Map<String, Object> format)
            throws OllamaBaseException, IOException, InterruptedException {
        URI uri = URI.create(this.host + "/api/generate");
 
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", model);
        requestBody.put("prompt", prompt);
        requestBody.put("stream", false);
        requestBody.put("format", format);
 
        String jsonData = Utils.getObjectMapper().writeValueAsString(requestBody);
        HttpClient httpClient = HttpClient.newHttpClient();
 
        HttpRequest request = HttpRequest.newBuilder(uri)
                .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) {
            OllamaStructuredResult structuredResult = Utils.getObjectMapper().readValue(responseBody,
                    OllamaStructuredResult.class);
            OllamaResult ollamaResult = new OllamaResult(structuredResult.getResponse(),
                    structuredResult.getResponseTime(), statusCode);
            return ollamaResult;
        } else {
            throw new OllamaBaseException(statusCode + " - " + responseBody);
        }
    }

    public OllamaResult generate(String model, String prompt, boolean raw, Options options)
            throws OllamaBaseException, IOException, InterruptedException {
        return generate(model, prompt, raw, options, null);
    }

    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<>();
 
        if (!prompt.startsWith("[AVAILABLE_TOOLS]")) {
            final Tools.PromptBuilder promptBuilder = new Tools.PromptBuilder();
            for (Tools.ToolSpecification spec : toolRegistry.getRegisteredSpecs()) {
                promptBuilder.withToolSpecification(spec);
            }
            promptBuilder.withPrompt(prompt);
            prompt = promptBuilder.build();
        }
 
        OllamaResult result = generate(model, prompt, raw, options, null);
        toolResult.setModelResult(result);
 
        String toolsResponse = result.getResponse();
        if (toolsResponse.contains("[TOOL_CALLS]")) {
            toolsResponse = toolsResponse.replace("[TOOL_CALLS]", "");
        }
 
        List<ToolFunctionCallSpec> toolFunctionCallSpecs = new ArrayList<>();
        ObjectMapper objectMapper = Utils.getObjectMapper();
 
        if (!toolsResponse.isEmpty()) {
            try {
                // Try to parse the string to see if it's a valid JSON
                JsonNode jsonNode = objectMapper.readTree(toolsResponse);
            } catch (JsonParseException e) {
                logger.warn("Response from model does not contain any tool calls. Returning the response as is.");
                return toolResult;
            }
            toolFunctionCallSpecs = objectMapper.readValue(
                    toolsResponse,
                    objectMapper.getTypeFactory().constructCollectionType(List.class, ToolFunctionCallSpec.class));
        }
        for (ToolFunctionCallSpec toolFunctionCallSpec : toolFunctionCallSpecs) {
            toolResults.put(toolFunctionCallSpec, invokeTool(toolFunctionCallSpec));
        }
        toolResult.setToolResults(toolResults);
        return toolResult;
    }

    public OllamaAsyncResultStreamer generateAsync(String model, String prompt, boolean raw) {
        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.start();
        return ollamaAsyncResultStreamer;
    }

    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));
        }
        OllamaGenerateRequest ollamaRequestModel = new OllamaGenerateRequest(model, prompt, images);
        ollamaRequestModel.setOptions(options.getOptionsMap());
        return generateSyncForOllamaRequestModel(ollamaRequestModel, streamHandler);
    }

    public OllamaResult generateWithImageFiles(String model, String prompt, List<File> imageFiles, Options options)
            throws OllamaBaseException, IOException, InterruptedException {
        return generateWithImageFiles(model, prompt, imageFiles, options, null);
    }

    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)));
        }
        OllamaGenerateRequest ollamaRequestModel = new OllamaGenerateRequest(model, prompt, images);
        ollamaRequestModel.setOptions(options.getOptionsMap());
        return generateSyncForOllamaRequestModel(ollamaRequestModel, streamHandler);
    }

    public OllamaResult generateWithImageURLs(String model, String prompt, List<String> imageURLs, Options options)
            throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
        return generateWithImageURLs(model, prompt, imageURLs, options, null);
    }

    public OllamaResult generateWithImages(String model, String prompt, List<byte[]> images, Options options, OllamaStreamHandler streamHandler) throws OllamaBaseException, IOException, InterruptedException {
        List<String> encodedImages = new ArrayList<>();
        for (byte[] image : images) {
            encodedImages.add(encodeByteArrayToBase64(image));
        }
        OllamaGenerateRequest ollamaRequestModel = new OllamaGenerateRequest(model, prompt, encodedImages);
        ollamaRequestModel.setOptions(options.getOptionsMap());
        return generateSyncForOllamaRequestModel(ollamaRequestModel, streamHandler);
    }

    public OllamaResult generateWithImages(String model, String prompt, List<byte[]> images, Options options) throws OllamaBaseException, IOException, InterruptedException {
        return generateWithImages(model, prompt, images, options, null);
    }

    public OllamaChatResult chat(String model, List<OllamaChatMessage> messages)
            throws OllamaBaseException, IOException, InterruptedException, ToolInvocationException {
        OllamaChatRequestBuilder builder = OllamaChatRequestBuilder.getInstance(model);
        return chat(builder.withMessages(messages).build());
    }

    public OllamaChatResult chat(OllamaChatRequest request)
            throws OllamaBaseException, IOException, InterruptedException, ToolInvocationException {
        return chat(request, null);
    }

    public OllamaChatResult chat(OllamaChatRequest request, OllamaStreamHandler streamHandler)
            throws OllamaBaseException, IOException, InterruptedException, ToolInvocationException {
        return chatStreaming(request, new OllamaChatStreamObserver(streamHandler));
    }

    public OllamaChatResult chatStreaming(OllamaChatRequest request, OllamaTokenHandler tokenHandler)
            throws OllamaBaseException, IOException, InterruptedException, ToolInvocationException {
        OllamaChatEndpointCaller requestCaller = new OllamaChatEndpointCaller(host, auth, requestTimeoutSeconds,
                verbose);
        OllamaChatResult result;
 
        // add all registered tools to Request
        request.setTools(toolRegistry.getRegisteredSpecs().stream().map(Tools.ToolSpecification::getToolPrompt)
                .collect(Collectors.toList()));
 
        if (tokenHandler != null) {
            request.setStream(true);
            result = requestCaller.call(request, tokenHandler);
        } else {
            result = requestCaller.callSync(request);
        }
 
        // check if toolCallIsWanted
        List<OllamaChatToolCalls> toolCalls = result.getResponseModel().getMessage().getToolCalls();
        int toolCallTries = 0;
        while (toolCalls != null && !toolCalls.isEmpty() && toolCallTries < maxChatToolCallRetries) {
            for (OllamaChatToolCalls toolCall : toolCalls) {
                String toolName = toolCall.getFunction().getName();
                ToolFunction toolFunction = toolRegistry.getToolFunction(toolName);
                if (toolFunction == null) {
                    throw new ToolInvocationException("Tool function not found: " + toolName);
                }
                Map<String, Object> arguments = toolCall.getFunction().getArguments();
                Object res = toolFunction.apply(arguments);
                request.getMessages().add(new OllamaChatMessage(OllamaChatMessageRole.TOOL,
                        "[TOOL_RESULTS]" + toolName + "(" + arguments.keySet() + ") : " + res + "[/TOOL_RESULTS]"));
            }
 
            if (tokenHandler != null) {
                result = requestCaller.call(request, tokenHandler);
            } else {
                result = requestCaller.callSync(request);
            }
            toolCalls = result.getResponseModel().getMessage().getToolCalls();
            toolCallTries++;
        }
 
        return result;
    }

    public void registerTool(Tools.ToolSpecification toolSpecification) {
        toolRegistry.addTool(toolSpecification.getFunctionName(), toolSpecification);
        if (this.verbose) {
            logger.debug("Registered tool: {}", toolSpecification.getFunctionName());
        }
    }

    public void registerTools(List<Tools.ToolSpecification> toolSpecifications) {
        for (Tools.ToolSpecification toolSpecification : toolSpecifications) {
            toolRegistry.addTool(toolSpecification.getFunctionName(), toolSpecification);
        }
    }

    public void registerAnnotatedTools() {
        try {
            Class<?> callerClass = null;
            try {
                callerClass = Class.forName(Thread.currentThread().getStackTrace()[2].getClassName());
            } catch (ClassNotFoundException e) {
                throw new RuntimeException(e);
            }
 
            OllamaToolService ollamaToolServiceAnnotation = callerClass.getDeclaredAnnotation(OllamaToolService.class);
            if (ollamaToolServiceAnnotation == null) {
                throw new IllegalStateException(callerClass + " is not annotated as " + OllamaToolService.class);
            }
 
            Class<?>[] providers = ollamaToolServiceAnnotation.providers();
            for (Class<?> provider : providers) {
                registerAnnotatedTools(provider.getDeclaredConstructor().newInstance());
            }
        } catch (InstantiationException | NoSuchMethodException | IllegalAccessException
                | InvocationTargetException e) {
            throw new RuntimeException(e);
        }
    }

    public void registerAnnotatedTools(Object object) {
        Class<?> objectClass = object.getClass();
        Method[] methods = objectClass.getMethods();
        for (Method m : methods) {
            ToolSpec toolSpec = m.getDeclaredAnnotation(ToolSpec.class);
            if (toolSpec == null) {
                continue;
            }
            String operationName = !toolSpec.name().isBlank() ? toolSpec.name() : m.getName();
            String operationDesc = !toolSpec.desc().isBlank() ? toolSpec.desc() : operationName;
 
            final Tools.PropsBuilder propsBuilder = new Tools.PropsBuilder();
            LinkedHashMap<String, String> methodParams = new LinkedHashMap<>();
            for (Parameter parameter : m.getParameters()) {
                final ToolProperty toolPropertyAnn = parameter.getDeclaredAnnotation(ToolProperty.class);
                String propType = parameter.getType().getTypeName();
                if (toolPropertyAnn == null) {
                    methodParams.put(parameter.getName(), null);
                    continue;
                }
                String propName = !toolPropertyAnn.name().isBlank() ? toolPropertyAnn.name() : parameter.getName();
                methodParams.put(propName, propType);
                propsBuilder.withProperty(propName, Tools.PromptFuncDefinition.Property.builder().type(propType)
                        .description(toolPropertyAnn.desc()).required(toolPropertyAnn.required()).build());
            }
            final Map<String, Tools.PromptFuncDefinition.Property> params = propsBuilder.build();
            List<String> reqProps = params.entrySet().stream().filter(e -> e.getValue().isRequired())
                    .map(Map.Entry::getKey).collect(Collectors.toList());
 
            Tools.ToolSpecification toolSpecification = Tools.ToolSpecification.builder().functionName(operationName)
                    .functionDescription(operationDesc)
                    .toolPrompt(Tools.PromptFuncDefinition.builder().type("function")
                            .function(Tools.PromptFuncDefinition.PromptFuncSpec.builder().name(operationName)
                                    .description(operationDesc).parameters(Tools.PromptFuncDefinition.Parameters
                                            .builder().type("object").properties(params).required(reqProps).build())
                                    .build())
                            .build())
                    .build();
 
            ReflectionalToolFunction reflectionalToolFunction = new ReflectionalToolFunction(object, m, methodParams);
            toolSpecification.setToolFunction(reflectionalToolFunction);
            toolRegistry.addTool(toolSpecification.getFunctionName(), toolSpecification);
        }
 
    }

    public OllamaChatMessageRole addCustomRole(String roleName) {
        return OllamaChatMessageRole.newCustomRole(roleName);
    }

    public List<OllamaChatMessageRole> listRoles() {
        return OllamaChatMessageRole.getRoles();
    }

    public OllamaChatMessageRole getRole(String roleName) throws RoleNotFoundException {
        return OllamaChatMessageRole.getRole(roleName);
    }
 
    // technical private methods //

    private static String encodeFileToBase64(File file) throws IOException {
        return Base64.getEncoder().encodeToString(Files.readAllBytes(file.toPath()));
    }

    private static String encodeByteArrayToBase64(byte[] bytes) {
        return Base64.getEncoder().encodeToString(bytes);
    }

    private OllamaResult generateSyncForOllamaRequestModel(OllamaGenerateRequest ollamaRequestModel,
            OllamaStreamHandler streamHandler) throws OllamaBaseException, IOException, InterruptedException {
        OllamaGenerateEndpointCaller requestCaller = new OllamaGenerateEndpointCaller(host, auth, requestTimeoutSeconds,
                verbose);
        OllamaResult result;
        if (streamHandler != null) {
            ollamaRequestModel.setStream(true);
            result = requestCaller.call(ollamaRequestModel, streamHandler);
        } else {
            result = requestCaller.callSync(ollamaRequestModel);
        }
        return result;
    }

    private HttpRequest.Builder getRequestBuilderDefault(URI uri) {
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder(uri).header("Content-Type", "application/json")
                .timeout(Duration.ofSeconds(requestTimeoutSeconds));
        if (isBasicAuthCredentialsSet()) {
            requestBuilder.header("Authorization", auth.getAuthHeaderValue());
        }
        return requestBuilder;
    }

    private boolean isBasicAuthCredentialsSet() {
        return auth != null;
    }
 
    private Object invokeTool(ToolFunctionCallSpec toolFunctionCallSpec) throws ToolInvocationException {
        try {
            String methodName = toolFunctionCallSpec.getName();
            Map<String, Object> arguments = toolFunctionCallSpec.getArguments();
            ToolFunction function = toolRegistry.getToolFunction(methodName);
            if (verbose) {
                logger.debug("Invoking function {} with arguments {}", methodName, arguments);
            }
            if (function == null) {
                throw new ToolNotFoundException(
                        "No such tool: " + methodName + ". Please register the tool before invoking it.");
            }
            return function.apply(arguments);
        } catch (Exception e) {
            throw new ToolInvocationException("Failed to invoke tool: " + toolFunctionCallSpec.getName(), e);
        }
    }
}