Add Ollama module (testcontainers#8369)

Co-authored-by: Kevin Wittek <[email protected]>

Author: 2 people

.github/ISSUE_TEMPLATE/bug_report.yaml

@@ -43,6 +43,7 @@ body:
         - Neo4j
         - NGINX
         - OceanBase
+        - Ollama
         - OpenFGA
         - Oracle Free
         - Oracle XE

.github/ISSUE_TEMPLATE/enhancement.yaml

@@ -43,6 +43,7 @@ body:
         - Neo4j
         - NGINX
         - OceanBase
+        - Ollama
         - OpenFGA
         - Oracle Free
         - Oracle XE

.github/ISSUE_TEMPLATE/feature.yaml

@@ -43,6 +43,7 @@ body:
         - Neo4j
         - NGINX
         - OceanBase
+        - Ollama
         - OpenFGA
         - Oracle Free
         - Oracle XE

.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:

.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:

.github/settings.yml

@@ -196,6 +196,9 @@ labels:
   - name: modules/oceanbase
     color: '#006b75'
 
+  - name: modules/ollama
+    color: '#006b75'
+
   - name: modules/openfga
     color: '#006b75'
 

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:
+
+<!--codeinclude-->
+[Ollama container](../../modules/ollama/src/test/java/org/testcontainers/ollama/OllamaContainerTest.java) inside_block:container
+<!--/codeinclude-->
+
+## 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
+    <dependency>
+        <groupId>org.testcontainers</groupId>
+        <artifactId>ollama</artifactId>
+        <version>{{latest_version}}</version>
+        <scope>test</scope>
+    </dependency>
+    ```

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

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'
+}

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.
+ * <p>
+ * Supported image: {@code ollama/ollama}
+ * <p>
+ * Exposed ports: 11434
+ */
+public class OllamaContainer extends GenericContainer<OllamaContainer> {
+
+    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<String, RuntimeInfo> 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<Image> 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);
+    }
+}