github.com/instill-ai/component@v0.16.0-beta/pkg/base/execution.go (about)

     1  package base
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"strconv"
     7  	"strings"
     8  
     9  	"github.com/santhosh-tekuri/jsonschema/v5"
    10  	"go.uber.org/zap"
    11  	"google.golang.org/protobuf/encoding/protojson"
    12  	"google.golang.org/protobuf/types/known/structpb"
    13  )
    14  
    15  type ExecutionWrapper struct {
    16  	Execution IExecution
    17  }
    18  
    19  type IExecution interface {
    20  	GetTask() string
    21  	GetLogger() *zap.Logger
    22  	GetUsageHandler() UsageHandler
    23  	GetTaskInputSchema() string
    24  	GetTaskOutputSchema() string
    25  
    26  	Execute([]*structpb.Struct) ([]*structpb.Struct, error)
    27  }
    28  
    29  func FormatErrors(inputPath string, e jsonschema.Detailed, errors *[]string) {
    30  
    31  	path := inputPath + e.InstanceLocation
    32  
    33  	pathItems := strings.Split(path, "/")
    34  	formatedPath := pathItems[0]
    35  	for _, pathItem := range pathItems[1:] {
    36  		if _, err := strconv.Atoi(pathItem); err == nil {
    37  			formatedPath += fmt.Sprintf("[%s]", pathItem)
    38  		} else {
    39  			formatedPath += fmt.Sprintf(".%s", pathItem)
    40  		}
    41  
    42  	}
    43  	*errors = append(*errors, fmt.Sprintf("%s: %s", formatedPath, e.Error))
    44  
    45  }
    46  
    47  // Validate the input and output format
    48  func Validate(data []*structpb.Struct, jsonSchema string, target string) error {
    49  
    50  	schStruct := &structpb.Struct{}
    51  	err := protojson.Unmarshal([]byte(jsonSchema), schStruct)
    52  	if err != nil {
    53  		return err
    54  	}
    55  
    56  	err = CompileInstillAcceptFormats(schStruct)
    57  	if err != nil {
    58  		return err
    59  	}
    60  
    61  	schStr, err := protojson.Marshal(schStruct)
    62  	if err != nil {
    63  		return err
    64  	}
    65  
    66  	c := jsonschema.NewCompiler()
    67  	c.RegisterExtension("instillAcceptFormats", InstillAcceptFormatsMeta, InstillAcceptFormatsCompiler{})
    68  	c.RegisterExtension("instillFormat", InstillFormatMeta, InstillFormatCompiler{})
    69  	if err := c.AddResource("schema.json", strings.NewReader(string(schStr))); err != nil {
    70  		return err
    71  	}
    72  	sch, err := c.Compile("schema.json")
    73  	if err != nil {
    74  		return err
    75  	}
    76  	errors := []string{}
    77  
    78  	for idx := range data {
    79  		var v interface{}
    80  		jsonData, err := protojson.Marshal(data[idx])
    81  		if err != nil {
    82  			errors = append(errors, fmt.Sprintf("%s[%d]: data error", target, idx))
    83  			continue
    84  		}
    85  
    86  		if err := json.Unmarshal(jsonData, &v); err != nil {
    87  			errors = append(errors, fmt.Sprintf("%s[%d]: data error", target, idx))
    88  			continue
    89  		}
    90  
    91  		if err = sch.Validate(v); err != nil {
    92  			e := err.(*jsonschema.ValidationError)
    93  
    94  			for _, valErr := range e.DetailedOutput().Errors {
    95  				inputPath := fmt.Sprintf("%s/%d", target, idx)
    96  				FormatErrors(inputPath, valErr, &errors)
    97  				for _, subValErr := range valErr.Errors {
    98  					FormatErrors(inputPath, subValErr, &errors)
    99  				}
   100  			}
   101  		}
   102  	}
   103  
   104  	if len(errors) > 0 {
   105  		return fmt.Errorf("%s", strings.Join(errors, "; "))
   106  	}
   107  
   108  	return nil
   109  }
   110  
   111  // ExecuteWithValidation executes the execution with validation
   112  func (e *ExecutionWrapper) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
   113  
   114  	if err := Validate(inputs, e.Execution.GetTaskInputSchema(), "inputs"); err != nil {
   115  		return nil, err
   116  	}
   117  
   118  	if e.Execution.GetUsageHandler() != nil {
   119  		if err := e.Execution.GetUsageHandler().Check(); err != nil {
   120  			return nil, err
   121  		}
   122  	}
   123  
   124  	outputs, err := e.Execution.Execute(inputs)
   125  	if err != nil {
   126  		return nil, err
   127  	}
   128  
   129  	if err := Validate(outputs, e.Execution.GetTaskOutputSchema(), "outputs"); err != nil {
   130  		return nil, err
   131  	}
   132  
   133  	if e.Execution.GetUsageHandler() != nil {
   134  		if err := e.Execution.GetUsageHandler().Collect(); err != nil {
   135  			return nil, err
   136  		}
   137  	}
   138  
   139  	return outputs, err
   140  }