diff --git a/pom.xml b/pom.xml index 6aaa5a05a..c5ab8868d 100644 --- a/pom.xml +++ b/pom.xml @@ -16,6 +16,7 @@ spring-ai-core spring-ai-openai spring-ai-azure-openai + spring-ai-ollama spring-ai-spring-boot-autoconfigure spring-ai-spring-boot-starters/spring-ai-starter-openai spring-ai-spring-boot-starters/spring-ai-starter-azure-openai diff --git a/spring-ai-ollama/README.md b/spring-ai-ollama/README.md new file mode 100644 index 000000000..21a76267c --- /dev/null +++ b/spring-ai-ollama/README.md @@ -0,0 +1,9 @@ +## Ollama + +Ollama lets you ge tup an running with large language models locally + +Refer to the official [README](https://github.com/jmorganca/ollama) to get started. + +Note, installing `ollama run llama2` will download a 4GB docker image. + +You can run the disabled test in `OllamaClientTests.java` to kick the tires. \ No newline at end of file diff --git a/spring-ai-ollama/pom.xml b/spring-ai-ollama/pom.xml new file mode 100644 index 000000000..1277901aa --- /dev/null +++ b/spring-ai-ollama/pom.xml @@ -0,0 +1,42 @@ + + + 4.0.0 + + org.springframework.experimental.ai + spring-ai + 0.7.0-SNAPSHOT + + + spring-ai-ollama + jar + Spring AI Ollama + Ollama support + + + 17 + 17 + UTF-8 + + + + + org.springframework.experimental.ai + spring-ai-core + ${project.parent.version} + + + + org.springframework.boot + spring-boot-starter-logging + + + + + org.springframework.boot + spring-boot-starter-test + test + + + \ No newline at end of file diff --git a/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/client/OllamaClient.java b/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/client/OllamaClient.java new file mode 100644 index 000000000..0f61c1567 --- /dev/null +++ b/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/client/OllamaClient.java @@ -0,0 +1,248 @@ +package org.springframework.ai.ollama.client; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.client.AiClient; +import org.springframework.ai.client.AiResponse; +import org.springframework.ai.client.Generation; +import org.springframework.ai.prompt.Prompt; +import org.springframework.util.CollectionUtils; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; +import java.util.*; +import java.util.function.Consumer; +import java.util.stream.Collectors; + +/** + * A client implementation for interacting with Ollama Service. This class acts as an + * interface between the application and the Ollama AI Service, handling request creation, + * communication, and response processing. + * + * @author nullptr + */ +public class OllamaClient implements AiClient { + + /** Logger for logging the events and messages. */ + private static final Logger log = LoggerFactory.getLogger(OllamaClient.class); + + /** Mapper for JSON serialization and deserialization. */ + private static final ObjectMapper jsonMapper = new ObjectMapper(); + + /** HTTP client for making asynchronous calls to the Ollama Service. */ + private static final HttpClient httpClient = HttpClient.newBuilder().build(); + + /** Base URL of the Ollama Service. */ + private final String baseUrl; + + /** Name of the model to be used for the AI service. */ + private final String model; + + /** Optional callback to handle individual generation results. */ + private Consumer simpleCallback; + + /** + * Constructs an OllamaClient with the specified base URL and model. + * @param baseUrl Base URL of the Ollama Service. + * @param model Model specification for the AI service. + */ + public OllamaClient(String baseUrl, String model) { + this.baseUrl = baseUrl; + this.model = model; + } + + /** + * Constructs an OllamaClient with the specified base URL, model, and a callback. + * @param baseUrl Base URL of the Ollama Service. + * @param model Model specification for the AI service. + * @param simpleCallback Callback to handle individual generation results. + */ + public OllamaClient(String baseUrl, String model, Consumer simpleCallback) { + this(baseUrl, model); + this.simpleCallback = simpleCallback; + } + + @Override + public AiResponse generate(Prompt prompt) { + validatePrompt(prompt); + + HttpRequest request = buildHttpRequest(prompt); + var response = sendRequest(request); + + List results = readGenerateResults(response.body()); + return getAiResponse(results); + } + + /** + * Validates the provided prompt. + * @param prompt The prompt to validate. + */ + protected void validatePrompt(Prompt prompt) { + if (CollectionUtils.isEmpty(prompt.getMessages())) { + throw new RuntimeException("The prompt message cannot be empty."); + } + + if (prompt.getMessages().size() > 1) { + log.warn("Only the first prompt message will be used; subsequent messages will be ignored."); + } + } + + /** + * Constructs an HTTP request for the provided prompt. + * @param prompt The prompt for which the request needs to be built. + * @return The constructed HttpRequest. + */ + protected HttpRequest buildHttpRequest(Prompt prompt) { + String requestBody = getGenerateRequestBody(prompt.getMessages().get(0).getContent()); + + // remove the suffix '/' if necessary + String url = !this.baseUrl.endsWith("/") ? this.baseUrl : this.baseUrl.substring(0, this.baseUrl.length() - 1); + + return HttpRequest.newBuilder() + .uri(URI.create("%s/api/generate".formatted(url))) + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .timeout(Duration.ofMinutes(5L)) + .build(); + } + + /** + * Sends the constructed HttpRequest and retrieves the HttpResponse. + * @param request The HttpRequest to be sent. + * @return HttpResponse containing the response data. + */ + protected HttpResponse sendRequest(HttpRequest request) { + var response = httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofInputStream()).join(); + if (response.statusCode() != 200) { + throw new RuntimeException("Ollama call returned an unexpected status: " + response.statusCode()); + } + return response; + } + + /** + * Serializes the prompt into a request body for the Ollama API call. + * @param prompt The prompt to be serialized. + * @return Serialized request body as a String. + */ + private String getGenerateRequestBody(String prompt) { + var data = Map.of("model", model, "prompt", prompt); + try { + return jsonMapper.writeValueAsString(data); + } + catch (JsonProcessingException ex) { + throw new RuntimeException("Failed to serialize the prompt to JSON", ex); + } + + } + + /** + * Reads and processes the results from the InputStream provided by the Ollama + * Service. + * @param inputStream InputStream containing the results from the Ollama Service. + * @return List of OllamaGenerateResult. + */ + protected List readGenerateResults(InputStream inputStream) { + try (BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream))) { + var results = new ArrayList(); + String line; + while ((line = bufferedReader.readLine()) != null) { + processResponseLine(line, results); + } + return results; + } + catch (IOException e) { + throw new RuntimeException("Error parsing Ollama generation response.", e); + } + } + + /** + * Processes a single line from the Ollama response. + * @param line The line to be processed. + * @param results List to which parsed results will be added. + */ + protected void processResponseLine(String line, List results) { + if (line.isBlank()) + return; + + log.debug("Received ollama generate response: {}", line); + + OllamaGenerateResult result; + try { + result = jsonMapper.readValue(line, OllamaGenerateResult.class); + } + catch (IOException e) { + throw new RuntimeException("Error parsing response line from Ollama.", e); + } + + if (result.getModel() == null || result.getDone() == null) { + throw new IllegalStateException("Received invalid data from Ollama. Model = " + result.getModel() + + " , Done = " + result.getDone()); + + } + + if (simpleCallback != null) { + simpleCallback.accept(result); + } + + results.add(result); + } + + /** + * Converts the list of OllamaGenerateResult into a structured AiResponse. + * @param results List of OllamaGenerateResult. + * @return Formulated AiResponse. + */ + protected AiResponse getAiResponse(List results) { + var ollamaResponse = results.stream() + .filter(Objects::nonNull) + .filter(it -> it.getResponse() != null && !it.getResponse().isBlank()) + .filter(it -> it.getDone() != null) + .map(OllamaGenerateResult::getResponse) + .collect(Collectors.joining("")); + + var generation = new Generation(ollamaResponse); + + // TODO investigate mapping of additional metadata/runtime info to the response. + // Determine if should be top + // level map vs. nested map + return new AiResponse(Collections.singletonList(generation), Map.of("ollama-generate-results", results)); + } + + /** + * @return Model name for the AI service. + */ + public String getModel() { + return model; + } + + /** + * @return Base URL of the Ollama Service. + */ + public String getBaseUrl() { + return baseUrl; + } + + /** + * @return Callback that handles individual generation results. + */ + public Consumer getSimpleCallback() { + return simpleCallback; + } + + /** + * Sets the callback that handles individual generation results. + * @param simpleCallback The callback to be set. + */ + public void setSimpleCallback(Consumer simpleCallback) { + this.simpleCallback = simpleCallback; + } + +} diff --git a/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/client/OllamaGenerateResult.java b/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/client/OllamaGenerateResult.java new file mode 100644 index 000000000..f6f997d64 --- /dev/null +++ b/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/client/OllamaGenerateResult.java @@ -0,0 +1,146 @@ +package org.springframework.ai.ollama.client; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.List; + +/** + * Ollama generate a completion api response model + * + * @author nullptr + */ +@JsonIgnoreProperties(ignoreUnknown = true) +public class OllamaGenerateResult { + + @JsonProperty("model") + private String model; + + @JsonProperty("created_at") + private String createdAt; + + @JsonProperty("response") + private String response; + + @JsonProperty("done") + private Boolean done; + + @JsonProperty("context") + private List context; + + @JsonProperty("total_duration") + private Long totalDuration; + + @JsonProperty("load_duration") + private Long loadDuration; + + @JsonProperty("prompt_eval_count") + private Long promptEvalCount; + + @JsonProperty("prompt_eval_duration") + private Long promptEvalDuration; + + @JsonProperty("eval_count") + private Long evalCount; + + @JsonProperty("eval_duration") + private Long evalDuration; + + public String getModel() { + return model; + } + + public void setModel(String model) { + this.model = model; + } + + public String getCreatedAt() { + return createdAt; + } + + public void setCreatedAt(String createdAt) { + this.createdAt = createdAt; + } + + public String getResponse() { + return response; + } + + public void setResponse(String response) { + this.response = response; + } + + public Boolean getDone() { + return done; + } + + public void setDone(Boolean done) { + this.done = done; + } + + public List getContext() { + return context; + } + + public void setContext(List context) { + this.context = context; + } + + public Long getTotalDuration() { + return totalDuration; + } + + public void setTotalDuration(Long totalDuration) { + this.totalDuration = totalDuration; + } + + public Long getLoadDuration() { + return loadDuration; + } + + public void setLoadDuration(Long loadDuration) { + this.loadDuration = loadDuration; + } + + public Long getPromptEvalCount() { + return promptEvalCount; + } + + public void setPromptEvalCount(Long promptEvalCount) { + this.promptEvalCount = promptEvalCount; + } + + public Long getPromptEvalDuration() { + return promptEvalDuration; + } + + public void setPromptEvalDuration(Long promptEvalDuration) { + this.promptEvalDuration = promptEvalDuration; + } + + public Long getEvalCount() { + return evalCount; + } + + public void setEvalCount(Long evalCount) { + this.evalCount = evalCount; + } + + public Long getEvalDuration() { + return evalDuration; + } + + public void setEvalDuration(Long evalDuration) { + this.evalDuration = evalDuration; + } + + @Override + public String toString() { + return "OllamaGenerateResult{" + "model='" + model + '\'' + ", createdAt='" + createdAt + '\'' + ", response='" + + response + '\'' + ", done='" + done + '\'' + ", context=" + context + ", totalDuration=" + + totalDuration + ", loadDuration=" + loadDuration + ", promptEvalCount=" + promptEvalCount + + ", promptEvalDuration=" + promptEvalDuration + ", evalCount=" + evalCount + ", evalDuration=" + + evalDuration + '}'; + } + +} diff --git a/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/client/OllamaClientTests.java b/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/client/OllamaClientTests.java new file mode 100644 index 000000000..3f04271a5 --- /dev/null +++ b/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/client/OllamaClientTests.java @@ -0,0 +1,45 @@ +package org.springframework.ai.ollama.client; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; +import org.springframework.ai.client.AiResponse; +import org.springframework.ai.prompt.Prompt; +import org.springframework.util.CollectionUtils; + +import java.util.function.Consumer; + +public class OllamaClientTests { + + @Test + @Disabled("For manual smoke testing only.") + public void smokeTest() { + OllamaClient ollama2 = getOllamaClient(); + + Prompt prompt = new Prompt("Hello"); + AiResponse aiResponse = ollama2.generate(prompt); + Assertions.assertNotNull(aiResponse); + Assertions.assertFalse(CollectionUtils.isEmpty(aiResponse.getGenerations())); + Assertions.assertFalse(CollectionUtils.isEmpty(aiResponse.getProviderOutput())); + Assertions.assertNotNull(aiResponse.getProviderOutput().get("ollama-generate-results")); + Assertions.assertNotNull(aiResponse.getGeneration()); + Assertions.assertNotNull(aiResponse.getGeneration().getText()); + } + + private static OllamaClient getOllamaClient() { + Consumer ollamaGenerateResultConsumer = it -> { + if (it.getDone()) { + System.out.println(); + System.out.printf("Total duration: %dms%n", it.getTotalDuration() / 1000 / 1000); + System.out.printf("Prompt tokens: %d%n", it.getPromptEvalCount()); + System.out.printf("Completion tokens: %d%n", it.getEvalCount()); + } + else { + System.out.print(it.getResponse()); + } + }; + + return new OllamaClient("http://127.0.0.1:11434", "llama2", ollamaGenerateResultConsumer); + } + +}