feat(pinecone): Add rerank task for Pinecone component (#773)

Because

- Rerank API is introduced from Pinecone
- Resolves instill-ai/instill-core#1114

This commit

- Added Pinecone Rerank Support. 
- Inputs: Documents and Query input. We can restrict output to n
documents by setting optional field (top-n)
   - Output: Reranked Documents and Scores
- Made URL in Pinecone Setup Optional as for rerank input URL is not
required. For others, this will still need to be input.

---------

Co-authored-by: chuang8511 <a0961192312@gmail.com>

Author: MubeenKodvavi

pkg/component/data/pinecone/v0/README.mdx

@@ -9,6 +9,7 @@ The Pinecone component is a data component that allows users to build and search
 It can carry out the following tasks:
 - [Query](#query)
 - [Upsert](#upsert)
+- [Rerank](#rerank)
 
 
 
@@ -40,7 +41,7 @@ ${connection.<my-connection-id>}`.
 | Field | Field ID | Type | Note |
 | :--- | :--- | :--- | :--- |
 | API Key (required) | `api-key` | string | Fill in your Pinecone AI API key. You can create an api key in Pinecone Console.  |
-| Pinecone Base URL (required) | `url` | string | Fill in your Pinecone base URL. It is in the form.  |
+| Pinecone Index URL | `url` | string | Fill in your Pinecone index URL. It is in the form.  |
 
 </div>
 
@@ -127,4 +128,32 @@ Writes vectors into a namespace. If a new value is upserted for an existing vect
 </div>
 
 
+### Rerank
+
+Rerank documents, such as text passages, according to their relevance to a query. The input is a list of documents and a query. The output is a list of documents, sorted by relevance to the query.
+
+<div class="markdown-col-no-wrap" data-col-1 data-col-2>
+
+| Input | ID | Type | Description |
+| :--- | :--- | :--- | :--- |
+| Task ID (required) | `task` | string | `TASK_RERANK` |
+| Query (required) | `query` | string | The query to rerank the documents. |
+| Documents (required) | `documents` | array[string] | The documents to rerank. |
+| Top N | `top-n` | integer | The number of results to return sorted by relevance. Defaults to the number of inputs. |
+</div>
+
+
+
+
+
+
+<div class="markdown-col-no-wrap" data-col-1 data-col-2>
+
+| Output | ID | Type | Description |
+| :--- | :--- | :--- | :--- |
+| Reranked Documents. | `documents` | array[string] | Reranked documents. |
+| Scores | `scores` | array[number] | The relevance score of the documents normalized between 0 and 1. |
+</div>
+
+
 

pkg/component/data/pinecone/v0/config/definition.json

@@ -1,7 +1,8 @@
 {
   "availableTasks": [
     "TASK_QUERY",
-    "TASK_UPSERT"
+    "TASK_UPSERT",
+    "TASK_RERANK"
   ],
   "custom": false,
   "documentationUrl": "https://www.instill.tech/docs/component/data/pinecone",

pkg/component/data/pinecone/v0/config/setup.json

@@ -16,7 +16,7 @@
       "type": "string"
     },
     "url": {
-      "description": "Fill in your Pinecone base URL. It is in the form.",
+      "description": "Fill in your Pinecone index URL. It is in the form.",
       "instillUpstreamTypes": [
         "value"
       ],
@@ -25,17 +25,15 @@
       ],
       "instillSecret": false,
       "instillUIOrder": 1,
-      "title": "Pinecone Base URL",
+      "title": "Pinecone Index URL",
       "type": "string"
     }
   },
   "required": [
-    "api-key",
-    "url"
+    "api-key"
   ],
   "instillEditOnNodeFields": [
-    "api-key",
-    "url"
+    "api-key"
   ],
   "title": "Pinecone Connection",
   "type": "object"

pkg/component/data/pinecone/v0/config/tasks.json

@@ -291,5 +291,97 @@
       "title": "Output",
       "type": "object"
     }
+  },
+  "TASK_RERANK": {
+    "instillShortDescription": "Rerank documents, such as text passages, according to their relevance to a query.",
+    "description": "Rerank documents, such as text passages, according to their relevance to a query. The input is a list of documents and a query. The output is a list of documents, sorted by relevance to the query.",
+    "input": {
+      "instillUIOrder": 0,
+      "properties": {
+        "query": {
+          "description": "The query to rerank the documents.",
+          "instillAcceptFormats": [
+            "string"
+          ],
+          "instullUIMultiline": false,
+          "instillUIOrder": 0,
+          "instillUpstreamTypes": [
+            "value",
+            "reference",
+            "template"
+          ],
+          "title": "Query",
+          "type": "string"
+        },
+        "documents": {
+          "description": "The documents to rerank.",
+          "instillUIOrder": 1,
+          "instillUpstreamTypes": [
+            "value",
+            "reference"
+          ],
+          "items": {
+            "type": "string"
+          },
+          "minItems": 1,
+          "title": "Documents",
+          "type": "array"
+        },
+        "top-n": {
+          "description": "The number of results to return sorted by relevance. Defaults to the number of inputs.",
+          "instillAcceptFormats": [
+            "integer"
+          ],
+          "instillUIOrder": 2,
+          "instillUpstreamTypes": [
+            "value",
+            "reference"
+          ],
+          "title": "Top N",
+          "type": "integer"
+        }
+      },
+      "required": [
+        "query",
+        "documents"
+      ],
+      "title": "Input",
+      "type": "object"
+    },
+    "output": {
+      "instillUIOrder": 0,
+      "properties": {
+        "documents": {
+          "description": "Reranked documents.",
+          "instillFormat": "array:string",
+          "items": {
+            "instillFormat": "string",
+            "title": "Documents",
+            "type": "string"
+          },
+          "instillUIOrder": 0,
+          "title": "Reranked documents.",
+          "type": "array"
+        },
+        "scores": {
+          "description": "The relevance score of the documents normalized between 0 and 1.",
+          "instillFormat": "array:number",
+          "items": {
+            "instillFormat": "number",
+            "title": "Score",
+            "type": "number"
+          },
+          "instillUIOrder": 1,
+          "title": "Scores",
+          "type": "array"
+        }
+      },
+      "required": [
+        "documents",
+        "scores"
+      ],
+      "title": "Output",
+      "type": "object"
+    }
   }
 }

pkg/component/data/pinecone/v0/main.go

@@ -17,9 +17,11 @@ import (
 const (
 	taskQuery  = "TASK_QUERY"
 	taskUpsert = "TASK_UPSERT"
+	taskRerank = "TASK_RERANK"
 
 	upsertPath = "/vectors/upsert"
 	queryPath  = "/query"
+	rerankPath = "/rerank"
 )
 
 //go:embed config/definition.json
@@ -59,7 +61,8 @@ func (c *component) CreateExecution(x base.ComponentExecution) (base.IExecution,
 	}, nil
 }
 
-func newClient(setup *structpb.Struct, logger *zap.Logger) *httpclient.Client {
+// newIndexClient creates a new httpclient.Client with the index URL provided in setup
+func newIndexClient(setup *structpb.Struct, logger *zap.Logger) *httpclient.Client {
 	c := httpclient.New("Pinecone", getURL(setup),
 		httpclient.WithLogger(logger),
 		httpclient.WithEndUserError(new(errBody)),
@@ -71,6 +74,23 @@ func newClient(setup *structpb.Struct, logger *zap.Logger) *httpclient.Client {
 	return c
 }
 
+// newBaseClient creates a new httpclient.Client with the default Pinecone API URL.
+func newBaseClient(setup *structpb.Struct, logger *zap.Logger) *httpclient.Client {
+	c := httpclient.New("Pinecone", "https://api.pinecone.io",
+		httpclient.WithLogger(logger),
+		httpclient.WithEndUserError(new(errBody)),
+	)
+
+	c.SetHeader("Api-Key", getAPIKey(setup))
+	c.SetHeader("User-Agent", "source_tag=instillai")
+
+	// Currently, by default Pinecone API redirects request to OLDEST stable version i.e. 2024-04 right now and does not support Rerank
+	// It is recommended by Pinecone to specify API version to use: https://docs.pinecone.io/reference/api/versioning#specify-an-api-version
+	c.SetHeader("X-Pinecone-API-Version", "2024-10")
+
+	return c
+}
+
 func getAPIKey(setup *structpb.Struct) string {
 	return setup.GetFields()["api-key"].GetStringValue()
 }
@@ -81,8 +101,6 @@ func getURL(setup *structpb.Struct) string {
 
 func (e *execution) Execute(ctx context.Context, jobs []*base.Job) error {
 
-	req := newClient(e.Setup, e.GetLogger()).R()
-
 	for _, job := range jobs {
 		input, err := job.Input.Read(ctx)
 		if err != nil {
@@ -92,6 +110,8 @@ func (e *execution) Execute(ctx context.Context, jobs []*base.Job) error {
 		var output *structpb.Struct
 		switch e.Task {
 		case taskQuery:
+			req := newIndexClient(e.Setup, e.GetLogger()).R()
+
 			inputStruct := queryInput{}
 			err := base.ConvertFromStructpb(input, &inputStruct)
 			if err != nil {
@@ -122,6 +142,8 @@ func (e *execution) Execute(ctx context.Context, jobs []*base.Job) error {
 				continue
 			}
 		case taskUpsert:
+			req := newIndexClient(e.Setup, e.GetLogger()).R()
+
 			v := upsertInput{}
 			err := base.ConvertFromStructpb(input, &v)
 			if err != nil {
@@ -145,7 +167,34 @@ func (e *execution) Execute(ctx context.Context, jobs []*base.Job) error {
 				job.Error.Error(ctx, err)
 				continue
 			}
+		case taskRerank:
+			// rerank task does not need index URL, so using the base client with the default pinecone API URL.
+			req := newBaseClient(e.Setup, e.GetLogger()).R()
+
+			// parse input struct
+			inputStruct := rerankInput{}
+			err := base.ConvertFromStructpb(input, &inputStruct)
+			if err != nil {
+				job.Error.Error(ctx, err)
+				continue
+			}
+
+			// make API request to rerank task
+			resp := rerankResp{}
+			req.SetResult(&resp).SetBody(inputStruct.asRequest())
+			if _, err := req.Post(rerankPath); err != nil {
+				job.Error.Error(ctx, httpclient.WrapURLError(err))
+				continue
+			}
+
+			// convert response to output struct
+			output, err = base.ConvertToStructpb(resp.toOutput())
+			if err != nil {
+				job.Error.Error(ctx, err)
+				continue
+			}
 		}
+
 		err = job.Output.Write(ctx, output)
 		if err != nil {
 			job.Error.Error(ctx, err)

pkg/component/data/pinecone/v0/structs.go

@@ -83,6 +83,66 @@ type upsertOutput struct {
 	RecordsUpserted int64 `json:"upserted-count"`
 }
 
+type Document struct {
+	Text string `json:"text"`
+}
+
+type rerankInput struct {
+	// not taking model as input for now as only one model is supported for rerank task: https://docs.pinecone.io/guides/inference/understanding-inference#models
+	//ModelName string   `json:"model-name"`
+	Query     string   `json:"query"`
+	Documents []string `json:"documents"`
+	TopN      int      `json:"top-n"`
+}
+
+func (r *rerankInput) asRequest() *rerankReq {
+	reqDocuments := make([]Document, 0, len(r.Documents))
+	for _, doc := range r.Documents {
+		reqDocuments = append(reqDocuments, Document{Text: doc})
+	}
+
+	// TODO: make model configurable in tasks.json
+	return &rerankReq{
+		Model:     "bge-reranker-v2-m3",
+		Query:     r.Query,
+		TopN:      r.TopN,
+		Documents: reqDocuments,
+	}
+}
+
+type rerankReq struct {
+	Model     string     `json:"model"`
+	Query     string     `json:"query"`
+	TopN      int        `json:"top_n,omitempty"`
+	Documents []Document `json:"documents"`
+}
+
+type rerankResp struct {
+	Data []struct {
+		Index    int      `json:"index"`
+		Document Document `json:"document"`
+		Score    float64  `json:"score"`
+	} `json:"data"`
+}
+
+func (r *rerankResp) toOutput() rerankOutput {
+	documents := make([]string, 0, len(r.Data))
+	scores := make([]float64, 0, len(r.Data))
+	for _, d := range r.Data {
+		documents = append(documents, d.Document.Text)
+		scores = append(scores, d.Score)
+	}
+	return rerankOutput{
+		Documents: documents,
+		Scores:    scores,
+	}
+}
+
+type rerankOutput struct {
+	Documents []string  `json:"documents"`
+	Scores    []float64 `json:"scores"`
+}
+
 type errBody struct {
 	Msg string `json:"message"`
 }