Add chat completions (#90)
This commit is contained in:
17
api_test.go
17
api_test.go
@@ -53,6 +53,23 @@ func TestAPI(t *testing.T) {
|
||||
t.Fatalf("Embedding error: %v", err)
|
||||
}
|
||||
|
||||
_, err = c.CreateChatCompletion(
|
||||
ctx,
|
||||
ChatCompletionRequest{
|
||||
Model: GPT3Dot5Turbo,
|
||||
Messages: []ChatCompletionMessage{
|
||||
{
|
||||
Role: "user",
|
||||
Content: "Hello!",
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
t.Errorf("CreateChatCompletion returned error: %v", err)
|
||||
}
|
||||
|
||||
stream, err := c.CreateCompletionStream(ctx, CompletionRequest{
|
||||
Prompt: "Ex falso quodlibet",
|
||||
Model: GPT3Ada,
|
||||
|
||||
68
chat.go
Normal file
68
chat.go
Normal file
@@ -0,0 +1,68 @@
|
||||
package gogpt
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrChatCompletionInvalidModel = errors.New("currently, only gpt-3.5-turbo and gpt-3.5-turbo-0301 are supported")
|
||||
)
|
||||
|
||||
type ChatCompletionMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// ChatCompletionRequest represents a request structure for chat completion API.
|
||||
type ChatCompletionRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []ChatCompletionMessage `json:"messages"`
|
||||
}
|
||||
|
||||
type ChatCompletionChoice struct {
|
||||
Index int `json:"index"`
|
||||
Message ChatCompletionMessage `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
// ChatCompletionResponse represents a response structure for chat completion API.
|
||||
type ChatCompletionResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChatCompletionChoice `json:"choices"`
|
||||
Usage Usage `json:"usage"`
|
||||
}
|
||||
|
||||
// CreateChatCompletion — API call to Creates a completion for the chat message.
|
||||
func (c *Client) CreateChatCompletion(
|
||||
ctx context.Context,
|
||||
request ChatCompletionRequest,
|
||||
) (response ChatCompletionResponse, err error) {
|
||||
model := request.Model
|
||||
if model != GPT3Dot5Turbo0301 && model != GPT3Dot5Turbo {
|
||||
err = ErrChatCompletionInvalidModel
|
||||
return
|
||||
}
|
||||
|
||||
var reqBytes []byte
|
||||
reqBytes, err = json.Marshal(request)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
urlSuffix := "/chat/completions"
|
||||
req, err := http.NewRequest("POST", c.fullURL(urlSuffix), bytes.NewBuffer(reqBytes))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
req = req.WithContext(ctx)
|
||||
err = c.sendRequest(req, &response)
|
||||
return
|
||||
}
|
||||
@@ -12,6 +12,8 @@ import (
|
||||
// GPT3 Models are designed for text-based tasks. For code-specific
|
||||
// tasks, please refer to the Codex series of models.
|
||||
const (
|
||||
GPT3Dot5Turbo0301 = "gpt-3.5-turbo-0301"
|
||||
GPT3Dot5Turbo = "gpt-3.5-turbo"
|
||||
GPT3TextDavinci003 = "text-davinci-003"
|
||||
GPT3TextDavinci002 = "text-davinci-002"
|
||||
GPT3TextCurie001 = "text-curie-001"
|
||||
|
||||
Reference in New Issue
Block a user