feat(chroma): update to chroma 1.x

docs/modules/chroma.md

@@ -47,7 +47,7 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom
 #### Image
 
 Use the second argument in the `Run` function to set a valid Docker image.
-In example: `Run(context.Background(), "chromadb/chroma:0.4.24")`.
+In example: `Run(context.Background(), "chromadb/chroma:1.4.0")`.
 
 ### Container Options
 
@@ -75,12 +75,11 @@ This method returns the REST endpoint of the Chroma container, using the default
 
 The following example demonstrates how to create a Chroma client using the Chroma module.
 
-First of all, you need to import the Chroma module and the Swagger client:
+First of all, you need to import the Chroma module and the chroma-go client:
 
 ```golang
 import (
-    chromago "github.com/amikos-tech/chroma-go"
-    "github.com/amikos-tech/chroma-go/types"
+    chromago "github.com/amikos-tech/chroma-go/pkg/api/v2"
 )
 ```
 

modules/chroma/chroma.go

@@ -16,7 +16,7 @@ type ChromaContainer struct {
 // Deprecated: use Run instead
 // RunContainer creates an instance of the Chroma container type
 func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomizer) (*ChromaContainer, error) {
-	return Run(ctx, "chromadb/chroma:0.4.24", opts...)
+	return Run(ctx, "chromadb/chroma:1.4.0", opts...)
 }
 
 // Run creates an instance of the Chroma container type
@@ -25,8 +25,7 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom
 		testcontainers.WithExposedPorts("8000/tcp"),
 		testcontainers.WithWaitStrategy(
 			wait.ForListeningPort("8000/tcp"),
-			wait.ForLog("Application startup complete"),
-			wait.ForHTTP("/api/v1/heartbeat").WithStatusCodeMatcher(func(status int) bool {
+			wait.ForHTTP("/api/v2/heartbeat").WithStatusCodeMatcher(func(status int) bool {
 				return status == 200
 			}),
 		),

modules/chroma/chroma_test.go

@@ -5,7 +5,7 @@ import (
 	"net/http"
 	"testing"
 
-	chromago "github.com/amikos-tech/chroma-go"
+	chromago "github.com/amikos-tech/chroma-go/pkg/api/v2"
 	"github.com/stretchr/testify/require"
 
 	"github.com/testcontainers/testcontainers-go"
@@ -15,7 +15,7 @@ import (
 func TestChroma(t *testing.T) {
 	ctx := context.Background()
 
-	ctr, err := chroma.Run(ctx, "chromadb/chroma:0.4.24")
+	ctr, err := chroma.Run(ctx, "chromadb/chroma:1.4.0")
 	testcontainers.CleanupContainer(t, ctr)
 	require.NoError(t, err)
 
@@ -37,12 +37,14 @@ func TestChroma(t *testing.T) {
 		// restEndpoint {
 		endpoint, err := ctr.RESTEndpoint(context.Background())
 		require.NoErrorf(tt, err, "failed to get REST endpoint")
-		chromaClient, err := chromago.NewClient(endpoint)
+		chromaClient, err := chromago.NewHTTPClient(chromago.WithBaseURL(endpoint))
 		// }
 		require.NoErrorf(tt, err, "failed to create client")
+		t.Cleanup(func() {
+			defer chromaClient.Close()
+		})
 
-		hb, err := chromaClient.Heartbeat(context.TODO())
+		err = chromaClient.Heartbeat(context.TODO())
 		require.NoError(tt, err)
-		require.NotNil(tt, hb["nanosecond heartbeat"])
 	})
 }

modules/chroma/examples_test.go

@@ -4,9 +4,10 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"math"
+	"path/filepath"
 
-	chromago "github.com/amikos-tech/chroma-go"
-	"github.com/amikos-tech/chroma-go/types"
+	chromago "github.com/amikos-tech/chroma-go/pkg/api/v2"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/modules/chroma"
@@ -16,7 +17,7 @@ func ExampleRun() {
 	// runChromaContainer {
 	ctx := context.Background()
 
-	chromaContainer, err := chroma.Run(ctx, "chromadb/chroma:0.4.24")
+	chromaContainer, err := chroma.Run(ctx, "chromadb/chroma:1.4.0")
 	defer func() {
 		if err := testcontainers.TerminateContainer(chromaContainer); err != nil {
 			log.Printf("failed to terminate container: %s", err)
@@ -44,7 +45,7 @@ func ExampleChromaContainer_connectWithClient() {
 	// createClient {
 	ctx := context.Background()
 
-	chromaContainer, err := chroma.Run(ctx, "chromadb/chroma:0.4.24")
+	chromaContainer, err := chroma.Run(ctx, "chromadb/chroma:1.4.0")
 	defer func() {
 		if err := testcontainers.TerminateContainer(chromaContainer); err != nil {
 			log.Printf("failed to terminate container: %s", err)
@@ -60,29 +61,38 @@ func ExampleChromaContainer_connectWithClient() {
 		log.Printf("failed to get REST endpoint: %s", err)
 		return
 	}
-	chromaClient, err := chromago.NewClient(endpoint)
+	chromaClient, err := chromago.NewHTTPClient(chromago.WithBaseURL(endpoint))
 	if err != nil {
 		log.Printf("failed to get client: %s", err)
 		return
 	}
 
-	hbs, errHb := chromaClient.Heartbeat(context.Background())
+	errHb := chromaClient.Heartbeat(context.Background())
+	// }
+	fmt.Println(errHb) // error is only returned if the heartbeat fails
+	// closeClient {
+	// ensure all resources are freed, e.g. TCP connections or the default embedding function which runs locally
+	defer func() {
+		err = chromaClient.Close()
+		if err != nil {
+			log.Printf("failed to close client: %s", err)
+		}
+	}()
 	// }
-	if _, ok := hbs["nanosecond heartbeat"]; ok {
-		fmt.Println(ok)
-	}
-
-	fmt.Println(errHb)
-
 	// Output:
-	// true
 	// <nil>
 }
 
 func ExampleChromaContainer_collections() {
 	ctx := context.Background()
 
-	chromaContainer, err := chroma.Run(ctx, "chromadb/chroma:0.4.24", testcontainers.WithEnv(map[string]string{"ALLOW_RESET": "true"}))
+	chromaContainer, err := chroma.Run(ctx, "chromadb/chroma:1.4.0",
+		testcontainers.WithFiles(testcontainers.ContainerFile{
+			HostFilePath:      filepath.Join("testdata", "v1-config.yaml"),
+			ContainerFilePath: "/config.yaml",
+			FileMode:          0o644,
+		}),
+	)
 	defer func() {
 		if err := testcontainers.TerminateContainer(chromaContainer); err != nil {
 			log.Printf("failed to terminate container: %s", err)
@@ -100,64 +110,82 @@ func ExampleChromaContainer_collections() {
 		log.Printf("failed to get REST endpoint: %s", err)
 		return
 	}
-	chromaClient, err := chromago.NewClient(endpoint)
+	chromaClient, err := chromago.NewHTTPClient(chromago.WithBaseURL(endpoint))
 	// }
 	if err != nil {
 		log.Printf("failed to get client: %s", err)
 		return
 	}
+	defer func() {
+		if err := chromaClient.Close(); err != nil {
+			log.Printf("failed to close client: %s", err)
+		}
+	}()
 	// reset {
-	reset, err := chromaClient.Reset(context.Background())
+	err = chromaClient.Reset(context.Background())
 	// }
 	if err != nil {
 		log.Printf("failed to reset: %s", err)
 		return
 	}
-	fmt.Printf("Reset successful: %v\n", reset)
+	fmt.Printf("Reset successful\n")
 
 	// createCollection {
 	// for testing we use a dummy hashing function NewConsistentHashEmbeddingFunction
-	col, err := chromaClient.CreateCollection(context.Background(), "test-collection", map[string]any{}, true, types.NewConsistentHashEmbeddingFunction(), types.L2)
-	// }
+	col, err := chromaClient.GetOrCreateCollection(context.Background(), "test-collection",
+		chromago.WithCollectionMetadataCreate(
+			chromago.NewMetadata(
+				chromago.NewStringAttribute("str", "hello2"),
+				chromago.NewIntAttribute("int", 1),
+				chromago.NewFloatAttribute("float", 1.1),
+			),
+		),
+	)
 	if err != nil {
 		log.Printf("failed to create collection: %s", err)
 		return
 	}
 
-	fmt.Println("Collection created:", col.Name)
+	fmt.Println("Collection created:", col.Name())
 
 	// addData {
 	// verify it's possible to add data to the collection
-	col1, err := col.Add(
-		context.Background(),
-		nil,                                      // embeddings
-		[]map[string]any{},                       // metadata
-		[]string{"test-doc-1", "test-doc-2"},     // documents
-		[]string{"test-label-1", "test-label-2"}, // ids
-	)
+	err = col.Add(context.Background(),
+		// chroma.WithIDGenerator(chroma.NewULIDGenerator()),
+		chromago.WithIDs("1", "2"),
+		chromago.WithTexts("hello world", "goodbye world"),
+		chromago.WithMetadatas(
+			chromago.NewDocumentMetadata(chromago.NewIntAttribute("int", 1)),
+			chromago.NewDocumentMetadata(chromago.NewStringAttribute("str1", "hello2")),
+		))
 	// }
 	if err != nil {
 		log.Printf("failed to add data to collection: %s", err)
 		return
 	}
-
-	fmt.Println(col1.Count(context.Background()))
+	count, err := col.Count(context.Background())
+	if err != nil {
+		log.Printf("failed to count collection: %s", err)
+		return
+	}
+	fmt.Println(count)
 
 	// queryCollection {
 	// verify it's possible to query the collection
-	queryResults, err := col1.QueryWithOptions(
+	queryResults, err := col.Query(
 		context.Background(),
-		types.WithQueryTexts([]string{"test-doc-1"}),
-		types.WithInclude(types.IDocuments, types.IEmbeddings, types.IMetadatas),
-		types.WithNResults(1),
+		chromago.WithQueryTexts("say hello"),
+		// chromago.WithInclude(types.IDocuments, types.IEmbeddings, types.IMetadatas),
+		chromago.WithNResults(1),
 	)
 	// }
 	if err != nil {
 		log.Printf("failed to query collection: %s", err)
 		return
 	}
 
-	fmt.Printf("Result of query: %v\n", queryResults)
+	distance := queryResults.GetDistancesGroups()[0][0]
+	fmt.Printf("Result of query: %v %v distance_ok=%v\n", queryResults.GetIDGroups()[0][0], queryResults.GetDocumentsGroups()[0][0], math.Abs(float64(distance)-0.7525849) < 1e-3)
 
 	// listCollections {
 	cols, err := chromaClient.ListCollections(context.Background())
@@ -170,7 +198,7 @@ func ExampleChromaContainer_collections() {
 	fmt.Println(len(cols))
 
 	// deleteCollection {
-	_, err = chromaClient.DeleteCollection(context.Background(), "test-collection")
+	err = chromaClient.DeleteCollection(context.Background(), "test-collection")
 	// }
 	if err != nil {
 		log.Printf("failed to delete collection: %s", err)
@@ -180,10 +208,10 @@ func ExampleChromaContainer_collections() {
 	fmt.Println(err)
 
 	// Output:
-	// Reset successful: true
+	// Reset successful
 	// Collection created: test-collection
-	// 2 <nil>
-	// Result of query: &{[[test-doc-1]] [[test-label-1]] [[map[]]] []}
+	// 2
+	// Result of query: 1 hello world distance_ok=true
 	// 1
 	// <nil>
 }

modules/chroma/go.mod

@@ -1,38 +1,44 @@
 module github.com/testcontainers/testcontainers-go/modules/chroma
 
-go 1.24.0
-
-toolchain go1.24.7
+go 1.24.11
 
 require (
-	github.com/amikos-tech/chroma-go v0.1.2
+	github.com/amikos-tech/chroma-go v0.3.2
 	github.com/stretchr/testify v1.11.1
 	github.com/testcontainers/testcontainers-go v0.40.0
 )
 
 require (
 	dario.cat/mergo v1.0.2 // indirect
-	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
-	github.com/Masterminds/semver v1.5.0 // indirect
+	github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 // indirect
+	github.com/Masterminds/semver/v3 v3.4.0 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
+	github.com/amikos-tech/pure-tokenizers v0.1.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/containerd/errdefs v1.0.0 // indirect
 	github.com/containerd/errdefs/pkg v0.3.0 // indirect
 	github.com/containerd/log v0.1.0 // indirect
-	github.com/containerd/platforms v0.2.1 // indirect
+	github.com/containerd/platforms v1.0.0-rc.1 // indirect
 	github.com/cpuguy83/dockercfg v0.3.2 // indirect
+	github.com/creasty/defaults v1.8.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/docker v28.5.2+incompatible // indirect
 	github.com/docker/go-connections v0.6.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/ebitengine/purego v0.9.1 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
+	github.com/gabriel-vasile/mimetype v1.4.12 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-ole/go-ole v1.2.6 // indirect
+	github.com/go-playground/locales v0.14.1 // indirect
+	github.com/go-playground/universal-translator v0.18.1 // indirect
+	github.com/go-playground/validator/v10 v10.30.1 // indirect
+	github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/klauspost/compress v1.18.0 // indirect
+	github.com/leodido/go-urn v1.4.0 // indirect
 	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
 	github.com/magiconair/properties v1.8.10 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
@@ -53,16 +59,19 @@ require (
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/tklauser/go-sysconf v0.3.16 // indirect
 	github.com/tklauser/numcpus v0.11.0 // indirect
+	github.com/yalue/onnxruntime_go v1.22.0 // indirect
 	github.com/yusufpapurcu/wmi v1.2.4 // indirect
 	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.56.0 // indirect
 	go.opentelemetry.io/otel v1.35.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.19.0 // indirect
 	go.opentelemetry.io/otel/metric v1.35.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.31.0 // indirect
 	go.opentelemetry.io/otel/trace v1.35.0 // indirect
-	go.opentelemetry.io/proto/otlp v1.0.0 // indirect
-	golang.org/x/crypto v0.45.0 // indirect
+	go.uber.org/multierr v1.10.0 // indirect
+	go.uber.org/zap v1.27.0 // indirect
+	golang.org/x/crypto v0.46.0 // indirect
 	golang.org/x/sys v0.39.0 // indirect
+	golang.org/x/text v0.32.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )