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..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 @@ -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,72 @@ public void chatCompletionStreaming() { }); } + static class OllamaContainer extends GenericContainer { + + private final DockerImageName dockerImageName; + + OllamaContainer(DockerImageName image) { + super(image); + this.dockerImageName = image; + withExposedPorts(11434); + withImagePullPolicy(dockerImageName -> !dockerImageName.getUnversionedPart().startsWith(MODEL_NAME)); + } + + @Override + protected void containerIsStarted(InspectContainerResponse containerInfo) { + if (!this.dockerImageName.getVersionPart().endsWith(MODEL_NAME)) { + try { + execInContainer("ollama", "pull", MODEL_NAME); + } + 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..a3d01974a45 --- /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.23"; + +}