feat: add UsageHandler interface (#87)

Because - We want to collect the usage of each component. We'll provide a `UsageHandler` interface to allow the pipeline-backend or other services to pass the usage handler implementation here. This commit - Adds `UsageHandler` interface.
instill-ai · Apr 18, 2024 · b9d9645 · b9d9645
1 parent a6de70e
commit b9d9645
Show file tree

Hide file tree

Showing 32 changed files with 138 additions and 109 deletions.
diff --git a/pkg/base/component.go b/pkg/base/component.go
@@ -59,6 +59,9 @@ type IComponent interface {
 	GetTaskInputSchemas() map[string]string
 	// Get task output schemas
 	GetTaskOutputSchemas() map[string]string
+
+	// Get usage handler
+	GetUsageHandler() UsageHandler
 }
 
 // Component is the basic component struct
@@ -76,6 +79,8 @@ type Component struct {
 	taskInputSchemas  map[string]string
 	taskOutputSchemas map[string]string
 
+	UsageHandler UsageHandler
+
 	// Logger
 	Logger *zap.Logger
 }
@@ -543,6 +548,10 @@ func (comp *Component) getDefinitionByID(defID string) (interface{}, error) {
 	return val, nil
 }
 
+func (comp *Component) GetUsageHandler() UsageHandler {
+	return comp.UsageHandler
+}
+
 // ConvertFromStructpb converts from structpb.Struct to a struct
 func ConvertFromStructpb(from *structpb.Struct, to interface{}) error {
 	inputJSON, err := protojson.Marshal(from)

diff --git a/pkg/base/execution.go b/pkg/base/execution.go
@@ -159,6 +159,12 @@ func (e *Execution) ExecuteWithValidation(inputs []*structpb.Struct) ([]*structp
 		return nil, err
 	}
 
+	if e.Component.GetUsageHandler() != nil {
+		if err := e.Component.GetUsageHandler().Check(); err != nil {
+			return nil, err
+		}
+	}
+
 	outputs, err := e.ComponentExecution.Execute(inputs)
 	if err != nil {
 		return nil, err
@@ -167,6 +173,13 @@ func (e *Execution) ExecuteWithValidation(inputs []*structpb.Struct) ([]*structp
 	if err := e.Validate(outputs, e.Component.GetTaskOutputSchemas()[task], "outputs"); err != nil {
 		return nil, err
 	}
+
+	if e.Component.GetUsageHandler() != nil {
+		if err := e.Component.GetUsageHandler().Collect(); err != nil {
+			return nil, err
+		}
+	}
+
 	return outputs, err
 }
 

diff --git a/pkg/base/usage.go b/pkg/base/usage.go
@@ -0,0 +1,7 @@
+package base
+
+// TODO: add parameters for usage check and collection
+type UsageHandler interface {
+	Check() error
+	Collect() error
+}
diff --git a/pkg/connector/airbyte/v0/main.go b/pkg/connector/airbyte/v0/main.go
@@ -58,7 +58,7 @@ type Execution struct {
 	connector *Connector
 }
 
-func Init(logger *zap.Logger, options ConnectorOptions) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler, options ConnectorOptions) base.IConnector {
 	once.Do(func() {
 
 		dockerClient, err := dockerclient.NewClientWithOpts(dockerclient.FromEnv, dockerclient.WithAPIVersionNegotiation())
@@ -73,7 +73,7 @@ func Init(logger *zap.Logger, options ConnectorOptions) base.IConnector {
 
 		connector = &Connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 			dockerClient: dockerClient,
 			cache:        cache,
@@ -103,9 +103,9 @@ func Init(logger *zap.Logger, options ConnectorOptions) base.IConnector {
 	return connector
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &Execution{}
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 	e.connector = c
 	return e, nil
 }

diff --git a/pkg/connector/archetypeai/v0/connector_test.go b/pkg/connector/archetypeai/v0/connector_test.go
@@ -226,7 +226,7 @@ func TestConnector_Execute(t *testing.T) {
 	}
 
 	logger := zap.NewNop()
-	connector := Init(logger)
+	connector := Init(logger, nil)
 	defID := uuid.Must(uuid.NewV4())
 
 	for _, tc := range testcases {
@@ -287,7 +287,7 @@ func TestConnector_CreateExecution(t *testing.T) {
 	c := qt.New(t)
 
 	logger := zap.NewNop()
-	connector := Init(logger)
+	connector := Init(logger, nil)
 	defID := uuid.Must(uuid.NewV4())
 
 	c.Run("nok - unsupported task", func(c *qt.C) {
@@ -304,7 +304,7 @@ func TestConnector_Test(t *testing.T) {
 	c := qt.New(t)
 
 	logger := zap.NewNop()
-	connector := Init(logger)
+	connector := Init(logger, nil)
 	defID := uuid.Must(uuid.NewV4())
 
 	c.Run("ok - connected", func(c *qt.C) {

diff --git a/pkg/connector/archetypeai/v0/main.go b/pkg/connector/archetypeai/v0/main.go
@@ -46,11 +46,11 @@ type execution struct {
 
 // Init returns an implementation of IConnector that interacts with Archetype
 // AI.
-func Init(logger *zap.Logger) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler) base.IConnector {
 	once.Do(func() {
 		baseConn = &connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 		}
 		if err := baseConn.LoadConnectorDefinition(definitionJSON, tasksJSON, nil); err != nil {
@@ -62,9 +62,9 @@ func Init(logger *zap.Logger) base.IConnector {
 }
 
 // CreateExecution returns an IExecution that executes tasks in Archetype AI.
-func (c *connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &execution{
-		client: newClient(config, logger),
+		client: newClient(connection, logger),
 	}
 
 	switch task {
@@ -81,7 +81,7 @@ func (c *connector) CreateExecution(defUID uuid.UUID, task string, config *struc
 		)
 	}
 
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 
 	return e, nil
 }

diff --git a/pkg/connector/bigquery/v0/main.go b/pkg/connector/bigquery/v0/main.go
@@ -37,11 +37,11 @@ type Execution struct {
 	base.Execution
 }
 
-func Init(logger *zap.Logger) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler) base.IConnector {
 	once.Do(func() {
 		connector = &Connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 		}
 		err := connector.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
@@ -52,9 +52,9 @@ func Init(logger *zap.Logger) base.IConnector {
 	return connector
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &Execution{}
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 	return e, nil
 }
 

diff --git a/pkg/connector/googlecloudstorage/v0/main.go b/pkg/connector/googlecloudstorage/v0/main.go
@@ -36,11 +36,11 @@ type Execution struct {
 	base.Execution
 }
 
-func Init(logger *zap.Logger) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler) base.IConnector {
 	once.Do(func() {
 		connector = &Connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 		}
 		err := connector.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
@@ -51,9 +51,9 @@ func Init(logger *zap.Logger) base.IConnector {
 	return connector
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &Execution{}
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 	return e, nil
 }
 

diff --git a/pkg/connector/googlesearch/v0/main.go b/pkg/connector/googlesearch/v0/main.go
@@ -38,11 +38,11 @@ type Execution struct {
 	base.Execution
 }
 
-func Init(logger *zap.Logger) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler) base.IConnector {
 	once.Do(func() {
 		connector = &Connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 		}
 		err := connector.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
@@ -53,9 +53,9 @@ func Init(logger *zap.Logger) base.IConnector {
 	return connector
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &Execution{}
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 	return e, nil
 }
 

diff --git a/pkg/connector/huggingface/v0/connector_test.go b/pkg/connector/huggingface/v0/connector_test.go
@@ -222,7 +222,7 @@ func TestConnector_ExecuteSpeechRecognition(t *testing.T) {
 
 func testTask(c *qt.C, p taskParams) {
 	logger := zap.NewNop()
-	connector := Init(logger)
+	connector := Init(logger, nil)
 	defID := uuid.Must(uuid.NewV4())
 
 	c.Run("nok - HTTP client error - "+p.task, func(c *qt.C) {

diff --git a/pkg/connector/huggingface/v0/main.go b/pkg/connector/huggingface/v0/main.go
@@ -55,11 +55,11 @@ type Execution struct {
 	base.Execution
 }
 
-func Init(logger *zap.Logger) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler) base.IConnector {
 	once.Do(func() {
 		connector = &Connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 		}
 		err := connector.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
@@ -70,9 +70,9 @@ func Init(logger *zap.Logger) base.IConnector {
 	return connector
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &Execution{}
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 	return e, nil
 }
 

diff --git a/pkg/connector/instill/v0/main.go b/pkg/connector/instill/v0/main.go
@@ -42,11 +42,11 @@ type Execution struct {
 	base.Execution
 }
 
-func Init(logger *zap.Logger) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler) base.IConnector {
 	once.Do(func() {
 		connector = &Connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 		}
 		err := connector.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
@@ -57,9 +57,9 @@ func Init(logger *zap.Logger) base.IConnector {
 	return connector
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &Execution{}
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 	return e, nil
 }
 

diff --git a/pkg/connector/main.go b/pkg/connector/main.go
@@ -39,29 +39,29 @@ type ConnectorOptions struct {
 	Airbyte airbyte.ConnectorOptions
 }
 
-func Init(logger *zap.Logger, options ConnectorOptions) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler, options ConnectorOptions) base.IConnector {
 	once.Do(func() {
 
 		connector = &Connector{
-			Connector:       base.Connector{Component: base.Component{Logger: logger}},
+			Connector:       base.Connector{Component: base.Component{Logger: logger, UsageHandler: usageHandler}},
 			connectorUIDMap: map[uuid.UUID]base.IConnector{},
 			connectorIDMap:  map[string]base.IConnector{},
 		}
 
-		connector.(*Connector).ImportDefinitions(stabilityai.Init(logger))
-		connector.(*Connector).ImportDefinitions(instill.Init(logger))
-		connector.(*Connector).ImportDefinitions(huggingface.Init(logger))
-		connector.(*Connector).ImportDefinitions(openai.Init(logger))
-		connector.(*Connector).ImportDefinitions(archetypeai.Init(logger))
-		connector.(*Connector).ImportDefinitions(numbers.Init(logger))
-		connector.(*Connector).ImportDefinitions(airbyte.Init(logger, options.Airbyte))
-		connector.(*Connector).ImportDefinitions(bigquery.Init(logger))
-		connector.(*Connector).ImportDefinitions(googlecloudstorage.Init(logger))
-		connector.(*Connector).ImportDefinitions(googlesearch.Init(logger))
-		connector.(*Connector).ImportDefinitions(pinecone.Init(logger))
-		connector.(*Connector).ImportDefinitions(redis.Init(logger))
-		connector.(*Connector).ImportDefinitions(restapi.Init(logger))
-		connector.(*Connector).ImportDefinitions(website.Init(logger))
+		connector.(*Connector).ImportDefinitions(stabilityai.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(instill.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(huggingface.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(openai.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(archetypeai.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(numbers.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(airbyte.Init(logger, usageHandler, options.Airbyte))
+		connector.(*Connector).ImportDefinitions(bigquery.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(googlecloudstorage.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(googlesearch.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(pinecone.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(redis.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(restapi.Init(logger, usageHandler))
+		connector.(*Connector).ImportDefinitions(website.Init(logger, usageHandler))
 
 	})
 	return connector
@@ -77,8 +77,8 @@ func (c *Connector) ImportDefinitions(con base.IConnector) {
 	}
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
-	return c.connectorUIDMap[defUID].CreateExecution(defUID, task, config, logger)
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+	return c.connectorUIDMap[defUID].CreateExecution(defUID, task, connection, logger)
 }
 
 func (c *Connector) Test(defUID uuid.UUID, config *structpb.Struct, logger *zap.Logger) error {

diff --git a/pkg/connector/numbers/v0/main.go b/pkg/connector/numbers/v0/main.go
@@ -105,12 +105,12 @@ type Output struct {
 	AssetUrls []string `json:"asset_urls"`
 }
 
-func Init(logger *zap.Logger) base.IConnector {
+func Init(logger *zap.Logger, usageHandler base.UsageHandler) base.IConnector {
 	once.Do(func() {
 
 		connector = &Connector{
 			Connector: base.Connector{
-				Component: base.Component{Logger: logger},
+				Component: base.Component{Logger: logger, UsageHandler: usageHandler},
 			},
 		}
 		err := connector.LoadConnectorDefinition(definitionJSON, tasksJSON, nil)
@@ -201,9 +201,9 @@ func (e *Execution) registerAsset(data []byte, reg Register) (string, error) {
 	}
 }
 
-func (c *Connector) CreateExecution(defUID uuid.UUID, task string, config *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
+func (c *Connector) CreateExecution(defUID uuid.UUID, task string, connection *structpb.Struct, logger *zap.Logger) (base.IExecution, error) {
 	e := &Execution{}
-	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, config, logger)
+	e.Execution = base.CreateExecutionHelper(e, c, defUID, task, connection, logger)
 	return e, nil
 }
 

diff --git a/pkg/connector/openai/v0/connector_test.go b/pkg/connector/openai/v0/connector_test.go
@@ -30,7 +30,7 @@ func TestConnector_Execute(t *testing.T) {
 	c := qt.New(t)
 
 	logger := zap.NewNop()
-	connector := Init(logger)
+	connector := Init(logger, nil)
 	defID := uuid.Must(uuid.NewV4())
 
 	testcases := []struct {
@@ -130,7 +130,7 @@ func TestConnector_Test(t *testing.T) {
 	c := qt.New(t)
 
 	logger := zap.NewNop()
-	connector := Init(logger)
+	connector := Init(logger, nil)
 	defID := uuid.Must(uuid.NewV4())
 
 	c.Run("nok - error", func(c *qt.C) {