Skip to content

Commit bbb4382

Browse files
mdelapenyamdelapenya
authored
feat: add ollama module (#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
1 parent 6a7d02d commit bbb4382

File tree

13 files changed

+833
-2
lines changed

13 files changed

+833
-2
lines changed

‎.github/dependabot.yml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,13 @@ updates:
212212
day: sunday
213213
open-pull-requests-limit: 3
214214
rebase-strategy: disabled
215+
- package-ecosystem: gomod
216+
directory: /modules/ollama
217+
schedule:
218+
interval: monthly
219+
day: sunday
220+
open-pull-requests-limit: 3
221+
rebase-strategy: disabled
215222
- package-ecosystem: gomod
216223
directory: /modules/openldap
217224
schedule:

‎.github/workflows/ci.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ jobs:
104104
matrix:
105105
go-version: [1.21.x, 1.x]
106106
platform: [ubuntu-latest]
107-
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]
107+
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]
108108
uses: ./.github/workflows/ci-test-go.yml
109109
with:
110110
go-version: ${{ matrix.go-version }}

‎.vscode/.testcontainers-go.code-workspace

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,10 @@
109109
"name": "module / neo4j",
110110
"path": "../modules/neo4j"
111111
},
112+
{
113+
"name": "module / ollama",
114+
"path": "../modules/ollama"
115+
},
112116
{
113117
"name": "module / openldap",
114118
"path": "../modules/openldap"

‎docs/modules/ollama.md

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
# Ollama
2+
3+
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>
4+
5+
## Introduction
6+
7+
The Testcontainers module for Ollama.
8+
9+
## Adding this module to your project dependencies
10+
11+
Please run the following command to add the Ollama module to your Go dependencies:
12+
13+
```
14+
go get github.com/testcontainers/testcontainers-go/modules/ollama
15+
```
16+
17+
## Usage example
18+
19+
<!--codeinclude-->
20+
[Creating a Ollama container](../../modules/ollama/examples_test.go) inside_block:runOllamaContainer
21+
<!--/codeinclude-->
22+
23+
## Module reference
24+
25+
The Ollama module exposes one entrypoint function to create the Ollama container, and this function receives two parameters:
26+
27+
```golang
28+
func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomizer) (*OllamaContainer, error)
29+
```
30+
31+
- `context.Context`, the Go context.
32+
- `testcontainers.ContainerCustomizer`, a variadic argument for passing options.
33+
34+
### Container Options
35+
36+
When starting the Ollama container, you can pass options in a variadic way to configure it.
37+
38+
#### Image
39+
40+
If you need to set a different Ollama Docker image, you can use `testcontainers.WithImage` with a valid Docker image
41+
for Ollama. E.g. `testcontainers.WithImage("ollama/ollama:0.1.25")`.
42+
43+
{% include "../features/common_functional_options.md" %}
44+
45+
### Container Methods
46+
47+
The Ollama container exposes the following methods:
48+
49+
#### ConnectionString
50+
51+
This method returns the connection string to connect to the Ollama container, using the default `11434` port.
52+
53+
<!--codeinclude-->
54+
[Get connection string](../../modules/ollama/ollama_test.go) inside_block:connectionString
55+
<!--/codeinclude-->
56+
57+
#### Commit
58+
59+
This method commits the container to a new image, returning the new image ID.
60+
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,
61+
and eventually use it as the base image for a new container. That will speed up the execution of the following containers.
62+
63+
<!--codeinclude-->
64+
[Commit Ollama image](../../modules/ollama/ollama_test.go) inside_block:commitOllamaContainer
65+
<!--/codeinclude-->
66+
67+
## Examples
68+
69+
### Loading Models
70+
71+
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).
72+
73+
!!!warning
74+
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.
75+
76+
The following examples use the `llama2` model to connect to the Ollama container using HTTP and Langchain.
77+
78+
<!--codeinclude-->
79+
[Using HTTP](../../modules/ollama/examples_test.go) inside_block:withHTTPModelLlama2
80+
[Using Langchaingo](../../modules/ollama/examples_test.go) inside_block:withLangchainModelLlama2
81+
<!--/codeinclude-->

‎mkdocs.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ nav:
8686
- modules/mysql.md
8787
- modules/nats.md
8888
- modules/neo4j.md
89+
- modules/ollama.md
8990
- modules/openldap.md
9091
- modules/opensearch.md
9192
- modules/postgres.md

‎modules/ollama/Makefile

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
include ../../commons-test.mk
2+
3+
.PHONY: test
4+
test:
5+
$(MAKE) test-ollama

‎modules/ollama/examples_test.go

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
package ollama_test
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"log"
7+
"net/http"
8+
"strings"
9+
10+
"github.com/tmc/langchaingo/llms"
11+
langchainollama "github.com/tmc/langchaingo/llms/ollama"
12+
13+
"github.com/testcontainers/testcontainers-go"
14+
tcollama "github.com/testcontainers/testcontainers-go/modules/ollama"
15+
)
16+
17+
func ExampleRunContainer() {
18+
// runOllamaContainer {
19+
ctx := context.Background()
20+
21+
ollamaContainer, err := tcollama.RunContainer(ctx, testcontainers.WithImage("ollama/ollama:0.1.25"))
22+
if err != nil {
23+
log.Fatalf("failed to start container: %s", err)
24+
}
25+
26+
// Clean up the container
27+
defer func() {
28+
if err := ollamaContainer.Terminate(ctx); err != nil {
29+
log.Fatalf("failed to terminate container: %s", err) // nolint:gocritic
30+
}
31+
}()
32+
// }
33+
34+
state, err := ollamaContainer.State(ctx)
35+
if err != nil {
36+
log.Fatalf("failed to get container state: %s", err) // nolint:gocritic
37+
}
38+
39+
fmt.Println(state.Running)
40+
41+
// Output:
42+
// true
43+
}
44+
45+
func ExampleRunContainer_withModel_llama2_http() {
46+
// withHTTPModelLlama2 {
47+
ctx := context.Background()
48+
49+
ollamaContainer, err := tcollama.RunContainer(
50+
ctx,
51+
testcontainers.WithImage("ollama/ollama:0.1.25"),
52+
)
53+
if err != nil {
54+
log.Fatalf("failed to start container: %s", err)
55+
}
56+
defer func() {
57+
if err := ollamaContainer.Terminate(ctx); err != nil {
58+
log.Fatalf("failed to terminate container: %s", err) // nolint:gocritic
59+
}
60+
}()
61+
62+
model := "llama2"
63+
64+
_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "pull", model})
65+
if err != nil {
66+
log.Fatalf("failed to pull model %s: %s", model, err)
67+
}
68+
69+
_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "run", model})
70+
if err != nil {
71+
log.Fatalf("failed to run model %s: %s", model, err)
72+
}
73+
74+
connectionStr, err := ollamaContainer.ConnectionString(ctx)
75+
if err != nil {
76+
log.Fatalf("failed to get connection string: %s", err) // nolint:gocritic
77+
}
78+
79+
httpClient := &http.Client{}
80+
81+
// generate a response
82+
payload := `{
83+
"model": "llama2",
84+
"prompt":"Why is the sky blue?"
85+
}`
86+
87+
req, err := http.NewRequest("POST", fmt.Sprintf("%s/api/generate", connectionStr), strings.NewReader(payload))
88+
if err != nil {
89+
log.Fatalf("failed to create request: %s", err) // nolint:gocritic
90+
}
91+
92+
resp, err := httpClient.Do(req)
93+
if err != nil {
94+
log.Fatalf("failed to get response: %s", err) // nolint:gocritic
95+
}
96+
// }
97+
98+
fmt.Println(resp.StatusCode)
99+
100+
// Intentionally not asserting the output, as we don't want to run this example in the tests.
101+
}
102+
103+
func ExampleRunContainer_withModel_llama2_langchain() {
104+
// withLangchainModelLlama2 {
105+
ctx := context.Background()
106+
107+
ollamaContainer, err := tcollama.RunContainer(
108+
ctx,
109+
testcontainers.WithImage("ollama/ollama:0.1.25"),
110+
)
111+
if err != nil {
112+
log.Fatalf("failed to start container: %s", err)
113+
}
114+
defer func() {
115+
if err := ollamaContainer.Terminate(ctx); err != nil {
116+
log.Fatalf("failed to terminate container: %s", err) // nolint:gocritic
117+
}
118+
}()
119+
120+
model := "llama2"
121+
122+
_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "pull", model})
123+
if err != nil {
124+
log.Fatalf("failed to pull model %s: %s", model, err)
125+
}
126+
127+
_, _, err = ollamaContainer.Exec(ctx, []string{"ollama", "run", model})
128+
if err != nil {
129+
log.Fatalf("failed to run model %s: %s", model, err)
130+
}
131+
132+
connectionStr, err := ollamaContainer.ConnectionString(ctx)
133+
if err != nil {
134+
log.Fatalf("failed to get connection string: %s", err) // nolint:gocritic
135+
}
136+
137+
var llm *langchainollama.LLM
138+
if llm, err = langchainollama.New(
139+
langchainollama.WithModel(model),
140+
langchainollama.WithServerURL(connectionStr),
141+
); err != nil {
142+
log.Fatalf("failed to create langchain ollama: %s", err) // nolint:gocritic
143+
}
144+
145+
completion, err := llm.Call(
146+
context.Background(),
147+
"how can Testcontainers help with testing?",
148+
llms.WithSeed(42), // the lower the seed, the more deterministic the completion
149+
llms.WithTemperature(0.0), // the lower the temperature, the more creative the completion
150+
)
151+
if err != nil {
152+
log.Fatalf("failed to create langchain ollama: %s", err) // nolint:gocritic
153+
}
154+
155+
words := []string{
156+
"easy", "isolation", "consistency",
157+
}
158+
lwCompletion := strings.ToLower(completion)
159+
160+
for _, word := range words {
161+
if strings.Contains(lwCompletion, word) {
162+
fmt.Println(true)
163+
}
164+
}
165+
166+
// }
167+
168+
// Intentionally not asserting the output, as we don't want to run this example in the tests.
169+
}

‎modules/ollama/go.mod

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
module github.com/testcontainers/testcontainers-go/modules/ollama
2+
3+
go 1.21
4+
5+
require (
6+
github.com/docker/docker v25.0.3+incompatible
7+
github.com/google/uuid v1.6.0
8+
github.com/testcontainers/testcontainers-go v0.28.0
9+
github.com/tmc/langchaingo v0.1.4
10+
)
11+
12+
require (
13+
dario.cat/mergo v1.0.0 // indirect
14+
github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
15+
github.com/Microsoft/go-winio v0.6.1 // indirect
16+
github.com/Microsoft/hcsshim v0.11.4 // indirect
17+
github.com/cenkalti/backoff/v4 v4.2.1 // indirect
18+
github.com/containerd/containerd v1.7.12 // indirect
19+
github.com/containerd/log v0.1.0 // indirect
20+
github.com/cpuguy83/dockercfg v0.3.1 // indirect
21+
github.com/distribution/reference v0.5.0 // indirect
22+
github.com/dlclark/regexp2 v1.8.1 // indirect
23+
github.com/docker/go-connections v0.5.0 // indirect
24+
github.com/docker/go-units v0.5.0 // indirect
25+
github.com/felixge/httpsnoop v1.0.3 // indirect
26+
github.com/go-logr/logr v1.2.4 // indirect
27+
github.com/go-logr/stdr v1.2.2 // indirect
28+
github.com/go-ole/go-ole v1.2.6 // indirect
29+
github.com/gogo/protobuf v1.3.2 // indirect
30+
github.com/golang/protobuf v1.5.3 // indirect
31+
github.com/klauspost/compress v1.16.0 // indirect
32+
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
33+
github.com/magiconair/properties v1.8.7 // indirect
34+
github.com/moby/patternmatcher v0.6.0 // indirect
35+
github.com/moby/sys/sequential v0.5.0 // indirect
36+
github.com/moby/sys/user v0.1.0 // indirect
37+
github.com/moby/term v0.5.0 // indirect
38+
github.com/morikuni/aec v1.0.0 // indirect
39+
github.com/opencontainers/go-digest v1.0.0 // indirect
40+
github.com/opencontainers/image-spec v1.1.0 // indirect
41+
github.com/pkg/errors v0.9.1 // indirect
42+
github.com/pkoukk/tiktoken-go v0.1.2 // indirect
43+
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
44+
github.com/shirou/gopsutil/v3 v3.23.12 // indirect
45+
github.com/shoenig/go-m1cpu v0.1.6 // indirect
46+
github.com/sirupsen/logrus v1.9.3 // indirect
47+
github.com/tklauser/go-sysconf v0.3.12 // indirect
48+
github.com/tklauser/numcpus v0.6.1 // indirect
49+
github.com/yusufpapurcu/wmi v1.2.3 // indirect
50+
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.45.0 // indirect
51+
go.opentelemetry.io/otel v1.19.0 // indirect
52+
go.opentelemetry.io/otel/metric v1.19.0 // indirect
53+
go.opentelemetry.io/otel/trace v1.19.0 // indirect
54+
golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect
55+
golang.org/x/mod v0.13.0 // indirect
56+
golang.org/x/sys v0.16.0 // indirect
57+
golang.org/x/tools v0.14.0 // indirect
58+
google.golang.org/genproto/googleapis/rpc v0.0.0-20231211222908-989df2bf70f3 // indirect
59+
google.golang.org/grpc v1.60.0 // indirect
60+
google.golang.org/protobuf v1.31.0 // indirect
61+
)
62+
63+
replace github.com/testcontainers/testcontainers-go => ../..

0 commit comments

Comments
 (0)