From 93ca7cd7345145de315e1d3c60f8b1ad07442cb1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edd=C3=BA=20Mel=C3=A9ndez=20Gonzales?= Date: Fri, 1 Mar 2024 19:13:45 -0500 Subject: [PATCH] Add Ollama module (#8369) Co-authored-by: Kevin Wittek --- .github/ISSUE_TEMPLATE/bug_report.yaml | 1 + .github/ISSUE_TEMPLATE/enhancement.yaml | 1 + .github/ISSUE_TEMPLATE/feature.yaml | 1 + .github/dependabot.yml | 5 ++ .github/labeler.yml | 4 + .github/settings.yml | 3 + docs/modules/ollama.md | 30 +++++++ mkdocs.yml | 1 + modules/ollama/build.gradle | 8 ++ .../ollama/OllamaContainer.java | 80 +++++++++++++++++++ .../ollama/OllamaContainerTest.java | 56 +++++++++++++ .../src/test/resources/logback-test.xml | 16 ++++ 12 files changed, 206 insertions(+) create mode 100644 docs/modules/ollama.md create mode 100644 modules/ollama/build.gradle create mode 100644 modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java create mode 100644 modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerTest.java create mode 100644 modules/ollama/src/test/resources/logback-test.xml diff --git a/.github/ISSUE_TEMPLATE/bug_report.yaml b/.github/ISSUE_TEMPLATE/bug_report.yaml index 270f292dcf7..1458b89d2f7 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yaml +++ b/.github/ISSUE_TEMPLATE/bug_report.yaml @@ -43,6 +43,7 @@ body: - Neo4j - NGINX - OceanBase + - Ollama - OpenFGA - Oracle Free - Oracle XE diff --git a/.github/ISSUE_TEMPLATE/enhancement.yaml b/.github/ISSUE_TEMPLATE/enhancement.yaml index 97297c9ddb2..2baf2be99eb 100644 --- a/.github/ISSUE_TEMPLATE/enhancement.yaml +++ b/.github/ISSUE_TEMPLATE/enhancement.yaml @@ -43,6 +43,7 @@ body: - Neo4j - NGINX - OceanBase + - Ollama - OpenFGA - Oracle Free - Oracle XE diff --git a/.github/ISSUE_TEMPLATE/feature.yaml b/.github/ISSUE_TEMPLATE/feature.yaml index 80048f00fb0..fb5013a41d3 100644 --- a/.github/ISSUE_TEMPLATE/feature.yaml +++ b/.github/ISSUE_TEMPLATE/feature.yaml @@ -43,6 +43,7 @@ body: - Neo4j - NGINX - OceanBase + - Ollama - OpenFGA - Oracle Free - Oracle XE diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 22cbe2c2748..0bb82dce9f4 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -224,6 +224,11 @@ updates: schedule: interval: "weekly" open-pull-requests-limit: 10 + - package-ecosystem: "gradle" + directory: "/modules/ollama" + schedule: + interval: "weekly" + open-pull-requests-limit: 10 - package-ecosystem: "gradle" directory: "/modules/openfga" schedule: diff --git a/.github/labeler.yml b/.github/labeler.yml index 2b9dd3bc899..387b1644801 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -139,6 +139,10 @@ - changed-files: - any-glob-to-any-file: - modules/oceanbase/**/* +"modules/ollama": + - changed-files: + - any-glob-to-any-file: + - modules/ollama/**/* "modules/openfga": - changed-files: - any-glob-to-any-file: diff --git a/.github/settings.yml b/.github/settings.yml index 6ac8714a40c..9cf3d27d5d4 100644 --- a/.github/settings.yml +++ b/.github/settings.yml @@ -196,6 +196,9 @@ labels: - name: modules/oceanbase color: '#006b75' + - name: modules/ollama + color: '#006b75' + - name: modules/openfga color: '#006b75' diff --git a/docs/modules/ollama.md b/docs/modules/ollama.md new file mode 100644 index 00000000000..97e8865bdc1 --- /dev/null +++ b/docs/modules/ollama.md @@ -0,0 +1,30 @@ +# Ollama + +Testcontainers module for [Ollama](https://hub.docker.com/r/ollama/ollama) . + +## Ollama's usage examples + +You can start an Ollama container instance from any Java application by using: + + +[Ollama container](../../modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerTest.java) inside_block:container + + +## Adding this module to your project dependencies + +Add the following dependency to your `pom.xml`/`build.gradle` file: + +=== "Gradle" + ```groovy + testImplementation "org.testcontainers:ollama:{{latest_version}}" + ``` + +=== "Maven" + ```xml + + org.testcontainers + ollama + {{latest_version}} + test + + ``` diff --git a/mkdocs.yml b/mkdocs.yml index ce0c33dcf16..d5f5b275288 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -90,6 +90,7 @@ nav: - modules/minio.md - modules/mockserver.md - modules/nginx.md + - modules/ollama.md - modules/openfga.md - modules/pulsar.md - modules/qdrant.md diff --git a/modules/ollama/build.gradle b/modules/ollama/build.gradle new file mode 100644 index 00000000000..b64083062ec --- /dev/null +++ b/modules/ollama/build.gradle @@ -0,0 +1,8 @@ +description = "Testcontainers :: Ollama" + +dependencies { + api project(':testcontainers') + + testImplementation 'org.assertj:assertj-core:3.25.1' + testImplementation 'io.rest-assured:rest-assured:5.4.0' +} diff --git a/modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java b/modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java new file mode 100644 index 00000000000..1c739b95a6b --- /dev/null +++ b/modules/ollama/src/main/java/org/testcontainers/ollama/OllamaContainer.java @@ -0,0 +1,80 @@ +package org.testcontainers.ollama; + +import com.github.dockerjava.api.DockerClient; +import com.github.dockerjava.api.model.DeviceRequest; +import com.github.dockerjava.api.model.Image; +import com.github.dockerjava.api.model.Info; +import com.github.dockerjava.api.model.RuntimeInfo; +import org.testcontainers.DockerClientFactory; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.utility.DockerImageName; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * Testcontainers implementation for Ollama. + *

+ * Supported image: {@code ollama/ollama} + *

+ * Exposed ports: 11434 + */ +public class OllamaContainer extends GenericContainer { + + private static final DockerImageName DOCKER_IMAGE_NAME = DockerImageName.parse("ollama/ollama"); + + public OllamaContainer(String image) { + this(DockerImageName.parse(image)); + } + + public OllamaContainer(DockerImageName dockerImageName) { + super(dockerImageName); + dockerImageName.assertCompatibleWith(DOCKER_IMAGE_NAME); + + Info info = this.dockerClient.infoCmd().exec(); + Map runtimes = info.getRuntimes(); + if (runtimes != null) { + if (runtimes.containsKey("nvidia")) { + withCreateContainerCmdModifier(cmd -> { + cmd + .getHostConfig() + .withDeviceRequests( + Collections.singletonList( + new DeviceRequest() + .withCapabilities(Collections.singletonList(Collections.singletonList("gpu"))) + .withCount(-1) + ) + ); + }); + } + } + withExposedPorts(11434); + } + + /** + * Commits the current file system changes in the container into a new image. + * Should be used for creating an image that contains a loaded model. + * @param imageName the name of the new image + */ + public void commitToImage(String imageName) { + DockerImageName dockerImageName = DockerImageName.parse(getDockerImageName()); + if (!dockerImageName.equals(DockerImageName.parse(imageName))) { + DockerClient dockerClient = DockerClientFactory.instance().client(); + List images = dockerClient.listImagesCmd().withReferenceFilter(imageName).exec(); + if (images.isEmpty()) { + DockerImageName imageModel = DockerImageName.parse(imageName); + dockerClient + .commitCmd(getContainerId()) + .withRepository(imageModel.getUnversionedPart()) + .withLabels(Collections.singletonMap("org.testcontainers.sessionId", "")) + .withTag(imageModel.getVersionPart()) + .exec(); + } + } + } + + public String getEndpoint() { + return "http://" + getHost() + ":" + getMappedPort(11434); + } +} diff --git a/modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerTest.java b/modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerTest.java new file mode 100644 index 00000000000..d33660b40b0 --- /dev/null +++ b/modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerTest.java @@ -0,0 +1,56 @@ +package org.testcontainers.ollama; + +import org.junit.Test; +import org.testcontainers.utility.Base58; +import org.testcontainers.utility.DockerImageName; + +import java.io.IOException; + +import static io.restassured.RestAssured.given; +import static org.assertj.core.api.Assertions.assertThat; + +public class OllamaContainerTest { + + @Test + public void withDefaultConfig() { + try ( // container { + OllamaContainer ollama = new OllamaContainer("ollama/ollama:0.1.26") + // } + ) { + ollama.start(); + + String version = given().baseUri(ollama.getEndpoint()).get("/api/version").jsonPath().get("version"); + assertThat(version).isEqualTo("0.1.26"); + } + } + + @Test + public void downloadModelAndCommitToImage() throws IOException, InterruptedException { + String newImageName = "tc-ollama-allminilm-" + Base58.randomString(4).toLowerCase(); + try (OllamaContainer ollama = new OllamaContainer("ollama/ollama:0.1.26")) { + ollama.start(); + ollama.execInContainer("ollama", "pull", "all-minilm"); + + String modelName = given() + .baseUri(ollama.getEndpoint()) + .get("/api/tags") + .jsonPath() + .getString("models[0].name"); + assertThat(modelName).contains("all-minilm"); + ollama.commitToImage(newImageName); + } + try ( + OllamaContainer ollama = new OllamaContainer( + DockerImageName.parse(newImageName).asCompatibleSubstituteFor("ollama/ollama") + ) + ) { + ollama.start(); + String modelName = given() + .baseUri(ollama.getEndpoint()) + .get("/api/tags") + .jsonPath() + .getString("models[0].name"); + assertThat(modelName).contains("all-minilm"); + } + } +} diff --git a/modules/ollama/src/test/resources/logback-test.xml b/modules/ollama/src/test/resources/logback-test.xml new file mode 100644 index 00000000000..83ef7a1a3ef --- /dev/null +++ b/modules/ollama/src/test/resources/logback-test.xml @@ -0,0 +1,16 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger - %msg%n + + + + + + + + +