@@ -4,11 +4,9 @@ import (
44 "context"
55 "fmt"
66 "log"
7- "os"
87
98 chromago "github.com/amikos-tech/chroma-go"
10- "github.com/amikos-tech/chroma-go/openai"
11- chromaopenapi "github.com/amikos-tech/chroma-go/swagger"
9+ "github.com/amikos-tech/chroma-go/types"
1210
1311 "github.com/testcontainers/testcontainers-go"
1412 "github.com/testcontainers/testcontainers-go/modules/chroma"
@@ -18,7 +16,7 @@ func ExampleRunContainer() {
1816 // runChromaContainer {
1917 ctx := context .Background ()
2018
21- chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.22 " ))
19+ chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.24 " ))
2220 if err != nil {
2321 log .Fatalf ("failed to start container: %s" , err )
2422 }
@@ -46,7 +44,7 @@ func ExampleChromaContainer_connectWithClient() {
4644 // createClient {
4745 ctx := context .Background ()
4846
49- chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.22 " ))
47+ chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.24 " ))
5048 if err != nil {
5149 log .Fatalf ("failed to start container: %s" , err )
5250 }
@@ -58,21 +56,13 @@ func ExampleChromaContainer_connectWithClient() {
5856 }
5957 }()
6058
61- connectionStr , err := chromaContainer .RESTEndpoint (ctx )
59+ endpoint , err := chromaContainer .RESTEndpoint (context . Background () )
6260 if err != nil {
6361 log .Fatalf ("failed to get REST endpoint: %s" , err ) // nolint:gocritic
6462 }
65-
66- // create the client connection and confirm that we can access the server with it
67- configuration := chromaopenapi .NewConfiguration ()
68- configuration .Servers = chromaopenapi.ServerConfigurations {
69- {
70- URL : connectionStr ,
71- Description : "Chromadb server url for this store" ,
72- },
73- }
74- chromaClient := & chromago.Client {
75- ApiClient : chromaopenapi .NewAPIClient (configuration ),
63+ chromaClient , err := chromago .NewClient (endpoint )
64+ if err != nil {
65+ log .Fatalf ("failed to get client: %s" , err ) // nolint:gocritic
7666 }
7767
7868 hbs , errHb := chromaClient .Heartbeat (context .Background ())
@@ -91,7 +81,7 @@ func ExampleChromaContainer_connectWithClient() {
9181func ExampleChromaContainer_collections () {
9282 ctx := context .Background ()
9383
94- chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.22" ))
84+ chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.24" ), testcontainers . WithEnv ( map [ string ] string { "ALLOW_RESET" : "true" } ))
9585 if err != nil {
9686 log .Fatalf ("failed to start container: %s" , err )
9787 }
@@ -102,51 +92,66 @@ func ExampleChromaContainer_collections() {
10292 }
10393 }()
10494
105- connectionStr , err := chromaContainer .RESTEndpoint (ctx )
95+ // getClient {
96+ // create the client connection and confirm that we can access the server with it
97+ endpoint , err := chromaContainer .RESTEndpoint (context .Background ())
10698 if err != nil {
10799 log .Fatalf ("failed to get REST endpoint: %s" , err ) // nolint:gocritic
108100 }
109-
110- // create the client connection and confirm that we can access the server with it
111- configuration := chromaopenapi .NewConfiguration ()
112- configuration .Servers = chromaopenapi.ServerConfigurations {
113- {
114- URL : connectionStr ,
115- Description : "Chromadb server url for this store" ,
116- },
101+ chromaClient , err := chromago .NewClient (endpoint )
102+ // }
103+ if err != nil {
104+ log .Fatalf ("failed to get client: %s" , err ) // nolint:gocritic
117105 }
118- chromaClient := & chromago.Client {
119- ApiClient : chromaopenapi .NewAPIClient (configuration ),
106+ // reset {
107+ reset , err := chromaClient .Reset (context .Background ())
108+ // }
109+ if err != nil {
110+ log .Fatalf ("failed to reset: %s" , err ) // nolint:gocritic
120111 }
112+ fmt .Printf ("Reset successful: %v\n " , reset )
121113
122114 // createCollection {
123- // for testing purposes, the OPENAI_API_KEY environment variable can be empty
124- // therefore this test is expected to succeed even though the API key is not set.
125- embeddingFunction := openai .NewOpenAIEmbeddingFunction (os .Getenv ("OPENAI_API_KEY" ))
126- distanceFunction := chromago .L2
127-
128- col , err := chromaClient .CreateCollection (context .Background (), "test-collection" , map [string ]any {}, true , embeddingFunction , distanceFunction )
115+ // for testing we use a dummy hashing function NewConsistentHashEmbeddingFunction
116+ col , err := chromaClient .CreateCollection (context .Background (), "test-collection" , map [string ]any {}, true , types .NewConsistentHashEmbeddingFunction (), types .L2 )
129117 // }
130118 if err != nil {
131119 log .Fatalf ("failed to create collection: %s" , err ) // nolint:gocritic
132120 }
133121
134122 fmt .Println ("Collection created:" , col .Name )
135123
124+ // addData {
136125 // verify it's possible to add data to the collection
137126 col1 , err := col .Add (
138127 context .Background (),
139- [][] float32 {{ 1 , 2 , 3 }, { 4 , 5 , 6 }},
140- []map [string ]interface {}{},
141- []string {"test-doc-1" , "test-doc-2" },
142- []string {"test-label-1" , "test-label-2" },
128+ nil , // embeddings
129+ []map [string ]interface {}{}, // metadata
130+ []string {"test-doc-1" , "test-doc-2" }, // documents
131+ []string {"test-label-1" , "test-label-2" }, // ids
143132 )
133+ // }
144134 if err != nil {
145135 log .Fatalf ("failed to add data to collection: %s" , err ) // nolint:gocritic
146136 }
147137
148138 fmt .Println (col1 .Count (context .Background ()))
149139
140+ // queryCollection {
141+ // verify it's possible to query the collection
142+ queryResults , err := col1 .QueryWithOptions (
143+ context .Background (),
144+ types .WithQueryTexts ([]string {"test-doc-1" }),
145+ types .WithInclude (types .IDocuments , types .IEmbeddings , types .IMetadatas ),
146+ types .WithNResults (1 ),
147+ )
148+ // }
149+ if err != nil {
150+ log .Fatalf ("failed to query collection: %s" , err ) // nolint:gocritic
151+ }
152+
153+ fmt .Printf ("Result of query: %v\n " , queryResults )
154+
150155 // listCollections {
151156 cols , err := chromaClient .ListCollections (context .Background ())
152157 // }
@@ -166,8 +171,10 @@ func ExampleChromaContainer_collections() {
166171 fmt .Println (err )
167172
168173 // Output:
174+ // Reset successful: true
169175 // Collection created: test-collection
170176 // 2 <nil>
177+ // Result of query: &{[[test-doc-1]] [[test-label-1]] [[map[]]] []}
171178 // 1
172179 // <nil>
173180}
0 commit comments