upgrade to gtp-4-turbo

openai/client.go

@@ -8,6 +8,8 @@ import (
 	resty "github.com/go-resty/resty/v2"
 )
 
+const MAX_TOKENS = 250
+
 type Client struct {
 	HttpClient *resty.Client
 }
@@ -19,6 +21,34 @@ func NewClient(token string) *Client {
 	return &Client{HttpClient: httpclient}
 }
 
+func (c Client) Chat(messages []*ChatMessage) (*ChatResponse, error) {
+	body := &ChatPayload{
+		Messages:  messages,
+		Model:     "gpt-4-1106-preview",
+		MaxTokens: MAX_TOKENS,
+	}
+
+	b, err := json.Marshal(body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal the payload. %s", err.Error())
+	}
+
+	res, err := c.HttpClient.NewRequest().
+		SetHeader("Content-Type", "application/json").
+		SetBody(b).
+		Post("https://api.openai.com/v1/chat/completions")
+	if err != nil {
+		return nil, fmt.Errorf("failed to query openai. Error: %s", err.Error())
+	}
+
+	var r ChatResponse
+	if err := json.Unmarshal(res.Body(), &r); err != nil {
+		return nil, fmt.Errorf("failed to parse the openai response. Error: %s", err.Error())
+	}
+
+	return &r, nil
+}
+
 func (c Client) Completions(prompt string) (*CompletionResponse, error) {
 	body := &CompletionPayload{
 		Model:       "text-davinci-003",
@@ -54,6 +84,32 @@ func cleanPrompt(p string) string {
 	return regex.ReplaceAllString(p, "")
 }
 
+type ChatResponse struct {
+	ID      string               `json:"id"`
+	Choices []ChatResponseChoice `json:"choices"`
+}
+
+type ChatResponseChoice struct {
+	Message ChatResponseChoiceMessage `json:"message"`
+}
+
+type ChatResponseChoiceMessage struct {
+	Content string `json:"content"`
+}
+
+type ChatPayload struct {
+	Messages  []*ChatMessage `json:"messages"`
+	Model     string         `json:"string"`
+	MaxTokens int            `json:"max_tokens"`
+}
+
+type ChatMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+// -----
+
 type CompletionPayload struct {
 	Model       string  `json:"model"`
 	Prompt      string  `json:"prompt"`