-
Notifications
You must be signed in to change notification settings - Fork 254
/
libaiac.go
181 lines (149 loc) · 4.74 KB
/
libaiac.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
171
172
173
174
175
176
177
178
179
180
181
package libaiac
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"regexp"
"strings"
"github.com/ido50/requests"
)
// Version contains aiac's version string
var Version = "development"
const OpenAIBackend = "https://api.openai.com/v1"
// Client is a structure used to continuously generate IaC code via OpenAPI/ChatGPT
type Client struct {
*requests.HTTPClient
apiKey string
apiVersion string
}
var (
// ErrResultTruncated is returned when the OpenAI API returned a truncated
// result. The reason for the truncation will be appended to the error
// string.
ErrResultTruncated = errors.New("result was truncated")
// ErrNoResults is returned if the OpenAI API returned an empty result. This
// should not generally happen.
ErrNoResults = errors.New("no results return from API")
// ErrUnsupportedModel is returned if the SetModel method is provided with
// an unsupported model
ErrUnsupportedModel = errors.New("unsupported model")
// ErrUnexpectedStatus is returned when the OpenAI API returned a response
// with an unexpected status code
ErrUnexpectedStatus = errors.New("OpenAI returned unexpected response")
// ErrRequestFailed is returned when the OpenAI API returned an error for
// the request
ErrRequestFailed = errors.New("request failed")
)
type NewClientOptions struct {
// APIKey is the OpenAI API key to use for requests. This is required.
ApiKey string
// ChatGPTURL is the URL to use for ChatGPT requests. This is optional nd by default to openai backend.
URL string
// APIVersion is the version of the OpenAI API to use. This is optional and by default to non specified.
APIVersion string
}
// NewClient creates a new instance of the Client struct, with the provided
// input options. Neither the OpenAI API nor ChatGPT are yet contacted at this
// point.
func NewClient(opts *NewClientOptions) *Client {
if opts == nil {
return nil
}
if opts.ApiKey == "" {
return nil
}
if opts.URL == "" {
opts.URL = OpenAIBackend
}
var authHeaderKey string
var authHeaderVal string
if opts.URL == OpenAIBackend {
authHeaderKey = "Authorization"
authHeaderVal = fmt.Sprintf("Bearer %s", opts.ApiKey)
} else {
authHeaderKey = "api-key"
authHeaderVal = opts.ApiKey
}
cli := &Client{
apiKey: strings.TrimPrefix(opts.ApiKey, "Bearer "),
apiVersion: opts.APIVersion,
}
cli.HTTPClient = requests.NewClient(opts.URL).
Accept("application/json").
Header(authHeaderKey, authHeaderVal).
ErrorHandler(func(
httpStatus int,
contentType string,
body io.Reader,
) error {
var res struct {
Error struct {
Message string `json:"Message"`
Type string `json:"type"`
} `json:"error"`
}
err := json.NewDecoder(body).Decode(&res)
if err != nil {
return fmt.Errorf(
"%w %s",
ErrUnexpectedStatus,
http.StatusText(httpStatus),
)
}
return fmt.Errorf(
"%w: [%s]: %s",
ErrRequestFailed,
res.Error.Type,
res.Error.Message,
)
})
return cli
}
// Response is the struct returned from methods generating code via the OpenAI
// API.
type Response struct {
// FullOutput is the complete output returned by the API. This is generally
// a Markdown-formatted Message that contains the generated code, plus
// explanations, if any.
FullOutput string
// Code is the extracted code section from the complete output. If code was
// not found or extraction otherwise failed, this will be the same as
// FullOutput.
Code string
// APIKeyUsed is the API key used when making the request.
APIKeyUsed string
// TokensUsed is the number of tokens utilized by the request. This is
// the "usage.total_tokens" value returned from the API.
TokensUsed int64
}
// GenerateCode sends the provided prompt to the OpenAI API and returns a
// Response object. It is a convenience wrapper around client.Complete (for
// text completion models) and client.Chat.Send (for chat models).
func (client *Client) GenerateCode(
ctx context.Context,
model Model,
prompt string,
msgs ...Message,
) (res Response, err error) {
if model.Type == ModelTypeChat {
chat := client.Chat(model)
return chat.Send(ctx, prompt, msgs...)
}
return client.Complete(ctx, model, prompt)
}
var codeRegex = regexp.MustCompile("(?ms)^```(?:[^\n]*)\n(.*?)\n```$")
// ExtractCode receives the full output string from the OpenAI API and attempts
// to extract a code block from it. OpenAI code blocks are generally Markdown
// blocks surrounded by the ``` string on both sides. If successful, the code
// string will be returned together with a true value, otherwise an empty string
// is returned together with a false value.
func ExtractCode(output string) (string, bool) {
m := codeRegex.FindStringSubmatch(output)
if m == nil || m[1] == "" {
return "", false
}
return m[1], true
}