github.com/instill-ai/component@v0.16.0-beta/pkg/connector/pinecone/v0/main.go (about)

     1  //go:generate compogen readme --connector ./config ./README.mdx
     2  package pinecone
     3  
     4  import (
     5  	_ "embed"
     6  	"sync"
     7  
     8  	"go.uber.org/zap"
     9  	"google.golang.org/protobuf/types/known/structpb"
    10  
    11  	"github.com/instill-ai/component/pkg/base"
    12  	"github.com/instill-ai/component/pkg/connector/util/httpclient"
    13  )
    14  
    15  const (
    16  	taskQuery  = "TASK_QUERY"
    17  	taskUpsert = "TASK_UPSERT"
    18  
    19  	upsertPath = "/vectors/upsert"
    20  	queryPath  = "/query"
    21  )
    22  
    23  //go:embed config/definition.json
    24  var definitionJSON []byte
    25  
    26  //go:embed config/tasks.json
    27  var tasksJSON []byte
    28  
    29  var once sync.Once
    30  var con *connector
    31  
    32  type connector struct {
    33  	base.BaseConnector
    34  }
    35  
    36  type execution struct {
    37  	base.BaseConnectorExecution
    38  }
    39  
    40  func Init(l *zap.Logger, u base.UsageHandler) *connector {
    41  	once.Do(func() {
    42  		con = &connector{
    43  			BaseConnector: base.BaseConnector{
    44  				Logger:       l,
    45  				UsageHandler: u,
    46  			},
    47  		}
    48  		err := con.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
    49  		if err != nil {
    50  			panic(err)
    51  		}
    52  	})
    53  	return con
    54  }
    55  
    56  func (c *connector) CreateExecution(sysVars map[string]any, connection *structpb.Struct, task string) (*base.ExecutionWrapper, error) {
    57  	return &base.ExecutionWrapper{Execution: &execution{
    58  		BaseConnectorExecution: base.BaseConnectorExecution{Connector: c, SystemVariables: sysVars, Connection: connection, Task: task},
    59  	}}, nil
    60  }
    61  
    62  func newClient(config *structpb.Struct, logger *zap.Logger) *httpclient.Client {
    63  	c := httpclient.New("Pinecone", getURL(config),
    64  		httpclient.WithLogger(logger),
    65  		httpclient.WithEndUserError(new(errBody)),
    66  	)
    67  
    68  	c.SetHeader("Api-Key", getAPIKey(config))
    69  
    70  	return c
    71  }
    72  
    73  func getAPIKey(config *structpb.Struct) string {
    74  	return config.GetFields()["api_key"].GetStringValue()
    75  }
    76  
    77  func getURL(config *structpb.Struct) string {
    78  	return config.GetFields()["url"].GetStringValue()
    79  }
    80  
    81  func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
    82  	req := newClient(e.Connection, e.GetLogger()).R()
    83  	outputs := []*structpb.Struct{}
    84  
    85  	for _, input := range inputs {
    86  		var output *structpb.Struct
    87  		switch e.Task {
    88  		case taskQuery:
    89  			inputStruct := queryInput{}
    90  			err := base.ConvertFromStructpb(input, &inputStruct)
    91  			if err != nil {
    92  				return nil, err
    93  			}
    94  
    95  			// Each query request can contain only one of the parameters
    96  			// vector, or id.
    97  			// Ref: https://docs.pinecone.io/reference/query
    98  			if inputStruct.ID != "" {
    99  				inputStruct.Vector = nil
   100  			}
   101  
   102  			resp := queryResp{}
   103  			req.SetResult(&resp).SetBody(inputStruct.asRequest())
   104  
   105  			if _, err := req.Post(queryPath); err != nil {
   106  				return nil, httpclient.WrapURLError(err)
   107  			}
   108  
   109  			resp = resp.filterOutBelowThreshold(inputStruct.MinScore)
   110  
   111  			output, err = base.ConvertToStructpb(resp)
   112  			if err != nil {
   113  				return nil, err
   114  			}
   115  		case taskUpsert:
   116  			v := upsertInput{}
   117  			err := base.ConvertFromStructpb(input, &v)
   118  			if err != nil {
   119  				return nil, err
   120  			}
   121  
   122  			resp := upsertResp{}
   123  			req.SetResult(&resp).SetBody(upsertReq{
   124  				Vectors:   []vector{v.vector},
   125  				Namespace: v.Namespace,
   126  			})
   127  
   128  			if _, err := req.Post(upsertPath); err != nil {
   129  				return nil, httpclient.WrapURLError(err)
   130  			}
   131  
   132  			output, err = base.ConvertToStructpb(upsertOutput(resp))
   133  			if err != nil {
   134  				return nil, err
   135  			}
   136  		}
   137  		outputs = append(outputs, output)
   138  	}
   139  	return outputs, nil
   140  }
   141  
   142  func (c *connector) Test(sysVars map[string]any, connection *structpb.Struct) error {
   143  	//TODO: change this
   144  	return nil
   145  }