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:
Eddú Meléndez
2023-11-18 22:11:43 -06:00
committed by Christian Tzolov
parent 6de57d5f37
commit 246ba173fe
3 changed files with 96 additions and 5 deletions

View File

@@ -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 }}

View File

@@ -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);
}
}
}

View File

@@ -0,0 +1,7 @@
package org.springframework.ai.autoconfigure.ollama;
public class OllamaImage {
static final String IMAGE = "ollama/ollama:0.1.23";
}