Skip to content

Commit

Permalink
Merge branch 'main' into Audrey/readme-updates
Browse files Browse the repository at this point in the history
  • Loading branch information
aulorbe authored Jul 30, 2024
2 parents f24aee5 + a8bec31 commit 128405c
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 21 deletions.
25 changes: 17 additions & 8 deletions pinecone/index_connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -200,9 +200,11 @@ func (idx *IndexConnection) UpsertVectors(ctx context.Context, in []*Vector) (ui
// Fields:
// - Vectors: The vectors fetched.
// - Usage: The usage information for the request.
// - Namespace: The namespace from which the vectors were fetched.
type FetchVectorsResponse struct {
Vectors map[string]*Vector `json:"vectors,omitempty"`
Usage *Usage `json:"usage,omitempty"`
Vectors map[string]*Vector `json:"vectors,omitempty"`
Usage *Usage `json:"usage,omitempty"`
Namespace string `json:"namespace"`
}

// FetchVectors fetches vectors by ID from a Pinecone index.
Expand Down Expand Up @@ -269,8 +271,9 @@ func (idx *IndexConnection) FetchVectors(ctx context.Context, ids []string) (*Fe
}

return &FetchVectorsResponse{
Vectors: vectors,
Usage: toUsage(res.Usage),
Vectors: vectors,
Usage: toUsage(res.Usage),
Namespace: idx.Namespace,
}, nil
}

Expand All @@ -295,10 +298,12 @@ type ListVectorsRequest struct {
// - VectorIds: The unique IDs of the returned vectors.
// - Usage: The usage information for the request.
// - NextPaginationToken: The token for paginating through results.
// - Namespace: The namespace vector ids are listed from.
type ListVectorsResponse struct {
VectorIds []*string `json:"vector_ids,omitempty"`
Usage *Usage `json:"usage,omitempty"`
NextPaginationToken *string `json:"next_pagination_token,omitempty"`
Namespace string `json:"namespace"`
}

// ListVectors lists vectors in a Pinecone index. You can filter vectors by prefix,
Expand Down Expand Up @@ -378,6 +383,7 @@ func (idx *IndexConnection) ListVectors(ctx context.Context, in *ListVectorsRequ
VectorIds: vectorIds,
Usage: &Usage{ReadUnits: derefOrDefault(res.Usage.ReadUnits, 0)},
NextPaginationToken: toPaginationToken(res.Pagination),
Namespace: idx.Namespace,
}, nil
}

Expand Down Expand Up @@ -406,9 +412,11 @@ type QueryByVectorValuesRequest struct {
// Fields:
// - Matches: The vectors that are most similar to the query vector.
// - Usage: The usage information for the request.
// - Namespace: The namespace from which the vectors were queried.
type QueryVectorsResponse struct {
Matches []*ScoredVector `json:"matches,omitempty"`
Usage *Usage `json:"usage,omitempty"`
Matches []*ScoredVector `json:"matches,omitempty"`
Usage *Usage `json:"usage,omitempty"`
Namespace string `json:"namespace"`
}

// QueryByVectorValues queries a Pinecone index for vectors that are most similar to a provided query vector.
Expand Down Expand Up @@ -989,8 +997,9 @@ func (idx *IndexConnection) query(ctx context.Context, req *data.QueryRequest) (
}

return &QueryVectorsResponse{
Matches: matches,
Usage: toUsage(res.Usage),
Matches: matches,
Usage: toUsage(res.Usage),
Namespace: idx.Namespace,
}, nil
}

Expand Down
31 changes: 18 additions & 13 deletions pinecone/index_connection_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -330,22 +330,24 @@ func TestMarshalFetchVectorsResponseUnit(t *testing.T) {
"vec-1": {Id: "vec-1", Values: []float32{0.01, 0.01, 0.01}},
"vec-2": {Id: "vec-2", Values: []float32{0.02, 0.02, 0.02}},
},
Usage: &Usage{ReadUnits: 5},
Usage: &Usage{ReadUnits: 5},
Namespace: "test-namespace",
},
want: `{"vectors":{"vec-1":{"id":"vec-1","values":[0.01,0.01,0.01]},"vec-2":{"id":"vec-2","values":[0.02,0.02,0.02]}},"usage":{"read_units":5}}`,
want: `{"vectors":{"vec-1":{"id":"vec-1","values":[0.01,0.01,0.01]},"vec-2":{"id":"vec-2","values":[0.02,0.02,0.02]}},"usage":{"read_units":5},"namespace":"test-namespace"}`,
},
{
name: "Fields omitted",
input: FetchVectorsResponse{},
want: `{}`,
want: `{"namespace":""}`,
},
{
name: "Fields empty",
input: FetchVectorsResponse{
Vectors: nil,
Usage: nil,
Vectors: nil,
Usage: nil,
Namespace: "",
},
want: `{}`,
want: `{"namespace":""}`,
},
}

Expand Down Expand Up @@ -378,22 +380,24 @@ func TestMarshalListVectorsResponseUnit(t *testing.T) {
VectorIds: []*string{&vectorId1, &vectorId2},
Usage: &Usage{ReadUnits: 5},
NextPaginationToken: &paginationToken,
Namespace: "test-namespace",
},
want: `{"vector_ids":["vec-1","vec-2"],"usage":{"read_units":5},"next_pagination_token":"next-token"}`,
want: `{"vector_ids":["vec-1","vec-2"],"usage":{"read_units":5},"next_pagination_token":"next-token","namespace":"test-namespace"}`,
},
{
name: "Fields omitted",
input: ListVectorsResponse{},
want: `{}`,
want: `{"namespace":""}`,
},
{
name: "Fields empty",
input: ListVectorsResponse{
VectorIds: nil,
Usage: nil,
NextPaginationToken: nil,
Namespace: "",
},
want: `{}`,
want: `{"namespace":""}`,
},
}

Expand Down Expand Up @@ -424,19 +428,20 @@ func TestMarshalQueryVectorsResponseUnit(t *testing.T) {
{Vector: &Vector{Id: "vec-1", Values: []float32{0.01, 0.01, 0.01}}, Score: 0.1},
{Vector: &Vector{Id: "vec-2", Values: []float32{0.02, 0.02, 0.02}}, Score: 0.2},
},
Usage: &Usage{ReadUnits: 5},
Usage: &Usage{ReadUnits: 5},
Namespace: "test-namespace",
},
want: `{"matches":[{"vector":{"id":"vec-1","values":[0.01,0.01,0.01]},"score":0.1},{"vector":{"id":"vec-2","values":[0.02,0.02,0.02]},"score":0.2}],"usage":{"read_units":5}}`,
want: `{"matches":[{"vector":{"id":"vec-1","values":[0.01,0.01,0.01]},"score":0.1},{"vector":{"id":"vec-2","values":[0.02,0.02,0.02]},"score":0.2}],"usage":{"read_units":5},"namespace":"test-namespace"}`,
},
{
name: "Fields omitted",
input: QueryVectorsResponse{},
want: `{}`,
want: `{"namespace":""}`,
},
{
name: "Fields empty",
input: QueryVectorsResponse{Matches: nil, Usage: nil},
want: `{}`,
want: `{"namespace":""}`,
},
}

Expand Down

0 comments on commit 128405c

Please sign in to comment.