-
Notifications
You must be signed in to change notification settings - Fork 254
/
chat.go
109 lines (92 loc) · 2.59 KB
/
chat.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
package libaiac
import (
"context"
"fmt"
"strings"
)
// Conversation is a struct used to converse with an OpenAI chat model. It
// maintains all messages sent/received in order to maintain context just like
// using ChatGPT.
type Conversation struct {
client *Client
model Model
messages []Message
}
type Message struct {
Role string `json:"role"`
Content string `json:"content"`
}
type chatResponse struct {
Choices []struct {
Message Message `json:"message"`
Index int64 `json:"index"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
Usage struct {
TotalTokens int64 `json:"total_tokens"`
} `json:"usage"`
}
// Chat initiates a conversation with an OpenAI chat model. A conversation
// maintains context, allowing to send further instructions to modify the output
// from previous requests, just like using the ChatGPT website.
func (client *Client) Chat(model Model) *Conversation {
if model.Type != ModelTypeChat {
return nil
}
return &Conversation{
client: client,
model: model,
}
}
// Send sends the provided message to the API and returns a Response object.
// To maintain context, all previous messages (whether from you to the API or
// vice-versa) are sent as well, allowing you to ask the API to modify the
// code it already generated.
func (conv *Conversation) Send(ctx context.Context, prompt string, msgs ...Message) (
res Response,
err error,
) {
var answer chatResponse
if len(msgs) > 0 {
conv.messages = append(conv.messages, msgs...)
}
conv.messages = append(conv.messages, Message{
Role: "user",
Content: prompt,
})
var apiVersion string
if len(conv.client.apiVersion) > 0 {
apiVersion = fmt.Sprintf("?api-version=%s", conv.client.apiVersion)
}
err = conv.client.NewRequest("POST",
fmt.Sprintf("/chat/completions%s", apiVersion)).
JSONBody(map[string]interface{}{
"model": conv.model.Name,
"messages": conv.messages,
"temperature": 0.2,
}).
Into(&answer).
RunContext(ctx)
if err != nil {
return res, fmt.Errorf("failed sending prompt: %w", err)
}
if len(answer.Choices) == 0 {
return res, ErrNoResults
}
if answer.Choices[0].FinishReason != "stop" {
return res, fmt.Errorf(
"%w: %s",
ErrResultTruncated,
answer.Choices[0].FinishReason,
)
}
conv.messages = append(conv.messages, answer.Choices[0].Message)
res.FullOutput = strings.TrimSpace(answer.Choices[0].Message.Content)
res.APIKeyUsed = conv.client.apiKey
res.TokensUsed = answer.Usage.TotalTokens
var ok bool
if res.Code, ok = ExtractCode(res.FullOutput); !ok {
res.Code = res.FullOutput
}
return res, nil
}