Rebase & Implement Rerank

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,6 +1343,125 @@ 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 calling [InferenceService.Rerank] and reranking documents
+// by a specified query and model.
+//
+// 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.
+//   - Parameters: (Optional) Additional model-specific parameters for the reranker
+//
+// [model]: https://docs.pinecone.io/guides/inference/understanding-inference#models
+type RerankRequest struct {
+	Model           string
+	Query           string
+	Documents       []Document
+	RankFields      *[]string
+	ReturnDocuments *bool
+	TopN            *int
+	Parameters      *map[string]string
+}
+
+// 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  []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
+// to the provided query using the specified model.
+//
+// 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)
+//	     }
+//
+//	     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 {
+		convertedDocuments[i] = inference.Document(doc)
+	}
+	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 rerank: ")
+	}
+	return decodeRerankResponse(res.Body)
+}
+
 func (c *Client) extractAuthHeader() map[string]string {
 	possibleAuthKeys := []string{
 		"api-key",
@@ -1423,6 +1542,16 @@ func decodeEmbeddingsList(resBody io.ReadCloser) (*inference.EmbeddingsList, err
 	return &embeddingsList, nil
 }
 
+func decodeRerankResponse(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

pinecone/client_test.go

@@ -313,6 +313,29 @@ 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"
+	topN := 4
+	retunDocuments := true
+	ranking, err := ts.client.Inference.Rerank(ctx, &RerankRequest{
+		Model:           rerankModel,
+		Query:           "i love cats",
+		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."},
+			{"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")

pinecone/models.go

@@ -158,6 +158,13 @@ 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"`
+}
+
 // MetadataFilter represents the [metadata filters] attached to a Pinecone request.
 // These optional metadata filters are applied to query and deletion requests.
 //