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
This commit is contained in:
committed by
Christian Tzolov
parent
6de57d5f37
commit
246ba173fe
5
.github/workflows/continuous-integration.yml
vendored
5
.github/workflows/continuous-integration.yml
vendored
@@ -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 }}
|
||||
|
||||
@@ -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.ai.chat.messages.UserMessage;
|
||||
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 @@ import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @author Eddú Meléndez
|
||||
* @since 0.8.0
|
||||
*/
|
||||
@Disabled("For manual smoke testing only.")
|
||||
@Testcontainers
|
||||
public class OllamaChatAutoConfigurationIT {
|
||||
|
||||
@@ -56,8 +60,15 @@ public class OllamaChatAutoConfigurationIT {
|
||||
|
||||
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;
|
||||
|
||||
@@ -138,4 +149,72 @@ public class OllamaChatAutoConfigurationIT {
|
||||
});
|
||||
}
|
||||
|
||||
static class OllamaContainer extends GenericContainer<OllamaContainer> {
|
||||
|
||||
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<Image> 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);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
package org.springframework.ai.autoconfigure.ollama;
|
||||
|
||||
public class OllamaImage {
|
||||
|
||||
static final String IMAGE = "ollama/ollama:0.1.23";
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user