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

     1  //go:generate compogen readme --connector ./config ./README.mdx
     2  package redis
     3  
     4  import (
     5  	"context"
     6  	_ "embed"
     7  	"fmt"
     8  	"sync"
     9  
    10  	"go.uber.org/zap"
    11  	"google.golang.org/protobuf/types/known/structpb"
    12  
    13  	"github.com/instill-ai/component/pkg/base"
    14  )
    15  
    16  const (
    17  	taskWriteChatMessage           = "TASK_WRITE_CHAT_MESSAGE"
    18  	taskWriteMultiModalChatMessage = "TASK_WRITE_MULTI_MODAL_CHAT_MESSAGE"
    19  	taskRetrieveChatHistory        = "TASK_RETRIEVE_CHAT_HISTORY"
    20  )
    21  
    22  var (
    23  	//go:embed config/definition.json
    24  	definitionJSON []byte
    25  	//go:embed config/tasks.json
    26  	tasksJSON []byte
    27  
    28  	once sync.Once
    29  	con  *connector
    30  )
    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 (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
    63  	outputs := []*structpb.Struct{}
    64  
    65  	client, err := NewClient(e.Connection)
    66  	if err != nil {
    67  		return outputs, err
    68  	}
    69  	defer client.Close()
    70  
    71  	for _, input := range inputs {
    72  		var output *structpb.Struct
    73  		switch e.Task {
    74  		case taskWriteChatMessage:
    75  			inputStruct := ChatMessageWriteInput{}
    76  			err := base.ConvertFromStructpb(input, &inputStruct)
    77  			if err != nil {
    78  				return nil, err
    79  			}
    80  			outputStruct := WriteMessage(client, inputStruct)
    81  			output, err = base.ConvertToStructpb(outputStruct)
    82  			if err != nil {
    83  				return nil, err
    84  			}
    85  		case taskWriteMultiModalChatMessage:
    86  			inputStruct := ChatMultiModalMessageWriteInput{}
    87  			err := base.ConvertFromStructpb(input, &inputStruct)
    88  			if err != nil {
    89  				return nil, err
    90  			}
    91  			outputStruct := WriteMultiModelMessage(client, inputStruct)
    92  			output, err = base.ConvertToStructpb(outputStruct)
    93  			if err != nil {
    94  				return nil, err
    95  			}
    96  		case taskRetrieveChatHistory:
    97  			inputStruct := ChatHistoryRetrieveInput{}
    98  			err := base.ConvertFromStructpb(input, &inputStruct)
    99  			if err != nil {
   100  				return nil, err
   101  			}
   102  			outputStruct := RetrieveSessionMessages(client, inputStruct)
   103  			output, err = base.ConvertToStructpb(outputStruct)
   104  			if err != nil {
   105  				return nil, err
   106  			}
   107  		default:
   108  			return nil, fmt.Errorf("unsupported task: %s", e.Task)
   109  		}
   110  		outputs = append(outputs, output)
   111  	}
   112  	return outputs, nil
   113  }
   114  
   115  func (c *connector) Test(sysVars map[string]any, connection *structpb.Struct) error {
   116  	client, err := NewClient(connection)
   117  	if err != nil {
   118  		return err
   119  	}
   120  	defer client.Close()
   121  
   122  	// Ping the Redis server to check the connection
   123  	_, err = client.Ping(context.Background()).Result()
   124  	if err != nil {
   125  		return err
   126  	}
   127  	return nil
   128  }