diff --git a/README.md b/README.md index f4c28e2..0a86b03 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,16 @@ # Upstash Vector Go Client +[![Go Reference](https://pkg.go.dev/badge/github.com/upstash/vector-go.svg)](https://pkg.go.dev/github.com/upstash/vector-go) + > [!NOTE] > **This project is in GA Stage.** > -> The Upstash Professional Support fully covers this project. It receives regular updates, and bug fixes. The Upstash team is committed to maintaining and improving its functionality. +> The Upstash Professional Support fully covers this project. It receives regular updates, and bug fixes. +> The Upstash team is committed to maintaining and improving its functionality. + +[Upstash](https://upstash.com/) Vector is a serverless vector database designed for working with vector embeddings. -This is the Go client for [Upstash](https://upstash.com/) Vector. +This is the HTTP-based Go client for [Upstash](https://upstash.com/) Vector. ## Documentation @@ -13,16 +18,25 @@ This is the Go client for [Upstash](https://upstash.com/) Vector. ## Installation +Use `go get` to install the Upstash Vector package: ```bash go get github.com/upstash/vector-go ``` +Import the Upstash Vector package in your project: + +```go +import "github.com/upstash/vector-go" +``` + ## Usage +In order to use this client, head out to [Upstash Console](https://console.upstash.com) and create a vector database. + ### Initializing the client -There are two pieces of configuration required to use the Upstash Vector index client: an REST token and REST URL. -Find your configuration values in the console dashboard at [https://console.upstash.com/](https://console.upstash.com/). +The REST token and REST URL configurations are required to initialize an Upstash Vector index client. +Find your configuration values in the console dashboard at [Upstash Console](https://console.upstash.com/). ```go import ( @@ -80,6 +94,35 @@ func main() { Upstash vector indexes support operations for working with vector data using operations such as upsert, query, fetch, and delete. +```go +import ( + "github.com/upstash/vector-go" +) + +func main() { + index := vector.NewIndex("", "") +} +``` + +Upstash Vector allows you to partition a single index into multiple isolated namespaces. + +You can specify a namespace for an index client with `Namespace(ns string)` function. +When you create a `Namespace` client, all index operations executed through this client become associated with the specified namespace. + +By default, the `Index` client is associated with the default namespace. + +```go +import ( + "github.com/upstash/vector-go" +) + +func main() { + index := vector.NewIndex("", "") + + // Returns a new Namespace client associated with the given namespace + ns := index.Namespace("") +``` + ### Upserting Vectors All vectors upserted to index must have the same dimensions. @@ -87,7 +130,7 @@ All vectors upserted to index must have the same dimensions. Upsert can be used to insert new vectors into index or to update existing vectors. -#### Upsert many +#### Upsert Many ```go upserts := []vector.Upsert{ @@ -287,3 +330,23 @@ err := index.Reset() ```go info, err := index.Info() ``` + +### List Namespaces + +All the names of active namespaces can be listed. + +```go +namespaces, err := index.ListNamespaces() +for _, ns : range namespaces { + fmt.Println(ns) +} +``` + +### Delete Namespaces + +A namespace can be deleted entirely if it exists. +The default namespaces cannot be deleted. + +```go +err := index.Namespace("ns").DeleteNamespace() +``` \ No newline at end of file diff --git a/delete.go b/delete.go index d9c9559..0a5aa38 100644 --- a/delete.go +++ b/delete.go @@ -2,11 +2,19 @@ package vector const deletePath = "/delete" -// Delete deletes the vector with the given id and reports -// whether the vector is deleted. If a vector with the given -// id is not found, Delete returns false. +// Delete deletes the vector with the given id in the default namespace and reports whether the vector is deleted. +// If a vector with the given id is not found, Delete returns false. func (ix *Index) Delete(id string) (ok bool, err error) { - data, err := ix.sendBytes(deletePath, []byte(id)) + return ix.deleteInternal(id, defaultNamespace) +} + +// DeleteMany deletes the vectors with the given ids in the default namespace and reports how many of them are deleted. +func (ix *Index) DeleteMany(ids []string) (count int, err error) { + return ix.deleteManyInternal(ids, defaultNamespace) +} + +func (ix *Index) deleteInternal(id string, ns string) (ok bool, err error) { + data, err := ix.sendBytes(buildPath(deletePath, ns), []byte(id)) if err != nil { return } @@ -16,10 +24,8 @@ func (ix *Index) Delete(id string) (ok bool, err error) { return } -// DeleteMany deletes the vectors with the given ids and reports -// how many of them are deleted. -func (ix *Index) DeleteMany(ids []string) (count int, err error) { - data, err := ix.sendJson(deletePath, ids) +func (ix *Index) deleteManyInternal(ids []string, ns string) (count int, err error) { + data, err := ix.sendJson(buildPath(deletePath, ns), ids) if err != nil { return } diff --git a/delete_test.go b/delete_test.go index 35d1c93..12b2d34 100644 --- a/delete_test.go +++ b/delete_test.go @@ -15,49 +15,56 @@ func randomString() string { } func TestDelete(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) - - id := randomString() - err = client.Upsert(Upsert{ - Id: id, - Vector: []float32{0, 1}, - }) - require.NoError(t, err) - - t.Run("existing id", func(t *testing.T) { - ok, err := client.Delete(id) - require.NoError(t, err) - require.True(t, ok) - }) - - t.Run("non existing id", func(t *testing.T) { - ok, err := client.Delete(randomString()) - require.NoError(t, err) - require.False(t, ok) - }) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClientWithNamespace(ns) + require.NoError(t, err) + id := randomString() + err = client.Upsert(Upsert{ + Id: id, + Vector: []float32{0, 1}, + }) + require.NoError(t, err) + + t.Run("existing id", func(t *testing.T) { + ok, err := client.Delete(id) + require.NoError(t, err) + require.True(t, ok) + }) + + t.Run("non existing id", func(t *testing.T) { + ok, err := client.Delete(randomString()) + require.NoError(t, err) + require.False(t, ok) + }) + }) + } } func TestDeleteMany(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) - - id0 := randomString() - id1 := randomString() - id2 := randomString() - err = client.UpsertMany([]Upsert{ - { - Id: id0, - Vector: []float32{0, 1}, - }, - { - Id: id1, - Vector: []float32{5, 10}, - }, - }) - require.NoError(t, err) - - count, err := client.DeleteMany([]string{id0, id1, id2}) - require.NoError(t, err) - require.Equal(t, 2, count) + for _, ns := range namespaces { + t.Run("namespace "+ns, func(t *testing.T) { + client, err := newTestClientWithNamespace(ns) + require.NoError(t, err) + + id0 := randomString() + id1 := randomString() + id2 := randomString() + err = client.UpsertMany([]Upsert{ + { + Id: id0, + Vector: []float32{0, 1}, + }, + { + Id: id1, + Vector: []float32{5, 10}, + }, + }) + require.NoError(t, err) + + count, err := client.DeleteMany([]string{id0, id1, id2}) + require.NoError(t, err) + require.Equal(t, 2, count) + }) + } } diff --git a/embedding_test.go b/embedding_test.go index 6fc5e98..2e2bfa4 100644 --- a/embedding_test.go +++ b/embedding_test.go @@ -1,84 +1,88 @@ package vector import ( + "github.com/stretchr/testify/require" "testing" "time" - - "github.com/stretchr/testify/require" ) func TestEmbedding(t *testing.T) { - client, err := newEmbeddingTestClient() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newEmbeddingTestClient() + require.NoError(t, err) - id0 := "tr" - id1 := "jp" - id2 := "uk" - id3 := "fr" - err = client.UpsertDataMany([]UpsertData{ - { - Id: id0, - Data: "Capital of Türkiye is Ankara.", - Metadata: map[string]any{"country": id0, "capital": "Ankara"}, - }, - { - Id: id1, - Data: "Capital of Japan is Tokyo.", - Metadata: map[string]any{"country": id1, "capital": "Tokyo"}, - }, - { - Id: id2, - Data: "Capital of England is London.", - Metadata: map[string]any{"country": id2, "capital": "London"}, - }, - { - Id: id3, - Data: "Capital of France is Paris.", - Metadata: map[string]any{"country": id3, "capital": "Paris"}, - }, - }) - require.NoError(t, err) + namespace := client.Namespace(ns) + id0 := "tr" + id1 := "jp" + id2 := "uk" + id3 := "fr" + err = namespace.UpsertDataMany([]UpsertData{ + { + Id: id0, + Data: "Capital of Türkiye is Ankara.", + Metadata: map[string]any{"country": id0, "capital": "Ankara"}, + }, + { + Id: id1, + Data: "Capital of Japan is Tokyo.", + Metadata: map[string]any{"country": id1, "capital": "Tokyo"}, + }, + { + Id: id2, + Data: "Capital of England is London.", + Metadata: map[string]any{"country": id2, "capital": "London"}, + }, + { + Id: id3, + Data: "Capital of France is Paris.", + Metadata: map[string]any{"country": id3, "capital": "Paris"}, + }, + }) + require.NoError(t, err) - require.Eventually(t, func() bool { - info, err := client.Info() - require.NoError(t, err) - return info.PendingVectorCount == 0 - }, 10*time.Second, 1*time.Second) + require.Eventually(t, func() bool { + info, err := client.Info() + require.NoError(t, err) + return info.PendingVectorCount == 0 + }, 10*time.Second, 1*time.Second) - t.Run("score", func(t *testing.T) { - scores, err := client.QueryData(QueryData{ - Data: "where is the capital of Japan?", - TopK: 1, - }) - require.NoError(t, err) - require.Equal(t, 1, len(scores)) - require.Equal(t, id1, scores[0].Id) - }) + t.Run("score", func(t *testing.T) { + scores, err := namespace.QueryData(QueryData{ + Data: "where is the capital of Japan?", + TopK: 1, + }) + require.NoError(t, err) + require.Equal(t, 1, len(scores)) + require.Equal(t, id1, scores[0].Id) + }) - t.Run("with metadata", func(t *testing.T) { - scores, err := client.QueryData(QueryData{ - Data: "Which country's capital is Ankara?", - TopK: 1, - IncludeMetadata: true, - }) - require.NoError(t, err) - require.Equal(t, 1, len(scores)) - require.Equal(t, id0, scores[0].Id) - require.Equal(t, map[string]any{"country": "tr", "capital": "Ankara"}, scores[0].Metadata) - }) + t.Run("with metadata", func(t *testing.T) { + scores, err := namespace.QueryData(QueryData{ + Data: "Which country's capital is Ankara?", + TopK: 1, + IncludeMetadata: true, + }) + require.NoError(t, err) + require.Equal(t, 1, len(scores)) + require.Equal(t, id0, scores[0].Id) + require.Equal(t, map[string]any{"country": "tr", "capital": "Ankara"}, scores[0].Metadata) + }) - t.Run("with metadata filtering", func(t *testing.T) { - query := QueryData{ - Data: "Where is the capital of France?", - TopK: 1, - IncludeMetadata: true, - Filter: `country = 'fr'`, - } + t.Run("with metadata filtering", func(t *testing.T) { + query := QueryData{ + Data: "Where is the capital of France?", + TopK: 1, + IncludeMetadata: true, + Filter: `country = 'fr'`, + } - scores, err := client.QueryData(query) - require.NoError(t, err) - require.Equal(t, 1, len(scores)) - require.Equal(t, id3, scores[0].Id) - require.Equal(t, map[string]any{"country": "fr", "capital": "Paris"}, scores[0].Metadata) - }) + scores, err := namespace.QueryData(query) + require.NoError(t, err) + require.Equal(t, 1, len(scores)) + require.Equal(t, id3, scores[0].Id) + require.Equal(t, map[string]any{"country": "fr", "capital": "Paris"}, scores[0].Metadata) + }) + }) + } } diff --git a/fetch.go b/fetch.go index 403a29c..19dd526 100644 --- a/fetch.go +++ b/fetch.go @@ -2,12 +2,15 @@ package vector const fetchPath = "/fetch" -// Fetch fetches one or more vectors with the ids passed into f. -// When f.IncludeVectors is true, values of the vectors are also -// returned. When f.IncludeMetadata is true, metadata of the vectors -// are also returned, if any. +// Fetch fetches one or more vectors in the default namespace with the ids passed into f. +// If IncludeVectors is set to true, the vector values are also returned. +// If IncludeMetadata is set to true, any associated metadata of the vectors is also returned, if any. func (ix *Index) Fetch(f Fetch) (vectors []Vector, err error) { - data, err := ix.sendJson(fetchPath, f) + return ix.fetchInternal(f, defaultNamespace) +} + +func (ix *Index) fetchInternal(f Fetch, ns string) (vectors []Vector, err error) { + data, err := ix.sendJson(buildPath(fetchPath, ns), f) if err != nil { return } diff --git a/fetch_test.go b/fetch_test.go index 891ba51..65c2ff6 100644 --- a/fetch_test.go +++ b/fetch_test.go @@ -7,69 +7,73 @@ import ( ) func TestFetch(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClientWithNamespace(ns) + require.NoError(t, err) - id0 := randomString() - id1 := randomString() - err = client.UpsertMany([]Upsert{ - { - Id: id0, - Vector: []float32{0, 1}, - Metadata: map[string]any{"foo": "bar"}, - }, - { - Id: id1, - Vector: []float32{5, 10}, - }, - }) - require.NoError(t, err) + id0 := randomString() + id1 := randomString() + err = client.UpsertMany([]Upsert{ + { + Id: id0, + Vector: []float32{0, 1}, + Metadata: map[string]any{"foo": "bar"}, + }, + { + Id: id1, + Vector: []float32{5, 10}, + }, + }) + require.NoError(t, err) - t.Run("single", func(t *testing.T) { - vectors, err := client.Fetch(Fetch{ - Ids: []string{id0}, - }) - require.NoError(t, err) + t.Run("single", func(t *testing.T) { + vectors, err := client.Fetch(Fetch{ + Ids: []string{id0}, + }) + require.NoError(t, err) - require.Equal(t, 1, len(vectors)) - require.Equal(t, id0, vectors[0].Id) - }) + require.Equal(t, 1, len(vectors)) + require.Equal(t, id0, vectors[0].Id) + }) - t.Run("many", func(t *testing.T) { - vectors, err := client.Fetch(Fetch{ - Ids: []string{id0, id1}, - }) - require.NoError(t, err) + t.Run("many", func(t *testing.T) { + vectors, err := client.Fetch(Fetch{ + Ids: []string{id0, id1}, + }) + require.NoError(t, err) - require.Equal(t, 2, len(vectors)) - require.Equal(t, id0, vectors[0].Id) - require.Equal(t, id1, vectors[1].Id) - }) + require.Equal(t, 2, len(vectors)) + require.Equal(t, id0, vectors[0].Id) + require.Equal(t, id1, vectors[1].Id) + }) - t.Run("non existing id", func(t *testing.T) { - vectors, err := client.Fetch(Fetch{ - Ids: []string{randomString()}, - }) - require.NoError(t, err) + t.Run("non existing id", func(t *testing.T) { + vectors, err := client.Fetch(Fetch{ + Ids: []string{randomString()}, + }) + require.NoError(t, err) - require.Equal(t, 1, len(vectors)) - require.Equal(t, "", vectors[0].Id) - }) + require.Equal(t, 1, len(vectors)) + require.Equal(t, "", vectors[0].Id) + }) - t.Run("with metadata and vectors", func(t *testing.T) { - vectors, err := client.Fetch(Fetch{ - Ids: []string{id0, id1}, - IncludeMetadata: true, - IncludeVectors: true, - }) - require.NoError(t, err) + t.Run("with metadata and vectors", func(t *testing.T) { + vectors, err := client.Fetch(Fetch{ + Ids: []string{id0, id1}, + IncludeMetadata: true, + IncludeVectors: true, + }) + require.NoError(t, err) - require.Equal(t, 2, len(vectors)) - require.Equal(t, id0, vectors[0].Id) - require.Equal(t, map[string]any{"foo": "bar"}, vectors[0].Metadata) - require.Equal(t, []float32{0, 1}, vectors[0].Vector) - require.Equal(t, id1, vectors[1].Id) - require.Nil(t, vectors[1].Metadata) // was upserted with nil metadata - require.Equal(t, []float32{5, 10}, vectors[1].Vector) - }) + require.Equal(t, 2, len(vectors)) + require.Equal(t, id0, vectors[0].Id) + require.Equal(t, map[string]any{"foo": "bar"}, vectors[0].Metadata) + require.Equal(t, []float32{0, 1}, vectors[0].Vector) + require.Equal(t, id1, vectors[1].Id) + require.Nil(t, vectors[1].Metadata) // was upserted with nil metadata + require.Equal(t, []float32{5, 10}, vectors[1].Vector) + }) + }) + } } diff --git a/index.go b/index.go index 9be20cb..2e69afc 100644 --- a/index.go +++ b/index.go @@ -4,14 +4,17 @@ import ( "bytes" "encoding/json" "errors" + "fmt" "io" "net/http" "os" + "runtime" ) const ( UrlEnvProperty = "UPSTASH_VECTOR_REST_URL" TokenEnvProperty = "UPSTASH_VECTOR_REST_TOKEN" + defaultNamespace = "" ) type Options struct { @@ -59,18 +62,21 @@ func NewIndexFromEnv() *Index { // with the given options. func NewIndexWith(options Options) *Index { options.init() - return &Index{ + index := &Index{ url: options.Url, token: options.Token, client: options.Client, } + index.generateHeaders() + return index } // Index is a client for Upstash Vector index. type Index struct { - url string - token string - client *http.Client + url string + token string + client *http.Client + headers http.Header } func (ix *Index) sendJson(path string, obj any) (data []byte, err error) { @@ -85,11 +91,11 @@ func (ix *Index) sendBytes(path string, obj []byte) (data []byte, err error) { } func (ix *Index) send(path string, r io.Reader) (data []byte, err error) { - request, err := http.NewRequest("POST", ix.url+path, r) + request, err := http.NewRequest(http.MethodPost, ix.url+path, r) if err != nil { return } - request.Header.Add("Authorization", "Bearer "+ix.token) + request.Header = ix.headers response, err := ix.client.Do(request) if err != nil { return @@ -110,3 +116,26 @@ func parseResponse[T any](data []byte) (t T, err error) { } return } + +func (ix *Index) generateHeaders() { + headers := http.Header{} + headers.Add("Authorization", "Bearer "+ix.token) + headers.Add("Upstash-Telemetry-Runtime", fmt.Sprintf("vector-go@%s", runtime.Version())) + var platform string + if os.Getenv("VERCEL") != "" { + platform = "vercel" + } else if os.Getenv("AWS_REGION") != "" { + platform = "aws" + } else { + platform = "unknown" + } + headers.Add("Upstash-Telemetry-Platform", platform) + ix.headers = headers +} + +func buildPath(path string, ns string) string { + if ns == "" { + return path + } + return path + "/" + ns +} diff --git a/index_test.go b/index_test.go index 92cc150..9909d40 100644 --- a/index_test.go +++ b/index_test.go @@ -2,13 +2,16 @@ package vector import ( "errors" + "github.com/joho/godotenv" + "github.com/stretchr/testify/require" "io/fs" "net/http" "os" "testing" +) - "github.com/joho/godotenv" - "github.com/stretchr/testify/require" +var ( + namespaces = [...]string{defaultNamespace, "ns"} ) func init() { @@ -24,23 +27,43 @@ func newTestClient() (*Index, error) { os.Getenv(TokenEnvProperty), ) - err := client.Reset() - if err != nil { - return nil, err + for _, ns := range namespaces { + err := client.Namespace(ns).Reset() + if err != nil { + return nil, err + } } return client, nil } +func newTestClientWithNamespace(ns string) (*Namespace, error) { + client := NewIndex( + os.Getenv(UrlEnvProperty), + os.Getenv(TokenEnvProperty), + ) + + for _, ns := range namespaces { + err := client.Namespace(ns).Reset() + if err != nil { + return nil, err + } + } + + return client.Namespace(ns), nil +} + func newEmbeddingTestClient() (*Index, error) { client := NewIndex( os.Getenv("EMBEDDING_"+UrlEnvProperty), os.Getenv("EMBEDDING_"+TokenEnvProperty), ) - err := client.Reset() - if err != nil { - return nil, err + for _, ns := range namespaces { + err := client.Namespace(ns).Reset() + if err != nil { + return nil, err + } } return client, nil @@ -55,26 +78,36 @@ func newTestClientWith(client *http.Client) (*Index, error) { c := NewIndexWith(opts) - err := c.Reset() - if err != nil { - return nil, err + for _, ns := range namespaces { + err := c.Namespace(ns).Reset() + if err != nil { + return nil, err + } } return c, nil } func TestNewClient(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) - - _, err = client.Info() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClient() + require.NoError(t, err) + + _, err = client.Info() + require.NoError(t, err) + }) + } } func TestNewClientWith(t *testing.T) { - client, err := newTestClientWith(&http.Client{}) - require.NoError(t, err) - - _, err = client.Info() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClientWith(&http.Client{}) + require.NoError(t, err) + + _, err = client.Info() + require.NoError(t, err) + }) + } } diff --git a/info.go b/info.go index e626937..a71752a 100644 --- a/info.go +++ b/info.go @@ -2,10 +2,13 @@ package vector const infoPath = "/info" -// Info returns some information about the index; such as -// vector count, vectors pending for indexing, size of the -// index on disk in bytes, dimension of the index, and -// the name of the similarity function used. +// Info returns some information about the index, including: +// - Total number of vectors across all namespaces +// - Total number of vectors waiting to be indexed across all namespaces +// - Total size of the index on disk in bytes +// - Vector dimension +// - Similarity function used +// - per-namespace vector and pending vector counts func (ix *Index) Info() (info IndexInfo, err error) { data, err := ix.sendJson(infoPath, nil) if err != nil { diff --git a/info_test.go b/info_test.go index 48465fe..4869847 100644 --- a/info_test.go +++ b/info_test.go @@ -1,24 +1,42 @@ package vector import ( - "testing" - "github.com/stretchr/testify/require" + "testing" + "time" ) func TestInfo(t *testing.T) { client, err := newTestClient() require.NoError(t, err) - err = client.Upsert(Upsert{ - Id: randomString(), - Vector: []float32{0, 1}, - }) - require.NoError(t, err) + for _, ns := range namespaces { + createNamespace(t, client, ns) + } info, err := client.Info() require.NoError(t, err) - require.Greater(t, info.VectorCount, 0) + require.Equal(t, info.VectorCount, 0) require.Equal(t, 2, info.Dimension) require.Equal(t, "COSINE", info.SimilarityFunction) + require.Equal(t, len(namespaces), len(info.Namespaces)) + for _, ns := range namespaces { + require.Contains(t, info.Namespaces, ns) + require.Equal(t, 0, info.Namespaces[ns].VectorCount) + require.Equal(t, 0, info.Namespaces[ns].PendingVectorCount) + } + + for _, ns := range namespaces { + err = client.Namespace(ns).Upsert(Upsert{ + Id: randomString(), + Vector: []float32{0, 1}, + }) + require.NoError(t, err) + } + + require.Eventually(t, func() bool { + info, err := client.Info() + require.NoError(t, err) + return info.VectorCount == len(namespaces) + }, 10*time.Second, 1*time.Second) } diff --git a/namespace.go b/namespace.go new file mode 100644 index 0000000..03f68f0 --- /dev/null +++ b/namespace.go @@ -0,0 +1,112 @@ +package vector + +const deleteNamespacePath = "/delete-namespace" +const listNamespacesPath = "/list-namespaces" + +type Namespace struct { + index *Index + ns string +} + +// Namespace returns a new client associated with the given namespace. +func (ix *Index) Namespace(namespace string) (i *Namespace) { + return &Namespace{ + index: ix, + ns: namespace, + } +} + +// ListNamespaces returns the list of names of namespaces. +func (ix *Index) ListNamespaces() (namespaces []string, err error) { + data, err := ix.sendJson(listNamespacesPath, nil) + if err != nil { + return + } + namespaces, err = parseResponse[[]string](data) + return +} + +// DeleteNamespace deletes the given namespace of index if it exists. +func (ns *Namespace) DeleteNamespace() error { + _, err := ns.index.sendBytes(buildPath(deleteNamespacePath, ns.ns), nil) + return err +} + +// Upsert updates or inserts a vector to the namespace of the index. +// Additional metadata can also be provided while upserting the vector. +func (ns *Namespace) Upsert(u Upsert) (err error) { + return ns.index.upsertInternal(u, ns.ns) +} + +// UpsertMany updates or inserts some vectors to the default namespace of the index. +// Additional metadata can also be provided for each vector. +func (ns *Namespace) UpsertMany(u []Upsert) (err error) { + return ns.index.upsertManyInternal(u, ns.ns) +} + +// UpsertData updates or inserts a vector to the namespace of the index +// by converting given raw data to an embedding on the server. +// Additional metadata can also be provided while upserting the vector. +func (ns *Namespace) UpsertData(u UpsertData) (err error) { + return ns.index.upsertDataInternal(u, ns.ns) +} + +// UpsertDataMany updates or inserts some vectors to the default namespace of the index +// by converting given raw data to an embedding on the server. +// Additional metadata can also be provided for each vector. +func (ns *Namespace) UpsertDataMany(u []UpsertData) (err error) { + return ns.index.upsertDataManyInternal(u, ns.ns) +} + +// Fetch fetches one or more vectors in the namespace with the ids passed into f. +// If IncludeVectors is set to true, the vector values are also returned. +// If IncludeMetadata is set to true, any associated metadata of the vectors is also returned, if any. +func (ns *Namespace) Fetch(f Fetch) (vectors []Vector, err error) { + return ns.index.fetchInternal(f, ns.ns) +} + +// QueryData returns the result of the query for the given data by converting it to an embedding on the server. +// When q.TopK is specified, the result will contain at most q.TopK many vectors. +// The returned list will contain vectors sorted in descending order of score, +// which correlates with the similarity of the vectors to the given query vector. +// When q.IncludeVectors is true, values of the vectors are also returned. +// When q.IncludeMetadata is true, metadata of the vectors are also returned, if any. +func (ns *Namespace) QueryData(q QueryData) (scores []VectorScore, err error) { + return ns.index.queryDataInternal(q, ns.ns) +} + +// Query returns the result of the query for the given vector in the namespace. +// When q.TopK is specified, the result will contain at most q.TopK many vectors. +// The returned list will contain vectors sorted in descending order of score, +// which correlates with the similarity of the vectors to the given query vector. +// When q.IncludeVectors is true, values of the vectors are also returned. +// When q.IncludeMetadata is true, metadata of the vectors are also returned, if any. +func (ns *Namespace) Query(q Query) (scores []VectorScore, err error) { + return ns.index.queryInternal(q, ns.ns) +} + +// Range returns a range of vectors, starting with r.Cursor (inclusive), +// until the end of the vectors in the index or until the given q.Limit. +// The initial cursor should be set to "0", and subsequent calls to +// Range might use the next cursor returned in the response. +// When r.IncludeVectors is true, values of the vectors are also returned. +// When r.IncludeMetadata is true, metadata of the vectors are also returned, if any. +func (ns *Namespace) Range(r Range) (vectors RangeVectors, err error) { + return ns.index.rangeInternal(r, ns.ns) +} + +// Delete deletes the vector with the given id in the namespace and reports whether the vector is deleted. +// If a vector with the given id is not found, Delete returns false. +func (ns *Namespace) Delete(id string) (ok bool, err error) { + return ns.index.deleteInternal(id, ns.ns) +} + +// DeleteMany deletes the vectors with the given ids in the namespace and reports how many of them are deleted. +func (ns *Namespace) DeleteMany(ids []string) (count int, err error) { + return ns.index.deleteManyInternal(ids, ns.ns) +} + +// Reset deletes all the vectors in the namespace of the index and resets it to initial state. +func (ns *Namespace) Reset() (err error) { + return ns.index.resetInternal(ns.ns) +} diff --git a/namespace_test.go b/namespace_test.go new file mode 100644 index 0000000..cf87ebd --- /dev/null +++ b/namespace_test.go @@ -0,0 +1,67 @@ +package vector + +import ( + "github.com/stretchr/testify/require" + "testing" + "time" +) + +func TestNamespace(t *testing.T) { + + t.Run("list namespaces", func(t *testing.T) { + client, err := newTestClient() + require.NoError(t, err) + + for _, ns := range namespaces { + createNamespace(t, client, ns) + } + + ns, err := client.ListNamespaces() + require.NoError(t, err) + require.Exactly(t, []string{"", "ns"}, ns) + }) + + t.Run("delete namespaces", func(t *testing.T) { + client, err := newTestClient() + require.NoError(t, err) + + for _, ns := range namespaces { + createNamespace(t, client, ns) + } + + for _, ns := range namespaces { + if ns == "" { + continue + } + + err := client.Namespace(ns).DeleteNamespace() + require.NoError(t, err) + } + + info, err := client.Info() + require.NoError(t, err) + require.Exactly(t, 1, len(info.Namespaces)) + }) +} + +func createNamespace(t *testing.T, client *Index, ns string) { + if ns == "" { + return + } + + namespace := client.Namespace(ns) + + err := namespace.Upsert(Upsert{ + Vector: []float32{0.1, 0.1}, + }) + require.NoError(t, err) + + require.Eventually(t, func() bool { + info, err := client.Info() + require.NoError(t, err) + return info.PendingVectorCount == 0 + }, 10*time.Second, 1*time.Second) + + err = namespace.Reset() + require.NoError(t, err) +} diff --git a/query.go b/query.go index 7cc8015..0b20f04 100644 --- a/query.go +++ b/query.go @@ -2,15 +2,18 @@ package vector const queryPath = "/query" -// Query returns the result of the query for the given vector -// When q.TopK is specified, the result will contain at most q.TopK -// many vectors. The returned list will contain vectors sorted in descending -// order of score, which correlates with the similarity of the vectors to the -// given query vector. When q.IncludeVectors is true, values of the vectors are -// also returned. When q.IncludeMetadata is true, metadata of the vectors are -// also returned, if any. +// Query returns the result of the query for the given vector in the default namespace. +// When q.TopK is specified, the result will contain at most q.TopK many vectors. +// The returned list will contain vectors sorted in descending order of score, +// which correlates with the similarity of the vectors to the given query vector. +// When q.IncludeVectors is true, values of the vectors are also returned. +// When q.IncludeMetadata is true, metadata of the vectors are also returned, if any. func (ix *Index) Query(q Query) (scores []VectorScore, err error) { - data, err := ix.sendJson(queryPath, q) + return ix.queryInternal(q, defaultNamespace) +} + +func (ix *Index) queryInternal(q Query, ns string) (scores []VectorScore, err error) { + data, err := ix.sendJson(buildPath(queryPath, ns), q) if err != nil { return } diff --git a/query_data.go b/query_data.go index 078d577..5ccdd3f 100644 --- a/query_data.go +++ b/query_data.go @@ -2,16 +2,18 @@ package vector const queryDataPath = "/query-data" -// QueryData returns the result of the query for the given data -// by converting it to an embedding on the server. -// When q.TopK is specified, the result will contain at most q.TopK -// many vectors. The returned list will contain vectors sorted in descending -// order of score, which correlates with the similarity of the vectors to the -// given query vector. When q.IncludeVectors is true, values of the vectors are -// also returned. When q.IncludeMetadata is true, metadata of the vectors are -// also returned, if any. +// QueryData returns the result of the query for the given data by converting it to an embedding on the server. +// When q.TopK is specified, the result will contain at most q.TopK many vectors. +// The returned list will contain vectors sorted in descending order of score, +// which correlates with the similarity of the vectors to the given query vector. +// When q.IncludeVectors is true, values of the vectors are also returned. +// When q.IncludeMetadata is true, metadata of the vectors are also returned, if any. func (ix *Index) QueryData(q QueryData) (scores []VectorScore, err error) { - data, err := ix.sendJson(queryDataPath, q) + return ix.queryDataInternal(q, defaultNamespace) +} + +func (ix *Index) queryDataInternal(q QueryData, ns string) (scores []VectorScore, err error) { + data, err := ix.sendJson(buildPath(queryDataPath, ns), q) if err != nil { return } diff --git a/query_test.go b/query_test.go index eb7e755..63b17b5 100644 --- a/query_test.go +++ b/query_test.go @@ -1,96 +1,100 @@ package vector import ( + "github.com/stretchr/testify/require" "testing" "time" - - "github.com/stretchr/testify/require" ) func TestQuery(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClient() + require.NoError(t, err) - id0 := randomString() - id1 := randomString() - id2 := randomString() - err = client.UpsertMany([]Upsert{ - { - Id: id0, - Vector: []float32{0, 1}, - Metadata: map[string]any{"foo": "bar"}, - }, - { - Id: id1, - Vector: []float32{5, 10}, - }, - { - Id: id2, - Vector: []float32{0.01, 1.01}, - Metadata: map[string]any{"foo": "nay"}, - }, - }) - require.NoError(t, err) + namespace := client.Namespace(ns) + id0 := randomString() + id1 := randomString() + id2 := randomString() + err = namespace.UpsertMany([]Upsert{ + { + Id: id0, + Vector: []float32{0, 1}, + Metadata: map[string]any{"foo": "bar"}, + }, + { + Id: id1, + Vector: []float32{5, 10}, + }, + { + Id: id2, + Vector: []float32{0.01, 1.01}, + Metadata: map[string]any{"foo": "nay"}, + }, + }) + require.NoError(t, err) - require.Eventually(t, func() bool { - info, err := client.Info() - require.NoError(t, err) - return info.PendingVectorCount == 0 - }, 10*time.Second, 1*time.Second) + require.Eventually(t, func() bool { + info, err := client.Info() + require.NoError(t, err) + return info.PendingVectorCount == 0 + }, 10*time.Second, 1*time.Second) - t.Run("score", func(t *testing.T) { - scores, err := client.Query(Query{ - Vector: []float32{0, 1}, - TopK: 2, - }) - require.NoError(t, err) - require.Equal(t, 2, len(scores)) - require.Equal(t, id0, scores[0].Id) - require.Equal(t, float32(1.0), scores[0].Score) - require.Equal(t, id2, scores[1].Id) - }) + t.Run("score", func(t *testing.T) { + scores, err := namespace.Query(Query{ + Vector: []float32{0, 1}, + TopK: 2, + }) + require.NoError(t, err) + require.Equal(t, 2, len(scores)) + require.Equal(t, id0, scores[0].Id) + require.Equal(t, float32(1.0), scores[0].Score) + require.Equal(t, id2, scores[1].Id) + }) - t.Run("with metadata and vectors", func(t *testing.T) { - scores, err := client.Query(Query{ - Vector: []float32{0, 1}, - TopK: 2, - IncludeMetadata: true, - IncludeVectors: true, - }) - require.NoError(t, err) - require.Equal(t, 2, len(scores)) - require.Equal(t, id0, scores[0].Id) - require.Equal(t, float32(1.0), scores[0].Score) - require.Equal(t, map[string]any{"foo": "bar"}, scores[0].Metadata) - require.Equal(t, []float32{0, 1}, scores[0].Vector) + t.Run("with metadata and vectors", func(t *testing.T) { + scores, err := namespace.Query(Query{ + Vector: []float32{0, 1}, + TopK: 2, + IncludeMetadata: true, + IncludeVectors: true, + }) + require.NoError(t, err) + require.Equal(t, 2, len(scores)) + require.Equal(t, id0, scores[0].Id) + require.Equal(t, float32(1.0), scores[0].Score) + require.Equal(t, map[string]any{"foo": "bar"}, scores[0].Metadata) + require.Equal(t, []float32{0, 1}, scores[0].Vector) - require.Equal(t, id2, scores[1].Id) - require.Equal(t, []float32{0.01, 1.01}, scores[1].Vector) - }) + require.Equal(t, id2, scores[1].Id) + require.Equal(t, []float32{0.01, 1.01}, scores[1].Vector) + }) - t.Run("with metadata filtering", func(t *testing.T) { - query := Query{ - Vector: []float32{0, 1}, - TopK: 10, - IncludeMetadata: true, - IncludeVectors: true, - Filter: `foo = 'bar'`, - } + t.Run("with metadata filtering", func(t *testing.T) { + query := Query{ + Vector: []float32{0, 1}, + TopK: 10, + IncludeMetadata: true, + IncludeVectors: true, + Filter: `foo = 'bar'`, + } - scores, err := client.Query(query) - require.NoError(t, err) - require.Equal(t, 1, len(scores)) - require.Equal(t, id0, scores[0].Id) - require.Equal(t, float32(1.0), scores[0].Score) - require.Equal(t, map[string]any{"foo": "bar"}, scores[0].Metadata) - require.Equal(t, []float32{0, 1}, scores[0].Vector) + scores, err := namespace.Query(query) + require.NoError(t, err) + require.Equal(t, 1, len(scores)) + require.Equal(t, id0, scores[0].Id) + require.Equal(t, float32(1.0), scores[0].Score) + require.Equal(t, map[string]any{"foo": "bar"}, scores[0].Metadata) + require.Equal(t, []float32{0, 1}, scores[0].Vector) - query.Filter = `foo = 'nay'` - scores, err = client.Query(query) - require.NoError(t, err) - require.Equal(t, 1, len(scores)) - require.Equal(t, id2, scores[0].Id) - require.Equal(t, map[string]any{"foo": "nay"}, scores[0].Metadata) - require.Equal(t, []float32{0.01, 1.01}, scores[0].Vector) - }) + query.Filter = `foo = 'nay'` + scores, err = namespace.Query(query) + require.NoError(t, err) + require.Equal(t, 1, len(scores)) + require.Equal(t, id2, scores[0].Id) + require.Equal(t, map[string]any{"foo": "nay"}, scores[0].Metadata) + require.Equal(t, []float32{0.01, 1.01}, scores[0].Vector) + }) + }) + } } diff --git a/range.go b/range.go index 58e1ca2..f367b36 100644 --- a/range.go +++ b/range.go @@ -7,10 +7,13 @@ const rangePath = "/range" // The initial cursor should be set to "0", and subsequent calls to // Range might use the next cursor returned in the response. // When r.IncludeVectors is true, values of the vectors are also returned. -// When r.IncludeMetadata is true, metadata of the vectors are also returned, -// if any. +// When r.IncludeMetadata is true, metadata of the vectors are also returned, if any. func (ix *Index) Range(r Range) (vectors RangeVectors, err error) { - data, err := ix.sendJson(rangePath, r) + return ix.rangeInternal(r, defaultNamespace) +} + +func (ix *Index) rangeInternal(r Range, ns string) (vectors RangeVectors, err error) { + data, err := ix.sendJson(buildPath(rangePath, ns), r) if err != nil { return } diff --git a/range_test.go b/range_test.go index 4fe6be8..303a1d0 100644 --- a/range_test.go +++ b/range_test.go @@ -1,79 +1,82 @@ package vector import ( - "testing" - "github.com/stretchr/testify/require" + "testing" ) func TestRange(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClientWithNamespace(ns) + require.NoError(t, err) - id0 := randomString() - id1 := randomString() - err = client.UpsertMany([]Upsert{ - { - Id: id0, - Vector: []float32{0, 1}, - Metadata: map[string]any{"foo": "bar"}, - }, - { - Id: id1, - Vector: []float32{5, 10}, - }, - }) - require.NoError(t, err) + id0 := randomString() + id1 := randomString() + err = client.UpsertMany([]Upsert{ + { + Id: id0, + Vector: []float32{0, 1}, + Metadata: map[string]any{"foo": "bar"}, + }, + { + Id: id1, + Vector: []float32{5, 10}, + }, + }) + require.NoError(t, err) - t.Run("range all", func(t *testing.T) { - vectors, err := client.Range(Range{ - Cursor: "", - Limit: 2, - }) - require.NoError(t, err) + t.Run("range all", func(t *testing.T) { + vectors, err := client.Range(Range{ + Cursor: "", + Limit: 2, + }) + require.NoError(t, err) - require.Equal(t, 2, len(vectors.Vectors)) - require.Equal(t, id0, vectors.Vectors[0].Id) - require.Equal(t, id1, vectors.Vectors[1].Id) - require.Equal(t, "", vectors.NextCursor) - }) + require.Equal(t, 2, len(vectors.Vectors)) + require.Equal(t, id0, vectors.Vectors[0].Id) + require.Equal(t, id1, vectors.Vectors[1].Id) + require.Equal(t, "", vectors.NextCursor) + }) - t.Run("range part by part", func(t *testing.T) { - vectors, err := client.Range(Range{ - Cursor: "", - Limit: 1, - }) - require.NoError(t, err) + t.Run("range part by part", func(t *testing.T) { + vectors, err := client.Range(Range{ + Cursor: "", + Limit: 1, + }) + require.NoError(t, err) - require.Equal(t, 1, len(vectors.Vectors)) - require.Equal(t, id0, vectors.Vectors[0].Id) - require.Equal(t, "1", vectors.NextCursor) + require.Equal(t, 1, len(vectors.Vectors)) + require.Equal(t, id0, vectors.Vectors[0].Id) + require.Equal(t, "1", vectors.NextCursor) - vectors, err = client.Range(Range{ - Cursor: "1", - }) - require.NoError(t, err) + vectors, err = client.Range(Range{ + Cursor: "1", + }) + require.NoError(t, err) - require.Equal(t, 1, len(vectors.Vectors)) - require.Equal(t, id1, vectors.Vectors[0].Id) - require.Equal(t, "", vectors.NextCursor) - }) + require.Equal(t, 1, len(vectors.Vectors)) + require.Equal(t, id1, vectors.Vectors[0].Id) + require.Equal(t, "", vectors.NextCursor) + }) - t.Run("with metadata and vectors", func(t *testing.T) { - vectors, err := client.Range(Range{ - Limit: 2, - IncludeMetadata: true, - IncludeVectors: true, - }) - require.NoError(t, err) + t.Run("with metadata and vectors", func(t *testing.T) { + vectors, err := client.Range(Range{ + Limit: 2, + IncludeMetadata: true, + IncludeVectors: true, + }) + require.NoError(t, err) - require.Equal(t, 2, len(vectors.Vectors)) - require.Equal(t, id0, vectors.Vectors[0].Id) - require.Equal(t, map[string]any{"foo": "bar"}, vectors.Vectors[0].Metadata) - require.Equal(t, []float32{0, 1}, vectors.Vectors[0].Vector) - require.Equal(t, id1, vectors.Vectors[1].Id) - require.Nil(t, vectors.Vectors[1].Metadata) // was upserted with nil metadata - require.Equal(t, []float32{5, 10}, vectors.Vectors[1].Vector) - require.Equal(t, "", vectors.NextCursor) - }) + require.Equal(t, 2, len(vectors.Vectors)) + require.Equal(t, id0, vectors.Vectors[0].Id) + require.Equal(t, map[string]any{"foo": "bar"}, vectors.Vectors[0].Metadata) + require.Equal(t, []float32{0, 1}, vectors.Vectors[0].Vector) + require.Equal(t, id1, vectors.Vectors[1].Id) + require.Nil(t, vectors.Vectors[1].Metadata) // was upserted with nil metadata + require.Equal(t, []float32{5, 10}, vectors.Vectors[1].Vector) + require.Equal(t, "", vectors.NextCursor) + }) + }) + } } diff --git a/reset.go b/reset.go index aad9b17..85b97ca 100644 --- a/reset.go +++ b/reset.go @@ -2,10 +2,13 @@ package vector const resetPath = "/reset" -// Reset deletes all the vectors in the index and resets -// it to initial state. +// Reset deletes all the vectors in the default namespace of the index and resets it to initial state. func (ix *Index) Reset() (err error) { - data, err := ix.send(resetPath, nil) + return ix.resetInternal(defaultNamespace) +} + +func (ix *Index) resetInternal(ns string) (err error) { + data, err := ix.send(buildPath(resetPath, ns), nil) if err != nil { return } diff --git a/reset_test.go b/reset_test.go index 1e8db99..5edbb4a 100644 --- a/reset_test.go +++ b/reset_test.go @@ -1,26 +1,33 @@ package vector import ( - "testing" - "github.com/stretchr/testify/require" + "testing" + "time" ) func TestReset(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClient() + require.NoError(t, err) - id := randomString() - err = client.Upsert(Upsert{ - Id: id, - Vector: []float32{0, 1}, - }) - require.NoError(t, err) + namespace := client.Namespace(ns) + id := randomString() + err = namespace.Upsert(Upsert{ + Id: id, + Vector: []float32{0, 1}, + }) + require.NoError(t, err) - err = client.Reset() - require.NoError(t, err) + err = namespace.Reset() + require.NoError(t, err) - info, err := client.Info() - require.NoError(t, err) - require.Equal(t, 0, info.VectorCount) + require.Eventually(t, func() bool { + info, err := client.Info() + require.NoError(t, err) + return info.VectorCount == 0 + }, 10*time.Second, 1*time.Second) + }) + } } diff --git a/types.go b/types.go index 7cf587e..b2dfdeb 100644 --- a/types.go +++ b/types.go @@ -139,6 +139,17 @@ type IndexInfo struct { // Name of the similarity function used in indexing and queries. SimilarityFunction string `json:"similarityFunction"` + + // Per-namespace vector and pending vector counts + Namespaces map[string]NamespaceInfo `json:"namespaces"` +} + +type NamespaceInfo struct { + // The number of vectors in the namespace of the index. + VectorCount int `json:"vectorCount"` + + // The number of vectors that are pending to be indexed. + PendingVectorCount int `json:"pendingVectorCount"` } type response[T any] struct { diff --git a/upsert.go b/upsert.go index a1b5956..fd7b339 100644 --- a/upsert.go +++ b/upsert.go @@ -2,10 +2,20 @@ package vector const upsertPath = "/upsert" -// Upsert updates or inserts a vector to the index. +// Upsert updates or inserts a vector to the default namespace of the index. // Additional metadata can also be provided while upserting the vector. func (ix *Index) Upsert(u Upsert) (err error) { - data, err := ix.sendJson(upsertPath, u) + return ix.upsertInternal(u, defaultNamespace) +} + +// UpsertMany updates or inserts some vectors to the default namespace of the index. +// Additional metadata can also be provided for each vector. +func (ix *Index) UpsertMany(u []Upsert) (err error) { + return ix.upsertManyInternal(u, defaultNamespace) +} + +func (ix *Index) upsertInternal(u Upsert, ns string) (err error) { + data, err := ix.sendJson(buildPath(upsertPath, ns), u) if err != nil { return } @@ -13,10 +23,8 @@ func (ix *Index) Upsert(u Upsert) (err error) { return } -// UpsertMany updates or inserts some vectors to the index. -// Additional metadata can also be provided for each vector. -func (ix *Index) UpsertMany(u []Upsert) (err error) { - data, err := ix.sendJson(upsertPath, u) +func (ix *Index) upsertManyInternal(u []Upsert, ns string) (err error) { + data, err := ix.sendJson(buildPath(upsertPath, ns), u) if err != nil { return } diff --git a/upsert_data.go b/upsert_data.go index 38deff2..2e593b1 100644 --- a/upsert_data.go +++ b/upsert_data.go @@ -2,11 +2,22 @@ package vector const upsertDataPath = "/upsert-data" -// UpsertData updates or inserts a vector to the index +// UpsertData updates or inserts a vector to the default namespace of the index // by converting given raw data to an embedding on the server. // Additional metadata can also be provided while upserting the vector. func (ix *Index) UpsertData(u UpsertData) (err error) { - data, err := ix.sendJson(upsertDataPath, u) + return ix.upsertDataInternal(u, defaultNamespace) +} + +// UpsertDataMany updates or inserts some vectors to the default namespace of the index +// by converting given raw data to an embedding on the server. +// Additional metadata can also be provided for each vector. +func (ix *Index) UpsertDataMany(u []UpsertData) (err error) { + return ix.upsertDataManyInternal(u, defaultNamespace) +} + +func (ix *Index) upsertDataInternal(u UpsertData, ns string) (err error) { + data, err := ix.sendJson(buildPath(upsertDataPath, ns), u) if err != nil { return } @@ -14,11 +25,8 @@ func (ix *Index) UpsertData(u UpsertData) (err error) { return } -// UpsertDataMany updates or inserts some vectors to the index -// by converting given raw data to an embedding on the server. -// Additional metadata can also be provided for each vector. -func (ix *Index) UpsertDataMany(u []UpsertData) (err error) { - data, err := ix.sendJson(upsertDataPath, u) +func (ix *Index) upsertDataManyInternal(u []UpsertData, ns string) (err error) { + data, err := ix.sendJson(buildPath(upsertDataPath, ns), u) if err != nil { return } diff --git a/upsert_test.go b/upsert_test.go index 4add9c3..5f0af7d 100644 --- a/upsert_test.go +++ b/upsert_test.go @@ -1,73 +1,76 @@ package vector import ( - "testing" - "github.com/stretchr/testify/require" + "testing" ) func TestUpsert(t *testing.T) { - client, err := newTestClient() - require.NoError(t, err) + for _, ns := range namespaces { + t.Run("namespace_"+ns, func(t *testing.T) { + client, err := newTestClientWithNamespace(ns) + require.NoError(t, err) - t.Run("single", func(t *testing.T) { - id := randomString() - err := client.Upsert(Upsert{ - Id: id, - Vector: []float32{0, 1}, - }) - require.NoError(t, err) + t.Run("single", func(t *testing.T) { + id := randomString() + err := client.Upsert(Upsert{ + Id: id, + Vector: []float32{0, 1}, + }) + require.NoError(t, err) - vectors, err := client.Fetch(Fetch{ - Ids: []string{id}, - }) - require.NoError(t, err) - require.Equal(t, 1, len(vectors)) - require.Equal(t, id, vectors[0].Id) - }) + vectors, err := client.Fetch(Fetch{ + Ids: []string{id}, + }) + require.NoError(t, err) + require.Equal(t, 1, len(vectors)) + require.Equal(t, id, vectors[0].Id) + }) - t.Run("many", func(t *testing.T) { - id0 := randomString() - id1 := randomString() - err = client.UpsertMany([]Upsert{ - { - Id: id0, - Vector: []float32{0, 1}, - }, - { - Id: id1, - Vector: []float32{5, 10}, - }, - }) - require.NoError(t, err) + t.Run("many", func(t *testing.T) { + id0 := randomString() + id1 := randomString() + err = client.UpsertMany([]Upsert{ + { + Id: id0, + Vector: []float32{0, 1}, + }, + { + Id: id1, + Vector: []float32{5, 10}, + }, + }) + require.NoError(t, err) - vectors, err := client.Fetch(Fetch{ - Ids: []string{id0, id1}, - }) - require.NoError(t, err) - require.Equal(t, 2, len(vectors)) - require.Equal(t, id0, vectors[0].Id) - require.Equal(t, id1, vectors[1].Id) - }) + vectors, err := client.Fetch(Fetch{ + Ids: []string{id0, id1}, + }) + require.NoError(t, err) + require.Equal(t, 2, len(vectors)) + require.Equal(t, id0, vectors[0].Id) + require.Equal(t, id1, vectors[1].Id) + }) - t.Run("with metadata", func(t *testing.T) { - id := randomString() - err := client.Upsert(Upsert{ - Id: id, - Vector: []float32{0, 1}, - Metadata: map[string]any{"foo": "bar"}, - }) - require.NoError(t, err) + t.Run("with metadata", func(t *testing.T) { + id := randomString() + err := client.Upsert(Upsert{ + Id: id, + Vector: []float32{0, 1}, + Metadata: map[string]any{"foo": "bar"}, + }) + require.NoError(t, err) - vectors, err := client.Fetch(Fetch{ - Ids: []string{id}, - IncludeMetadata: true, - IncludeVectors: true, + vectors, err := client.Fetch(Fetch{ + Ids: []string{id}, + IncludeMetadata: true, + IncludeVectors: true, + }) + require.NoError(t, err) + require.Equal(t, 1, len(vectors)) + require.Equal(t, id, vectors[0].Id) + require.Equal(t, map[string]any{"foo": "bar"}, vectors[0].Metadata) + require.Equal(t, []float32{0, 1}, vectors[0].Vector) + }) }) - require.NoError(t, err) - require.Equal(t, 1, len(vectors)) - require.Equal(t, id, vectors[0].Id) - require.Equal(t, map[string]any{"foo": "bar"}, vectors[0].Metadata) - require.Equal(t, []float32{0, 1}, vectors[0].Vector) - }) + } }