redis缓存替换+pgvector向量替换

ai-chat-backend/pkg/config/config.go

@@ -0,0 +1,185 @@
+package config
+
+import (
+	"bufio"
+	"github.com/spf13/viper"
+	"log"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+)
+
+type Config struct {
+	Http struct {
+		IP   string
+		Port int
+	}
+	BasicAuthUser     string `mapstructure:"basic_auth_user"`
+	BasicAuthPassword string `mapstructure:"basic_auth_password"`
+	FrontendPath      string `mapstructure:"frontend_path"`
+	Log struct {
+		Level   string
+		LogPath string `mapstructure:"logPath"`
+	} `mapstructure:"log"`
+	Chat struct {
+		Model             string  `mapstructure:"model"`
+		MaxTokens         int     `mapstructure:"max_tokens"`
+		Temperature       float32 `mapstructure:"temperature"`
+		TopP              float32 `mapstructure:"top_p"`
+		PresencePenalty   float32 `mapstructure:"presence_penalty"`
+		FrequencyPenalty  float32 `mapstructure:"frequency_penalty"`
+		BotDesc           string  `mapstructure:"bot_desc"`
+		MinResponseTokens int     `mapstructure:"min_response_tokens"`
+		ContextTTL        int     `mapstructure:"context_ttl"`
+		ContextLen        int     `mapstructure:"context_len"`
+	}
+	DependOn struct {
+		AiChatService struct {
+			Address     string
+			AccessToken string
+		} `mapstructure:"ai-chat-service"`
+	}
+}
+
+var conf *Config
+
+func InitConfig(filePath string, typ ...string) {
+	loadProjectDotEnv(filePath)
+	v := viper.New()
+	v.SetConfigFile(filePath)
+	if len(typ) > 0 {
+		v.SetConfigType(typ[0])
+	}
+	err := v.ReadInConfig()
+	if err != nil {
+		log.Fatal(err)
+	}
+	conf = &Config{}
+	err = v.Unmarshal(conf)
+	if err != nil {
+		log.Fatal(err)
+	}
+	normalizeConfig(conf)
+	overrideConfigFromEnv(conf)
+
+}
+
+func GetConfig() *Config {
+	return conf
+}
+
+func normalizeConfig(conf *Config) {
+	if conf.Http.IP == "" {
+		conf.Http.IP = "0.0.0.0"
+	}
+	if conf.Http.Port == 0 {
+		conf.Http.Port = 7080
+	}
+	if conf.FrontendPath == "" {
+		conf.FrontendPath = "www"
+	}
+	if conf.Chat.Model == "" {
+		conf.Chat.Model = "kimi-k2.5"
+	}
+	if conf.Chat.MaxTokens == 0 {
+		conf.Chat.MaxTokens = 4096
+	}
+	if conf.Chat.Temperature == 0 {
+		conf.Chat.Temperature = 1.0
+	}
+	if conf.Chat.TopP == 0 {
+		conf.Chat.TopP = 1.0
+	}
+	if conf.Chat.MinResponseTokens == 0 {
+		conf.Chat.MinResponseTokens = 600
+	}
+	if conf.DependOn.AiChatService.Address == "" {
+		conf.DependOn.AiChatService.Address = "localhost:50055"
+	}
+}
+
+func overrideConfigFromEnv(conf *Config) {
+	overrideString(&conf.Http.IP, os.Getenv("SERVER_HOST"))
+	overrideInt(&conf.Http.Port, os.Getenv("SERVER_PORT"))
+	overrideString(&conf.BasicAuthUser, os.Getenv("BASIC_AUTH_USER"))
+	overrideString(&conf.BasicAuthPassword, os.Getenv("BASIC_AUTH_PASSWORD"))
+	overrideString(&conf.FrontendPath, os.Getenv("FRONTEND_PATH"))
+	overrideString(&conf.Chat.Model, os.Getenv("OPENAI_MODEL"))
+	overrideInt(&conf.Chat.MaxTokens, os.Getenv("OPENAI_MAX_TOKENS"))
+	overrideScaledFloat32(&conf.Chat.Temperature, os.Getenv("OPENAI_TEMPERATURE"))
+	overrideScaledFloat32(&conf.Chat.PresencePenalty, os.Getenv("OPENAI_PRESENCE_PENALTY"))
+	overrideScaledFloat32(&conf.Chat.FrequencyPenalty, os.Getenv("OPENAI_FREQUENCY_PENALTY"))
+	overrideInt(&conf.Chat.MinResponseTokens, os.Getenv("CHAT_MIN_RESPONSE_TOKENS"))
+	overrideString(&conf.DependOn.AiChatService.Address, firstNonEmpty(
+		os.Getenv("AI_CHAT_SERVICE_ADDRESS"),
+		os.Getenv("AI_CHAT_SERVICE_ADDR"),
+	))
+	overrideString(&conf.DependOn.AiChatService.AccessToken, os.Getenv("AI_CHAT_SERVICE_ACCESS_TOKEN"))
+}
+
+func loadProjectDotEnv(configFilePath string) {
+	projectRoot := filepath.Dir(filepath.Dir(configFilePath))
+	envPath := filepath.Join(projectRoot, ".env")
+	file, err := os.Open(envPath)
+	if err != nil {
+		return
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+		key, value, ok := strings.Cut(line, "=")
+		if !ok {
+			continue
+		}
+		key = strings.TrimSpace(key)
+		value = strings.TrimSpace(value)
+		if key == "" {
+			continue
+		}
+		if _, exists := os.LookupEnv(key); exists {
+			continue
+		}
+		_ = os.Setenv(key, value)
+	}
+}
+
+func firstNonEmpty(values ...string) string {
+	for _, value := range values {
+		if strings.TrimSpace(value) != "" {
+			return value
+		}
+	}
+	return ""
+}
+
+func overrideString(target *string, value string) {
+	if strings.TrimSpace(value) != "" {
+		*target = strings.TrimSpace(value)
+	}
+}
+
+func overrideInt(target *int, value string) {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return
+	}
+	if parsed, err := strconv.Atoi(value); err == nil {
+		*target = parsed
+	}
+}
+
+func overrideScaledFloat32(target *float32, value string) {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return
+	}
+	if parsed, err := strconv.ParseFloat(value, 32); err == nil {
+		*target = float32(parsed / 100.0)
+	}
+}

ai-chat-backend/pkg/controllers/chat.go

@@ -0,0 +1,192 @@
+package controllers
+
+import (
+	"ai-chat-backend/pkg/config"
+	"ai-chat-backend/pkg/log"
+	"ai-chat-backend/services"
+	ai_chat_service "ai-chat-backend/services/ai-chat-service"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+	"time"
+
+	ai_chat_service_proto "ai-chat-backend/services/ai-chat-service/proto"
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	openai "github.com/sashabaranov/go-openai"
+	"k8s.io/klog/v2"
+)
+
+type ChatService struct {
+	config *config.Config
+	log    log.ILogger
+}
+
+type ChatCompletionParams struct {
+	Model                 string        `json:"model"`
+	MaxTokens             int           `json:"max_tokens,omitempty"`
+	Temperature           float32       `json:"temperature,omitempty"`
+	PresencePenalty       float32       `json:"presence_penalty,omitempty"`
+	FrequencyPenalty      float32       `json:"frequency_penalty,omitempty"`
+	ChatSessionTTL        time.Duration `json:"chat_session_ttl"`
+	ChatMinResponseTokens int           `json:"chat_min_response_tokens"`
+}
+
+type ChatMessageRequest struct {
+	Prompt  string                    `json:"prompt"`
+	Options ChatMessageRequestOptions `json:"options"`
+}
+
+type ChatMessageRequestOptions struct {
+	Name            string `json:"name"`
+	ParentMessageId string `json:"parentMessageId"`
+}
+
+type ChatMessage struct {
+	ID              string                                              `json:"id"`
+	Text            string                                              `json:"text"`
+	Role            string                                              `json:"role"`
+	Name            string                                              `json:"name"`
+	Delta           string                                              `json:"delta"`
+	Detail          *ai_chat_service_proto.ChatCompletionStreamResponse `json:"detail"`
+	TokenCount      int                                                 `json:"tokenCount"`
+	ParentMessageId string                                              `json:"parentMessageId"`
+}
+
+func NewChatService(config *config.Config, log log.ILogger) (*ChatService, error) {
+	return &ChatService{
+		config: config,
+		log:    log,
+	}, nil
+}
+
+func (chat *ChatService) ChatProcess(ctx *gin.Context) {
+	payload := ChatMessageRequest{}
+	if err := ctx.BindJSON(&payload); err != nil {
+		klog.Error(err)
+		ctx.JSON(200, gin.H{
+			"status":  "Fail",
+			"message": fmt.Sprintf("%v", err),
+			"data":    nil,
+		})
+		return
+	}
+
+	messageID := uuid.New().String()
+
+	result := ChatMessage{
+		ID:              uuid.New().String(),
+		Role:            openai.ChatMessageRoleAssistant,
+		Text:            "",
+		ParentMessageId: messageID,
+	}
+
+	aiChatServicePool := ai_chat_service.GetAiChatServiceClientPool()
+	conn := aiChatServicePool.Get()
+	defer aiChatServicePool.Put(conn)
+	ctx1 := services.AppendBearerTokenToContext(context.Background(), chat.config.DependOn.AiChatService.AccessToken)
+	in := &ai_chat_service_proto.ChatCompletionRequest{
+		Id:            messageID,
+		Message:       payload.Prompt,
+		Pid:           payload.Options.ParentMessageId,
+		EnableContext: false,
+		ChatParam: &ai_chat_service_proto.ChatParam{
+			Model:             chat.config.Chat.Model,
+			MaxTokens:         int32(chat.config.Chat.MaxTokens),
+			Temperature:       chat.config.Chat.Temperature,
+			TopP:              chat.topP(),
+			PresencePenalty:   chat.config.Chat.PresencePenalty,
+			FrequencyPenalty:  chat.config.Chat.FrequencyPenalty,
+			BotDesc:           chat.config.Chat.BotDesc,
+			ContextTTL:        int32(chat.config.Chat.ContextTTL),
+			ContextLen:        int32(chat.config.Chat.ContextLen),
+			MinResponseTokens: int32(chat.config.Chat.MinResponseTokens),
+		},
+	}
+	if in.Pid != "" {
+		in.EnableContext = true
+	}
+
+	aiChatServiceClient := ai_chat_service_proto.NewChatClient(conn)
+	stream, err := aiChatServiceClient.ChatCompletionStream(ctx1, in)
+	if err != nil {
+		chat.log.Error(err)
+		ctx.JSON(200, gin.H{
+			"status":  "Fail",
+			"message": fmt.Sprintf("%v", err),
+			"data":    nil,
+		})
+		return
+	}
+	defer stream.CloseSend()
+
+	firstChunk := true
+	ctx.Header("Content-type", "application/octet-stream")
+	for {
+		rsp, err := stream.Recv()
+		if errors.Is(err, io.EOF) {
+			return
+		}
+
+		if err != nil {
+			klog.Error(err)
+			ctx.JSON(200, gin.H{
+				"status":  "Fail",
+				"message": fmt.Sprintf("OpenAI Event Error %v", err),
+				"data":    nil,
+			})
+			return
+		}
+
+		if rsp.Id != "" {
+			result.ID = rsp.Id
+		}
+
+		if len(rsp.Choices) > 0 {
+			content := rsp.Choices[0].Delta.Content
+			result.Delta = content
+			if len(content) > 0 {
+				result.Text += content
+			}
+			result.Detail = rsp
+		}
+
+		bts, err := json.Marshal(result)
+		if err != nil {
+			klog.Error(err)
+			ctx.JSON(200, gin.H{
+				"status":  "Fail",
+				"message": fmt.Sprintf("OpenAI Event Marshal Error %v", err),
+				"data":    nil,
+			})
+			return
+		}
+
+		if !firstChunk {
+			ctx.Writer.Write([]byte("\n"))
+		} else {
+			firstChunk = false
+		}
+
+		if _, err := ctx.Writer.Write(bts); err != nil {
+			klog.Error(err)
+			return
+		}
+
+		ctx.Writer.Flush()
+	}
+}
+
+func (chat *ChatService) topP() float32 {
+	model := strings.ToLower(chat.config.Chat.Model)
+	if strings.HasPrefix(model, "kimi-") || strings.HasPrefix(model, "moonshot-") {
+		return 0.95
+	}
+	if chat.config.Chat.TopP > 0 {
+		return chat.config.Chat.TopP
+	}
+	return 1
+}

ai-chat-backend/pkg/log/README.md

@@ -0,0 +1,3 @@
+# 日志框架
+1. 可通过包调用日志打印，也可以通过对象调用日志打印
+2. 可以自动切分日志文件

ai-chat-backend/pkg/log/hook.go

@@ -0,0 +1,19 @@
+package log
+
+import "github.com/sirupsen/logrus"
+import nativeLog "log"
+
+type errorHook struct {
+}
+
+func (*errorHook) Levels() []logrus.Level {
+	return []logrus.Level{
+		logrus.PanicLevel,
+		logrus.FatalLevel,
+		logrus.ErrorLevel,
+	}
+}
+func (*errorHook) Fire(entry *logrus.Entry) error {
+	nativeLog.Println(entry.Message, entry.Data)
+	return nil
+}

ai-chat-backend/pkg/log/log.go

@@ -0,0 +1,233 @@
+package log
+
+import (
+	"errors"
+	"fmt"
+	"github.com/sirupsen/logrus"
+	"io"
+	"runtime"
+)
+
+type ILogger interface {
+	SetLevel(lvl string)
+	SetOutput(writer io.Writer)
+	SetPrintCaller(bool)
+	SetCaller(caller func() (file string, line int, funcName string, err error))
+	Trace(args ...interface{})
+	Debug(args ...interface{})
+	Info(args ...interface{})
+	Warning(args ...interface{})
+	Error(args ...interface{})
+	Fatal(args ...interface{})
+	Panic(args ...interface{})
+	TraceF(format string, args ...interface{})
+	DebugF(format string, args ...interface{})
+	InfoF(format string, args ...interface{})
+	WarningF(format string, args ...interface{})
+	ErrorF(format string, args ...interface{})
+	FatalF(format string, args ...interface{})
+	PanicF(format string, args ...interface{})
+	WithFields(fields map[string]interface{}) ILogger
+}
+type Logger struct {
+	entry *logrus.Entry
+	// panic,fatal,error,warn,warning,info,debug,trace
+	level       string
+	printCaller bool
+	caller      func() (file string, line int, funcName string, err error)
+}
+
+// 设置日志打印级别
+func (l *Logger) SetLevel(lvl string) {
+	if lvl == "" {
+		return
+	}
+	level, err := logrus.ParseLevel(lvl)
+	if err == nil {
+		l.level = lvl
+		l.entry.Logger.Level = level
+	}
+}
+
+// 设置日志输出位置
+func (l *Logger) SetOutput(writer io.Writer) {
+	l.entry.Logger.SetOutput(writer)
+}
+
+// 设置是否打印调用信息
+func (l *Logger) SetPrintCaller(printCaller bool) {
+	l.printCaller = printCaller
+}
+func (l *Logger) SetCaller(caller func() (file string, line int, funcName string, err error)) {
+	l.caller = caller
+}
+
+// 获取caller信息
+func (l *Logger) getCallerInfo(level logrus.Level) map[string]interface{} {
+	mp := make(map[string]interface{})
+	if l.printCaller == true || level != logrus.InfoLevel {
+		file, line, funcName, err := l.caller()
+		if err == nil {
+			mp["file"] = fmt.Sprintf("%s:%d", file, line)
+			mp["func"] = funcName
+		}
+	}
+	return mp
+}
+
+func (l *Logger) log(level logrus.Level, args ...interface{}) {
+	l.entry.WithFields(l.getCallerInfo(level)).Log(level, args...)
+}
+func (l *Logger) logf(level logrus.Level, format string, args ...interface{}) {
+	l.entry.WithFields(l.getCallerInfo(level)).Logf(level, format, args...)
+}
+func (l *Logger) Trace(args ...interface{}) {
+	l.log(logrus.TraceLevel, args...)
+}
+func (l *Logger) Debug(args ...interface{}) {
+	l.log(logrus.DebugLevel, args...)
+}
+func (l *Logger) Info(args ...interface{}) {
+	l.log(logrus.InfoLevel, args...)
+}
+func (l *Logger) Warning(args ...interface{}) {
+	l.log(logrus.WarnLevel, args...)
+}
+func (l *Logger) Error(args ...interface{}) {
+	l.log(logrus.ErrorLevel, args...)
+}
+func (l *Logger) Fatal(args ...interface{}) {
+	l.log(logrus.FatalLevel, args...)
+}
+func (l *Logger) Panic(args ...interface{}) {
+	l.log(logrus.PanicLevel, args...)
+}
+func (l *Logger) TraceF(format string, args ...interface{}) {
+	l.logf(logrus.TraceLevel, format, args...)
+}
+func (l *Logger) DebugF(format string, args ...interface{}) {
+	l.logf(logrus.DebugLevel, format, args...)
+}
+func (l *Logger) InfoF(format string, args ...interface{}) {
+	l.logf(logrus.InfoLevel, format, args...)
+}
+func (l *Logger) WarningF(format string, args ...interface{}) {
+	l.logf(logrus.WarnLevel, format, args...)
+}
+func (l *Logger) ErrorF(format string, args ...interface{}) {
+	l.logf(logrus.ErrorLevel, format, args...)
+}
+func (l *Logger) FatalF(format string, args ...interface{}) {
+	l.logf(logrus.FatalLevel, format, args...)
+}
+func (l *Logger) PanicF(format string, args ...interface{}) {
+	l.logf(logrus.PanicLevel, format, args...)
+}
+func (l *Logger) WithFields(fields map[string]interface{}) ILogger {
+	entry := l.entry.WithFields(fields)
+	return &Logger{entry: entry, level: l.level, printCaller: l.printCaller, caller: l.caller}
+}
+
+var log *Logger
+
+func NewLogger() ILogger {
+	return newLogger()
+}
+func newLogger() *Logger {
+	log := logrus.New()
+	log.SetLevel(logrus.InfoLevel)
+	log.AddHook(&errorHook{})
+	logger := &Logger{
+		entry:  logrus.NewEntry(log),
+		caller: defaultCaller,
+	}
+	return logger
+}
+
+func init() {
+	log = newLogger()
+}
+
+// 设置日志打印级别
+func SetLevel(lvl string) {
+	if lvl == "" {
+		return
+	}
+	level, err := logrus.ParseLevel(lvl)
+	if err == nil {
+		log.level = lvl
+		log.entry.Logger.Level = level
+	}
+}
+
+// 设置日志的输出位置
+func SetOutput(writer io.Writer) {
+	log.entry.Logger.SetOutput(writer)
+}
+
+// 设置是否打印调用信息
+func SetPrintCaller(printCaller bool) {
+	log.printCaller = printCaller
+}
+
+func SetCaller(caller func() (file string, line int, funcName string, err error)) {
+	log.caller = caller
+}
+
+func defaultCaller() (file string, line int, funcName string, err error) {
+	pc, f, l, ok := runtime.Caller(4)
+	if !ok {
+		err = errors.New("caller failure")
+		return
+	}
+	funcName = runtime.FuncForPC(pc).Name()
+	file, line = f, l
+	return
+}
+
+func Trace(args ...interface{}) {
+	log.log(logrus.TraceLevel, args...)
+}
+func Debug(args ...interface{}) {
+	log.log(logrus.DebugLevel, args...)
+}
+func Info(args ...interface{}) {
+	log.log(logrus.InfoLevel, args...)
+}
+func Warning(args ...interface{}) {
+	log.log(logrus.WarnLevel, args...)
+}
+func Error(args ...interface{}) {
+	log.log(logrus.ErrorLevel, args...)
+}
+func Fatal(args ...interface{}) {
+	log.log(logrus.FatalLevel, args...)
+}
+func Panic(args ...interface{}) {
+	log.log(logrus.PanicLevel, args...)
+}
+func TraceF(format string, args ...interface{}) {
+	log.logf(logrus.TraceLevel, format, args...)
+}
+func DebugF(format string, args ...interface{}) {
+	log.logf(logrus.DebugLevel, format, args...)
+}
+func InfoF(format string, args ...interface{}) {
+	log.logf(logrus.InfoLevel, format, args...)
+}
+func WarningF(format string, args ...interface{}) {
+	log.logf(logrus.WarnLevel, format, args...)
+}
+func ErrorF(format string, args ...interface{}) {
+	log.logf(logrus.ErrorLevel, format, args...)
+}
+func FatalF(format string, args ...interface{}) {
+	log.logf(logrus.FatalLevel, format, args...)
+}
+func PanicF(format string, args ...interface{}) {
+	log.logf(logrus.PanicLevel, format, args...)
+}
+func WithFields(fields map[string]interface{}) *Logger {
+	entry := log.entry.WithFields(fields)
+	return &Logger{entry: entry, level: log.level, printCaller: log.printCaller, caller: log.caller}
+}

ai-chat-backend/pkg/log/rotate_writer.go

@@ -0,0 +1,58 @@
+package log
+
+import (
+	"gopkg.in/natefinch/lumberjack.v2"
+	"io"
+	"sync"
+)
+
+type fileRotateWriter struct {
+	data map[string]io.Writer
+	sync.RWMutex
+}
+
+func (frw *fileRotateWriter) getWriter(logPath string) io.Writer {
+	frw.RLock()
+	defer frw.RUnlock()
+	w, ok := frw.data[logPath]
+	if !ok {
+		return nil
+	}
+	return w
+}
+func (frw *fileRotateWriter) setWriter(logPath string, w io.Writer) io.Writer {
+	frw.Lock()
+	defer frw.Unlock()
+	frw.data[logPath] = w
+	return w
+}
+
+var _fileRotateWriter *fileRotateWriter
+
+func init() {
+	_fileRotateWriter = &fileRotateWriter{
+		data: map[string]io.Writer{},
+	}
+}
+
+func GetRotateWriter(logPath string) io.Writer {
+	if logPath == "" {
+		panic("日志文件路径不能为空")
+	}
+	writer := _fileRotateWriter.getWriter(logPath)
+	if writer != nil {
+		return writer
+	}
+	writer = &lumberjack.Logger{
+		//文件名
+		Filename: logPath,
+		//单个文件大小单位MB
+		MaxSize: 1,
+		//最多保留文件数
+		MaxBackups: 15,
+		//最长保留时间（天）
+		MaxAge:    7,
+		LocalTime: true,
+	}
+	return _fileRotateWriter.setWriter(logPath, writer)
+}

ai-chat-backend/pkg/middlewares/cors.go

@@ -0,0 +1,24 @@
+package middlewares
+
+import (
+	"github.com/gin-gonic/gin"
+	"net/http"
+)
+
+func Cors() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		method := c.Request.Method
+		origin := c.Request.Header.Get("Origin")
+		if origin != "" {
+			c.Header("Access-Control-Allow-Origin", "*") // 可将将 * 替换为指定的域名
+			c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
+			c.Header("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept, Authorization")
+			c.Header("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers, Cache-Control, Content-Language, Content-Type")
+			c.Header("Access-Control-Allow-Credentials", "true")
+		}
+		if method == "OPTIONS" {
+			c.AbortWithStatus(http.StatusNoContent)
+		}
+		c.Next()
+	}
+}

ai-chat-backend/pkg/middlewares/rate_limit.go

@@ -0,0 +1,23 @@
+package middlewares
+
+import (
+	"github.com/gin-gonic/gin"
+	"golang.org/x/time/rate"
+)
+
+func RateLimitMiddleware(r rate.Limit, b int) gin.HandlerFunc {
+	limiter := rate.NewLimiter(r, b)
+	return func(c *gin.Context) {
+		if !limiter.Allow() {
+			// 请求被限制，返回错误信息
+			c.JSON(429, gin.H{
+				"status":  "Fail",
+				"message": "Too many requests, please try again later",
+				"data":    nil,
+			})
+			c.Abort()
+			return
+		}
+		c.Next()
+	}
+}

ai-chat-backend/pkg/tokenizer/tokenizer.go

@@ -0,0 +1,43 @@
+package tokenizer
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+
+	"github.com/sashabaranov/go-openai"
+)
+
+type tokenInfo struct {
+	Code  int    `json:"code"`
+	Count int    `json:"num_tokens"`
+	Msg   string `json:"msg"`
+}
+
+func GetTokenCount(message openai.ChatCompletionMessage, model string) (int, error) {
+	url := fmt.Sprintf("http://192.168.239.161:3002/tokenizer/%s", model)
+	info := tokenInfo{}
+	if err := postJSON(url, &message, &info); err != nil {
+		return 0, err
+	}
+	if info.Code != 200 {
+		return 0, fmt.Errorf("%v", info.Msg)
+	}
+	return info.Count, nil
+}
+
+func postJSON(url string, requestData *openai.ChatCompletionMessage, responseData *tokenInfo) error {
+	requestBody, err := json.Marshal(requestData)
+	if err != nil {
+		return err
+	}
+
+	resp, err := http.Post(url, "application/json", bytes.NewBuffer(requestBody))
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	return json.NewDecoder(resp.Body).Decode(responseData)
+}

ai-chat-backend/pkg/utils/utils.go

@@ -0,0 +1,65 @@
+package utils
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"k8s.io/klog/v2"
+)
+
+func ReplaceInFile(filePath string, targetStr string, replaceStr string) error {
+	// 读取文件内容
+	content, err := ioutil.ReadFile(filePath)
+	if err != nil {
+		return err
+	}
+	// 判断是否需要替换
+	if !strings.Contains(string(content), targetStr) {
+		return nil
+	}
+
+	// 替换字符串并写回文件，保持原有 filemode
+	info, err := os.Stat(filePath)
+	if err != nil {
+		return err
+	}
+	newContent := strings.ReplaceAll(string(content), targetStr, replaceStr)
+	err = ioutil.WriteFile(filePath, []byte(newContent), info.Mode())
+	if err != nil {
+		return err
+	}
+	klog.Infof("Replaced in file: %s\n", filePath)
+	return nil
+}
+
+func ReplaceFiles(rootDir string, replacePairs map[string]string) error {
+	err := filepath.Walk(rootDir, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		// 如果当前路径是目录，则继续遍历
+		if info.IsDir() {
+			return nil
+		}
+		// 处理文件
+		for targetStr, replaceStr := range replacePairs {
+			err = ReplaceInFile(path, targetStr, replaceStr)
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func Reverse[S ~[]E, E any](s S) {
+	for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
+		s[i], s[j] = s[j], s[i]
+	}
+}