@@ -3,12 +3,13 @@ package chroma_test
33import (
44 "context"
55 "fmt"
6+ "log"
7+
68 chromago "github.com/amikos-tech/chroma-go"
7- chromaopenapi "github.com/amikos-tech/chroma-go/swagger"
89 "github.com/amikos-tech/chroma-go/types"
10+
911 "github.com/testcontainers/testcontainers-go"
1012 "github.com/testcontainers/testcontainers-go/modules/chroma"
11- "log"
1213)
1314
1415func ExampleRunContainer () {
@@ -55,21 +56,9 @@ func ExampleChromaContainer_connectWithClient() {
5556 }
5657 }()
5758
58- connectionStr , err := chromaContainer .RESTEndpoint ( ctx )
59+ chromaClient , err := chromaContainer .GetClient ( )
5960 if err != nil {
60- log .Fatalf ("failed to get REST endpoint: %s" , err ) // nolint:gocritic
61- }
62-
63- // create the client connection and confirm that we can access the server with it
64- configuration := chromaopenapi .NewConfiguration ()
65- configuration .Servers = chromaopenapi.ServerConfigurations {
66- {
67- URL : connectionStr ,
68- Description : "Chromadb server url for this store" ,
69- },
70- }
71- chromaClient := & chromago.Client {
72- ApiClient : chromaopenapi .NewAPIClient (configuration ),
61+ log .Fatalf ("failed to get client: %s" , err ) // nolint:gocritic
7362 }
7463
7564 hbs , errHb := chromaClient .Heartbeat (context .Background ())
@@ -88,7 +77,7 @@ func ExampleChromaContainer_connectWithClient() {
8877func ExampleChromaContainer_collections () {
8978 ctx := context .Background ()
9079
91- chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.24" ))
80+ chromaContainer , err := chroma .RunContainer (ctx , testcontainers .WithImage ("chromadb/chroma:0.4.24" ), testcontainers . WithEnv ( map [ string ] string { "ALLOW_RESET" : "true" }) )
9281 if err != nil {
9382 log .Fatalf ("failed to start container: %s" , err )
9483 }
@@ -99,8 +88,20 @@ func ExampleChromaContainer_collections() {
9988 }
10089 }()
10190
91+ // getClient {
10292 // create the client connection and confirm that we can access the server with it
10393 chromaClient , err := chromaContainer .GetClient (chromago .WithDebug (true ))
94+ // }
95+ if err != nil {
96+ log .Fatalf ("failed to get client: %s" , err ) // nolint:gocritic
97+ }
98+ // reset {
99+ reset , err := chromaClient .Reset (context .Background ())
100+ // }
101+ if err != nil {
102+ log .Fatalf ("failed to reset: %s" , err ) // nolint:gocritic
103+ }
104+ fmt .Printf ("Reset successful: %v\n " , reset )
104105
105106 // createCollection {
106107 // for testing we use a dummy hashing function NewConsistentHashEmbeddingFunction
@@ -116,10 +117,10 @@ func ExampleChromaContainer_collections() {
116117 // verify it's possible to add data to the collection
117118 col1 , err := col .Add (
118119 context .Background (),
119- types .NewEmbeddingsFromFloat32 ([][]float32 {{1 , 2 , 3 }, {4 , 5 , 6 }}), // or set this to nil to use the EmbeddingFunction
120- []map [string ]interface {}{}, // metadata
121- []string {"test-doc-1" , "test-doc-2" }, // documents
122- []string {"test-label-1" , "test-label-2" }, // ids
120+ nil , // or set the embedding types.NewEmbeddingsFromFloat32([][]float32{{1, 2, 3}, {4, 5, 6}})
121+ []map [string ]interface {}{}, // metadata
122+ []string {"test-doc-1" , "test-doc-2" }, // documents
123+ []string {"test-label-1" , "test-label-2" }, // ids
123124 )
124125 // }
125126 if err != nil {
@@ -140,7 +141,6 @@ func ExampleChromaContainer_collections() {
140141 if err != nil {
141142 log .Fatalf ("failed to query collection: %s" , err ) // nolint:gocritic
142143 }
143- // }
144144
145145 fmt .Printf ("Result of query: %v\n " , queryResults )
146146
@@ -163,8 +163,10 @@ func ExampleChromaContainer_collections() {
163163 fmt .Println (err )
164164
165165 // Output:
166+ // Reset successful: true
166167 // Collection created: test-collection
167168 // 2 <nil>
169+ // Result of query: &{[[test-doc-1]] [[test-label-1]] [[map[]]] []}
168170 // 1
169171 // <nil>
170172}
0 commit comments