fix openai request

This commit is contained in:
unintendedfraud
2023-05-06 11:50:58 +02:00
parent a028fb8b06
commit 31c44edc0f
+20 -6
View File
@@ -3,6 +3,7 @@ package openai
import ( import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"regexp"
"github.com/go-resty/resty/v2" "github.com/go-resty/resty/v2"
) )
@@ -19,10 +20,11 @@ func NewClient(token string) *Client {
} }
func (c Client) ChatCompletions(prompt string) (*CompletionResponse, error) { func (c Client) ChatCompletions(prompt string) (*CompletionResponse, error) {
fmt.Println("using prompt: ", prompt) fmt.Println("using prompt: ", prompt, cleanPrompt(prompt))
body := &ChatCompletionPayload{ body := &ChatCompletionPayload{
Model: "gpt-3.5-turbo", Model: "gpt-3.5-turbo",
Prompt: prompt, Messages: []CompletionPayloadMessage{},
Prompt: cleanPrompt(prompt),
MaxTokens: 50, MaxTokens: 50,
Temperature: 0.2, Temperature: 0.2,
N: 1, N: 1,
@@ -57,8 +59,14 @@ func (c Client) ChatCompletions(prompt string) (*CompletionResponse, error) {
return &r, nil return &r, nil
} }
func cleanPrompt(p string) string {
regex := regexp.MustCompile("(<@\d+>)")
return regex.ReplaceAllString(p, "")
}
type ChatCompletionPayload struct { type ChatCompletionPayload struct {
Model string `json:"model"` Model string `json:"model"`
Messages []CompletionPayloadMessage `json:"messages"`
Prompt string `json:"prompt"` Prompt string `json:"prompt"`
MaxTokens int `json:"max_tokens"` MaxTokens int `json:"max_tokens"`
Temperature float32 `json:"temperature"` Temperature float32 `json:"temperature"`
@@ -66,11 +74,17 @@ type ChatCompletionPayload struct {
N int `json:"n"` N int `json:"n"`
} }
type CompletionPayloadMessage struct {
Role string `json:"role"`
Content string `json:"content"`
Name string `json:"name"`
}
type CompletionResponse struct { type CompletionResponse struct {
ID string `json:"id` ID string `json:"id"`
Object string `json:"object` Object string `json:"object"`
Created int64 `json:"created` Created int64 `json:"created"`
Model string `json:"model` Model string `json:"model"`
Usage CompletionUsage `json:"usage"` Usage CompletionUsage `json:"usage"`
Choices []CompletionChoice `json:"choices"` Choices []CompletionChoice `json:"choices"`
} }