From d29451b5330bb44f9c1ce84227351edd28780412 Mon Sep 17 00:00:00 2001 From: Sam Ayo Date: Wed, 11 Sep 2024 23:45:23 +0100 Subject: [PATCH 1/7] implemented rerankdocument api --- pinecone/client.go | 116 +++++++++++++++++++++++++++++++++++++++++++++ pinecone/models.go | 4 ++ 2 files changed, 120 insertions(+) diff --git a/pinecone/client.go b/pinecone/client.go index 01e1c06..ff3845c 100644 --- a/pinecone/client.go +++ b/pinecone/client.go @@ -1343,6 +1343,112 @@ func (i *InferenceService) Embed(ctx context.Context, in *EmbedRequest) (*infere return decodeEmbeddingsList(res.Body) } +// Document represents a single document to be reranked. +// +// Fields: +// - ID: (Required) The unique identifier of the document. +// - Text: (Required) The text content of the document. +type Document struct { + Id string + Text string +} + +// RerankRequest holds the parameters for reranking a list of documents based on a query. +// +// Fields: +// - Model: (Required) The model to use for reranking. +// - Query: (Required) The query string to use for reranking the documents. +// - ReturnDocuments: (Required) A boolean indicating whether to return the documents in the response. +// - TopN: (Required) The number of top documents to return after reranking. +// - Documents: (Required) A list of Document objects to be reranked. +type RerankRequest struct { + Model string + Query string + ReturnDocuments bool + TopN int + Documents []Document +} + +type RankResult struct { + Document *Document `json:"document,omitempty"` + Index IntegrationTests `json:"index"` + Score float32 `json:"score"` +} + +type RerankResponse struct { + Data []RankResult `json:"data,omitempty"` + Model string `json:"model,omitempty"` + Usage RerankUsage `json:"usage"` +} + +// Rerank processes the reranking of documents based on the provided query and model. +// +// Fields: +// - Model: The model to use for reranking. +// - Query: The query to base the reranking on. +// - ReturnDocuments: A boolean indicating whether to return the documents. +// - TopN: The number of top documents to return. +// - Documents: A list of documents to be reranked. +// +// Example: +// +// ctx := context.Background() +// +// clientParams := pinecone.NewClientParams{ +// ApiKey: "YOUR_API_KEY", +// SourceTag: "your_source_identifier", // optional +// } +// +// pc, err := pinecone.NewClient(clientParams) +// +// if err != nil { +// log.Fatalf("Failed to create Client: %v", err) +// } else { +// fmt.Println("Successfully created a new Client object!") +// } +// +// in := pinecone.RerankRequest{ +// Model: "your_model", +// Query: "your_query", +// ReturnDocuments: true, +// TopN: 5, +// Documents: []pinecone.Document{{Id: "doc1", Text: "text1"}, {Id: "doc2", Text: "text2"}}, +// } +// +// res, err := pc.Inference.Rerank(ctx, in) +// if err != nil { +// log.Fatalf("Failed to rerank: %v", err) +// } else { +// fmt.Printf("Successfully reranked documents: %+v", res) +// } +func (i *InferenceService) Rerank(ctx context.Context, in *RerankRequest) (*RerankResponse, error) { + + if len(in.Documents) <= 2 { + return nil, fmt.Errorf("documents must contain at least more than two value") + } + // Convert documents to the expected type + convertedDocuments := make([]inference.Document, len(in.Documents)) + for i, doc := range in.Documents { + convertedDocuments[i] = inference.Document{"Id": doc.Id, "Text": doc.Text} + } + req := inference.RerankJSONRequestBody{ + Model: in.Model, + Query: in.Query, + ReturnDocuments: &in.ReturnDocuments, + TopN: &in.TopN, + Documents: convertedDocuments, + } + res, err := i.client.Rerank(ctx, req) + if err != nil { + return nil, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return nil, handleErrorResponseBody(res, "failed to embed: ") + } + return decodeRerankList(res.Body) +} + func (c *Client) extractAuthHeader() map[string]string { possibleAuthKeys := []string{ "api-key", @@ -1423,6 +1529,16 @@ func decodeEmbeddingsList(resBody io.ReadCloser) (*inference.EmbeddingsList, err return &embeddingsList, nil } +func decodeRerankList(resBody io.ReadCloser) (*RerankResponse, error) { + var rerankResponse RerankResponse + err := json.NewDecoder(resBody).Decode(&rerankResponse) + if err != nil { + return nil, fmt.Errorf("failed to decode rerank response: %w", err) + } + + return &rerankResponse, nil +} + func toCollection(cm *db_control.CollectionModel) *Collection { if cm == nil { return nil diff --git a/pinecone/models.go b/pinecone/models.go index a29cf9d..aa43236 100644 --- a/pinecone/models.go +++ b/pinecone/models.go @@ -158,6 +158,10 @@ type Usage struct { ReadUnits uint32 `json:"read_units"` } +type RerankUsage struct { + RerankUnits *int `json:"rerank_units,omitempty"` +} + // MetadataFilter represents the [metadata filters] attached to a Pinecone request. // These optional metadata filters are applied to query and deletion requests. // From 343b9d4414d0c7c532cf6b8f221d4b66c00754c2 Mon Sep 17 00:00:00 2001 From: Sam Ayo Date: Wed, 11 Sep 2024 23:56:55 +0100 Subject: [PATCH 2/7] Update client_test.go --- pinecone/client_test.go | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/pinecone/client_test.go b/pinecone/client_test.go index 7771849..f931b5a 100644 --- a/pinecone/client_test.go +++ b/pinecone/client_test.go @@ -313,6 +313,26 @@ func (ts *IntegrationTests) TestGenerateEmbeddingsInvalidInputs() { require.Contains(ts.T(), err.Error(), "TextInputs must contain at least one value") } +func (ts *IntegrationTests) TestRerankDocument() { + ctx := context.Background() + rerankModel := "bge-reranker-v2-m3" + ranking, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ + Model: rerankModel, + Query: "The tech company Apple is known for its innovative products like the iPhone.", + ReturnDocuments: true, + TopN: 4, + Documents: []Document{ + {Id: "vec1", Text: "Apple is a popular fruit known for its sweetness and crisp texture."}, + {Id: "vec2", Text: "Many people enjoy eating apples as a healthy snack."}, + {Id: "vec3", Text: "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, + {Id: "vec4", Text: "An apple a day keeps the doctor away, as the saying goes."}, + }}) + + require.NoError(ts.T(), err) + require.NotNil(ts.T(), ranking, "Expected reranking result to be non-nil") + require.Equal(ts.T(), 4, len(*ranking.Data), "Expected 4 rankings") +} + // Unit tests: func TestExtractAuthHeaderUnit(t *testing.T) { globalApiKey := os.Getenv("PINECONE_API_KEY") From d76c4aabbdb8c238ce6774251fde43abced699c9 Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Wed, 16 Oct 2024 01:28:20 -0400 Subject: [PATCH 3/7] rebase external changes for rerank against release-candidate/2024 branch, refactor code to use generated inference module, correct a few issues with Documents argument shape and rerank test --- codegen/apis | 2 +- pinecone/client.go | 37 +++++++++++++------------------------ pinecone/client_test.go | 18 ++++++++++-------- 3 files changed, 24 insertions(+), 33 deletions(-) diff --git a/codegen/apis b/codegen/apis index 3002f1e..9d3a41f 160000 --- a/codegen/apis +++ b/codegen/apis @@ -1 +1 @@ -Subproject commit 3002f1e62b895d2e64fba53e346ad84eb4719934 +Subproject commit 9d3a41f6ae657d9b0b818065dcdc3adf76bdb7fe diff --git a/pinecone/client.go b/pinecone/client.go index ff3845c..4b32cf3 100644 --- a/pinecone/client.go +++ b/pinecone/client.go @@ -1343,15 +1343,7 @@ func (i *InferenceService) Embed(ctx context.Context, in *EmbedRequest) (*infere return decodeEmbeddingsList(res.Body) } -// Document represents a single document to be reranked. -// -// Fields: -// - ID: (Required) The unique identifier of the document. -// - Text: (Required) The text content of the document. -type Document struct { - Id string - Text string -} +type Document map[string]string // RerankRequest holds the parameters for reranking a list of documents based on a query. // @@ -1362,23 +1354,25 @@ type Document struct { // - TopN: (Required) The number of top documents to return after reranking. // - Documents: (Required) A list of Document objects to be reranked. type RerankRequest struct { + Documents []Document Model string Query string - ReturnDocuments bool - TopN int - Documents []Document + Parameters *map[string]string + RankFields *[]string + ReturnDocuments *bool + TopN *int } -type RankResult struct { +type RerankResult struct { Document *Document `json:"document,omitempty"` Index IntegrationTests `json:"index"` Score float32 `json:"score"` } type RerankResponse struct { - Data []RankResult `json:"data,omitempty"` - Model string `json:"model,omitempty"` - Usage RerankUsage `json:"usage"` + Data []RerankResult `json:"data,omitempty"` + Model string `json:"model,omitempty"` + Usage RerankUsage `json:"usage"` } // Rerank processes the reranking of documents based on the provided query and model. @@ -1422,20 +1416,15 @@ type RerankResponse struct { // fmt.Printf("Successfully reranked documents: %+v", res) // } func (i *InferenceService) Rerank(ctx context.Context, in *RerankRequest) (*RerankResponse, error) { - - if len(in.Documents) <= 2 { - return nil, fmt.Errorf("documents must contain at least more than two value") - } - // Convert documents to the expected type convertedDocuments := make([]inference.Document, len(in.Documents)) for i, doc := range in.Documents { - convertedDocuments[i] = inference.Document{"Id": doc.Id, "Text": doc.Text} + convertedDocuments[i] = inference.Document(doc) } req := inference.RerankJSONRequestBody{ Model: in.Model, Query: in.Query, - ReturnDocuments: &in.ReturnDocuments, - TopN: &in.TopN, + ReturnDocuments: in.ReturnDocuments, + TopN: in.TopN, Documents: convertedDocuments, } res, err := i.client.Rerank(ctx, req) diff --git a/pinecone/client_test.go b/pinecone/client_test.go index f931b5a..b0a9604 100644 --- a/pinecone/client_test.go +++ b/pinecone/client_test.go @@ -316,21 +316,23 @@ func (ts *IntegrationTests) TestGenerateEmbeddingsInvalidInputs() { func (ts *IntegrationTests) TestRerankDocument() { ctx := context.Background() rerankModel := "bge-reranker-v2-m3" + topN := 4 ranking, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ Model: rerankModel, - Query: "The tech company Apple is known for its innovative products like the iPhone.", - ReturnDocuments: true, - TopN: 4, + Query: "i love cats", + ReturnDocuments: nil, + TopN: &topN, + RankFields: &[]string{"text"}, Documents: []Document{ - {Id: "vec1", Text: "Apple is a popular fruit known for its sweetness and crisp texture."}, - {Id: "vec2", Text: "Many people enjoy eating apples as a healthy snack."}, - {Id: "vec3", Text: "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, - {Id: "vec4", Text: "An apple a day keeps the doctor away, as the saying goes."}, + {"id": "vec1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, + {"id": "vec2", "text": "Many people enjoy eating apples as a healthy snack."}, + {"id": "vec3", "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, + {"id": "vec4", "text": "An apple a day keeps the doctor away, as the saying goes."}, }}) require.NoError(ts.T(), err) require.NotNil(ts.T(), ranking, "Expected reranking result to be non-nil") - require.Equal(ts.T(), 4, len(*ranking.Data), "Expected 4 rankings") + require.Equal(ts.T(), 4, len(ranking.Data), "Expected 4 rankings") } // Unit tests: From 7df61e9eee2a840010050bbf18700a58bbc54b63 Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Wed, 16 Oct 2024 18:09:18 -0400 Subject: [PATCH 4/7] add and reformat doc comments for reranking --- pinecone/client.go | 132 ++++++++++++++++++++++++---------------- pinecone/client_test.go | 3 +- pinecone/models.go | 3 + 3 files changed, 83 insertions(+), 55 deletions(-) diff --git a/pinecone/client.go b/pinecone/client.go index 4b32cf3..f833de5 100644 --- a/pinecone/client.go +++ b/pinecone/client.go @@ -1261,7 +1261,7 @@ type InferenceService struct { // Parameters: // - ctx: A context.Context object controls the request's lifetime, allowing for the request // to be canceled or to timeout according to the context's deadline. -// - in: A pointer to an EmbedRequest object that contains the model t4o use for embedding generation, the +// - in: A pointer to an EmbedRequest object that contains the model to use for embedding generation, the // list of input strings to generate embeddings for, and any additional parameters to use for generation. // // Returns a pointer to an EmbeddingsList object or an error. @@ -1343,78 +1343,102 @@ func (i *InferenceService) Embed(ctx context.Context, in *EmbedRequest) (*infere return decodeEmbeddingsList(res.Body) } +// Document is a map representing the document to be reranked. type Document map[string]string -// RerankRequest holds the parameters for reranking a list of documents based on a query. +// RerankRequest holds the parameters for calling [InferenceService.Rerank] and reranking documents +// by a specified query and model. // // Fields: -// - Model: (Required) The model to use for reranking. -// - Query: (Required) The query string to use for reranking the documents. -// - ReturnDocuments: (Required) A boolean indicating whether to return the documents in the response. -// - TopN: (Required) The number of top documents to return after reranking. +// - Model: "The [model] to use for reranking. +// - Query: (Required) The query to compare with documents. // - Documents: (Required) A list of Document objects to be reranked. +// - RankFields: (Optional) A list of document fields to use for ranking. +// - ReturnDocuments: (Optional) Whether to include documents in the response. Defaults to true. +// - TopN: (Optional) How many documents to return. Defaults to the length of input Documents. +// - Parameters: (Optional) Additional model-specific parameters for the reranker +// +// [model]: https://docs.pinecone.io/guides/inference/understanding-inference#models type RerankRequest struct { - Documents []Document Model string Query string - Parameters *map[string]string + Documents []Document RankFields *[]string ReturnDocuments *bool TopN *int + Parameters *map[string]string } -type RerankResult struct { - Document *Document `json:"document,omitempty"` - Index IntegrationTests `json:"index"` - Score float32 `json:"score"` +// Represents a ranked document with a relevance score and an index position. +// +// Fields: +// - Document: The [Document]. +// - Index: The index position of the document from the original request. This can be used +// to locate the position of the document relative to others described in the request. +// - Score: The relevance score of the document indicating how closely it matches the query. +type RankedDocument struct { + Document *Document `json:"document,omitempty"` + Index int `json:"index"` + Score float32 `json:"score"` } +// RerankResponse is the result of a reranking operation. +// +// Fields: +// - Data: A list of documents which have been reranked. The documents are sorted in order of relevance, +// with the first being the most relevant. +// - Model: The model used to rerank documents. +// - Usage: Usage statistics for the reranking operation. type RerankResponse struct { - Data []RerankResult `json:"data,omitempty"` - Model string `json:"model,omitempty"` - Usage RerankUsage `json:"usage"` + Data []RankedDocument `json:"data,omitempty"` + Model string `json:"model"` + Usage RerankUsage `json:"usage"` } -// Rerank processes the reranking of documents based on the provided query and model. +// Rerank documents with associated relevance scores that represent the relevance of each document +// to the provided query using the specified model. // -// Fields: -// - Model: The model to use for reranking. -// - Query: The query to base the reranking on. -// - ReturnDocuments: A boolean indicating whether to return the documents. -// - TopN: The number of top documents to return. -// - Documents: A list of documents to be reranked. +// Parameters: +// - ctx: A context.Context object controls the request's lifetime, allowing for the request +// to be canceled or to timeout according to the context's deadline. +// - in: A pointer to a [RerankRequest] object that contains the model, query, and documents to use for reranking. // // Example: // -// ctx := context.Background() -// -// clientParams := pinecone.NewClientParams{ -// ApiKey: "YOUR_API_KEY", -// SourceTag: "your_source_identifier", // optional -// } -// -// pc, err := pinecone.NewClient(clientParams) -// -// if err != nil { -// log.Fatalf("Failed to create Client: %v", err) -// } else { -// fmt.Println("Successfully created a new Client object!") -// } -// -// in := pinecone.RerankRequest{ -// Model: "your_model", -// Query: "your_query", -// ReturnDocuments: true, -// TopN: 5, -// Documents: []pinecone.Document{{Id: "doc1", Text: "text1"}, {Id: "doc2", Text: "text2"}}, -// } -// -// res, err := pc.Inference.Rerank(ctx, in) -// if err != nil { -// log.Fatalf("Failed to rerank: %v", err) -// } else { -// fmt.Printf("Successfully reranked documents: %+v", res) -// } +// ctx := context.Background() +// +// clientParams := pinecone.NewClientParams{ +// ApiKey: "YOUR_API_KEY", +// SourceTag: "your_source_identifier", // optional +// } +// +// pc, err := pinecone.NewClient(clientParams) +// if err != nil { +// log.Fatalf("Failed to create Client: %v", err) +// } +// +// rerankModel := "bge-reranker-v2-m3" +// topN := 2 +// retunDocuments := true +// documents := []pinecone.Document{ +// {"id": "vec1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, +// {"id": "vec2", "text": "Many people enjoy eating apples as a healthy snack."}, +// {"id": "vec3", "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, +// {"id": "vec4", "text": "An apple a day keeps the doctor away, as the saying goes."}, +// } +// +// ranking, err := pc.Inference.Rerank(ctx, &pinecone.RerankRequest{ +// Model: rerankModel, +// Query: "i love cats", +// ReturnDocuments: &retunDocuments, +// TopN: &topN, +// RankFields: &[]string{"text"}, +// Documents: documents, +// }) +// if err != nil { +// log.Fatalf("Failed to rerank: %v", err) +// } +// fmt.Printf("Rerank result: %+v\n", ranking) func (i *InferenceService) Rerank(ctx context.Context, in *RerankRequest) (*RerankResponse, error) { convertedDocuments := make([]inference.Document, len(in.Documents)) for i, doc := range in.Documents { @@ -1433,9 +1457,9 @@ func (i *InferenceService) Rerank(ctx context.Context, in *RerankRequest) (*Rera } defer res.Body.Close() if res.StatusCode != http.StatusOK { - return nil, handleErrorResponseBody(res, "failed to embed: ") + return nil, handleErrorResponseBody(res, "failed to rerank: ") } - return decodeRerankList(res.Body) + return decodeRerankResponse(res.Body) } func (c *Client) extractAuthHeader() map[string]string { @@ -1518,7 +1542,7 @@ func decodeEmbeddingsList(resBody io.ReadCloser) (*inference.EmbeddingsList, err return &embeddingsList, nil } -func decodeRerankList(resBody io.ReadCloser) (*RerankResponse, error) { +func decodeRerankResponse(resBody io.ReadCloser) (*RerankResponse, error) { var rerankResponse RerankResponse err := json.NewDecoder(resBody).Decode(&rerankResponse) if err != nil { diff --git a/pinecone/client_test.go b/pinecone/client_test.go index b0a9604..a2d45fd 100644 --- a/pinecone/client_test.go +++ b/pinecone/client_test.go @@ -317,10 +317,11 @@ func (ts *IntegrationTests) TestRerankDocument() { ctx := context.Background() rerankModel := "bge-reranker-v2-m3" topN := 4 + retunDocuments := true ranking, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ Model: rerankModel, Query: "i love cats", - ReturnDocuments: nil, + ReturnDocuments: &retunDocuments, TopN: &topN, RankFields: &[]string{"text"}, Documents: []Document{ diff --git a/pinecone/models.go b/pinecone/models.go index aa43236..00fa28a 100644 --- a/pinecone/models.go +++ b/pinecone/models.go @@ -158,6 +158,9 @@ type Usage struct { ReadUnits uint32 `json:"read_units"` } +// RerankUsage is the usage stats ([Rerank Units]) for a reranking request. +// +// [Read Units]: https://docs.pinecone.io/guides/organizations/manage-cost/understanding-cost#rerank type RerankUsage struct { RerankUnits *int `json:"rerank_units,omitempty"` } From be67550c2772a76f1a2e3ad8b36a3a38e01890f9 Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Fri, 18 Oct 2024 12:26:54 -0400 Subject: [PATCH 5/7] review feedback --- pinecone/client.go | 35 +++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/pinecone/client.go b/pinecone/client.go index f833de5..fa8a537 100644 --- a/pinecone/client.go +++ b/pinecone/client.go @@ -1351,11 +1351,12 @@ type Document map[string]string // // Fields: // - Model: "The [model] to use for reranking. -// - Query: (Required) The query to compare with documents. -// - Documents: (Required) A list of Document objects to be reranked. -// - RankFields: (Optional) A list of document fields to use for ranking. -// - ReturnDocuments: (Optional) Whether to include documents in the response. Defaults to true. -// - TopN: (Optional) How many documents to return. Defaults to the length of input Documents. +// - Query: (Required) The query to rerank Documents against. +// - Documents: (Required) A list of Document objects to be reranked. The default is "text", but you can +// specify this behavior with [RerankRequest.RankFields]. +// - RankFields: (Optional) The fields to rank the Documents by. If not provided, the default is "text". +// - ReturnDocuments: (Optional) Whether to include Documents in the response. Defaults to true. +// - TopN: (Optional) How many Documents to return. Defaults to the length of input Documents. // - Parameters: (Optional) Additional model-specific parameters for the reranker // // [model]: https://docs.pinecone.io/guides/inference/understanding-inference#models @@ -1373,9 +1374,9 @@ type RerankRequest struct { // // Fields: // - Document: The [Document]. -// - Index: The index position of the document from the original request. This can be used +// - Index: The index position of the Document from the original request. This can be used // to locate the position of the document relative to others described in the request. -// - Score: The relevance score of the document indicating how closely it matches the query. +// - Score: The relevance score of the Document indicating how closely it matches the query. type RankedDocument struct { Document *Document `json:"document,omitempty"` Index int `json:"index"` @@ -1385,17 +1386,19 @@ type RankedDocument struct { // RerankResponse is the result of a reranking operation. // // Fields: -// - Data: A list of documents which have been reranked. The documents are sorted in order of relevance, +// - Data: A list of Documents which have been reranked. The Documents are sorted in order of relevance, // with the first being the most relevant. -// - Model: The model used to rerank documents. -// - Usage: Usage statistics for the reranking operation. +// - Model: The model used to rerank Documents. +// - Usage: Usage statistics ([Rerank Units]) for the reranking operation. +// +// [Read Units]: https://docs.pinecone.io/guides/organizations/manage-cost/understanding-cost#rerank type RerankResponse struct { Data []RankedDocument `json:"data,omitempty"` Model string `json:"model"` Usage RerankUsage `json:"usage"` } -// Rerank documents with associated relevance scores that represent the relevance of each document +// Rerank Documents with associated relevance scores that represent the relevance of each Document // to the provided query using the specified model. // // Parameters: @@ -1421,15 +1424,15 @@ type RerankResponse struct { // topN := 2 // retunDocuments := true // documents := []pinecone.Document{ -// {"id": "vec1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, -// {"id": "vec2", "text": "Many people enjoy eating apples as a healthy snack."}, -// {"id": "vec3", "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, -// {"id": "vec4", "text": "An apple a day keeps the doctor away, as the saying goes."}, +// {"id": "doc1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, +// {"id": "doc2", "text": "Many people enjoy eating apples as a healthy snack."}, +// {"id": "doc3", "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, +// {"id": "doc4", "text": "An apple a day keeps the doctor away, as the saying goes."}, // } // // ranking, err := pc.Inference.Rerank(ctx, &pinecone.RerankRequest{ // Model: rerankModel, -// Query: "i love cats", +// Query: "i love to eat apples", // ReturnDocuments: &retunDocuments, // TopN: &topN, // RankFields: &[]string{"text"}, From 825349453bf4e3ed2be66a1d8f634da38f5a9caa Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Fri, 18 Oct 2024 15:01:29 -0400 Subject: [PATCH 6/7] add additional rerank integration tests --- pinecone/client.go | 4 +- pinecone/client_test.go | 152 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 149 insertions(+), 7 deletions(-) diff --git a/pinecone/client.go b/pinecone/client.go index fa8a537..89f8865 100644 --- a/pinecone/client.go +++ b/pinecone/client.go @@ -1450,9 +1450,11 @@ func (i *InferenceService) Rerank(ctx context.Context, in *RerankRequest) (*Rera req := inference.RerankJSONRequestBody{ Model: in.Model, Query: in.Query, + Documents: convertedDocuments, + RankFields: in.RankFields, ReturnDocuments: in.ReturnDocuments, TopN: in.TopN, - Documents: convertedDocuments, + Parameters: in.Parameters, } res, err := i.client.Rerank(ctx, req) if err != nil { diff --git a/pinecone/client_test.go b/pinecone/client_test.go index a2d45fd..eb54fa2 100644 --- a/pinecone/client_test.go +++ b/pinecone/client_test.go @@ -13,11 +13,11 @@ import ( "time" "github.com/google/go-cmp/cmp" + "github.com/google/uuid" "github.com/pinecone-io/go-pinecone/internal/gen" "github.com/pinecone-io/go-pinecone/internal/gen/db_control" "github.com/pinecone-io/go-pinecone/internal/provider" - "github.com/google/uuid" "github.com/pinecone-io/go-pinecone/internal/utils" "github.com/stretchr/testify/assert" @@ -277,6 +277,11 @@ func (ts *IntegrationTests) TestConfigureIndexHitPodLimit() { } func (ts *IntegrationTests) TestGenerateEmbeddings() { + // Run Embed tests once rather than duplicating across serverless & pods + if ts.indexType == "pod" { + ts.T().Skip("Skipping Embed tests for pods") + } + ctx := context.Background() embeddingModel := "multilingual-e5-large" embeddings, err := ts.client.Inference.Embed(ctx, &EmbedRequest{ @@ -313,17 +318,21 @@ func (ts *IntegrationTests) TestGenerateEmbeddingsInvalidInputs() { require.Contains(ts.T(), err.Error(), "TextInputs must contain at least one value") } -func (ts *IntegrationTests) TestRerankDocument() { +func (ts *IntegrationTests) TestRerankDocumentDefaultField() { + // Run Rerank tests once rather than duplicating across serverless & pods + if ts.indexType == "pod" { + ts.T().Skip("Skipping Rerank tests for pods") + } + ctx := context.Background() rerankModel := "bge-reranker-v2-m3" - topN := 4 + topN := 2 retunDocuments := true ranking, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ Model: rerankModel, - Query: "i love cats", + Query: "i love apples", ReturnDocuments: &retunDocuments, TopN: &topN, - RankFields: &[]string{"text"}, Documents: []Document{ {"id": "vec1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, {"id": "vec2", "text": "Many people enjoy eating apples as a healthy snack."}, @@ -333,7 +342,138 @@ func (ts *IntegrationTests) TestRerankDocument() { require.NoError(ts.T(), err) require.NotNil(ts.T(), ranking, "Expected reranking result to be non-nil") - require.Equal(ts.T(), 4, len(ranking.Data), "Expected 4 rankings") + require.Equal(ts.T(), topN, len(ranking.Data), "Expected %v rankings", topN) + + doc := *ranking.Data[0].Document + _, exists := doc["text"] + require.True(ts.T(), exists, "Expected '%s' to exist in Document map", "text") + _, exists = doc["id"] + require.True(ts.T(), exists, "Expected '%s' to exist in Document map", "id") +} + +func (ts *IntegrationTests) TestRerankDocumentCustomField() { + // Run Rerank tests once rather than duplicating across serverless & pods + if ts.indexType == "pod" { + ts.T().Skip("Skipping Rerank tests for pods") + } + + ctx := context.Background() + rerankModel := "bge-reranker-v2-m3" + topN := 2 + retunDocuments := true + ranking, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ + Model: rerankModel, + Query: "i love apples", + ReturnDocuments: &retunDocuments, + TopN: &topN, + RankFields: &[]string{"customField"}, + Documents: []Document{ + {"id": "vec1", "customField": "Apple is a popular fruit known for its sweetness and crisp texture."}, + {"id": "vec2", "customField": "Many people enjoy eating apples as a healthy snack."}, + {"id": "vec3", "customField": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, + {"id": "vec4", "customField": "An apple a day keeps the doctor away, as the saying goes."}, + }}) + + require.NoError(ts.T(), err) + require.NotNil(ts.T(), ranking, "Expected reranking result to be non-nil") + require.Equal(ts.T(), topN, len(ranking.Data), "Expected %v rankings", topN) + + doc := *ranking.Data[0].Document + _, exists := doc["customField"] + require.True(ts.T(), exists, "Expected '%s' to exist in Document map", "customField") + _, exists = doc["id"] + require.True(ts.T(), exists, "Expected '%s' to exist in Document map", "id") +} + +func (ts *IntegrationTests) TestRerankDocumentAllDefaults() { + // Run Rerank tests once rather than duplicating across serverless & pods + if ts.indexType == "pod" { + ts.T().Skip("Skipping Rerank tests for pods") + } + + ctx := context.Background() + rerankModel := "bge-reranker-v2-m3" + ranking, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ + Model: rerankModel, + Query: "i love apples", + Documents: []Document{ + {"id": "vec1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, + {"id": "vec2", "text": "Many people enjoy eating apples as a healthy snack."}, + {"id": "vec3", "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, + {"id": "vec4", "text": "An apple a day keeps the doctor away, as the saying goes."}, + }}) + + require.NoError(ts.T(), err) + require.NotNil(ts.T(), ranking, "Expected reranking result to be non-nil") + require.Equal(ts.T(), 4, len(ranking.Data), "Expected %v rankings", 4) + + doc := *ranking.Data[0].Document + _, exists := doc["text"] + require.True(ts.T(), exists, "Expected '%s' to exist in Document map", "text") + _, exists = doc["id"] + require.True(ts.T(), exists, "Expected '%s' to exist in Document map", "id") +} + +func (ts *IntegrationTests) TestRerankDocumentsMultipleRankFields() { + // Run Rerank tests once rather than duplicating across serverless & pods + if ts.indexType == "pod" { + ts.T().Skip("Skipping Rerank tests for pods") + } + + ctx := context.Background() + rerankModel := "bge-reranker-v2-m3" + _, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ + Model: rerankModel, + Query: "i love apples", + RankFields: &[]string{"text", "custom-field"}, + Documents: []Document{ + { + "id": "vec1", + "text": "Apple is a popular fruit known for its sweetness and crisp texture.", + "custom-field": "another field", + }, + { + "id": "vec2", + "text": "Many people enjoy eating apples as a healthy snack.", + "custom-field": "another field", + }, + { + "id": "vec3", + "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces.", + "custom-field": "another field", + }, + { + "id": "vec4", + "text": "An apple a day keeps the doctor away, as the saying goes.", + "custom-field": "another field", + }, + }}) + + require.Error(ts.T(), err) + require.Contains(ts.T(), err.Error(), "Only one rank field is supported for model") +} + +func (ts *IntegrationTests) TestRerankDocumentFieldError() { + // Run Rerank tests once rather than duplicating across serverless & pods + if ts.indexType == "pod" { + ts.T().Skip("Skipping Rerank tests for pods") + } + + ctx := context.Background() + rerankModel := "bge-reranker-v2-m3" + _, err := ts.client.Inference.Rerank(ctx, &RerankRequest{ + Model: rerankModel, + Query: "i love apples", + RankFields: &[]string{"custom-field"}, + Documents: []Document{ + {"id": "vec1", "text": "Apple is a popular fruit known for its sweetness and crisp texture."}, + {"id": "vec2", "text": "Many people enjoy eating apples as a healthy snack."}, + {"id": "vec3", "text": "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."}, + {"id": "vec4", "text": "An apple a day keeps the doctor away, as the saying goes."}, + }}) + + require.Error(ts.T(), err) + require.Contains(ts.T(), err.Error(), "field 'custom-field' not found in document") } // Unit tests: From e1c3343045437182787b3f39b5c009b0c75071b8 Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Fri, 18 Oct 2024 16:56:25 -0400 Subject: [PATCH 7/7] fix link in doc comment --- pinecone/models.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pinecone/models.go b/pinecone/models.go index 00fa28a..43d71c6 100644 --- a/pinecone/models.go +++ b/pinecone/models.go @@ -160,7 +160,7 @@ type Usage struct { // RerankUsage is the usage stats ([Rerank Units]) for a reranking request. // -// [Read Units]: https://docs.pinecone.io/guides/organizations/manage-cost/understanding-cost#rerank +// [Rerank Units]: https://docs.pinecone.io/guides/organizations/manage-cost/understanding-cost#rerank type RerankUsage struct { RerankUnits *int `json:"rerank_units,omitempty"` }