From 4294b8a62a89bc7a65451fb035ce0f2c43edfdf1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edd=C3=BA=20Mel=C3=A9ndez?= Date: Sat, 18 Nov 2023 22:11:43 -0600 Subject: [PATCH 1/2] Enable Ollama integration test Currently, test is disable because in every run the image should be downloaded and then pull the model. Now, taking advantage of Testcontainers and a GHA, a new image is created on-the-fly with the model in it and then cached, so, next executions will reuse the cached image instead. Fixes #121 --- .github/workflows/continuous-integration.yml | 5 + .../ollama/OllamaAutoConfigurationIT.java | 91 ++++++++++++++++++- .../ai/autoconfigure/ollama/OllamaImage.java | 7 ++ 3 files changed, 98 insertions(+), 5 deletions(-) create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml index 1f0dca2e8ec..afd5a487e3b 100644 --- a/.github/workflows/continuous-integration.yml +++ b/.github/workflows/continuous-integration.yml @@ -19,6 +19,11 @@ jobs: distribution: 'temurin' cache: 'maven' + - name: Cache Docker images. + uses: ScribeMD/docker-cache@0.3.6 + with: + key: docker-${{ runner.os }}-${{ hashFiles('**/OllamaImage.java') }} + - name: Build with Maven and deploy to Artifactory env: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java index a00a036f77f..0e0caa665a7 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java @@ -16,10 +16,12 @@ package org.springframework.ai.autoconfigure.ollama; +import com.github.dockerjava.api.DockerClient; +import com.github.dockerjava.api.command.InspectContainerResponse; +import com.github.dockerjava.api.model.Image; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.ChatResponse; import org.springframework.ai.chat.Generation; @@ -32,12 +34,14 @@ import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.testcontainers.DockerClientFactory; import org.testcontainers.containers.GenericContainer; -import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; import reactor.core.publisher.Flux; import java.io.IOException; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.stream.Collectors; @@ -46,9 +50,9 @@ /** * @author Christian Tzolov + * @author Eddú Meléndez * @since 0.8.0 */ -@Disabled("For manual smoke testing only.") @Testcontainers public class OllamaAutoConfigurationIT { @@ -56,8 +60,15 @@ public class OllamaAutoConfigurationIT { private static String MODEL_NAME = "mistral"; - @Container - static GenericContainer ollamaContainer = new GenericContainer<>("ollama/ollama:0.1.23").withExposedPorts(11434); + private static final String OLLAMA_WITH_MODEL = "%s-%s".formatted(MODEL_NAME, OllamaImage.IMAGE); + + private static final OllamaContainer ollamaContainer; + + static { + ollamaContainer = new OllamaContainer(OllamaDockerImageName.image()); + ollamaContainer.start(); + createImage(ollamaContainer, OLLAMA_WITH_MODEL); + } static String baseUrl; @@ -120,4 +131,74 @@ public void chatCompletionStreaming() { }); } + static class OllamaContainer extends GenericContainer { + + private static final String MODEL = "orca-mini"; + + private final DockerImageName dockerImageName; + + OllamaContainer(DockerImageName image) { + super(image); + this.dockerImageName = image; + withExposedPorts(11434); + withImagePullPolicy(dockerImageName -> !dockerImageName.getVersionPart().endsWith(MODEL)); + } + + @Override + protected void containerIsStarted(InspectContainerResponse containerInfo) { + if (!this.dockerImageName.getVersionPart().endsWith(MODEL)) { + try { + execInContainer("ollama", "pull", MODEL); + } + catch (IOException | InterruptedException e) { + throw new RuntimeException("Error pulling orca-mini model", e); + } + } + } + + } + + static void createImage(GenericContainer container, String localImageName) { + DockerImageName dockerImageName = DockerImageName.parse(container.getDockerImageName()); + if (!dockerImageName.equals(DockerImageName.parse(localImageName))) { + DockerClient dockerClient = DockerClientFactory.instance().client(); + List images = dockerClient.listImagesCmd().withReferenceFilter(localImageName).exec(); + if (images.isEmpty()) { + DockerImageName imageModel = DockerImageName.parse(localImageName); + dockerClient.commitCmd(container.getContainerId()) + .withRepository(imageModel.getUnversionedPart()) + .withLabels(Collections.singletonMap("org.testcontainers.sessionId", "")) + .withTag(imageModel.getVersionPart()) + .exec(); + } + } + } + + static class OllamaDockerImageName { + + private final String baseImage; + + private final String localImageName; + + OllamaDockerImageName(String baseImage, String localImageName) { + this.baseImage = baseImage; + this.localImageName = localImageName; + } + + static DockerImageName image() { + return new OllamaDockerImageName(OllamaImage.IMAGE, OLLAMA_WITH_MODEL).resolve(); + } + + private DockerImageName resolve() { + var dockerImageName = DockerImageName.parse(this.baseImage); + var dockerClient = DockerClientFactory.instance().client(); + var images = dockerClient.listImagesCmd().withReferenceFilter(this.localImageName).exec(); + if (images.isEmpty()) { + return dockerImageName; + } + return DockerImageName.parse(this.localImageName); + } + + } + } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java new file mode 100644 index 00000000000..a617dc74779 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java @@ -0,0 +1,7 @@ +package org.springframework.ai.autoconfigure.ollama; + +public class OllamaImage { + + static final String IMAGE = "ollama/ollama:0.1.10"; + +} From 0ee042febe78bd98cacb6e03387d65b7312665c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edd=C3=BA=20Mel=C3=A9ndez?= Date: Tue, 27 Feb 2024 12:54:31 -0600 Subject: [PATCH 2/2] Fix --- .../autoconfigure/ollama/OllamaAutoConfigurationIT.java | 8 +++----- .../ai/autoconfigure/ollama/OllamaImage.java | 2 +- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java index 0e0caa665a7..f0c0f175450 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java @@ -133,22 +133,20 @@ public void chatCompletionStreaming() { static class OllamaContainer extends GenericContainer { - private static final String MODEL = "orca-mini"; - private final DockerImageName dockerImageName; OllamaContainer(DockerImageName image) { super(image); this.dockerImageName = image; withExposedPorts(11434); - withImagePullPolicy(dockerImageName -> !dockerImageName.getVersionPart().endsWith(MODEL)); + withImagePullPolicy(dockerImageName -> !dockerImageName.getUnversionedPart().startsWith(MODEL_NAME)); } @Override protected void containerIsStarted(InspectContainerResponse containerInfo) { - if (!this.dockerImageName.getVersionPart().endsWith(MODEL)) { + if (!this.dockerImageName.getVersionPart().endsWith(MODEL_NAME)) { try { - execInContainer("ollama", "pull", MODEL); + execInContainer("ollama", "pull", MODEL_NAME); } catch (IOException | InterruptedException e) { throw new RuntimeException("Error pulling orca-mini model", e); diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java index a617dc74779..a3d01974a45 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java @@ -2,6 +2,6 @@ public class OllamaImage { - static final String IMAGE = "ollama/ollama:0.1.10"; + static final String IMAGE = "ollama/ollama:0.1.23"; }