Merge pull request #6 from CsterKuroi/main

support proxy and context

chatgpt/chatgpt.go

@@ -2,25 +2,36 @@ package chatgpt
 
 import (
 	"github.com/poorjobless/wechatbot/config"
-	gpt35 "github.com/AlmazDelDiablo/gpt3-5-turbo-go"
+	gpt35 "github.com/poorjobless/wechatbot/gpt-client"
 )
 
-func Completions(msg string) (string, error) {
-	c := gpt35.NewClient(config.LoadConfig().ApiKey)
+var Cache CacheInterface
+
+func init() {
+	Cache = GetSessionCache()
+}
+
+func Completions(session, msg string) (string, error) {
+	ms := Cache.GetMsg(session)
+	message := &gpt35.Message{
+		Role:    "user",
+		Content: msg,
+	}
+	ms = append(ms, *message)
+	c := gpt35.NewClient(config.LoadConfig().ApiKey, config.LoadConfig().Proxy)
 	req := &gpt35.Request{
-		Model: gpt35.ModelGpt35Turbo,
-		Messages: []*gpt35.Message{
-			{
-				Role:    gpt35.RoleUser,
-				Content: msg,
-			},
-		},
+		Model:       gpt35.ModelGpt35Turbo,
+		Messages:    ms,
+		MaxTokens:   1000,
+		Temperature: 0.7,
 	}
 
 	resp, err := c.GetChat(req)
 	if err != nil {
 		panic(err)
 	}
+	ms = append(ms, resp.Choices[0].Message)
+	Cache.SetMsg(session, ms)
 
 	return resp.Choices[0].Message.Content, err
 }

chatgpt/sessionCache.go

@@ -0,0 +1,74 @@
+package chatgpt
+
+import (
+	"encoding/json"
+	"time"
+
+	"github.com/patrickmn/go-cache"
+
+	gpt35 "github.com/poorjobless/wechatbot/gpt-client"
+)
+
+type SessionService struct {
+	cache *cache.Cache
+}
+
+type SessionMeta struct {
+	Msg []gpt35.Message `json:"msg,omitempty"`
+}
+
+type CacheInterface interface {
+	GetMsg(sessionId string) []gpt35.Message
+	SetMsg(sessionId string, msg []gpt35.Message)
+	Clear(sessionId string)
+}
+
+var sessionServices *SessionService
+
+func getLength(strPool []gpt35.Message) int {
+	var total int
+	for _, v := range strPool {
+		bytes, _ := json.Marshal(v)
+		total += len(string(bytes))
+	}
+	return total
+}
+
+func (s *SessionService) GetMsg(sessionId string) (msg []gpt35.Message) {
+	sessionContext, ok := s.cache.Get(sessionId)
+	if !ok {
+		return nil
+	}
+	sessionMeta := sessionContext.(*SessionMeta)
+	return sessionMeta.Msg
+}
+
+func (s *SessionService) SetMsg(sessionId string, msg []gpt35.Message) {
+	maxLength := 4096
+	maxCacheTime := time.Hour * 12
+
+	for getLength(msg) > maxLength {
+		msg = append(msg[:1], msg[2:]...)
+	}
+
+	sessionContext, ok := s.cache.Get(sessionId)
+	if !ok {
+		sessionMeta := &SessionMeta{Msg: msg}
+		s.cache.Set(sessionId, sessionMeta, maxCacheTime)
+		return
+	}
+	sessionMeta := sessionContext.(*SessionMeta)
+	sessionMeta.Msg = msg
+	s.cache.Set(sessionId, sessionMeta, maxCacheTime)
+}
+
+func (s *SessionService) Clear(sessionId string) {
+	s.cache.Delete(sessionId)
+}
+
+func GetSessionCache() CacheInterface {
+	if sessionServices == nil {
+		sessionServices = &SessionService{cache: cache.New(time.Hour*12, time.Hour*1)}
+	}
+	return sessionServices
+}

config/config.go

@@ -11,6 +11,8 @@ import (
 type Configuration struct {
 	// gtp apikey
 	ApiKey string `json:"api_key"`
+	// proxy
+	Proxy string `json:"proxy"`
 	// 自动通过好友
 	AutoPass bool `json:"auto_pass"`
 }

go.mod

@@ -3,5 +3,6 @@ module github.com/poorjobless/wechatbot
 go 1.16
 
 require (
-	github.com/eatmoreapple/openwechat v1.3.0
+	github.com/eatmoreapple/openwechat v1.4.1
+	github.com/patrickmn/go-cache v2.1.0+incompatible
 )

go.sum

@@ -1,4 +1,4 @@
-github.com/eatmoreapple/openwechat v1.2.1 h1:ez4oqF/Y2NSEX/DbPV8lvj7JlfkYqvieeo4awx5lzfU=
-github.com/eatmoreapple/openwechat v1.2.1/go.mod h1:61HOzTyvLobGdgWhL68jfGNwTJEv0mhQ1miCXQrvWU8=
-github.com/eatmoreapple/openwechat v1.3.0 h1:HlAx7emjaTnvAJ+fcPSc3J3PyF/tULBDWk0KoizARs8=
-github.com/eatmoreapple/openwechat v1.3.0/go.mod h1:61HOzTyvLobGdgWhL68jfGNwTJEv0mhQ1miCXQrvWU8=
+github.com/eatmoreapple/openwechat v1.4.1 h1:hIVEr2Xaj+r1SXzdTigqhIXiuu6TZd+NPWdEVVt/qeM=
+github.com/eatmoreapple/openwechat v1.4.1/go.mod h1:ZxMcq7IpVWVU9JG7ERjExnm5M8/AQ6yZTtX30K3rwRQ=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=

gpt-client/client.go

@@ -0,0 +1,87 @@
+package gpt35
+
+import (
+	"bytes"
+	"encoding/json"
+	"log"
+	"net/http"
+	"net/url"
+	"time"
+)
+
+const ModelGpt35Turbo = "gpt-3.5-turbo"
+
+const MaxTokensGpt35Turbo = 4096
+
+const (
+	RoleUser      RoleType = "user"
+	RoleAssistant RoleType = "assistant"
+	RoleSystem    RoleType = "system"
+)
+
+const DefaultUrl = "https://api.openai.com/v1/chat/completions"
+
+type Client struct {
+	transport *http.Client
+	apiKey    string
+	url       string
+}
+
+func NewClient(apiKey, proxy string) *Client {
+	proxyUrl, err := url.Parse(proxy)
+	if err != nil {
+		log.Fatal(err)
+	}
+	transport := &http.Transport{
+		Proxy: http.ProxyURL(proxyUrl),
+	}
+
+	return &Client{
+		transport: &http.Client{
+			Transport: transport,
+			Timeout:   110 * time.Second,
+		},
+		apiKey: apiKey,
+		url:    DefaultUrl,
+	}
+}
+
+func NewClientCustomUrl(apiKey string, url string) *Client {
+	return &Client{
+		transport: http.DefaultClient,
+		apiKey:    apiKey,
+		url:       url,
+	}
+}
+
+func (c *Client) GetChat(r *Request) (*Response, error) {
+	jsonData, err := json.Marshal(r)
+	if err != nil {
+		return nil, err
+	}
+
+	req, err := http.NewRequest("POST", c.url, bytes.NewBuffer(jsonData))
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+c.apiKey)
+
+	httpResp, err := c.transport.Do(req)
+
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		_ = httpResp.Body.Close()
+	}()
+
+	var resp Response
+	err = json.NewDecoder(httpResp.Body).Decode(&resp)
+	if err != nil {
+		return nil, err
+	}
+
+	return &resp, nil
+}

gpt-client/types.go

@@ -0,0 +1,51 @@
+package gpt35
+
+type RoleType string
+type ModelType string
+
+type Request struct {
+	Model       string    `json:"model"`
+	Messages    []Message `json:"messages"`
+	MaxTokens   int       `json:"max_tokens"`
+	Temperature float64   `json:"temperature"`
+}
+
+type Response struct {
+	Id      string `json:"id"`
+	Object  string `json:"object"`
+	Created int64  `json:"created"`
+	Model   string `json:"model"`
+	Choices []struct {
+		Message      Message `json:"message"`
+		FinishReason string  `json:"finish_reason"`
+	} `json:"choices"`
+	Usage struct {
+		PromptTokens    int `json:"prompt_tokens"`
+		CompletionTokes int `json:"completion_tokens"`
+		TotalTokens     int `json:"total_tokens"`
+	}
+}
+
+type Message struct {
+	Role    RoleType `json:"role"`
+	Content string   `json:"content"`
+}
+
+type Choice struct {
+	Index        int     `json:"index"`
+	Message      Message `json:"message"`
+	FinishReason string  `json:"finish_reason"`
+}
+
+type Usage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
+}
+
+type Error struct {
+	Message string `json:"message"`
+	Type    string `json:"type"`
+	Param   string `json:"param"`
+	Code    string `json:"code"`
+}

handlers/group_msg_handler.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/eatmoreapple/openwechat"
+
 	"github.com/poorjobless/wechatbot/chatgpt"
 )
 
@@ -39,10 +40,15 @@ func (g *GroupMessageHandler) ReplyText(msg *openwechat.Message) error {
 		return nil
 	}
 
+	if strings.Contains(msg.Content, "clear") {
+		chatgpt.Cache.Clear(group.ID())
+		log.Println("clear", group.ID())
+		return nil
+	}
 	// 替换掉@文本，然后向GPT发起请求
-	replaceText := "@" + sender.Self.NickName
+	replaceText := "@" + sender.NickName
 	requestText := strings.TrimSpace(strings.ReplaceAll(msg.Content, replaceText, ""))
-	reply, err := chatgpt.Completions(requestText)
+	reply, err := chatgpt.Completions(group.ID(), requestText)
 	if err != nil {
 		log.Printf("chatgpt request error: %v \n", err)
 		msg.ReplyText("机器人神了，我一会发现了就去修。")
@@ -62,7 +68,7 @@ func (g *GroupMessageHandler) ReplyText(msg *openwechat.Message) error {
 	// 回复@我的用户
 	reply = strings.TrimSpace(reply)
 	reply = strings.Trim(reply, "\n")
-	atText := "@" + groupSender.NickName + ": "
+	atText := "@" + groupSender.NickName + "  "
 	replyText := atText + reply
 	_, err = msg.ReplyText(replyText)
 	if err != nil {

handlers/handler.go

@@ -4,6 +4,7 @@ import (
 	"log"
 
 	"github.com/eatmoreapple/openwechat"
+
 	"github.com/poorjobless/wechatbot/config"
 )
 
@@ -41,7 +42,7 @@ func Handler(msg *openwechat.Message) {
 	// 好友申请
 	if msg.IsFriendAdd() {
 		if config.LoadConfig().AutoPass {
-			_, err := msg.Agree("你好我是基于chatGPT引擎开发的微信机器人，你可以向我提问任何问题。")
+			_, err := msg.Agree("你好我是基于GPT3.5引擎的微信机器人，你可以向我提问任何问题。")
 			if err != nil {
 				log.Fatalf("add friend agree error : %v", err)
 				return

handlers/user_msg_handler.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/eatmoreapple/openwechat"
+
 	"github.com/poorjobless/wechatbot/chatgpt"
 )
 
@@ -33,10 +34,16 @@ func (g *UserMessageHandler) ReplyText(msg *openwechat.Message) error {
 	sender, err := msg.Sender()
 	log.Printf("Received User %v Text Msg : %v", sender.NickName, msg.Content)
 
+	if strings.Contains(msg.Content, "clear") {
+		chatgpt.Cache.Clear(sender.ID())
+		log.Println("clear", sender.ID())
+		return nil
+	}
+
 	// 向GPT发起请求
 	requestText := strings.TrimSpace(msg.Content)
 	requestText = strings.Trim(msg.Content, "\n")
-	reply, err := chatgpt.Completions(requestText)
+	reply, err := chatgpt.Completions(sender.ID(), requestText)
 	if err != nil {
 		log.Printf("chatgpt request error: %v \n", err)
 		msg.ReplyText("机器人神了，我一会发现了就去修。")