tokenizer

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

@@ -0,0 +1,90 @@
+package config
+
+import (
+	"github.com/spf13/viper"
+	"log"
+)
+
+type Config struct {
+	Server struct {
+		IP          string
+		Port        int
+		AccessToken string
+	}
+	Log struct {
+		Level   string
+		LogPath string `mapstructure:"logPath"`
+	} `mapstructure:"log"`
+	Chat struct {
+		ApiKey            string  `mapstructure:"api_key"`
+		BaseUrl           string  `mapstructure:"base_url"`
+		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"`
+	}
+	Mysql struct {
+		DSN         string
+		MaxLifeTime int
+		MaxOpenConn int
+		MaxIdleConn int
+	}
+	Redis struct {
+		Host string
+		Port int
+		Pwd  string `mapstructure:"pwd"`
+	}
+	DependOn struct {
+		Sensitive struct {
+			Address     string
+			AccessToken string
+		}
+		Keywords struct {
+			Address     string
+			AccessToken string
+		}
+		Tokenizer struct {
+			Address string
+		}
+	}
+	VectorDB struct {
+		Url                string
+		Username           string
+		Pwd                string
+		Database           string
+		Timeout            int
+		MaxIdleConnPerHost int
+		ReadConsistency    string
+		IdleConnTimeout    int
+	}
+}
+
+var conf *Config
+
+func InitConfig(filePath string, typ ...string) {
+	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)
+	}
+
+}
+
+func GetConfig() *Config {
+	return conf
+}

ai-chat-service/pkg/db/mysql/mysql.go

@@ -0,0 +1,28 @@
+package mysql
+
+import (
+	"ai-chat-service/pkg/config"
+	"database/sql"
+	_ "github.com/go-sql-driver/mysql"
+	"time"
+)
+
+var db *sql.DB
+
+func InitMysql(cnf *config.Config) {
+	var err error
+	if cnf.Mysql.DSN == "" {
+		panic("数据库连接字符串不能为空")
+	}
+	db, err = sql.Open("mysql", cnf.Mysql.DSN)
+	if err != nil {
+		panic(err)
+	}
+	db.SetMaxOpenConns(cnf.Mysql.MaxOpenConn)
+	db.SetMaxIdleConns(cnf.Mysql.MaxIdleConn)
+	db.SetConnMaxLifetime(time.Second * time.Duration(cnf.Mysql.MaxLifeTime))
+}
+
+func GetDB() *sql.DB {
+	return db
+}

ai-chat-service/pkg/db/redis/prefix.go

@@ -0,0 +1,14 @@
+package redis
+
+import "strings"
+
+const ServicePrefix = "ai_chat_service_"
+
+func GetKey(key string, parts ...string) string {
+	key = ServicePrefix + key
+	if len(parts) == 0 {
+		return key
+	}
+	key += "_" + strings.Join(parts, "_")
+	return key
+}

ai-chat-service/pkg/db/redis/redis.go

@@ -0,0 +1,57 @@
+package redis
+
+import (
+	"ai-chat-service/pkg/config"
+	"context"
+	"fmt"
+	redis "github.com/redis/go-redis/v9"
+	"sync"
+)
+
+type RedisPool interface {
+	Get() *redis.Client
+	Put(client *redis.Client)
+}
+
+var pool RedisPool
+
+type redisPool struct {
+	pool sync.Pool
+}
+
+func (p *redisPool) Get() *redis.Client {
+	client := p.pool.Get().(*redis.Client)
+	if client.Ping(context.Background()).Err() != nil {
+		client = p.pool.New().(*redis.Client)
+	}
+	return client
+}
+
+func (p *redisPool) Put(client *redis.Client) {
+	if client.Ping(context.Background()).Err() != nil {
+		return
+	}
+	p.pool.Put(client)
+}
+
+func getPool(cnf *config.Config) RedisPool {
+	return &redisPool{
+		pool: sync.Pool{
+			New: func() any {
+				rdb := redis.NewClient(&redis.Options{
+					Addr:     fmt.Sprintf("%s:%d", cnf.Redis.Host, cnf.Redis.Port),
+					Password: cnf.Redis.Pwd,
+				})
+				return rdb
+			},
+		},
+	}
+}
+
+func InitRedisPool(cnf *config.Config) {
+	pool = getPool(cnf)
+}
+
+func GetPool() RedisPool {
+	return pool
+}

ai-chat-service/pkg/db/vector/vector.go

@@ -0,0 +1,29 @@
+package vector
+
+import (
+	"ai-chat-service/pkg/config"
+	"ai-chat-service/pkg/log"
+	"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
+	"time"
+)
+
+var vdb *tcvectordb.Client
+
+func InitDB(config *config.Config) {
+	var defaultOption = &tcvectordb.ClientOption{
+		Timeout:            time.Second * time.Duration(config.VectorDB.Timeout),
+		MaxIdldConnPerHost: config.VectorDB.MaxIdleConnPerHost,
+		IdleConnTimeout:    time.Second * time.Duration(config.VectorDB.IdleConnTimeout),
+		ReadConsistency:    tcvectordb.ReadConsistency(config.VectorDB.ReadConsistency),
+	}
+	var err error
+	vdb, err = tcvectordb.NewClient(config.VectorDB.Url, config.VectorDB.Username, config.VectorDB.Pwd, defaultOption)
+	if err != nil {
+		log.Error(err)
+		return
+	}
+}
+
+func GetVdb() *tcvectordb.Client {
+	return vdb
+}

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

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

ai-chat-service/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-service/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-service/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-service/pkg/zerror/error_code.go

@@ -0,0 +1,14 @@
+package zerror
+
+type ZErrorCode string
+
+func getErrMsg(errCode ZErrorCode) string {
+	msg, ok := errorMsgs[errCode]
+	if ok {
+		return msg
+	}
+	return ""
+}
+
+// 错误码与之对应的错误消息
+var errorMsgs = map[ZErrorCode]string{}

ai-chat-service/pkg/zerror/zerror.go

@@ -0,0 +1,101 @@
+package zerror
+
+import (
+	"errors"
+	"fmt"
+)
+
+type ZError struct {
+	ErrCode ZErrorCode `json:"err_code,omitempty"`
+	ErrMsg  string     `json:"err_msg,omitempty"`
+	errs    []error
+}
+
+func (e *ZError) Error() string {
+	if e == nil {
+		return ""
+	}
+	if e.ErrMsg != "" {
+		return fmt.Sprintf("ErrCode:%s; ErrMsg:%s;", e.ErrCode, e.ErrMsg)
+	}
+	res := ""
+	if e.errs == nil || len(e.errs) == 0 {
+		return res
+	}
+	var first = true
+	for _, err := range e.errs {
+		if first {
+			res = err.Error()
+			first = false
+		} else {
+			res += ";" + err.Error()
+		}
+	}
+	return res
+}
+func (e *ZError) Errors() []error {
+	if e == nil {
+		return nil
+	}
+	return e.errs
+}
+func (e *ZError) Append(err error) {
+	if e == nil || err == nil {
+		return
+	}
+	ze, ok := err.(*ZError)
+	if ok {
+		e.errs = append(e.errs, ze.errs...)
+	} else {
+		e.errs = append(e.errs, err)
+	}
+}
+
+func NewByErr(err ...error) error {
+	res := &ZError{
+		errs: make([]error, 0),
+	}
+	for i, e := range err {
+		if e == nil {
+			continue
+		}
+		ze, ok := err[i].(*ZError)
+		if ok {
+			res.errs = append(res.errs, ze.errs...)
+		} else {
+			res.errs = append(res.errs, err[i])
+		}
+	}
+	if len(res.errs) > 0 {
+		return res
+	}
+	return nil
+}
+func NewByCode(errCode ZErrorCode, errMsg ...string) error {
+	msg := ""
+	if len(errMsg) > 0 {
+		msg = errMsg[0]
+	} else {
+		msg = getErrMsg(errCode)
+	}
+	return &ZError{
+		ErrCode: errCode,
+		ErrMsg:  msg,
+	}
+}
+func NewByMsg(msg string) error {
+	err := errors.New(msg)
+	return NewByErr(err)
+}
+func Errors(err error) []error {
+	if err == nil {
+		return nil
+	}
+	ze, ok := err.(*ZError)
+	if !ok {
+		return []error{err}
+	}
+	//将一个空的[]error切片与ze.Errors()返回的错误切片合并在一起
+	//返回一个新切片
+	return append(([]error)(nil), ze.Errors()...)
+}