From c8261c5638c031109c068571dce64ee9749c25f0 Mon Sep 17 00:00:00 2001
From: sashabaranov <677093+sashabaranov@users.noreply.github.com>
Date: Thu, 2 Mar 2023 00:47:25 +0400
Subject: [PATCH] Add chat completions (#90)

---
 api_test.go   | 17 +++++++++++++
 chat.go       | 68 +++++++++++++++++++++++++++++++++++++++++++++++++++
 completion.go |  2 ++
 3 files changed, 87 insertions(+)
 create mode 100644 chat.go

diff --git a/api_test.go b/api_test.go
index e7b6813..e5c3be4 100644
--- a/api_test.go
+++ b/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,
diff --git a/chat.go b/chat.go
new file mode 100644
index 0000000..a5ad321
--- /dev/null
+++ b/chat.go
@@ -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
+}
diff --git a/completion.go b/completion.go
index fdf2a27..b40b9cd 100644
--- a/completion.go
+++ b/completion.go
@@ -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"