Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add ollama module #2265

Merged
merged 24 commits into from
Mar 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
ceb0570
chore: bootstrap ollama module
mdelapenya Feb 21, 2024
969de4c
feat: enrich Ollama container with opts
mdelapenya Feb 23, 2024
fb4e249
chore: automatically detect GPU support
mdelapenya Feb 23, 2024
6cf76f8
chore: skip example on CI
mdelapenya Feb 23, 2024
d7f5329
chore: mod tidy
mdelapenya Feb 23, 2024
db9a5b9
fix: lint
mdelapenya Feb 23, 2024
b9c3fb8
fix: handle error
mdelapenya Feb 23, 2024
bdaed03
fix: first pull the model
mdelapenya Feb 27, 2024
c2b161e
chore: remove prompt capabilities by now
mdelapenya Feb 27, 2024
552d40c
fix: wait for the container more time
mdelapenya Feb 29, 2024
201e4f7
Merge branch 'main' into ollama-module
mdelapenya Mar 1, 2024
39a9c9c
Merge branch 'main' into ollama-module
mdelapenya Mar 4, 2024
29567a8
chore: run mod tidy
mdelapenya Mar 4, 2024
152183f
feat: support commiting the image for speeding up the containers
mdelapenya Mar 4, 2024
eb84b92
Merge branch 'main' into ollama-module
mdelapenya Mar 4, 2024
002adff
fix: typo
mdelapenya Mar 4, 2024
09b00c6
chore: bump docker/docker
mdelapenya Mar 4, 2024
4da53f5
chore: delegate the creation of the target image to the user
mdelapenya Mar 4, 2024
dbf7e1c
chore: test the new container
mdelapenya Mar 4, 2024
850c31b
chore: extract assertions to a function
mdelapenya Mar 4, 2024
c710fa9
Merge branch 'main' into ollama-module
mdelapenya Mar 4, 2024
f766ecf
chore: go mod tidy
mdelapenya Mar 4, 2024
e2fdda2
Merge branch 'main' into ollama-module
mdelapenya Mar 4, 2024
3431940
chore: do not couple implementation with the ollama CLI
mdelapenya Mar 4, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
7 changes: 7 additions & 0 deletions .github/dependabot.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
Expand Down
4 changes: 4 additions & 0 deletions .vscode/.testcontainers-go.code-workspace
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,10 @@
"name": "module / neo4j",
"path": "../modules/neo4j"
},
{
"name": "module / ollama",
"path": "../modules/ollama"
},
{
"name": "module / openldap",
"path": "../modules/openldap"
Expand Down
81 changes: 81 additions & 0 deletions docs/modules/ollama.md
Original file line number Diff line number Diff line change
@@ -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-->
1 change: 1 addition & 0 deletions mkdocs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions modules/ollama/Makefile
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
include ../../commons-test.mk

.PHONY: test
test:
$(MAKE) test-ollama
169 changes: 169 additions & 0 deletions modules/ollama/examples_test.go
Original file line number Diff line number Diff line change
@@ -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.
}
63 changes: 63 additions & 0 deletions modules/ollama/go.mod
Original file line number Diff line number Diff line change
@@ -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 => ../..