VaalaCat/go-openai
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: allow more input types to functions, fix tests (#377)

Browse Source 
* feat: use json.rawMessage, test functions

* chore: lint

* fix: tests

the ChatCompletion mock server doesn't actually run otherwise. N=0
is the default request but the server will treat it as n=1

* fix: tests should default to n=1 completions

* chore: add back removed interfaces, custom marshal

* chore: lint

* chore: lint

* chore: add some tests

* chore: appease lint

* clean up JSON schema + tests

* chore: lint

* feat: remove backwards compatible functions

for illustrative purposes

* fix: revert params change

* chore: use interface{}

* chore: add test

* chore: add back FunctionDefine

* chore: /s/interface{}/any

* chore: add back jsonschemadefinition

* chore: testcov

* chore: lint

* chore: remove pointers

* chore: update comment

* chore: address CR

added test for compatibility as well

---------

Co-authored-by: James <jmacwhyte@MacBooger-II.local>
This commit is contained in:
Chris Hua
2023-06-21 08:58:27 -04:00
committed by GitHub
parent e948150829
commit f22da8a7ed
3 changed files with 180 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
chat_test.go
View File

@@ -67,6 +67,130 @@ func TestChatCompletions(t *testing.T) {
checks.NoError(t, err, "CreateChatCompletion error")
}
// TestChatCompletionsFunctions tests including a function call.
func TestChatCompletionsFunctions(t *testing.T) {
client, server, teardown := setupOpenAITestServer()
defer teardown()
server.RegisterHandler("/v1/chat/completions", handleChatCompletionEndpoint)
t.Run("bytes", func(t *testing.T) {
//nolint:lll
msg := json.RawMessage(`{"properties":{"count":{"type":"integer","description":"total number of words in sentence"},"words":{"items":{"type":"string"},"type":"array","description":"list of words in sentence"}},"type":"object","required":["count","words"]}`)
_, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
MaxTokens: 5,
Model: GPT3Dot5Turbo0613,
Messages: []ChatCompletionMessage{
{
Role: ChatMessageRoleUser,
Content: "Hello!",
},
},
Functions: []FunctionDefine{{
Name: "test",
Parameters: &msg,
}},
})
checks.NoError(t, err, "CreateChatCompletion with functions error")
})
t.Run("struct", func(t *testing.T) {
type testMessage struct {
Count int `json:"count"`
Words []string `json:"words"`
}
msg := testMessage{
Count: 2,
Words: []string{"hello", "world"},
}
_, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
MaxTokens: 5,
Model: GPT3Dot5Turbo0613,
Messages: []ChatCompletionMessage{
{
Role: ChatMessageRoleUser,
Content: "Hello!",
},
},
Functions: []FunctionDefinition{{
Name: "test",
Parameters: &msg,
}},
})
checks.NoError(t, err, "CreateChatCompletion with functions error")
})
t.Run("JSONSchemaDefine", func(t *testing.T) {
_, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
MaxTokens: 5,
Model: GPT3Dot5Turbo0613,
Messages: []ChatCompletionMessage{
{
Role: ChatMessageRoleUser,
Content: "Hello!",
},
},
Functions: []FunctionDefinition{{
Name: "test",
Parameters: &JSONSchemaDefinition{
Type: JSONSchemaTypeObject,
Properties: map[string]JSONSchemaDefinition{
"count": {
Type: JSONSchemaTypeNumber,
Description: "total number of words in sentence",
},
"words": {
Type: JSONSchemaTypeArray,
Description: "list of words in sentence",
Items: &JSONSchemaDefinition{
Type: JSONSchemaTypeString,
},
},
"enumTest": {
Type: JSONSchemaTypeString,
Enum: []string{"hello", "world"},
},
},
},
}},
})
checks.NoError(t, err, "CreateChatCompletion with functions error")
})
t.Run("JSONSchemaDefineWithFunctionDefine", func(t *testing.T) {
// this is a compatibility check
_, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
MaxTokens: 5,
Model: GPT3Dot5Turbo0613,
Messages: []ChatCompletionMessage{
{
Role: ChatMessageRoleUser,
Content: "Hello!",
},
},
Functions: []FunctionDefine{{
Name: "test",
Parameters: &JSONSchemaDefine{
Type: JSONSchemaTypeObject,
Properties: map[string]JSONSchemaDefine{
"count": {
Type: JSONSchemaTypeNumber,
Description: "total number of words in sentence",
},
"words": {
Type: JSONSchemaTypeArray,
Description: "list of words in sentence",
Items: &JSONSchemaDefine{
Type: JSONSchemaTypeString,
},
},
"enumTest": {
Type: JSONSchemaTypeString,
Enum: []string{"hello", "world"},
},
},
},
}},
})
checks.NoError(t, err, "CreateChatCompletion with functions error")
})
}
func TestAzureChatCompletions(t *testing.T) {
client, server, teardown := setupAzureTestServer()
defer teardown()
@@ -109,7 +233,34 @@ func handleChatCompletionEndpoint(w http.ResponseWriter, r *http.Request) {
Model: completionReq.Model,
}
// create completions
for i := 0; i < completionReq.N; i++ {
n := completionReq.N
if n == 0 {
n = 1
}
for i := 0; i < n; i++ {
// if there are functions, include them
if len(completionReq.Functions) > 0 {
var fcb []byte
b := completionReq.Functions[0].Parameters
fcb, err = json.Marshal(b)
if err != nil {
http.Error(w, "could not marshal function parameters", http.StatusInternalServerError)
return
}
res.Choices = append(res.Choices, ChatCompletionChoice{
Message: ChatCompletionMessage{
Role: ChatMessageRoleFunction,
// this is valid json so it should be fine
FunctionCall: &FunctionCall{
Name: completionReq.Functions[0].Name,
Arguments: string(fcb),
},
},
Index: i,
})
continue
}
// generate a random string of length completionReq.Length
completionStr := strings.Repeat("a", completionReq.MaxTokens)
@@ -121,8 +272,8 @@ func handleChatCompletionEndpoint(w http.ResponseWriter, r *http.Request) {
Index: i,
})
}
inputTokens := numTokens(completionReq.Messages[0].Content) * completionReq.N
completionTokens := completionReq.MaxTokens * completionReq.N
inputTokens := numTokens(completionReq.Messages[0].Content) * n
completionTokens := completionReq.MaxTokens * n
res.Usage = Usage{
PromptTokens: inputTokens,
CompletionTokens: completionTokens,