go-bot/test/openai.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"go-bot/tools"

	"os"
	"strings"
	"time"

	"github.com/redis/go-redis/v9"
	"github.com/sashabaranov/go-openai"
)

var ctx = context.Background()

// 初始化 Redis 客户端
func initRedis() *redis.Client {
	rdb := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379", // Redis 地址
		Password: "",               // no password set
		DB:       0,                // use default DB
	})
	return rdb
}

// 检查请求频率，10秒内只能请求一次
func checkRequestFrequency(rdb *redis.Client, groupID string, qqID string) bool {
	key := fmt.Sprintf("last_request:%s:%s", groupID, qqID)
	lastRequest, err := rdb.Get(ctx, key).Int64()
	println("key:", key)
	now := time.Now().Unix()
	println("now:", now)
	println("lastRequest:", lastRequest)

	if err == redis.Nil {
		// 键不存在，这是第一次请求
		rdb.Set(ctx, key, now, 10*time.Second)

		return true
	} else if err != nil {
		fmt.Println("获取上次请求时间时出错:", err)
		return false
	}
	currentRequest := getContext(rdb, key)
	fmt.Println("currentRequest:", currentRequest)
	if now-lastRequest < 10 {
		fmt.Printf("请求过于频繁。距离上次请求还有 %d 秒。\n", 10-(now-lastRequest))
		return false // 频率超限，拒绝请求
	}

	// 更新最后请求时间，并设置 10 秒的过期时间
	rdb.Set(ctx, key, now, 10*time.Second)
	return true
}

// 添加对话到上下文，超过5条则删除所有上下文
func addToContext(rdb *redis.Client, groupID string, qqID string, message string, context int64) {
	key := fmt.Sprintf("context:%s:%s", groupID, qqID)

	// 如果上下文超过5条，删除所有上下文
	listLength := rdb.LLen(ctx, key).Val()
	if listLength > context {
		rdb.Del(ctx, key) // 删除该用户的所有上下文
	}
	rdb.RPush(ctx, key, message) // 添加新消息到列表
}

// 获取当前上下文
func getContext(rdb *redis.Client, key string) []string {
	// key :=
	context, err := rdb.LRange(ctx, key, 0, -1).Result()
	if err != nil {
		fmt.Println("Error fetching context:", err)
		return nil
	}
	return context
}

func main() {
	// ctx := context.Background()
	groupID := "12345"
	qqID := "67890"
	rdb := initRedis()

	OPENAI_API_KEY, OPENAI_BaseURL, MODEL, PROMPT, CONTEXT := tools.GetOAIConfig()
	oaiconfig := openai.DefaultConfig(OPENAI_API_KEY)
	oaiconfig.BaseURL = OPENAI_BaseURL
	client := openai.NewClientWithConfig(oaiconfig)

	messages := make([]openai.ChatCompletionMessage, 0)
	messages = append(messages, openai.ChatCompletionMessage{
		Role:    openai.ChatMessageRoleUser,
		Content: PROMPT,
	})
	reader := bufio.NewReader(os.Stdin)
	fmt.Println("Conversation")
	fmt.Println("---------------------")

	for {

		fmt.Print("-> ")
		text, _ := reader.ReadString('\n')
		// convert CRLF to LF
		text = strings.Replace(text, "\n", "", -1)

		messages = append(messages, openai.ChatCompletionMessage{
			Role:    openai.ChatMessageRoleUser,
			Content: text,
		})
		// 输出 messages 的内容
		for i, msg := range messages {
			fmt.Printf("Message %d: Role: %s, Content: %s\n", i, msg.Role, msg.Content)
		}

		resp, err := client.CreateChatCompletion(
			context.Background(),
			openai.ChatCompletionRequest{
				Model:    MODEL,
				Messages: messages,
			},
		)

		if err != nil {
			fmt.Printf("ChatCompletion error: %v\n", err)
			continue
		}

		content := resp.Choices[0].Message.Content
		messages = append(messages, openai.ChatCompletionMessage{
			Role:    openai.ChatMessageRoleAssistant,
			Content: content,
		})
		// 添加新消息到上下文
		addToContext(rdb, groupID, qqID, text, CONTEXT)

		fmt.Println(content)
		fmt.Println("---------------------")
		// 检查请求频率
		if !checkRequestFrequency(rdb, groupID, qqID) {
			// fmt.Println("请求太频繁，请稍后再试。")
			return
		}
		fmt.Println(getContext(rdb, "context:"+groupID+":"+qqID))
	}
}