tokenizer

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
+}