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

     1  //go:generate compogen readme --connector ./config ./README.mdx
     2  package googlesearch
     3  
     4  import (
     5  	"context"
     6  	_ "embed"
     7  	"encoding/json"
     8  	"fmt"
     9  	"sync"
    10  
    11  	"go.uber.org/zap"
    12  	"google.golang.org/api/customsearch/v1"
    13  	"google.golang.org/api/option"
    14  	"google.golang.org/protobuf/types/known/structpb"
    15  
    16  	"github.com/instill-ai/component/pkg/base"
    17  )
    18  
    19  const (
    20  	taskSearch = "TASK_SEARCH"
    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  // NewService creates a Google custom search service
    63  func NewService(apiKey string) (*customsearch.Service, error) {
    64  	return customsearch.NewService(context.Background(), option.WithAPIKey(apiKey))
    65  }
    66  
    67  func getAPIKey(config *structpb.Struct) string {
    68  	return config.GetFields()["api_key"].GetStringValue()
    69  }
    70  
    71  func getSearchEngineID(config *structpb.Struct) string {
    72  	return config.GetFields()["cse_id"].GetStringValue()
    73  }
    74  
    75  func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
    76  
    77  	service, err := NewService(getAPIKey(e.Connection))
    78  	if err != nil || service == nil {
    79  		return nil, fmt.Errorf("error creating Google custom search service: %v", err)
    80  	}
    81  	cseListCall := service.Cse.List().Cx(getSearchEngineID(e.Connection))
    82  
    83  	outputs := []*structpb.Struct{}
    84  
    85  	for _, input := range inputs {
    86  		switch e.Task {
    87  		case taskSearch:
    88  
    89  			inputStruct := SearchInput{}
    90  			err := base.ConvertFromStructpb(input, &inputStruct)
    91  			if err != nil {
    92  				return nil, err
    93  			}
    94  
    95  			// Make the search request
    96  			outputStruct, err := search(cseListCall, inputStruct)
    97  
    98  			if err != nil {
    99  				return nil, err
   100  			}
   101  
   102  			outputJSON, err := json.Marshal(outputStruct)
   103  			if err != nil {
   104  				return nil, err
   105  			}
   106  			output := structpb.Struct{}
   107  			err = json.Unmarshal(outputJSON, &output)
   108  			if err != nil {
   109  				return nil, err
   110  			}
   111  			outputs = append(outputs, &output)
   112  
   113  		default:
   114  			return nil, fmt.Errorf("not supported task: %s", e.Task)
   115  		}
   116  	}
   117  
   118  	return outputs, nil
   119  }
   120  
   121  func (c *connector) Test(sysVars map[string]any, connection *structpb.Struct) error {
   122  
   123  	service, err := NewService(getAPIKey(connection))
   124  	if err != nil || service == nil {
   125  		return fmt.Errorf("error creating Google custom search service: %v", err)
   126  	}
   127  	if service == nil {
   128  		return fmt.Errorf("error creating Google custom search service: %v", err)
   129  	}
   130  	return nil
   131  }