add get-models api

Author: idoubi

chat.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"log"
 	"net/http"
 	"strings"
 
@@ -16,9 +17,11 @@ import (
 // ChatText chat reply with text format
 type ChatText struct {
 	data           string // event data
-	ConversationID string // conversation context id
-	MessageID      string // current message id, can used as next chat's parent_message_id
-	Content        string // text content
+	ConversationID string `json:"conversation_id"` // conversation context id
+	MessageID      string `json:"message_id"`      // current message id, can used as next chat's parent_message_id
+	Content        string `json:"content"`         // text content
+	Model          string `json:"model"`           // chat model
+	CreatedAt      int64  `json:"created_at"`      // message create_time
 }
 
 // ChatStream chat reply with sream
@@ -27,9 +30,16 @@ type ChatStream struct {
 	Err    error          // error message
 }
 
+// ChatText raw data
+func (c *ChatText) Raw() string {
+	return c.data
+}
+
 // ChatText format to string
 func (c *ChatText) String() string {
-	return c.data
+	b, _ := json.Marshal(c)
+
+	return string(b)
 }
 
 // GetChatText will return text message
@@ -72,6 +82,19 @@ func (c *Client) GetChatStream(message string, args ...string) (*ChatStream, err
 		return nil, fmt.Errorf("send message failed: %v", err)
 	}
 
+	contentType := resp.Header.Get("Content-Type")
+	// not event-strem response
+	if !strings.HasPrefix(contentType, "text/event-stream") {
+		defer resp.Body.Close()
+
+		body, _ := ioutil.ReadAll(resp.Body)
+		if c.opts.Debug {
+			log.Printf("http response info: %s\n", body)
+		}
+
+		return nil, fmt.Errorf("response failed: [%s] %s", resp.Status, body)
+	}
+
 	chatStream := &ChatStream{
 		Stream: make(chan *ChatText),
 		Err:    nil,
@@ -119,6 +142,8 @@ func (c *Client) parseChatText(text string) (*ChatText, error) {
 	conversationID := res.Get("conversation_id").String()
 	messageID := res.Get("message.id").String()
 	content := res.Get("message.content.parts.0").String()
+	model := res.Get("message.metadata.model_slug").String()
+	createdAt := res.Get("message.create_time").Int()
 
 	if conversationID == "" || messageID == "" {
 		return nil, fmt.Errorf("invalid chat text")
@@ -129,6 +154,8 @@ func (c *Client) parseChatText(text string) (*ChatText, error) {
 		ConversationID: conversationID,
 		MessageID:      messageID,
 		Content:        content,
+		Model:          model,
+		CreatedAt:      createdAt,
 	}, nil
 }
 

client.go

@@ -12,6 +12,7 @@ const (
 	BASE_URI         = "https://chat.openai.com"
 	AUTH_SESSION_URI = "https://chat.openai.com/api/auth/session"
 	CONVERSATION_URI = "https://chat.openai.com/backend-api/conversation"
+	GET_MODELS_URI   = "https://chat.openai.com/backend-api/models"
 	USER_AGENT       = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36"
 	EOF_TEXT         = "[DONE]"
 )
@@ -81,9 +82,16 @@ func (c *Client) doRequest(req *http.Request) (*http.Response, error) {
 	resp, err := c.httpCli.Do(req)
 
 	if c.opts.Debug {
-		respInfo, _ := httputil.DumpResponse(resp, true)
+		respInfo, _ := httputil.DumpResponse(resp, false)
 		log.Printf("http response info: \n%s\n", respInfo)
 	}
 
 	return resp, err
 }
+
+// WithModel: set chat model
+func (c *Client) WithModel(model string) *Client {
+	c.opts.Model = model
+
+	return c
+}

examples/.gitignore

@@ -0,0 +1 @@
+init_test.go

examples/chat_test.go

@@ -6,6 +6,8 @@ import (
 )
 
 func ExampleClient_GetChatText() {
+	cli := getClient()
+
 	message := "say hello to me"
 
 	log.Printf("start get chat text")
@@ -23,6 +25,8 @@ func ExampleClient_GetChatText() {
 }
 
 func ExampleClient_GetContinuousChatText() {
+	cli := getClient()
+
 	message := "say hello to me"
 
 	log.Printf("start get chat text")
@@ -54,18 +58,20 @@ func ExampleClient_GetContinuousChatText() {
 }
 
 func ExampleClient_GetChatStream() {
+	cli := getClient()
+
 	message := "say hello to me"
 
 	log.Printf("start get chat stream")
 
-	stream, err := cli.GetChatStream(message)
+	stream, err := cli.WithModel("gpt-4").GetChatStream(message)
 	if err != nil {
 		log.Fatalf("get chat stream failed: %v\n", err)
 	}
 
 	var answer string
 	for text := range stream.Stream {
-		log.Printf("stream text: %s\n", text.Content)
+		log.Printf("stream text: %s\n", text)
 		answer = text.Content
 	}
 

examples/client_test.go

@@ -8,24 +8,25 @@ import (
 	chatgpt "github.com/chatgp/chatgpt-go"
 )
 
-// chatgpt client
-var cli *chatgpt.Client
+var (
+	debug        bool
+	accessToken  string
+	sessionToken string
+	cfValue      string
+	puid         string
+)
 
 func ExampleNewClient() {
-	fmt.Printf("%T", cli)
+	fmt.Printf("%T", getClient())
 
 	// Output: *chatgpt.Client
 }
 
-func init() {
-	token := `copy-from-cookies`
-	cfValue := "copy-from-cookies"
-	puid := "copy-from-cookies"
-
+func getClient() *chatgpt.Client {
 	cookies := []*http.Cookie{
 		{
 			Name:  "__Secure-next-auth.session-token",
-			Value: token,
+			Value: sessionToken,
 		},
 		{
 			Name:  "cf_clearance",
@@ -37,9 +38,12 @@ func init() {
 		},
 	}
 
-	cli = chatgpt.NewClient(
-		chatgpt.WithDebug(false),
+	cli := chatgpt.NewClient(
 		chatgpt.WithTimeout(60*time.Second),
+		chatgpt.WithDebug(debug),
+		chatgpt.WithAccessToken(accessToken),
 		chatgpt.WithCookies(cookies),
 	)
+
+	return cli
 }

examples/models_test.go

@@ -0,0 +1,27 @@
+package examples
+
+import (
+	"fmt"
+	"log"
+)
+
+func ExampleClient_GetModels() {
+	cli := getClient()
+
+	res, cookies, err := cli.GetModels()
+
+	if err != nil {
+		log.Fatalf("get models failed: %v\n", err)
+	}
+
+	for _, v := range res.Get("models").Array() {
+		log.Printf("model: %s\n", v.String())
+	}
+
+	for _, v := range cookies {
+		log.Printf("cookie: %s, %s, expires: %v\n", v.Name, v.Value, v.Expires)
+	}
+
+	fmt.Println("xxx")
+	// Output: xxx
+}

examples/session_test.go

@@ -0,0 +1,21 @@
+package examples
+
+import (
+	"fmt"
+	"log"
+)
+
+func ExampleClient_AuthSession() {
+	cli := getClient()
+
+	res, err := cli.AuthSession()
+
+	if err != nil {
+		log.Fatalf("auth session failed: %v\n", err)
+	}
+
+	log.Printf("session: %s\n", res)
+
+	fmt.Println("xxx")
+	// Output: xxx
+}

models.go

@@ -0,0 +1,46 @@
+package chatgpt
+
+import (
+	"fmt"
+	"io/ioutil"
+	"log"
+	"net/http"
+
+	"github.com/tidwall/gjson"
+)
+
+// GetModels to get all availabel model s
+func (c *Client) GetModels() (*gjson.Result, []*http.Cookie, error) {
+	req, err := http.NewRequest(http.MethodGet, GET_MODELS_URI, nil)
+	if err != nil {
+		return nil, nil, fmt.Errorf("new request failed: %v", err)
+	}
+
+	accessToken, err := c.getAccessToken()
+	if err != nil {
+		return nil, nil, fmt.Errorf("get accessToken failed: %v", err)
+	}
+
+	bearerToken := fmt.Sprintf("Bearer %s", accessToken)
+	req.Header.Set("Authorization", bearerToken)
+
+	resp, err := c.doRequest(req)
+
+	if err != nil {
+		return nil, nil, fmt.Errorf("do request failed: %v", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, nil, fmt.Errorf("read response body failed: %v", err)
+	}
+
+	if c.opts.Debug {
+		log.Printf("http response info: %s\n", body)
+	}
+
+	res := gjson.ParseBytes(body)
+
+	return &res, resp.Cookies(), nil
+}

session.go

@@ -3,13 +3,14 @@ package chatgpt
 import (
 	"fmt"
 	"io/ioutil"
+	"log"
 	"net/http"
 
 	"github.com/tidwall/gjson"
 )
 
-// authSession will check if session is expired and return a new accessToken
-func (c *Client) authSession() (*gjson.Result, error) {
+// AuthSession will check if session is expired and return a new accessToken
+func (c *Client) AuthSession() (*gjson.Result, error) {
 	req, err := http.NewRequest(http.MethodGet, AUTH_SESSION_URI, nil)
 	if err != nil {
 		return nil, fmt.Errorf("new request failed: %v", err)
@@ -27,10 +28,11 @@ func (c *Client) authSession() (*gjson.Result, error) {
 		return nil, fmt.Errorf("read response body failed: %v", err)
 	}
 
-	res := gjson.ParseBytes(body)
-	if res.String() == "" {
-		return nil, fmt.Errorf("parse response body failed")
+	if c.opts.Debug {
+		log.Printf("http response info: %s\n", body)
 	}
 
+	res := gjson.ParseBytes(body)
+
 	return &res, nil
 }

token.go

@@ -11,7 +11,7 @@ func (c *Client) getAccessToken() (string, error) {
 	}
 
 	// fetch new accessToken
-	res, err := c.authSession()
+	res, err := c.AuthSession()
 	if err != nil {
 		return "", fmt.Errorf("fetch new accessToken failed: %v", err)
 	}