From 9176c94c094a1d4fae802edc86f40ac5e980c6cd Mon Sep 17 00:00:00 2001 From: Mehmet Dogan Date: Wed, 30 Oct 2024 14:09:45 +0300 Subject: [PATCH] Add data support to upsert/query/fetch operations --- .github/workflows/build.yml | 2 +- README.md | 17 ++++++++++++++++- embedding_test.go | 7 ++++++- go.mod | 2 +- query_test.go | 12 +++++++++++- types.go | 21 +++++++++++++++++++++ upsert_test.go | 6 +++++- 7 files changed, 61 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index ee70128..dbc6fe6 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,7 +24,7 @@ jobs: uses: actions/checkout@v4 - name: Setup Go - uses: actions/setup-go@v4 + uses: actions/setup-go@v5 with: go-version-file: 'go.mod' cache-dependency-path: 'go.sum' diff --git a/README.md b/README.md index 0a86b03..aa9709b 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,7 @@ upserts := []vector.Upsert{ Id: "1", Vector: []float32{0.0, 1.0}, Metadata: map[string]any{"foo": "bar"}, // optional metadata + Data: "vector data", // optional data }, } @@ -205,12 +206,15 @@ When `IncludeVectors` is `true`, the response will contain the vector values. When `IncludeMetadata` is `true`, the response will contain the metadata of the vectors, if any. +When `IncludeData` is `true`, the response will contain the data of the vectors, if any. + ```go scores, err := index.Query(vector.Query{ Vector: []float32{0.0, 1.0}, TopK: 2, IncludeVectors: false, IncludeMetadata: false, + IncludeData: false, }) ``` @@ -224,6 +228,7 @@ scores, err := index.Query(vector.Query{ TopK: 2, IncludeVectors: false, IncludeMetadata: false, + IncludeData: false, Filter: `foo = 'bar'` }) ``` @@ -240,12 +245,15 @@ When `IncludeVectors` is `true`, the response will contain the vector values. When `IncludeMetadata` is `true`, the response will contain the metadata of the vectors, if any. +When `IncludeData` is `true`, the response will contain the data of the vectors, if any. + ```go scores, err := index.QueryData(vector.QueryData{ Data: "Where is the capital of Turkey?", TopK: 2, IncludeVectors: false, IncludeMetadata: false, + IncludeData: false, Filter: `foo = 'bar'` }) ``` @@ -259,11 +267,14 @@ When `IncludeVectors` is `true`, the response will contain the vector values. When `IncludeMetadata` is `true`, the response will contain the metadata of the vectors, if any. +When `IncludeData` is `true`, the response will contain the data of the vectors, if any. + ```go vectors, err := index.Fetch(vector.Fetch{ Ids: []string{"0", "1"}, IncludeVectors: false, IncludeMetadata: false, + IncludeData: false, }) ``` @@ -294,12 +305,15 @@ When `IncludeVectors` is `true`, the response will contain the vector values. When `IncludeMetadata` is `true`, the response will contain the metadata of the vectors, if any. +When `IncludeData` is `true`, the response will contain the data of the vectors, if any. + ```go vectors, err := index.Range(vector.Range{ Cursor: "0", Limit: 10, IncludeVectors: false, IncludeMetadata: false, + IncludeData: false, }) for vectors.NextCursor != "" { @@ -313,6 +327,7 @@ for vectors.NextCursor != "" { Limit: 10, IncludeVectors: false, IncludeMetadata: false, + IncludeData: false, }) } ``` @@ -349,4 +364,4 @@ The default namespaces cannot be deleted. ```go err := index.Namespace("ns").DeleteNamespace() -``` \ No newline at end of file +``` diff --git a/embedding_test.go b/embedding_test.go index 2e2bfa4..2ee90cd 100644 --- a/embedding_test.go +++ b/embedding_test.go @@ -1,9 +1,10 @@ package vector import ( - "github.com/stretchr/testify/require" "testing" "time" + + "github.com/stretchr/testify/require" ) func TestEmbedding(t *testing.T) { @@ -62,11 +63,13 @@ func TestEmbedding(t *testing.T) { Data: "Which country's capital is Ankara?", TopK: 1, IncludeMetadata: true, + IncludeData: 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) + require.Equal(t, "Capital of Türkiye is Ankara.", scores[0].Data) }) t.Run("with metadata filtering", func(t *testing.T) { @@ -74,6 +77,7 @@ func TestEmbedding(t *testing.T) { Data: "Where is the capital of France?", TopK: 1, IncludeMetadata: true, + IncludeData: true, Filter: `country = 'fr'`, } @@ -82,6 +86,7 @@ func TestEmbedding(t *testing.T) { 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) + require.Equal(t, "Capital of France is Paris.", scores[0].Data) }) }) } diff --git a/go.mod b/go.mod index e174c80..a716ee9 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/upstash/vector-go -go 1.20 +go 1.22 require ( github.com/joho/godotenv v1.5.1 diff --git a/query_test.go b/query_test.go index 63b17b5..0ea8d67 100644 --- a/query_test.go +++ b/query_test.go @@ -1,9 +1,10 @@ package vector import ( - "github.com/stretchr/testify/require" "testing" "time" + + "github.com/stretchr/testify/require" ) func TestQuery(t *testing.T) { @@ -21,10 +22,12 @@ func TestQuery(t *testing.T) { Id: id0, Vector: []float32{0, 1}, Metadata: map[string]any{"foo": "bar"}, + Data: "vector0 data", }, { Id: id1, Vector: []float32{5, 10}, + Data: "vector1 data", }, { Id: id2, @@ -58,16 +61,20 @@ func TestQuery(t *testing.T) { TopK: 2, IncludeMetadata: true, IncludeVectors: true, + IncludeData: 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, "vector0 data", scores[0].Data) 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, map[string]any{"foo": "nay"}, scores[1].Metadata) + require.Empty(t, scores[1].Data) }) t.Run("with metadata filtering", func(t *testing.T) { @@ -76,6 +83,7 @@ func TestQuery(t *testing.T) { TopK: 10, IncludeMetadata: true, IncludeVectors: true, + IncludeData: true, Filter: `foo = 'bar'`, } @@ -85,6 +93,7 @@ func TestQuery(t *testing.T) { 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, "vector0 data", scores[0].Data) require.Equal(t, []float32{0, 1}, scores[0].Vector) query.Filter = `foo = 'nay'` @@ -94,6 +103,7 @@ func TestQuery(t *testing.T) { 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) + require.Empty(t, scores[0].Data) }) }) } diff --git a/types.go b/types.go index b2dfdeb..269a4c1 100644 --- a/types.go +++ b/types.go @@ -7,6 +7,9 @@ type Upsert struct { // Vector values. Vector []float32 `json:"vector"` + // Optional data of the vector. + Data string `json:"data,omitempty"` + // Optional metadata of the vector. Metadata map[string]any `json:"metadata,omitempty"` } @@ -37,6 +40,9 @@ type Query struct { // Whether to include metadata in the query response, if any. IncludeMetadata bool `json:"includeMetadata,omitempty"` + // Whether to include data in the query response, if any. + IncludeData bool `json:"includeData,omitempty"` + // Query filter Filter any `json:"filter,omitempty"` } @@ -56,6 +62,9 @@ type QueryData struct { // Whether to include metadata in the query response, if any. IncludeMetadata bool `json:"includeMetadata,omitempty"` + // Whether to include data in the query response, if any. + IncludeData bool `json:"includeData,omitempty"` + // Query filter Filter any `json:"filter,omitempty"` } @@ -69,6 +78,9 @@ type Vector struct { // Optional metadata of the vector, if any. Metadata map[string]any `json:"metadata,omitempty"` + + // Optional data of the vector. + Data string `json:"data,omitempty"` } type VectorScore struct { @@ -84,6 +96,9 @@ type VectorScore struct { // Optional metadata of the vector, if any. Metadata map[string]any `json:"metadata,omitempty"` + + // Optional data of the vector. + Data string `json:"data,omitempty"` } type Fetch struct { @@ -95,6 +110,9 @@ type Fetch struct { // Whether to include metadata in the fetch response, if any. IncludeMetadata bool `json:"includeMetadata,omitempty"` + + // Whether to include data in the query response, if any. + IncludeData bool `json:"includeData,omitempty"` } type Range struct { @@ -110,6 +128,9 @@ type Range struct { // Whether to include metadata in the fetch response, if any. IncludeMetadata bool `json:"includeMetadata,omitempty"` + + // Whether to include data in the query response, if any. + IncludeData bool `json:"includeData,omitempty"` } type RangeVectors struct { diff --git a/upsert_test.go b/upsert_test.go index 5f0af7d..b3d70ce 100644 --- a/upsert_test.go +++ b/upsert_test.go @@ -1,8 +1,9 @@ package vector import ( - "github.com/stretchr/testify/require" "testing" + + "github.com/stretchr/testify/require" ) func TestUpsert(t *testing.T) { @@ -57,6 +58,7 @@ func TestUpsert(t *testing.T) { Id: id, Vector: []float32{0, 1}, Metadata: map[string]any{"foo": "bar"}, + Data: "some data", }) require.NoError(t, err) @@ -64,12 +66,14 @@ func TestUpsert(t *testing.T) { Ids: []string{id}, IncludeMetadata: true, IncludeVectors: true, + IncludeData: 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.Equal(t, "some data", vectors[0].Data) }) }) }