forked from sashabaranov/go-openai
-
Notifications
You must be signed in to change notification settings - Fork 0
/
api_integration_test.go
170 lines (149 loc) · 4.03 KB
/
api_integration_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
//go:build integration
package openai_test
import (
"context"
"errors"
"io"
"os"
"testing"
"github.com/sashabaranov/go-openai"
"github.com/sashabaranov/go-openai/internal/test/checks"
"github.com/sashabaranov/go-openai/jsonschema"
)
func TestAPI(t *testing.T) {
apiToken := os.Getenv("OPENAI_TOKEN")
if apiToken == "" {
t.Skip("Skipping testing against production OpenAI API. Set OPENAI_TOKEN environment variable to enable it.")
}
var err error
c := openai.NewClient(apiToken)
ctx := context.Background()
_, err = c.ListEngines(ctx)
checks.NoError(t, err, "ListEngines error")
_, err = c.GetEngine(ctx, "davinci")
checks.NoError(t, err, "GetEngine error")
fileRes, err := c.ListFiles(ctx)
checks.NoError(t, err, "ListFiles error")
if len(fileRes.Files) > 0 {
_, err = c.GetFile(ctx, fileRes.Files[0].ID)
checks.NoError(t, err, "GetFile error")
} // else skip
embeddingReq := openai.EmbeddingRequest{
Input: []string{
"The food was delicious and the waiter",
"Other examples of embedding request",
},
Model: openai.AdaSearchQuery,
}
_, err = c.CreateEmbeddings(ctx, embeddingReq)
checks.NoError(t, err, "Embedding error")
_, err = c.CreateChatCompletion(
ctx,
openai.ChatCompletionRequest{
Model: openai.GPT3Dot5Turbo,
Messages: []openai.ChatCompletionMessage{
{
Role: openai.ChatMessageRoleUser,
Content: "Hello!",
},
},
},
)
checks.NoError(t, err, "CreateChatCompletion (without name) returned error")
_, err = c.CreateChatCompletion(
ctx,
openai.ChatCompletionRequest{
Model: openai.GPT3Dot5Turbo,
Messages: []openai.ChatCompletionMessage{
{
Role: openai.ChatMessageRoleUser,
Name: "John_Doe",
Content: "Hello!",
},
},
},
)
checks.NoError(t, err, "CreateChatCompletion (with name) returned error")
stream, err := c.CreateCompletionStream(ctx, openai.CompletionRequest{
Prompt: "Ex falso quodlibet",
Model: openai.GPT3Ada,
MaxTokens: 5,
Stream: true,
})
checks.NoError(t, err, "CreateCompletionStream returned error")
defer stream.Close()
counter := 0
for {
_, err = stream.Recv()
if err != nil {
if errors.Is(err, io.EOF) {
break
}
t.Errorf("Stream error: %v", err)
} else {
counter++
}
}
if counter == 0 {
t.Error("Stream did not return any responses")
}
_, err = c.CreateChatCompletion(
context.Background(),
openai.ChatCompletionRequest{
Model: openai.GPT3Dot5Turbo,
Messages: []openai.ChatCompletionMessage{
{
Role: openai.ChatMessageRoleUser,
Content: "What is the weather like in Boston?",
},
},
Functions: []openai.FunctionDefinition{{
Name: "get_current_weather",
Parameters: jsonschema.Definition{
Type: jsonschema.Object,
Properties: map[string]jsonschema.Definition{
"location": {
Type: jsonschema.String,
Description: "The city and state, e.g. San Francisco, CA",
},
"unit": {
Type: jsonschema.String,
Enum: []string{"celsius", "fahrenheit"},
},
},
Required: []string{"location"},
},
}},
},
)
checks.NoError(t, err, "CreateChatCompletion (with functions) returned error")
}
func TestAPIError(t *testing.T) {
apiToken := os.Getenv("OPENAI_TOKEN")
if apiToken == "" {
t.Skip("Skipping testing against production OpenAI API. Set OPENAI_TOKEN environment variable to enable it.")
}
var err error
c := openai.NewClient(apiToken + "_invalid")
ctx := context.Background()
_, err = c.ListEngines(ctx)
checks.HasError(t, err, "ListEngines should fail with an invalid key")
var apiErr *openai.APIError
if !errors.As(err, &apiErr) {
t.Fatalf("Error is not an APIError: %+v", err)
}
if apiErr.HTTPStatusCode != 401 {
t.Fatalf("Unexpected API error status code: %d", apiErr.HTTPStatusCode)
}
switch v := apiErr.Code.(type) {
case string:
if v != "invalid_api_key" {
t.Fatalf("Unexpected API error code: %s", v)
}
default:
t.Fatalf("Unexpected API error code type: %T", v)
}
if apiErr.Error() == "" {
t.Fatal("Empty error message occurred")
}
}