Skip to content

Commit

Permalink
feat(universalai): add instill credit function for universal ai compo…
Browse files Browse the repository at this point in the history 
…nent (#680)

This is a port of instill-ai/component#362

Because

- we will support universal ai to execute vendors' ai models

This commit

- import the vendors' key to component and execution

---------

Co-authored-by: chuang8511 <[email protected]>
  • Loading branch information
@jvallesm @chuang8511
jvallesm and chuang8511 authored Oct 7, 2024
1 parent 49aed28 commit 9ce10b5
Show file tree
Hide file tree
Showing 8 changed files with 263 additions and 106 deletions.
2 changes: 0 additions & 2 deletions pkg/component/ai/universalai/v0/component_test.go
View file

This file was deleted.

2 changes: 1 addition & 1 deletion pkg/component/ai/universalai/v0/config/definition.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
"type": "COMPONENT_TYPE_AI",
"uid": "7656cb11-d504-4ca0-b481-6ef80964f2c9",
"vendorAttributes": {},
"version": "0.1.0",
"version": "0.2.0",
"sourceUrl": "https://github.com/instill-ai/pipeline-backend/blob/main/pkg/component/ai/universalai/v0",
"releaseStage": "RELEASE_STAGE_ALPHA"
}
62 changes: 59 additions & 3 deletions pkg/component/ai/universalai/v0/config/setup.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,60 @@
"$schema": "http://json-schema.org/draft-07/schema#",
"additionalProperties": true,
"properties": {
"model": {
"description": "The model to be used. Now, it only supports OpenAI model, and will support more models in the future.",
"instillShortDescription": "The model to be used.",
"instillAcceptFormats": [
"string"
],
"enum": [
"o1-preview",
"o1-mini",
"gpt-4o-mini",
"gpt-4o",
"gpt-4o-2024-05-13",
"gpt-4o-2024-08-06",
"gpt-4-turbo",
"gpt-4-turbo-2024-04-09",
"gpt-4-0125-preview",
"gpt-4-turbo-preview",
"gpt-4-1106-preview",
"gpt-4-vision-preview",
"gpt-4",
"gpt-4-0314",
"gpt-4-0613",
"gpt-4-32k",
"gpt-4-32k-0314",
"gpt-4-32k-0613",
"gpt-3.5-turbo",
"gpt-3.5-turbo-16k",
"gpt-3.5-turbo-0301",
"gpt-3.5-turbo-0613",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo-0125",
"gpt-3.5-turbo-16k-0613"
],
"instillCredentialMap": {
"values": [
"o1-preview",
"o1-mini",
"gpt-4o",
"gpt-4o-2024-08-06",
"gpt-4-turbo",
"gpt-4-vision-preview",
"gpt-4",
"gpt-4-32k",
"gpt-3.5-turbo",
"gpt-4o-mini"
],
"targets": [
"setup.api-key"
]
},
"instillUIOrder": 0,
"title": "Model Name",
"type": "string"
},
"api-key": {
"description": "Fill in your API key from the vendor's platform.",
"instillUpstreamTypes": [
Expand All @@ -12,7 +66,7 @@
],
"instillSecret": true,
"instillCredential": true,
"instillUIOrder": 0,
"instillUIOrder": 1,
"title": "API Key",
"type": "string"
},
Expand All @@ -24,12 +78,14 @@
"instillAcceptFormats": [
"string"
],
"instillUIOrder": 1,
"instillUIOrder": 2,
"title": "Organization ID",
"type": "string"
}
},
"required": [],
"required": [
"model"
],
"instillEditOnNodeFields": [
"api-key"
],
Expand Down
57 changes: 1 addition & 56 deletions pkg/component/ai/universalai/v0/config/tasks.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,60 +15,6 @@
"instillShortDescription": "Input data",
"type": "object",
"properties": {
"model": {
"description": "The model to be used. Now, it only supports OpenAI model, and will support more models in the future.",
"instillShortDescription": "The model to be used.",
"instillAcceptFormats": [
"string"
],
"enum": [
"o1-preview",
"o1-mini",
"gpt-4o-mini",
"gpt-4o",
"gpt-4o-2024-05-13",
"gpt-4o-2024-08-06",
"gpt-4-turbo",
"gpt-4-turbo-2024-04-09",
"gpt-4-0125-preview",
"gpt-4-turbo-preview",
"gpt-4-1106-preview",
"gpt-4-vision-preview",
"gpt-4",
"gpt-4-0314",
"gpt-4-0613",
"gpt-4-32k",
"gpt-4-32k-0314",
"gpt-4-32k-0613",
"gpt-3.5-turbo",
"gpt-3.5-turbo-16k",
"gpt-3.5-turbo-0301",
"gpt-3.5-turbo-0613",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo-0125",
"gpt-3.5-turbo-16k-0613"
],
"instillCredentialMap": {
"values": [
"o1-preview",
"o1-mini",
"gpt-4o",
"gpt-4o-2024-08-06",
"gpt-4-turbo",
"gpt-4-vision-preview",
"gpt-4",
"gpt-4-32k",
"gpt-3.5-turbo",
"gpt-4o-mini"
],
"targets": [
"setup.api-key"
]
},
"instillUIOrder": 0,
"title": "Model Name",
"type": "string"
},
"messages": {
"title": "Chat Messages",
"type": "array",
Expand Down Expand Up @@ -203,12 +149,11 @@
"role"
]
},
"instillUIOrder": 1,
"instillUIOrder": 0,
"description": "List of chat messages"
}
},
"required": [
"model",
"messages"
],
"instillUIOrder": 0
Expand Down
57 changes: 49 additions & 8 deletions pkg/component/ai/universalai/v0/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,11 +36,12 @@ var (
type component struct {
base.Component

instillAPIKey string
instillAPIKey map[string]string
}

type execution struct {
base.ComponentExecution

usesInstillCredentials bool
execute func(*structpb.Struct, *base.Job, context.Context) (*structpb.Struct, error)
}
Expand All @@ -59,7 +60,12 @@ func Init(bc base.Component) *component {
}

func (c *component) CreateExecution(x base.ComponentExecution) (base.IExecution, error) {
resolvedSetup, resolved, err := c.resolveSetup(x.Setup)

model := getModel(x.GetSetup())
vendor := modelVendorMap[model]

resolvedSetup, resolved, err := c.resolveSetup(vendor, x.Setup)

if err != nil {
return nil, err
}
Expand All @@ -73,7 +79,7 @@ func (c *component) CreateExecution(x base.ComponentExecution) (base.IExecution,

switch x.Task {
case TextChatTask:
e.execute = e.ExecuteTextChat
e.execute = e.executeTextChat
default:
return nil, fmt.Errorf("unknown task: %s", x.Task)
}
Expand All @@ -92,15 +98,18 @@ func (e *execution) UsesInstillCredentials() bool {
// WithInstillCredentials loads Instill credentials into the component, which
// can be used to configure it with globally defined parameters instead of with
// user-defined credential values.
func (c *component) WithInstillCredentials(s map[string]any) *component {
c.instillAPIKey = base.ReadFromGlobalConfig(cfgAPIKey, s)
func (c *component) WithInstillCredentials(vendor string, s map[string]any) *component {
if c.instillAPIKey == nil {
c.instillAPIKey = make(map[string]string)
}
c.instillAPIKey[vendor] = base.ReadFromGlobalConfig(cfgAPIKey, s)
return c
}

// resolveSetup checks whether the component is configured to use the Instill
// credentials injected during initialization and, if so, returns a new setup
// with the secret credential values.
func (c *component) resolveSetup(setup *structpb.Struct) (*structpb.Struct, bool, error) {
func (c *component) resolveSetup(vendor string, setup *structpb.Struct) (*structpb.Struct, bool, error) {
if setup == nil || setup.Fields == nil {
setup = &structpb.Struct{Fields: map[string]*structpb.Value{}}
}
Expand All @@ -111,10 +120,42 @@ func (c *component) resolveSetup(setup *structpb.Struct) (*structpb.Struct, bool
}
}

if c.instillAPIKey == "" {
if c.instillAPIKey[vendor] == "" {
return nil, false, base.NewUnresolvedCredential(cfgAPIKey)
}

setup.GetFields()[cfgAPIKey] = structpb.NewStringValue(c.instillAPIKey)
setup.GetFields()[cfgAPIKey] = structpb.NewStringValue(c.instillAPIKey[vendor])
return setup, true, nil
}

func getModel(setup *structpb.Struct) string {
return setup.GetFields()["model"].GetStringValue()
}

var modelVendorMap = map[string]string{
"o1-preview": "openai",
"o1-mini": "openai",
"gpt-4o-mini": "openai",
"gpt-4o": "openai",
"gpt-4o-2024-05-13": "openai",
"gpt-4o-2024-08-06": "openai",
"gpt-4-turbo": "openai",
"gpt-4-turbo-2024-04-09": "openai",
"gpt-4-0125-preview": "openai",
"gpt-4-turbo-preview": "openai",
"gpt-4-1106-preview": "openai",
"gpt-4-vision-preview": "openai",
"gpt-4": "openai",
"gpt-4-0314": "openai",
"gpt-4-0613": "openai",
"gpt-4-32k": "openai",
"gpt-4-32k-0314": "openai",
"gpt-4-32k-0613": "openai",
"gpt-3.5-turbo": "openai",
"gpt-3.5-turbo-16k": "openai",
"gpt-3.5-turbo-0301": "openai",
"gpt-3.5-turbo-0613": "openai",
"gpt-3.5-turbo-1106": "openai",
"gpt-3.5-turbo-0125": "openai",
"gpt-3.5-turbo-16k-0613": "openai",
}
67 changes: 36 additions & 31 deletions pkg/component/ai/universalai/v0/text_chat_task.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,20 +13,33 @@ import (
openaiv1 "github.com/instill-ai/pipeline-backend/pkg/component/ai/openai/v1"
)

func (e *execution) ExecuteTextChat(input *structpb.Struct, job *base.Job, ctx context.Context) (*structpb.Struct, error) {
func (e *execution) executeTextChat(input *structpb.Struct, job *base.Job, ctx context.Context) (*structpb.Struct, error) {

x := e.ComponentExecution
model := getModel(x.GetSetup())

err := insertModel(input, model)

if err != nil {
return nil, err
}

inputStruct := ai.TextChatInput{}

if err := base.ConvertFromStructpb(input, &inputStruct); err != nil {
return nil, fmt.Errorf("failed to convert input to TextChatInput: %w", err)
return nil, err
}

x := e.ComponentExecution
vendor := ModelVendorMap[inputStruct.Data.Model]
vendor, ok := modelVendorMap[model]

if !ok {
return nil, fmt.Errorf("unsupported vendor for model: %s", model)
}

client, err := newClient(x.GetSetup(), x.GetLogger(), vendor)

if err != nil {
return nil, fmt.Errorf("failed to create client: %w", err)
return nil, err
}

switch vendor {
Expand All @@ -37,30 +50,22 @@ func (e *execution) ExecuteTextChat(input *structpb.Struct, job *base.Job, ctx c
}
}

var ModelVendorMap = map[string]string{
"o1-preview": "openai",
"o1-mini": "openai",
"gpt-4o-mini": "openai",
"gpt-4o": "openai",
"gpt-4o-2024-05-13": "openai",
"gpt-4o-2024-08-06": "openai",
"gpt-4-turbo": "openai",
"gpt-4-turbo-2024-04-09": "openai",
"gpt-4-0125-preview": "openai",
"gpt-4-turbo-preview": "openai",
"gpt-4-1106-preview": "openai",
"gpt-4-vision-preview": "openai",
"gpt-4": "openai",
"gpt-4-0314": "openai",
"gpt-4-0613": "openai",
"gpt-4-32k": "openai",
"gpt-4-32k-0314": "openai",
"gpt-4-32k-0613": "openai",
"gpt-3.5-turbo": "openai",
"gpt-3.5-turbo-16k": "openai",
"gpt-3.5-turbo-0301": "openai",
"gpt-3.5-turbo-0613": "openai",
"gpt-3.5-turbo-1106": "openai",
"gpt-3.5-turbo-0125": "openai",
"gpt-3.5-turbo-16k-0613": "openai",
// In the implementation, the model is more like the input of execution than the setup of the component.
// However, we should set the model in setup to be able to resolve the setup for the key in the vendor map.
// To avoid users inputting the model in the setup and params, we insert the model into input data.
func insertModel(input *structpb.Struct, model string) error {

inputData, ok := input.Fields["data"]
if !ok {
return fmt.Errorf("failed to get data from input: no 'data' field found")
}

dataStruct, ok := inputData.GetKind().(*structpb.Value_StructValue)
if !ok {
return fmt.Errorf("data field is not a struct")
}

dataStruct.StructValue.Fields["model"] = structpb.NewStringValue(model)

return nil
}
Loading

0 comments on commit 9ce10b5

Please sign in to comment.