feat: add ollama module (testcontainers#2265)

* chore: bootstrap ollama module

* feat: enrich Ollama container with opts

* chore: automatically detect GPU support

* chore: skip example on CI

* chore: mod tidy

* fix: lint

* fix: handle error

* fix: first pull the model

* chore: remove prompt capabilities by now

* fix: wait for the container more time

* chore: run mod tidy

* feat: support commiting the image for speeding up the containers

* fix: typo

* chore: bump docker/docker

* chore: delegate the creation of the target image to the user

* chore: test the new container

* chore: extract assertions to a function

* chore: go mod tidy

* chore: do not couple implementation with the ollama CLI

.github/dependabot.yml

@@ -212,6 +212,13 @@ updates:
         day: sunday
       open-pull-requests-limit: 3
       rebase-strategy: disabled
+    - package-ecosystem: gomod
+      directory: /modules/ollama
+      schedule:
+        interval: monthly
+        day: sunday
+      open-pull-requests-limit: 3
+      rebase-strategy: disabled
     - package-ecosystem: gomod
       directory: /modules/openldap
       schedule:

.github/workflows/ci.yml

@@ -104,7 +104,7 @@ jobs:
       matrix:
         go-version: [1.21.x, 1.x]
         platform: [ubuntu-latest]
-        module: [artemis, cassandra, chroma, clickhouse, cockroachdb, compose, consul, couchbase, elasticsearch, gcloud, inbucket, k3s, k6, kafka, localstack, mariadb, milvus, minio, mockserver, mongodb, mssql, mysql, nats, neo4j, openldap, opensearch, postgres, pulsar, qdrant, rabbitmq, redis, redpanda, surrealdb, vault, weaviate]
+        module: [artemis, cassandra, chroma, clickhouse, cockroachdb, compose, consul, couchbase, elasticsearch, gcloud, inbucket, k3s, k6, kafka, localstack, mariadb, milvus, minio, mockserver, mongodb, mssql, mysql, nats, neo4j, ollama, openldap, opensearch, postgres, pulsar, qdrant, rabbitmq, redis, redpanda, surrealdb, vault, weaviate]
     uses: ./.github/workflows/ci-test-go.yml
     with:
       go-version: ${{ matrix.go-version }}

.vscode/.testcontainers-go.code-workspace

@@ -109,6 +109,10 @@
             "name": "module / neo4j",
             "path": "../modules/neo4j"
         },
+        {
+            "name": "module / ollama",
+            "path": "../modules/ollama"
+        },
         {
             "name": "module / openldap",
             "path": "../modules/openldap"

docs/modules/ollama.md

@@ -0,0 +1,81 @@
+# Ollama
+
+Not available until the next release of testcontainers-go <a href="https://github.com/testcontainers/testcontainers-go"><span class="tc-version">:material-tag: main</span></a>
+
+## Introduction
+
+The Testcontainers module for Ollama.
+
+## Adding this module to your project dependencies
+
+Please run the following command to add the Ollama module to your Go dependencies:
+
+```
+go get github.com/testcontainers/testcontainers-go/modules/ollama
+```
+
+## Usage example
+
+<!--codeinclude-->
+[Creating a Ollama container](../../modules/ollama/examples_test.go) inside_block:runOllamaContainer
+<!--/codeinclude-->
+
+## Module reference
+
+The Ollama module exposes one entrypoint function to create the Ollama container, and this function receives two parameters:
+
+```golang
+func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomizer) (*OllamaContainer, error)
+```
+
+- `context.Context`, the Go context.
+- `testcontainers.ContainerCustomizer`, a variadic argument for passing options.
+
+### Container Options
+
+When starting the Ollama container, you can pass options in a variadic way to configure it.
+
+#### Image
+
+If you need to set a different Ollama Docker image, you can use `testcontainers.WithImage` with a valid Docker image
+for Ollama. E.g. `testcontainers.WithImage("ollama/ollama:0.1.25")`.
+
+{% include "../features/common_functional_options.md" %}
+
+### Container Methods
+
+The Ollama container exposes the following methods:
+
+#### ConnectionString
+
+This method returns the connection string to connect to the Ollama container, using the default `11434` port.
+
+<!--codeinclude-->
+[Get connection string](../../modules/ollama/ollama_test.go) inside_block:connectionString
+<!--/codeinclude-->
+
+#### Commit
+
+This method commits the container to a new image, returning the new image ID.
+It should be used after a model has been pulled and loaded into the container in order to create a new image with the model,
+and eventually use it as the base image for a new container. That will speed up the execution of the following containers.
+
+<!--codeinclude-->
+[Commit Ollama image](../../modules/ollama/ollama_test.go) inside_block:commitOllamaContainer
+<!--/codeinclude-->
+
+## Examples
+
+### Loading Models
+
+It's possible to initialise the Ollama container with a specific model passed as parameter. The supported models are described in the Ollama project: [https://github.com/ollama/ollama?tab=readme-ov-file](https://github.com/ollama/ollama?tab=readme-ov-file) and [https://ollama.com/library](https://ollama.com/library).
+
+!!!warning
+    At the moment you use one of those models, the Ollama image will load the model and could take longer to start because of that.
+
+The following examples use the `llama2` model to connect to the Ollama container using HTTP and Langchain.
+
+<!--codeinclude-->
+[Using HTTP](../../modules/ollama/examples_test.go) inside_block:withHTTPModelLlama2
+[Using Langchaingo](../../modules/ollama/examples_test.go) inside_block:withLangchainModelLlama2
+<!--/codeinclude-->

mkdocs.yml

@@ -86,6 +86,7 @@ nav:
         - modules/mysql.md
         - modules/nats.md
         - modules/neo4j.md
+        - modules/ollama.md
         - modules/openldap.md
         - modules/opensearch.md
         - modules/postgres.md

modules/ollama/Makefile

@@ -0,0 +1,5 @@
+include ../../commons-test.mk
+
+.PHONY: test
+test:
+	$(MAKE) test-ollama

modules/ollama/examples_test.go

@@ -0,0 +1,169 @@
+package ollama_test
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net/http"
+	"strings"
+
+	"github.com/tmc/langchaingo/llms"
+	langchainollama "github.com/tmc/langchaingo/llms/ollama"
+
+	"github.com/testcontainers/testcontainers-go"
+	tcollama "github.com/testcontainers/testcontainers-go/modules/ollama"
+)
+
+func ExampleRunContainer() {
+	// runOllamaContainer {
+	ctx := context.Background()
+
+	ollamaContainer, err := tcollama.RunContainer(ctx, testcontainers.WithImage("ollama/ollama:0.1.25"))
+	if err != nil {
+		log.Fatalf("failed to start container: %s", err)
+	}
+
+	// Clean up the container
+	defer func() {
+		if err := ollamaContainer.Terminate(ctx); err != nil {
+			log.Fatalf("failed to terminate container: %s", err) // nolint:gocritic
+		}
+	}()
+	// }
+
+	state, err := ollamaContainer.State(ctx)
+	if err != nil {
+		log.Fatalf("failed to get container state: %s", err) // nolint:gocritic
+	}
+
+	fmt.Println(state.Running)
+
+	// Output:
+	// true
+}
+
+func ExampleRunContainer_withModel_llama2_http() {
+	// withHTTPModelLlama2 {
+	ctx := context.Background()
+
+	ollamaContainer, err := tcollama.RunContainer(
+		ctx,
+		testcontainers.WithImage("ollama/ollama:0.1.25"),
+	)
+	if err != nil {
+		log.Fatalf("failed to start container: %s", err)
+	}
+	defer func() {
+		if err := ollamaContainer.Terminate(ctx); err != nil {
+			log.Fatalf("failed to terminate container: %s", err) // nolint:gocritic
+		}
+	}()
+
+	model := "llama2"
+
+	_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "pull", model})
+	if err != nil {
+		log.Fatalf("failed to pull model %s: %s", model, err)
+	}
+
+	_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "run", model})
+	if err != nil {
+		log.Fatalf("failed to run model %s: %s", model, err)
+	}
+
+	connectionStr, err := ollamaContainer.ConnectionString(ctx)
+	if err != nil {
+		log.Fatalf("failed to get connection string: %s", err) // nolint:gocritic
+	}
+
+	httpClient := &http.Client{}
+
+	// generate a response
+	payload := `{
+	"model": "llama2",
+	"prompt":"Why is the sky blue?"
+}`
+
+	req, err := http.NewRequest("POST", fmt.Sprintf("%s/api/generate", connectionStr), strings.NewReader(payload))
+	if err != nil {
+		log.Fatalf("failed to create request: %s", err) // nolint:gocritic
+	}
+
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		log.Fatalf("failed to get response: %s", err) // nolint:gocritic
+	}
+	// }
+
+	fmt.Println(resp.StatusCode)
+
+	// Intentionally not asserting the output, as we don't want to run this example in the tests.
+}
+
+func ExampleRunContainer_withModel_llama2_langchain() {
+	// withLangchainModelLlama2 {
+	ctx := context.Background()
+
+	ollamaContainer, err := tcollama.RunContainer(
+		ctx,
+		testcontainers.WithImage("ollama/ollama:0.1.25"),
+	)
+	if err != nil {
+		log.Fatalf("failed to start container: %s", err)
+	}
+	defer func() {
+		if err := ollamaContainer.Terminate(ctx); err != nil {
+			log.Fatalf("failed to terminate container: %s", err) // nolint:gocritic
+		}
+	}()
+
+	model := "llama2"
+
+	_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "pull", model})
+	if err != nil {
+		log.Fatalf("failed to pull model %s: %s", model, err)
+	}
+
+	_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "run", model})
+	if err != nil {
+		log.Fatalf("failed to run model %s: %s", model, err)
+	}
+
+	connectionStr, err := ollamaContainer.ConnectionString(ctx)
+	if err != nil {
+		log.Fatalf("failed to get connection string: %s", err) // nolint:gocritic
+	}
+
+	var llm *langchainollama.LLM
+	if llm, err = langchainollama.New(
+		langchainollama.WithModel(model),
+		langchainollama.WithServerURL(connectionStr),
+	); err != nil {
+		log.Fatalf("failed to create langchain ollama: %s", err) // nolint:gocritic
+	}
+
+	completion, err := llm.Call(
+		context.Background(),
+		"how can Testcontainers help with testing?",
+		llms.WithSeed(42),         // the lower the seed, the more deterministic the completion
+		llms.WithTemperature(0.0), // the lower the temperature, the more creative the completion
+	)
+	if err != nil {
+		log.Fatalf("failed to create langchain ollama: %s", err) // nolint:gocritic
+	}
+
+	words := []string{
+		"easy", "isolation", "consistency",
+	}
+	lwCompletion := strings.ToLower(completion)
+
+	for _, word := range words {
+		if strings.Contains(lwCompletion, word) {
+			fmt.Println(true)
+		}
+	}
+
+	// }
+
+	// Intentionally not asserting the output, as we don't want to run this example in the tests.
+}

modules/ollama/go.mod

@@ -0,0 +1,63 @@
+module github.com/testcontainers/testcontainers-go/modules/ollama
+
+go 1.21
+
+require (
+	github.com/docker/docker v25.0.3+incompatible
+	github.com/google/uuid v1.6.0
+	github.com/testcontainers/testcontainers-go v0.28.0
+	github.com/tmc/langchaingo v0.1.4
+)
+
+require (
+	dario.cat/mergo v1.0.0 // indirect
+	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
+	github.com/Microsoft/go-winio v0.6.1 // indirect
+	github.com/Microsoft/hcsshim v0.11.4 // indirect
+	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
+	github.com/containerd/containerd v1.7.12 // indirect
+	github.com/containerd/log v0.1.0 // indirect
+	github.com/cpuguy83/dockercfg v0.3.1 // indirect
+	github.com/distribution/reference v0.5.0 // indirect
+	github.com/dlclark/regexp2 v1.8.1 // indirect
+	github.com/docker/go-connections v0.5.0 // indirect
+	github.com/docker/go-units v0.5.0 // indirect
+	github.com/felixge/httpsnoop v1.0.3 // indirect
+	github.com/go-logr/logr v1.2.4 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
+	github.com/go-ole/go-ole v1.2.6 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/klauspost/compress v1.16.0 // indirect
+	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
+	github.com/magiconair/properties v1.8.7 // indirect
+	github.com/moby/patternmatcher v0.6.0 // indirect
+	github.com/moby/sys/sequential v0.5.0 // indirect
+	github.com/moby/sys/user v0.1.0 // indirect
+	github.com/moby/term v0.5.0 // indirect
+	github.com/morikuni/aec v1.0.0 // indirect
+	github.com/opencontainers/go-digest v1.0.0 // indirect
+	github.com/opencontainers/image-spec v1.1.0 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.2 // indirect
+	github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
+	github.com/shirou/gopsutil/v3 v3.23.12 // indirect
+	github.com/shoenig/go-m1cpu v0.1.6 // indirect
+	github.com/sirupsen/logrus v1.9.3 // indirect
+	github.com/tklauser/go-sysconf v0.3.12 // indirect
+	github.com/tklauser/numcpus v0.6.1 // indirect
+	github.com/yusufpapurcu/wmi v1.2.3 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.45.0 // indirect
+	go.opentelemetry.io/otel v1.19.0 // indirect
+	go.opentelemetry.io/otel/metric v1.19.0 // indirect
+	go.opentelemetry.io/otel/trace v1.19.0 // indirect
+	golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect
+	golang.org/x/mod v0.13.0 // indirect
+	golang.org/x/sys v0.16.0 // indirect
+	golang.org/x/tools v0.14.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20231211222908-989df2bf70f3 // indirect
+	google.golang.org/grpc v1.60.0 // indirect
+	google.golang.org/protobuf v1.31.0 // indirect
+)
+
+replace github.com/testcontainers/testcontainers-go => ../..