GPT with memories 2 levels

Add the MaxTokens config
rust-kotlin · Mar 11, 2023 · b5c1dbf · b5c1dbf
1 parent 0fc207c
commit b5c1dbf
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 13 deletions.
diff --git a/chat/chat.go b/chat/chat.go
@@ -5,31 +5,42 @@ import (
 	openai "github.com/sashabaranov/go-openai"
 )
 
-func createMessage(content string) []openai.ChatCompletionMessage {
+func createMessage(content string, contents []string) []openai.ChatCompletionMessage {
 	return []openai.ChatCompletionMessage{
 		{
 			Role:    openai.ChatMessageRoleSystem,
 			Content: "猫娘是一种拟态化的任务,你来扮演猫娘,名字是'咪咪',你说的每一句话都要加上'喵',并且我将作为你的主人存在,在讲话的时候,你应该选择一个合适的情绪,并使用合适的颜文字表达你的情绪.",
 		},
+		{
+			Role:    openai.ChatMessageRoleUser,
+			Content: contents[0],
+		},
+		{
+			Role:    openai.ChatMessageRoleUser,
+			Content: contents[1],
+		},
 		{
 			Role:    openai.ChatMessageRoleUser,
 			Content: content,
-		}}
+		},
+	}
 }
 
-func CreateChat(aClient *openai.Client, model string, content string) (string, error) {
-	message := createMessage(content)
+func CreateChat(aClient *openai.Client, model string, maxTokens int, content string, contents []string) (string, error) {
+	message := createMessage(content, contents)
 	resp, err := aClient.CreateChatCompletion(
 		context.Background(),
 		openai.ChatCompletionRequest{
 			Model:     model,
 			Messages:  message,
-			MaxTokens: 1000,
+			MaxTokens: maxTokens,
 		},
 	)
 	if err != nil {
 		return "", err
 	}
 	result := resp.Choices[0].Message.Content
+	contents[0] = contents[1]
+	contents[1] = result
 	return result, nil
 }
diff --git a/connect/connect.go b/connect/connect.go
@@ -32,9 +32,14 @@ func init() {
 	}
 	// 创建客户端
 	aClient := client.CreateClient(config.Apikey, config.Proxy)
-
+	// 创建一个哈希表
+	dict := make(map[int64][]string)
 	engine.OnMessage(zero.OnlyToMe).Handle(func(ctx *zero.Ctx) {
-		result, err := chat.CreateChat(aClient, config.Model, ctx.ExtractPlainText())
+		qq := ctx.Event.UserID
+		if _, ok := dict[qq]; !ok {
+			dict[qq] = make([]string, 2)
+		}
+		result, err := chat.CreateChat(aClient, config.Model, config.MaxTokens, ctx.ExtractPlainText(), dict[qq])
 		if err != nil {
 			fmt.Println("Error creating chat:", err)
 			return

diff --git a/json/json.go b/json/json.go
@@ -7,17 +7,19 @@ import (
 )
 
 type Config struct {
-	Apikey string `json:"apikey"`
-	Model  string `json:"model"`
-	Proxy  string `json:"proxy"`
+	Apikey    string `json:"apikey"`
+	Model     string `json:"model"`
+	Proxy     string `json:"proxy"`
+	MaxTokens int    `json:"max_tokens"`
 }
 
 // LoadConfig 从文件中读取配置，如果文件不存在则返回默认配置
 func LoadConfig(filename string) (*Config, error) {
 	config := &Config{
-		Apikey: "your api key",
-		Model:  "gpt-3.5-turbo",
-		Proxy:  "http://localhost:7890",
+		Apikey:    "your api key",
+		Model:     "gpt-3.5-turbo",
+		Proxy:     "http://localhost:7890",
+		MaxTokens: 200,
 	}
 
 	// 如果文件不存在，则创建文件并写入默认配置