From 538c6c063adcfdaaaeaab922e2c992e9e4cbf6b8 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 12 May 2023 21:47:18 +0800 Subject: [PATCH 1/2] [add] openai.* process --- openai/process.go | 95 +++++++++++++++++++++++++++++++++++++++++- openai/process_test.go | 68 +++++++++++++++++++++++++++++- 2 files changed, 160 insertions(+), 3 deletions(-) diff --git a/openai/process.go b/openai/process.go index bad5f911..0149e89b 100644 --- a/openai/process.go +++ b/openai/process.go @@ -1,12 +1,20 @@ package openai import ( + "context" + + "github.com/yaoapp/gou/http" "github.com/yaoapp/gou/process" + "github.com/yaoapp/gou/runtime/v8/bridge" "github.com/yaoapp/kun/exception" + "github.com/yaoapp/kun/log" ) func init() { - process.Register("yao.openai.tiktoken", ProcessTiktoken) + process.RegisterGroup("openai", map[string]process.Handler{ + "tiktoken": ProcessTiktoken, + "chat.completions": ProcessChatCompletions, + }) } // ProcessTiktoken get number of tokens @@ -20,3 +28,88 @@ func ProcessTiktoken(process *process.Process) interface{} { } return nums } + +// ProcessChatCompletions get number of tokens +func ProcessChatCompletions(process *process.Process) interface{} { + + process.ValidateArgNums(2) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + model := process.ArgsString(0) + messages := []map[string]interface{}{} + intput := process.ArgsArray(1) + for idx, v := range intput { + message, ok := v.(map[string]interface{}) + if !ok { + exception.New("ChatCompletions input must be array of map, index %d", 400, idx).Throw() + } + messages = append(messages, message) + } + + ai, err := New(model) + if err != nil { + exception.New("ChatCompletions error: %s", 400, err).Throw() + } + + options := map[string]interface{}{} + if process.NumOfArgs() > 2 { + if opts, ok := process.Args[2].(map[string]interface{}); ok { + options = opts + } + } + + if process.NumOfArgs() == 3 { + data, ex := ai.ChatCompletionsWith(ctx, messages, options, nil) + if ex != nil { + ex.Throw() + } + return data + } + + if process.NumOfArgs() == 4 { + + switch cb := process.Args[3].(type) { + case func(data []byte) int: + res, ex := ai.ChatCompletionsWith(ctx, messages, options, cb) + if ex != nil { + ex.Throw() + } + return res + + case bridge.FunctionT: + res, ex := ai.ChatCompletionsWith(ctx, messages, options, func(data []byte) int { + + v, err := cb.Call(string(data)) + if err != nil { + log.Error("Call callback function error: %s", err.Error()) + return http.HandlerReturnError + } + + ret, ok := v.(int) + if !ok { + log.Error("Callback function must return int") + return http.HandlerReturnError + } + + return ret + }) + + if ex != nil { + ex.Throw() + } + return res + + default: + exception.New("ChatCompletions error: invalid callback arguments", 400).Throw() + return nil + } + } + + res, ex := ai.ChatCompletionsWith(ctx, messages, options, nil) + if ex != nil { + ex.Throw() + } + return res + +} diff --git a/openai/process_test.go b/openai/process_test.go index 9ab42e07..2900e8ec 100644 --- a/openai/process_test.go +++ b/openai/process_test.go @@ -5,15 +5,79 @@ import ( "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/process" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/test" ) func TestProcessTiktoken(t *testing.T) { // Hash args := []interface{}{"gpt-3.5-turbo", "hello world"} - res := process.New("yao.openai.Tiktoken", args...).Run() + res := process.New("openai.Tiktoken", args...).Run() assert.Equal(t, 2, res) args = []interface{}{"gpt-3.5-turbo", "你好世界！"} - res = process.New("yao.openai.Tiktoken", args...).Run() + res = process.New("openai.Tiktoken", args...).Run() assert.Equal(t, 6, res) } + +func TestProcessChatCompletions(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + args := []interface{}{"gpt-3_5-turbo", []map[string]interface{}{{"role": "user", "content": "hello"}}} + res := process.New("openai.chat.Completions", args...).Run() + data, ok := res.(map[string]interface{}) + if !ok { + t.Fatalf("ChatCompletions return type error") + } + assert.NotEmpty(t, data["id"]) + + // With options + args = []interface{}{ + "gpt-3_5-turbo", + []map[string]interface{}{{"role": "user", "content": "hello"}}, + map[string]interface{}{"max_tokens": 2}, + } + res = process.New("openai.chat.Completions", args...).Run() + data, ok = res.(map[string]interface{}) + if !ok { + t.Fatalf("ChatCompletions return type error") + } + + usage, ok := data["usage"].(map[string]interface{}) + if !ok { + t.Fatalf("ChatCompletions return type error") + } + assert.Equal(t, 2, int(usage["completion_tokens"].(float64))) + + // With callback + content := []byte{} + args = []interface{}{ + "gpt-3_5-turbo", + []map[string]interface{}{{"role": "user", "content": "hello"}}, + nil, + func(data []byte) int { + + content = append(content, data...) + if len(data) == 0 { + res = append(content, []byte("\n")...) + } + + if string(data) == "data: [DONE]" { + return 0 + } + + return 1 + }, + } + res = process.New("openai.chat.Completions", args...).Run() + assert.Contains(t, string(content), "[DONE]") + + // With JS Callback + res, err := process.New("scripts.openai.TestProcessChatCompletions").Exec() + if err != nil { + t.Fatal(err) + } + + assert.Contains(t, res, "[DONE]") +} From 1d102f110bee2aeb63f6809b9376664e05c12f01 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 12 May 2023 22:11:31 +0800 Subject: [PATCH 2/2] [add] openai process --- openai/process.go | 56 +++++++++++++++++++++++++++++++++++++++--- openai/process_test.go | 22 +++++++++++++++++ 2 files changed, 74 insertions(+), 4 deletions(-) diff --git a/openai/process.go b/openai/process.go index 0149e89b..ce0e266f 100644 --- a/openai/process.go +++ b/openai/process.go @@ -12,12 +12,14 @@ import ( func init() { process.RegisterGroup("openai", map[string]process.Handler{ - "tiktoken": ProcessTiktoken, - "chat.completions": ProcessChatCompletions, + "tiktoken": ProcessTiktoken, + "embeddings": ProcessEmbeddings, + "chat.completions": ProcessChatCompletions, + "audio.transcriptions": ProcessAudioTranscriptions, }) } -// ProcessTiktoken get number of tokens +// ProcessTiktoken openai.Tiktoken func ProcessTiktoken(process *process.Process) interface{} { process.ValidateArgNums(2) model := process.ArgsString(0) @@ -29,7 +31,53 @@ func ProcessTiktoken(process *process.Process) interface{} { return nums } -// ProcessChatCompletions get number of tokens +// ProcessEmbeddings openai.Embeddings +func ProcessEmbeddings(process *process.Process) interface{} { + process.ValidateArgNums(2) + model := process.ArgsString(0) + input := process.Args[1] + user := "" + if process.NumOfArgs() > 2 { + user = process.ArgsString(2) + } + + ai, err := New(model) + if err != nil { + exception.New("ChatCompletions error: %s", 400, err).Throw() + } + + res, ex := ai.Embeddings(input, user) + if ex != nil { + ex.Throw() + } + return res +} + +// ProcessAudioTranscriptions openai.audio.Transcriptions +func ProcessAudioTranscriptions(process *process.Process) interface{} { + process.ValidateArgNums(2) + model := process.ArgsString(0) + dataBase64 := process.ArgsString(1) + options := map[string]interface{}{} + if process.NumOfArgs() > 2 { + if opts, ok := process.Args[2].(map[string]interface{}); ok { + options = opts + } + } + + ai, err := New(model) + if err != nil { + exception.New("ChatCompletions error: %s", 400, err).Throw() + } + + res, ex := ai.AudioTranscriptions(dataBase64, options) + if ex != nil { + ex.Throw() + } + return res +} + +// ProcessChatCompletions openai.chat.Completions func ProcessChatCompletions(process *process.Process) interface{} { process.ValidateArgNums(2) diff --git a/openai/process_test.go b/openai/process_test.go index 2900e8ec..262ffd31 100644 --- a/openai/process_test.go +++ b/openai/process_test.go @@ -20,6 +20,28 @@ func TestProcessTiktoken(t *testing.T) { assert.Equal(t, 6, res) } +func TestProcessEmbeddings(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + args := []interface{}{"text-embedding-ada-002", "hello world"} + data := process.New("openai.Embeddings", args...).Run() + assert.NotNil(t, data.(map[string]interface{})["data"]) + + args = []interface{}{"text-embedding-ada-002", []string{"The food was delicious and the waiter", "hello"}, "user-01"} + data = process.New("openai.Embeddings", args...).Run() + assert.NotNil(t, data.(map[string]interface{})["data"]) +} + +func TestProcessAudioTranscriptions(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + args := []interface{}{"whisper-1", audio(t)} + data := process.New("openai.audio.Transcriptions", args...).Run() + assert.Equal(t, "今晚打老虎", data.(map[string]interface{})["text"]) +} + func TestProcessChatCompletions(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean()