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()