From 73a1cedc9074ffb74741ef180639912509cf2bac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mehmet=20Tokg=C3=B6z?= Date: Tue, 7 May 2024 23:10:15 +0300 Subject: [PATCH 1/5] Add namespaces support for existing endpoints and implement new ListNamespaces and RemoveNamespace methods. --- delete.go | 4 +- delete_test.go | 91 ++++++++++++++------------ embedding_test.go | 141 ++++++++++++++++++++-------------------- fetch.go | 2 +- fetch_test.go | 116 +++++++++++++++++---------------- index.go | 57 ++++++++++++---- index_test.go | 79 ++++++++++++++++------- info.go | 2 +- info_test.go | 31 +++++---- namespace.go | 38 +++++++++++ namespace_test.go | 65 +++++++++++++++++++ query.go | 2 +- query_data.go | 2 +- query_test.go | 161 +++++++++++++++++++++++----------------------- range.go | 2 +- range_test.go | 127 ++++++++++++++++++------------------ reset.go | 7 +- reset_test.go | 35 +++++----- types.go | 10 +++ upsert.go | 7 +- upsert_data.go | 4 +- upsert_test.go | 119 +++++++++++++++++----------------- 22 files changed, 659 insertions(+), 443 deletions(-) create mode 100644 namespace.go create mode 100644 namespace_test.go diff --git a/delete.go b/delete.go index d9c9559..214156f 100644 --- a/delete.go +++ b/delete.go @@ -6,7 +6,7 @@ const deletePath = "/delete" // 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)) + data, err := ix.sendBytes(deletePath, []byte(id), true) if err != nil { return } @@ -19,7 +19,7 @@ func (ix *Index) Delete(id string) (ok bool, err error) { // 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) + data, err := ix.sendJson(deletePath, ids, true) 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..9f385d3 100644 --- a/embedding_test.go +++ b/embedding_test.go @@ -1,84 +1,87 @@ 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(ns) + 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) + 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) - 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 := 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("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 := 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 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 := 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) + }) + }) + } } diff --git a/fetch.go b/fetch.go index 403a29c..ceb9605 100644 --- a/fetch.go +++ b/fetch.go @@ -7,7 +7,7 @@ const fetchPath = "/fetch" // returned. When f.IncludeMetadata is true, metadata of the vectors // are also returned, if any. func (ix *Index) Fetch(f Fetch) (vectors []Vector, err error) { - data, err := ix.sendJson(fetchPath, f) + data, err := ix.sendJson(fetchPath, f, true) 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..672701e 100644 --- a/index.go +++ b/index.go @@ -4,9 +4,12 @@ import ( "bytes" "encoding/json" "errors" + "fmt" "io" "net/http" + "net/url" "os" + "runtime" ) const ( @@ -59,37 +62,45 @@ 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 + namespace string + headers http.Header } -func (ix *Index) sendJson(path string, obj any) (data []byte, err error) { +func (ix *Index) sendJson(path string, obj any, ns bool) (data []byte, err error) { if data, err = json.Marshal(obj); err != nil { return } - return ix.sendBytes(path, data) + return ix.sendBytes(path, data, ns) } -func (ix *Index) sendBytes(path string, obj []byte) (data []byte, err error) { - return ix.send(path, bytes.NewReader(obj)) +func (ix *Index) sendBytes(path string, obj []byte, ns bool) (data []byte, err error) { + return ix.send(path, bytes.NewReader(obj), ns) } -func (ix *Index) send(path string, r io.Reader) (data []byte, err error) { - request, err := http.NewRequest("POST", ix.url+path, r) +func (ix *Index) send(path string, r io.Reader, ns bool) (data []byte, err error) { + url, err := ix.getUrl(path, ns) + if err != nil { + return nil, err + } + request, err := http.NewRequest(http.MethodPost, url, 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 +121,27 @@ func parseResponse[T any](data []byte) (t T, err error) { } return } + +func (ix *Index) getUrl(path string, ns bool) (result string, err error) { + if ns { + return url.JoinPath(ix.url, path, ix.namespace) + } else { + return url.JoinPath(ix.url, path) + } +} + +func (ix *Index) generateHeaders() { + headers := http.Header{} + headers.Add("Authorization", "Bearer "+ix.token) + headers.Add("Upstash-Telemetry-Runtime", fmt.Sprintf("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 +} diff --git a/index_test.go b/index_test.go index 92cc150..14aff04 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{"", "ns"} ) func init() { @@ -24,26 +27,46 @@ 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 newEmbeddingTestClient() (*Index, error) { +func newTestClientWithNamespace(ns string) (*Index, 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(ns string) (*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 + return client.Namespace(ns), nil } func newTestClientWith(client *http.Client) (*Index, error) { @@ -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..0829321 100644 --- a/info.go +++ b/info.go @@ -7,7 +7,7 @@ const infoPath = "/info" // index on disk in bytes, dimension of the index, and // the name of the similarity function used. func (ix *Index) Info() (info IndexInfo, err error) { - data, err := ix.sendJson(infoPath, nil) + data, err := ix.sendJson(infoPath, nil, false) if err != nil { return } diff --git a/info_test.go b/info_test.go index 48465fe..c53956f 100644 --- a/info_test.go +++ b/info_test.go @@ -1,24 +1,27 @@ package vector import ( - "testing" - "github.com/stretchr/testify/require" + "testing" ) func TestInfo(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) - err = client.Upsert(Upsert{ - Id: randomString(), - Vector: []float32{0, 1}, - }) - require.NoError(t, err) + err = client.Upsert(Upsert{ + Id: randomString(), + Vector: []float32{0, 1}, + }) + require.NoError(t, err) - info, err := client.Info() - require.NoError(t, err) - require.Greater(t, info.VectorCount, 0) - require.Equal(t, 2, info.Dimension) - require.Equal(t, "COSINE", info.SimilarityFunction) + info, err := client.Info() + require.NoError(t, err) + require.Greater(t, info.VectorCount, 0) + require.Equal(t, 2, info.Dimension) + require.Equal(t, "COSINE", info.SimilarityFunction) + }) + } } diff --git a/namespace.go b/namespace.go new file mode 100644 index 0000000..c6a3252 --- /dev/null +++ b/namespace.go @@ -0,0 +1,38 @@ +package vector + +import "errors" + +const deleteNamespacePath = "/delete-namespace" +const listNamespacesPath = "/list-namespaces" + +func (ix *Index) Namespace(ns string) (i *Index) { + return &Index{ + url: ix.url, + token: ix.token, + client: ix.client, + namespace: ns, + headers: ix.headers, + } +} + +// DeleteNamespace deletes the given namespace of index if it exists. +func (ix *Index) DeleteNamespace(ns string) error { + if ns == "" { + return errors.New("cannot delete the default namespace") + } + pns := ix.namespace + ix.namespace = ns + _, err := ix.sendBytes(deleteNamespacePath, nil, true) + ix.namespace = pns + return err +} + +// ListNamespaces returns the list of names of namespaces for the index. +func (ix *Index) ListNamespaces() (namespaces []string, err error) { + data, err := ix.sendJson(listNamespacesPath, nil, false) + if err != nil { + return + } + namespaces, err = parseResponse[[]string](data) + return +} diff --git a/namespace_test.go b/namespace_test.go new file mode 100644 index 0000000..e11ead5 --- /dev/null +++ b/namespace_test.go @@ -0,0 +1,65 @@ +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 + } + + client.DeleteNamespace(ns) + } + + ns, err := client.ListNamespaces() + require.NoError(t, err) + require.Exactly(t, []string{""}, ns) + }) +} + +func createNamespace(t *testing.T, client *Index, ns string) { + if ns == "" { + return + } + + client = client.Namespace(ns) + + err := client.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) + + client.Reset() +} diff --git a/query.go b/query.go index 7cc8015..2bd4612 100644 --- a/query.go +++ b/query.go @@ -10,7 +10,7 @@ const queryPath = "/query" // 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) + data, err := ix.sendJson(queryPath, q, true) if err != nil { return } diff --git a/query_data.go b/query_data.go index 078d577..38a3742 100644 --- a/query_data.go +++ b/query_data.go @@ -11,7 +11,7 @@ const queryDataPath = "/query-data" // 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) + data, err := ix.sendJson(queryDataPath, q, true) if err != nil { return } diff --git a/query_test.go b/query_test.go index eb7e755..226a1aa 100644 --- a/query_test.go +++ b/query_test.go @@ -1,96 +1,99 @@ 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 := newTestClientWithNamespace(ns) + 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) + 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) - 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 := 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("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 := 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) - 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 := 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) - 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 = 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) + }) + }) + } } diff --git a/range.go b/range.go index 58e1ca2..d41d0fc 100644 --- a/range.go +++ b/range.go @@ -10,7 +10,7 @@ const rangePath = "/range" // 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) + data, err := ix.sendJson(rangePath, r, true) 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..5fb5a16 100644 --- a/reset.go +++ b/reset.go @@ -2,10 +2,11 @@ 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 a namespace of the index and resets it to initial state. +// When not specified, the default namespace is used. +// Use Namespace(ns string) method to specify a namespace for the client. func (ix *Index) Reset() (err error) { - data, err := ix.send(resetPath, nil) + data, err := ix.send(resetPath, nil, true) if err != nil { return } diff --git a/reset_test.go b/reset_test.go index 1e8db99..095123e 100644 --- a/reset_test.go +++ b/reset_test.go @@ -1,26 +1,31 @@ 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 := newTestClientWithNamespace(ns) + require.NoError(t, err) - id := randomString() - err = client.Upsert(Upsert{ - Id: id, - Vector: []float32{0, 1}, - }) - require.NoError(t, err) + id := randomString() + err = client.Upsert(Upsert{ + Id: id, + Vector: []float32{0, 1}, + }) + require.NoError(t, err) - err = client.Reset() - require.NoError(t, err) + err = client.Reset() - 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..cee4abf 100644 --- a/types.go +++ b/types.go @@ -139,6 +139,16 @@ type IndexInfo struct { // Name of the similarity function used in indexing and queries. SimilarityFunction string `json:"similarityFunction"` + + NamespaceInfo map[string]NamespaceInfo `json:"namespaces"` +} + +type NamespaceInfo struct { + // The number of vectors in 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..81cef05 100644 --- a/upsert.go +++ b/upsert.go @@ -4,8 +4,11 @@ const upsertPath = "/upsert" // Upsert updates or inserts a vector to the index. // Additional metadata can also be provided while upserting the vector. +// Also, data can be upserted into particular namespaces of the index +// by assigning a namespace with index.Namespace() method. +// When no namespace is provided, the default namespace is used. func (ix *Index) Upsert(u Upsert) (err error) { - data, err := ix.sendJson(upsertPath, u) + data, err := ix.sendJson(upsertPath, u, true) if err != nil { return } @@ -16,7 +19,7 @@ func (ix *Index) Upsert(u Upsert) (err error) { // 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) + data, err := ix.sendJson(upsertPath, u, true) if err != nil { return } diff --git a/upsert_data.go b/upsert_data.go index 38deff2..ed3ee7d 100644 --- a/upsert_data.go +++ b/upsert_data.go @@ -6,7 +6,7 @@ const upsertDataPath = "/upsert-data" // 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) + data, err := ix.sendJson(upsertDataPath, u, true) if err != nil { return } @@ -18,7 +18,7 @@ func (ix *Index) UpsertData(u UpsertData) (err error) { // 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) + data, err := ix.sendJson(upsertDataPath, u, true) 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) - }) + } } From 207f98de4949bd3ebf0402cbbc1b002a9d04b9ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mehmet=20Tokg=C3=B6z?= Date: Wed, 8 May 2024 09:51:16 +0300 Subject: [PATCH 2/5] update documentations. --- README.md | 50 +++++++++++++++++++++++++++++++++++++++++++---- delete.go | 10 +++++----- fetch.go | 8 ++++---- info.go | 11 +++++++---- info_test.go | 41 ++++++++++++++++++++++++++------------ namespace.go | 1 + namespace_test.go | 10 ++++++---- reset.go | 3 +-- types.go | 2 +- upsert.go | 9 ++++----- upsert_data.go | 6 ++++-- 11 files changed, 107 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index f4c28e2..3ccb37d 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,12 @@ # 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. This is the Go client for [Upstash](https://upstash.com/) Vector. @@ -19,10 +22,12 @@ go get 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 +85,23 @@ func main() { Upstash vector indexes support operations for working with vector data using operations such as upsert, query, fetch, and delete. +Index operations are associated with a namespace variable. +You can specify a namespace for an index client with `Namespace(ns string)` function. +If namespace is not specified, default namespace is used. + +```go +import ( + "github.com/upstash/vector-go" +) + +func main() { + index := vector.NewIndex("", "") + + // Returns a new index client associated with namespace "" + ns := index.Namespace("") +} +``` + ### Upserting Vectors All vectors upserted to index must have the same dimensions. @@ -87,7 +109,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 +309,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.DeleteNamespace("ns") +``` \ No newline at end of file diff --git a/delete.go b/delete.go index 214156f..5a06bc2 100644 --- a/delete.go +++ b/delete.go @@ -2,9 +2,9 @@ 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 in a namespace with the given id and reports whether the vector is deleted. +// If a vector with the given id is not found, Delete returns false. +// If namespace is not specified, the default namespace is used. func (ix *Index) Delete(id string) (ok bool, err error) { data, err := ix.sendBytes(deletePath, []byte(id), true) if err != nil { @@ -16,8 +16,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. +// DeleteMany deletes the vectors in a namespace with the given ids and reports how many of them are deleted. +// If namespace is not specified, the default namespace is used. func (ix *Index) DeleteMany(ids []string) (count int, err error) { data, err := ix.sendJson(deletePath, ids, true) if err != nil { diff --git a/fetch.go b/fetch.go index ceb9605..ac0689a 100644 --- a/fetch.go +++ b/fetch.go @@ -2,10 +2,10 @@ 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 a 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. +// If namespace is not specified, the default namespace is used. func (ix *Index) Fetch(f Fetch) (vectors []Vector, err error) { data, err := ix.sendJson(fetchPath, f, true) if err != nil { diff --git a/info.go b/info.go index 0829321..d081e2a 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, false) if err != nil { diff --git a/info_test.go b/info_test.go index c53956f..50c1daa 100644 --- a/info_test.go +++ b/info_test.go @@ -3,25 +3,40 @@ package vector import ( "github.com/stretchr/testify/require" "testing" + "time" ) func TestInfo(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) + createNamespace(t, client, ns) + } - err = client.Upsert(Upsert{ - Id: randomString(), - Vector: []float32{0, 1}, - }) - require.NoError(t, err) + info, err := client.Info() + require.NoError(t, err) + 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.NamespaceInfo)) + for _, ns := range namespaces { + require.Contains(t, info.NamespaceInfo, ns) + require.Equal(t, 0, info.NamespaceInfo[ns].VectorCount) + require.Equal(t, 0, info.NamespaceInfo[ns].PendingVectorCount) + } - info, err := client.Info() - require.NoError(t, err) - require.Greater(t, info.VectorCount, 0) - require.Equal(t, 2, info.Dimension) - require.Equal(t, "COSINE", info.SimilarityFunction) + 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 index c6a3252..51c6c20 100644 --- a/namespace.go +++ b/namespace.go @@ -5,6 +5,7 @@ import "errors" const deleteNamespacePath = "/delete-namespace" const listNamespacesPath = "/list-namespaces" +// Namespace returns a new Index client for a specific namespace. func (ix *Index) Namespace(ns string) (i *Index) { return &Index{ url: ix.url, diff --git a/namespace_test.go b/namespace_test.go index e11ead5..09dba34 100644 --- a/namespace_test.go +++ b/namespace_test.go @@ -34,12 +34,13 @@ func TestNamespace(t *testing.T) { continue } - client.DeleteNamespace(ns) + err := client.DeleteNamespace(ns) + require.NoError(t, err) } - ns, err := client.ListNamespaces() + info, err := client.Info() require.NoError(t, err) - require.Exactly(t, []string{""}, ns) + require.Exactly(t, 1, len(info.NamespaceInfo)) }) } @@ -61,5 +62,6 @@ func createNamespace(t *testing.T, client *Index, ns string) { return info.PendingVectorCount == 0 }, 10*time.Second, 1*time.Second) - client.Reset() + err = client.Reset() + require.NoError(t, err) } diff --git a/reset.go b/reset.go index 5fb5a16..a584132 100644 --- a/reset.go +++ b/reset.go @@ -3,8 +3,7 @@ package vector const resetPath = "/reset" // Reset deletes all the vectors in a namespace of the index and resets it to initial state. -// When not specified, the default namespace is used. -// Use Namespace(ns string) method to specify a namespace for the client. +// If namespace is not specified, the default namespace is used. func (ix *Index) Reset() (err error) { data, err := ix.send(resetPath, nil, true) if err != nil { diff --git a/types.go b/types.go index cee4abf..e36b867 100644 --- a/types.go +++ b/types.go @@ -144,7 +144,7 @@ type IndexInfo struct { } type NamespaceInfo struct { - // The number of vectors in the index. + // The number of vectors in the namespace of the index. VectorCount int `json:"vectorCount"` // The number of vectors that are pending to be indexed. diff --git a/upsert.go b/upsert.go index 81cef05..c4dffa3 100644 --- a/upsert.go +++ b/upsert.go @@ -2,11 +2,9 @@ package vector const upsertPath = "/upsert" -// Upsert updates or inserts a vector to the index. +// Upsert updates or inserts a vector to a namespace of the index. // Additional metadata can also be provided while upserting the vector. -// Also, data can be upserted into particular namespaces of the index -// by assigning a namespace with index.Namespace() method. -// When no namespace is provided, the default namespace is used. +// If namespace is not specified, the default namespace is used. func (ix *Index) Upsert(u Upsert) (err error) { data, err := ix.sendJson(upsertPath, u, true) if err != nil { @@ -16,8 +14,9 @@ func (ix *Index) Upsert(u Upsert) (err error) { return } -// UpsertMany updates or inserts some vectors to the index. +// UpsertMany updates or inserts some vectors to a namespace of the index. // Additional metadata can also be provided for each vector. +// If namespace is not specified, the default namespace is used. func (ix *Index) UpsertMany(u []Upsert) (err error) { data, err := ix.sendJson(upsertPath, u, true) if err != nil { diff --git a/upsert_data.go b/upsert_data.go index ed3ee7d..265f62d 100644 --- a/upsert_data.go +++ b/upsert_data.go @@ -2,9 +2,10 @@ package vector const upsertDataPath = "/upsert-data" -// UpsertData updates or inserts a vector to the index +// UpsertData updates or inserts a vector to a 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. +// If namespace is not specified, the default namespace is used. func (ix *Index) UpsertData(u UpsertData) (err error) { data, err := ix.sendJson(upsertDataPath, u, true) if err != nil { @@ -14,9 +15,10 @@ func (ix *Index) UpsertData(u UpsertData) (err error) { return } -// UpsertDataMany updates or inserts some vectors to the index +// UpsertDataMany updates or inserts some vectors to a namespace of the index // by converting given raw data to an embedding on the server. // Additional metadata can also be provided for each vector. +// If namespace is not specified, the default namespace is used. func (ix *Index) UpsertDataMany(u []UpsertData) (err error) { data, err := ix.sendJson(upsertDataPath, u, true) if err != nil { From 8dd878569c117ea9fa662211a9eaea113beaa863 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mehmet=20Tokg=C3=B6z?= Date: Wed, 8 May 2024 10:17:05 +0300 Subject: [PATCH 3/5] fix build error. --- reset_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/reset_test.go b/reset_test.go index 095123e..d37e2fa 100644 --- a/reset_test.go +++ b/reset_test.go @@ -20,6 +20,7 @@ func TestReset(t *testing.T) { require.NoError(t, err) err = client.Reset() + require.NoError(t, err) require.Eventually(t, func() bool { info, err := client.Info() From e9f80f834f2b13ebd391bd31bc013305e6dc9a59 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mehmet=20Tokg=C3=B6z?= Date: Fri, 10 May 2024 09:10:51 +0300 Subject: [PATCH 4/5] Create new Namespace client for per-ns operations. --- README.md | 33 ++++++++++--- delete.go | 20 +++++--- embedding_test.go | 11 +++-- fetch.go | 9 ++-- index.go | 34 ++++++-------- index_test.go | 6 +-- info.go | 2 +- info_test.go | 8 ++-- namespace.go | 117 +++++++++++++++++++++++++++++++++++++--------- namespace_test.go | 10 ++-- query.go | 19 ++++---- query_data.go | 20 ++++---- query_test.go | 13 +++--- range.go | 9 ++-- reset.go | 9 ++-- reset_test.go | 7 +-- types.go | 3 +- upsert.go | 22 +++++---- upsert_data.go | 24 ++++++---- 19 files changed, 250 insertions(+), 126 deletions(-) diff --git a/README.md b/README.md index 3ccb37d..0a86b03 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,9 @@ > 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. -This is the Go client for [Upstash](https://upstash.com/) Vector. +[Upstash](https://upstash.com/) Vector is a serverless vector database designed for working with vector embeddings. + +This is the HTTP-based Go client for [Upstash](https://upstash.com/) Vector. ## Documentation @@ -16,10 +18,17 @@ 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. @@ -85,9 +94,22 @@ func main() { Upstash vector indexes support operations for working with vector data using operations such as upsert, query, fetch, and delete. -Index operations are associated with a namespace variable. +```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. -If namespace is not specified, default namespace is used. +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 ( @@ -97,9 +119,8 @@ import ( func main() { index := vector.NewIndex("", "") - // Returns a new index client associated with namespace "" + // Returns a new Namespace client associated with the given namespace ns := index.Namespace("") -} ``` ### Upserting Vectors @@ -327,5 +348,5 @@ A namespace can be deleted entirely if it exists. The default namespaces cannot be deleted. ```go -err := index.DeleteNamespace("ns") +err := index.Namespace("ns").DeleteNamespace() ``` \ No newline at end of file diff --git a/delete.go b/delete.go index 5a06bc2..bb49352 100644 --- a/delete.go +++ b/delete.go @@ -2,11 +2,19 @@ package vector const deletePath = "/delete" -// Delete deletes the vector in a namespace with the given id and reports whether the vector is deleted. +// 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. -// If namespace is not specified, the default namespace is used. func (ix *Index) Delete(id string) (ok bool, err error) { - data, err := ix.sendBytes(deletePath, []byte(id), true) + return ix.deleteInternal(id, "") +} + +// 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, "") +} + +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 in a namespace with the given ids and reports how many of them are deleted. -// If namespace is not specified, the default namespace is used. -func (ix *Index) DeleteMany(ids []string) (count int, err error) { - data, err := ix.sendJson(deletePath, ids, true) +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/embedding_test.go b/embedding_test.go index 9f385d3..2e2bfa4 100644 --- a/embedding_test.go +++ b/embedding_test.go @@ -9,14 +9,15 @@ import ( func TestEmbedding(t *testing.T) { for _, ns := range namespaces { t.Run("namespace_"+ns, func(t *testing.T) { - client, err := newEmbeddingTestClient(ns) + client, err := newEmbeddingTestClient() require.NoError(t, err) + namespace := client.Namespace(ns) id0 := "tr" id1 := "jp" id2 := "uk" id3 := "fr" - err = client.UpsertDataMany([]UpsertData{ + err = namespace.UpsertDataMany([]UpsertData{ { Id: id0, Data: "Capital of Türkiye is Ankara.", @@ -47,7 +48,7 @@ func TestEmbedding(t *testing.T) { }, 10*time.Second, 1*time.Second) t.Run("score", func(t *testing.T) { - scores, err := client.QueryData(QueryData{ + scores, err := namespace.QueryData(QueryData{ Data: "where is the capital of Japan?", TopK: 1, }) @@ -57,7 +58,7 @@ func TestEmbedding(t *testing.T) { }) t.Run("with metadata", func(t *testing.T) { - scores, err := client.QueryData(QueryData{ + scores, err := namespace.QueryData(QueryData{ Data: "Which country's capital is Ankara?", TopK: 1, IncludeMetadata: true, @@ -76,7 +77,7 @@ func TestEmbedding(t *testing.T) { Filter: `country = 'fr'`, } - scores, err := client.QueryData(query) + scores, err := namespace.QueryData(query) require.NoError(t, err) require.Equal(t, 1, len(scores)) require.Equal(t, id3, scores[0].Id) diff --git a/fetch.go b/fetch.go index ac0689a..768712d 100644 --- a/fetch.go +++ b/fetch.go @@ -2,12 +2,15 @@ package vector const fetchPath = "/fetch" -// Fetch fetches one or more vectors in a namespace with the ids passed into f. +// 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. -// If namespace is not specified, the default namespace is used. func (ix *Index) Fetch(f Fetch) (vectors []Vector, err error) { - data, err := ix.sendJson(fetchPath, f, true) + return ix.fetchInternal(f, "") +} + +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/index.go b/index.go index 672701e..75cc4d9 100644 --- a/index.go +++ b/index.go @@ -7,7 +7,6 @@ import ( "fmt" "io" "net/http" - "net/url" "os" "runtime" ) @@ -80,23 +79,19 @@ type Index struct { headers http.Header } -func (ix *Index) sendJson(path string, obj any, ns bool) (data []byte, err error) { +func (ix *Index) sendJson(path string, obj any) (data []byte, err error) { if data, err = json.Marshal(obj); err != nil { return } - return ix.sendBytes(path, data, ns) + return ix.sendBytes(path, data) } -func (ix *Index) sendBytes(path string, obj []byte, ns bool) (data []byte, err error) { - return ix.send(path, bytes.NewReader(obj), ns) +func (ix *Index) sendBytes(path string, obj []byte) (data []byte, err error) { + return ix.send(path, bytes.NewReader(obj)) } -func (ix *Index) send(path string, r io.Reader, ns bool) (data []byte, err error) { - url, err := ix.getUrl(path, ns) - if err != nil { - return nil, err - } - request, err := http.NewRequest(http.MethodPost, url, r) +func (ix *Index) send(path string, r io.Reader) (data []byte, err error) { + request, err := http.NewRequest(http.MethodPost, ix.url+path, r) if err != nil { return } @@ -122,18 +117,10 @@ func parseResponse[T any](data []byte) (t T, err error) { return } -func (ix *Index) getUrl(path string, ns bool) (result string, err error) { - if ns { - return url.JoinPath(ix.url, path, ix.namespace) - } else { - return url.JoinPath(ix.url, path) - } -} - func (ix *Index) generateHeaders() { headers := http.Header{} headers.Add("Authorization", "Bearer "+ix.token) - headers.Add("Upstash-Telemetry-Runtime", fmt.Sprintf("go@%s", runtime.Version())) + headers.Add("Upstash-Telemetry-Runtime", fmt.Sprintf("vector-go@%s", runtime.Version())) var platform string if os.Getenv("VERCEL") != "" { platform = "vercel" @@ -145,3 +132,10 @@ func (ix *Index) generateHeaders() { 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 14aff04..b7f3ef3 100644 --- a/index_test.go +++ b/index_test.go @@ -37,7 +37,7 @@ func newTestClient() (*Index, error) { return client, nil } -func newTestClientWithNamespace(ns string) (*Index, error) { +func newTestClientWithNamespace(ns string) (*Namespace, error) { client := NewIndex( os.Getenv(UrlEnvProperty), os.Getenv(TokenEnvProperty), @@ -53,7 +53,7 @@ func newTestClientWithNamespace(ns string) (*Index, error) { return client.Namespace(ns), nil } -func newEmbeddingTestClient(ns string) (*Index, error) { +func newEmbeddingTestClient() (*Index, error) { client := NewIndex( os.Getenv("EMBEDDING_"+UrlEnvProperty), os.Getenv("EMBEDDING_"+TokenEnvProperty), @@ -66,7 +66,7 @@ func newEmbeddingTestClient(ns string) (*Index, error) { } } - return client.Namespace(ns), nil + return client, nil } func newTestClientWith(client *http.Client) (*Index, error) { diff --git a/info.go b/info.go index d081e2a..a71752a 100644 --- a/info.go +++ b/info.go @@ -10,7 +10,7 @@ const infoPath = "/info" // - Similarity function used // - per-namespace vector and pending vector counts func (ix *Index) Info() (info IndexInfo, err error) { - data, err := ix.sendJson(infoPath, nil, false) + data, err := ix.sendJson(infoPath, nil) if err != nil { return } diff --git a/info_test.go b/info_test.go index 50c1daa..4869847 100644 --- a/info_test.go +++ b/info_test.go @@ -19,11 +19,11 @@ func TestInfo(t *testing.T) { 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.NamespaceInfo)) + require.Equal(t, len(namespaces), len(info.Namespaces)) for _, ns := range namespaces { - require.Contains(t, info.NamespaceInfo, ns) - require.Equal(t, 0, info.NamespaceInfo[ns].VectorCount) - require.Equal(t, 0, info.NamespaceInfo[ns].PendingVectorCount) + 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 { diff --git a/namespace.go b/namespace.go index 51c6c20..03f68f0 100644 --- a/namespace.go +++ b/namespace.go @@ -1,39 +1,112 @@ package vector -import "errors" - const deleteNamespacePath = "/delete-namespace" const listNamespacesPath = "/list-namespaces" -// Namespace returns a new Index client for a specific namespace. -func (ix *Index) Namespace(ns string) (i *Index) { - return &Index{ - url: ix.url, - token: ix.token, - client: ix.client, - namespace: ns, - headers: ix.headers, - } +type Namespace struct { + index *Index + ns string } -// DeleteNamespace deletes the given namespace of index if it exists. -func (ix *Index) DeleteNamespace(ns string) error { - if ns == "" { - return errors.New("cannot delete the default namespace") +// Namespace returns a new client associated with the given namespace. +func (ix *Index) Namespace(namespace string) (i *Namespace) { + return &Namespace{ + index: ix, + ns: namespace, } - pns := ix.namespace - ix.namespace = ns - _, err := ix.sendBytes(deleteNamespacePath, nil, true) - ix.namespace = pns - return err } -// ListNamespaces returns the list of names of namespaces for the index. +// ListNamespaces returns the list of names of namespaces. func (ix *Index) ListNamespaces() (namespaces []string, err error) { - data, err := ix.sendJson(listNamespacesPath, nil, false) + 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 index 09dba34..cf87ebd 100644 --- a/namespace_test.go +++ b/namespace_test.go @@ -34,13 +34,13 @@ func TestNamespace(t *testing.T) { continue } - err := client.DeleteNamespace(ns) + err := client.Namespace(ns).DeleteNamespace() require.NoError(t, err) } info, err := client.Info() require.NoError(t, err) - require.Exactly(t, 1, len(info.NamespaceInfo)) + require.Exactly(t, 1, len(info.Namespaces)) }) } @@ -49,9 +49,9 @@ func createNamespace(t *testing.T, client *Index, ns string) { return } - client = client.Namespace(ns) + namespace := client.Namespace(ns) - err := client.Upsert(Upsert{ + err := namespace.Upsert(Upsert{ Vector: []float32{0.1, 0.1}, }) require.NoError(t, err) @@ -62,6 +62,6 @@ func createNamespace(t *testing.T, client *Index, ns string) { return info.PendingVectorCount == 0 }, 10*time.Second, 1*time.Second) - err = client.Reset() + err = namespace.Reset() require.NoError(t, err) } diff --git a/query.go b/query.go index 2bd4612..8651ecd 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, true) + return ix.queryInternal(q, "") +} + +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 38a3742..b3e8594 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, true) + return ix.queryDataInternal(q, "") +} + +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 226a1aa..63b17b5 100644 --- a/query_test.go +++ b/query_test.go @@ -9,13 +9,14 @@ import ( func TestQuery(t *testing.T) { for _, ns := range namespaces { t.Run("namespace_"+ns, func(t *testing.T) { - client, err := newTestClientWithNamespace(ns) + client, err := newTestClient() require.NoError(t, err) + namespace := client.Namespace(ns) id0 := randomString() id1 := randomString() id2 := randomString() - err = client.UpsertMany([]Upsert{ + err = namespace.UpsertMany([]Upsert{ { Id: id0, Vector: []float32{0, 1}, @@ -40,7 +41,7 @@ func TestQuery(t *testing.T) { }, 10*time.Second, 1*time.Second) t.Run("score", func(t *testing.T) { - scores, err := client.Query(Query{ + scores, err := namespace.Query(Query{ Vector: []float32{0, 1}, TopK: 2, }) @@ -52,7 +53,7 @@ func TestQuery(t *testing.T) { }) t.Run("with metadata and vectors", func(t *testing.T) { - scores, err := client.Query(Query{ + scores, err := namespace.Query(Query{ Vector: []float32{0, 1}, TopK: 2, IncludeMetadata: true, @@ -78,7 +79,7 @@ func TestQuery(t *testing.T) { Filter: `foo = 'bar'`, } - scores, err := client.Query(query) + scores, err := namespace.Query(query) require.NoError(t, err) require.Equal(t, 1, len(scores)) require.Equal(t, id0, scores[0].Id) @@ -87,7 +88,7 @@ func TestQuery(t *testing.T) { require.Equal(t, []float32{0, 1}, scores[0].Vector) query.Filter = `foo = 'nay'` - scores, err = client.Query(query) + scores, err = namespace.Query(query) require.NoError(t, err) require.Equal(t, 1, len(scores)) require.Equal(t, id2, scores[0].Id) diff --git a/range.go b/range.go index d41d0fc..8e03b36 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, true) + return ix.rangeInternal(r, "") +} + +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/reset.go b/reset.go index a584132..6fcf363 100644 --- a/reset.go +++ b/reset.go @@ -2,10 +2,13 @@ package vector const resetPath = "/reset" -// Reset deletes all the vectors in a namespace of the index and resets it to initial state. -// If namespace is not specified, the default namespace is used. +// 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, true) + return ix.resetInternal("") +} + +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 d37e2fa..5edbb4a 100644 --- a/reset_test.go +++ b/reset_test.go @@ -9,17 +9,18 @@ import ( func TestReset(t *testing.T) { for _, ns := range namespaces { t.Run("namespace_"+ns, func(t *testing.T) { - client, err := newTestClientWithNamespace(ns) + client, err := newTestClient() require.NoError(t, err) + namespace := client.Namespace(ns) id := randomString() - err = client.Upsert(Upsert{ + err = namespace.Upsert(Upsert{ Id: id, Vector: []float32{0, 1}, }) require.NoError(t, err) - err = client.Reset() + err = namespace.Reset() require.NoError(t, err) require.Eventually(t, func() bool { diff --git a/types.go b/types.go index e36b867..b2dfdeb 100644 --- a/types.go +++ b/types.go @@ -140,7 +140,8 @@ type IndexInfo struct { // Name of the similarity function used in indexing and queries. SimilarityFunction string `json:"similarityFunction"` - NamespaceInfo map[string]NamespaceInfo `json:"namespaces"` + // Per-namespace vector and pending vector counts + Namespaces map[string]NamespaceInfo `json:"namespaces"` } type NamespaceInfo struct { diff --git a/upsert.go b/upsert.go index c4dffa3..d31aed3 100644 --- a/upsert.go +++ b/upsert.go @@ -2,11 +2,20 @@ package vector const upsertPath = "/upsert" -// Upsert updates or inserts a vector to a namespace of 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. -// If namespace is not specified, the default namespace is used. func (ix *Index) Upsert(u Upsert) (err error) { - data, err := ix.sendJson(upsertPath, u, true) + return ix.upsertInternal(u, "") +} + +// 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, "") +} + +func (ix *Index) upsertInternal(u Upsert, ns string) (err error) { + data, err := ix.sendJson(buildPath(upsertPath, ns), u) if err != nil { return } @@ -14,11 +23,8 @@ func (ix *Index) Upsert(u Upsert) (err error) { return } -// UpsertMany updates or inserts some vectors to a namespace of the index. -// Additional metadata can also be provided for each vector. -// If namespace is not specified, the default namespace is used. -func (ix *Index) UpsertMany(u []Upsert) (err error) { - data, err := ix.sendJson(upsertPath, u, true) +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 265f62d..433e788 100644 --- a/upsert_data.go +++ b/upsert_data.go @@ -2,12 +2,22 @@ package vector const upsertDataPath = "/upsert-data" -// UpsertData updates or inserts a vector to a namespace of 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. -// If namespace is not specified, the default namespace is used. func (ix *Index) UpsertData(u UpsertData) (err error) { - data, err := ix.sendJson(upsertDataPath, u, true) + return ix.upsertDataInternal(u, "") +} + +// 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, "") +} + +func (ix *Index) upsertDataInternal(u UpsertData, ns string) (err error) { + data, err := ix.sendJson(buildPath(upsertDataPath, ns), u) if err != nil { return } @@ -15,12 +25,8 @@ func (ix *Index) UpsertData(u UpsertData) (err error) { return } -// UpsertDataMany updates or inserts some vectors to a namespace of the index -// by converting given raw data to an embedding on the server. -// Additional metadata can also be provided for each vector. -// If namespace is not specified, the default namespace is used. -func (ix *Index) UpsertDataMany(u []UpsertData) (err error) { - data, err := ix.sendJson(upsertDataPath, u, true) +func (ix *Index) upsertDataManyInternal(u []UpsertData, ns string) (err error) { + data, err := ix.sendJson(buildPath(upsertDataPath, ns), u) if err != nil { return } From b160897efcc9490a061fe2860eb011a8e0fe2487 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mehmet=20Tokg=C3=B6z?= Date: Fri, 10 May 2024 09:20:06 +0300 Subject: [PATCH 5/5] Remove namespace field on Index. --- delete.go | 4 ++-- fetch.go | 2 +- index.go | 10 +++++----- index_test.go | 2 +- query.go | 2 +- query_data.go | 2 +- range.go | 2 +- reset.go | 2 +- upsert.go | 4 ++-- upsert_data.go | 4 ++-- 10 files changed, 17 insertions(+), 17 deletions(-) diff --git a/delete.go b/delete.go index bb49352..0a5aa38 100644 --- a/delete.go +++ b/delete.go @@ -5,12 +5,12 @@ const deletePath = "/delete" // 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) { - return ix.deleteInternal(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, "") + return ix.deleteManyInternal(ids, defaultNamespace) } func (ix *Index) deleteInternal(id string, ns string) (ok bool, err error) { diff --git a/fetch.go b/fetch.go index 768712d..19dd526 100644 --- a/fetch.go +++ b/fetch.go @@ -6,7 +6,7 @@ const fetchPath = "/fetch" // 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) { - return ix.fetchInternal(f, "") + return ix.fetchInternal(f, defaultNamespace) } func (ix *Index) fetchInternal(f Fetch, ns string) (vectors []Vector, err error) { diff --git a/index.go b/index.go index 75cc4d9..2e69afc 100644 --- a/index.go +++ b/index.go @@ -14,6 +14,7 @@ import ( const ( UrlEnvProperty = "UPSTASH_VECTOR_REST_URL" TokenEnvProperty = "UPSTASH_VECTOR_REST_TOKEN" + defaultNamespace = "" ) type Options struct { @@ -72,11 +73,10 @@ func NewIndexWith(options Options) *Index { // Index is a client for Upstash Vector index. type Index struct { - url string - token string - client *http.Client - namespace string - headers http.Header + url string + token string + client *http.Client + headers http.Header } func (ix *Index) sendJson(path string, obj any) (data []byte, err error) { diff --git a/index_test.go b/index_test.go index b7f3ef3..9909d40 100644 --- a/index_test.go +++ b/index_test.go @@ -11,7 +11,7 @@ import ( ) var ( - namespaces = [...]string{"", "ns"} + namespaces = [...]string{defaultNamespace, "ns"} ) func init() { diff --git a/query.go b/query.go index 8651ecd..0b20f04 100644 --- a/query.go +++ b/query.go @@ -9,7 +9,7 @@ const queryPath = "/query" // 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) { - return ix.queryInternal(q, "") + return ix.queryInternal(q, defaultNamespace) } func (ix *Index) queryInternal(q Query, ns string) (scores []VectorScore, err error) { diff --git a/query_data.go b/query_data.go index b3e8594..5ccdd3f 100644 --- a/query_data.go +++ b/query_data.go @@ -9,7 +9,7 @@ const queryDataPath = "/query-data" // 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) { - return ix.queryDataInternal(q, "") + return ix.queryDataInternal(q, defaultNamespace) } func (ix *Index) queryDataInternal(q QueryData, ns string) (scores []VectorScore, err error) { diff --git a/range.go b/range.go index 8e03b36..f367b36 100644 --- a/range.go +++ b/range.go @@ -9,7 +9,7 @@ const rangePath = "/range" // 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 (ix *Index) Range(r Range) (vectors RangeVectors, err error) { - return ix.rangeInternal(r, "") + return ix.rangeInternal(r, defaultNamespace) } func (ix *Index) rangeInternal(r Range, ns string) (vectors RangeVectors, err error) { diff --git a/reset.go b/reset.go index 6fcf363..85b97ca 100644 --- a/reset.go +++ b/reset.go @@ -4,7 +4,7 @@ const resetPath = "/reset" // Reset deletes all the vectors in the default namespace of the index and resets it to initial state. func (ix *Index) Reset() (err error) { - return ix.resetInternal("") + return ix.resetInternal(defaultNamespace) } func (ix *Index) resetInternal(ns string) (err error) { diff --git a/upsert.go b/upsert.go index d31aed3..fd7b339 100644 --- a/upsert.go +++ b/upsert.go @@ -5,13 +5,13 @@ const upsertPath = "/upsert" // 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) { - return ix.upsertInternal(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, "") + return ix.upsertManyInternal(u, defaultNamespace) } func (ix *Index) upsertInternal(u Upsert, ns string) (err error) { diff --git a/upsert_data.go b/upsert_data.go index 433e788..2e593b1 100644 --- a/upsert_data.go +++ b/upsert_data.go @@ -6,14 +6,14 @@ const upsertDataPath = "/upsert-data" // 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) { - return ix.upsertDataInternal(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, "") + return ix.upsertDataManyInternal(u, defaultNamespace) } func (ix *Index) upsertDataInternal(u UpsertData, ns string) (err error) {