diff --git a/pkg/component/ai/universalai/v0/component_test.go b/pkg/component/ai/universalai/v0/component_test.go deleted file mode 100644 index 19a54bdcc..000000000 --- a/pkg/component/ai/universalai/v0/component_test.go +++ /dev/null @@ -1,2 +0,0 @@ -// TODO: chuang8511 -package universalai diff --git a/pkg/component/ai/universalai/v0/config/definition.json b/pkg/component/ai/universalai/v0/config/definition.json index 831fa2bab..5be3b99b1 100644 --- a/pkg/component/ai/universalai/v0/config/definition.json +++ b/pkg/component/ai/universalai/v0/config/definition.json @@ -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" } diff --git a/pkg/component/ai/universalai/v0/config/setup.json b/pkg/component/ai/universalai/v0/config/setup.json index 5e9bc6f22..e44618ced 100644 --- a/pkg/component/ai/universalai/v0/config/setup.json +++ b/pkg/component/ai/universalai/v0/config/setup.json @@ -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": [ @@ -12,7 +66,7 @@ ], "instillSecret": true, "instillCredential": true, - "instillUIOrder": 0, + "instillUIOrder": 1, "title": "API Key", "type": "string" }, @@ -24,12 +78,14 @@ "instillAcceptFormats": [ "string" ], - "instillUIOrder": 1, + "instillUIOrder": 2, "title": "Organization ID", "type": "string" } }, - "required": [], + "required": [ + "model" + ], "instillEditOnNodeFields": [ "api-key" ], diff --git a/pkg/component/ai/universalai/v0/config/tasks.json b/pkg/component/ai/universalai/v0/config/tasks.json index 365fcc871..e88557c89 100644 --- a/pkg/component/ai/universalai/v0/config/tasks.json +++ b/pkg/component/ai/universalai/v0/config/tasks.json @@ -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", @@ -203,12 +149,11 @@ "role" ] }, - "instillUIOrder": 1, + "instillUIOrder": 0, "description": "List of chat messages" } }, "required": [ - "model", "messages" ], "instillUIOrder": 0 diff --git a/pkg/component/ai/universalai/v0/main.go b/pkg/component/ai/universalai/v0/main.go index 14b39728e..02e02193e 100644 --- a/pkg/component/ai/universalai/v0/main.go +++ b/pkg/component/ai/universalai/v0/main.go @@ -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) } @@ -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 } @@ -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) } @@ -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{}} } @@ -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", +} diff --git a/pkg/component/ai/universalai/v0/text_chat_task.go b/pkg/component/ai/universalai/v0/text_chat_task.go index 3703db607..6883ed3b4 100644 --- a/pkg/component/ai/universalai/v0/text_chat_task.go +++ b/pkg/component/ai/universalai/v0/text_chat_task.go @@ -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 { @@ -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 } diff --git a/pkg/component/ai/universalai/v0/text_chat_task_test.go b/pkg/component/ai/universalai/v0/text_chat_task_test.go new file mode 100644 index 000000000..b9589d9b9 --- /dev/null +++ b/pkg/component/ai/universalai/v0/text_chat_task_test.go @@ -0,0 +1,110 @@ +package universalai + +import ( + "context" + "testing" + + "google.golang.org/protobuf/types/known/structpb" + + qt "github.com/frankban/quicktest" + + "github.com/instill-ai/pipeline-backend/pkg/component/base" +) + +func TestExecuteTextChat(t *testing.T) { + c := qt.New(t) + + testCases := []struct { + name string + modelName string + fakeAPIKey string + wantErrMsg string + }{ + // This test case validate that the model is not supported. + { + name: "Unsupported Model", + modelName: "testModel", + fakeAPIKey: "testAPIKey", + wantErrMsg: "unsupported vendor for model: testModel", + }, + // This test case validate that the request is sent from UniversalAI to OpenAI. + // The other cases should be supported in openai/v1/text_chat_task_test.go. + { + name: "OpenAI Model", + modelName: "gpt-4", + fakeAPIKey: "", + wantErrMsg: "send request to openai error with error code: 401", + }, + } + + component := Init(base.Component{}) + component.instillAPIKey = map[string]string{ + "openai": "testAPIKey", + } + + for _, tc := range testCases { + c.Run(tc.name, func(t *qt.C) { + setup, err := structpb.NewStruct(map[string]any{ + "model": tc.modelName, + "api-key": tc.fakeAPIKey, + }) + + c.Assert(err, qt.IsNil) + + x := base.ComponentExecution{ + Component: component, + Setup: setup, + Task: "TASK_CHAT", + } + + e, err := component.CreateExecution(x) + + c.Assert(err, qt.IsNil) + + ctx := context.Background() + job := &base.Job{} + + input, err := structpb.NewStruct(map[string]interface{}{ + "data": map[string]interface{}{ + "messages": []interface{}{ + map[string]interface{}{ + "role": "user", + "name": "John", + "content": []interface{}{ + map[string]interface{}{ + "type": "text", + "text": "Hello, how can I help you?", + }, + }, + }, + map[string]interface{}{ + "role": "assistant", + "content": []interface{}{ + map[string]interface{}{ + "type": "text", + "text": "I'm here to assist you.", + }, + }, + }, + }, + }, + "parameter": map[string]interface{}{ + "max-tokens": 100, + "temperature": 0.7, + "top-p": 0.9, + "stream": false, + }, + }) + + c.Assert(err, qt.IsNil) + + execution := e.(*execution) + + _, err = execution.executeTextChat(input, job, ctx) + + c.Assert(err.Error(), qt.Contains, tc.wantErrMsg) + + }) + } + +} diff --git a/pkg/component/store/store.go b/pkg/component/store/store.go index 8f25f0f5f..47d7d80e6 100644 --- a/pkg/component/store/store.go +++ b/pkg/component/store/store.go @@ -19,6 +19,7 @@ import ( "github.com/instill-ai/pipeline-backend/pkg/component/ai/ollama/v0" "github.com/instill-ai/pipeline-backend/pkg/component/ai/openai/v0" "github.com/instill-ai/pipeline-backend/pkg/component/ai/stabilityai/v0" + "github.com/instill-ai/pipeline-backend/pkg/component/ai/universalai/v0" "github.com/instill-ai/pipeline-backend/pkg/component/application/asana/v0" "github.com/instill-ai/pipeline-backend/pkg/component/application/email/v0" "github.com/instill-ai/pipeline-backend/pkg/component/application/freshdesk/v0" @@ -123,11 +124,12 @@ func Init( conn = conn.WithInstillCredentials(secrets[conn.GetDefinitionID()]) compStore.Import(conn) } - // { - // conn := universalai.Init(baseComp) - // conn = conn.WithInstillCredentials(secrets[conn.GetDefinitionID()]) - // compStore.Import(conn) - // } + { + conn := universalai.Init(baseComp) + // Please apply more keys when we add more vendors + conn = conn.WithInstillCredentials("openai", secrets["openai"]) + compStore.Import(conn) + } { // Anthropic conn := anthropic.Init(baseComp)