redis缓存替换+pgvector向量替换

ai-chat-service/chat-server/main.go

@@ -10,7 +10,6 @@ import (
 	"ai-chat-service/pkg/config"
 	"ai-chat-service/pkg/db/mysql"
 	"ai-chat-service/pkg/db/redis"
-	"ai-chat-service/pkg/db/vector"
 	"ai-chat-service/pkg/log"
 	"ai-chat-service/proto"
 	"flag"
@@ -57,17 +56,19 @@ func main() {
 	mysql.InitMysql(cnf)
 	// 初始化redis
 	redis.InitRedisPool(cnf)
-	// 初始化向量数据库
-	vector.InitDB(cnf)
 
 	recordsData := data.NewChatRecordsData(mysql.GetDB())
+	vectorRecordsData, err := vector_data.NewChatRecordsData(cnf)
+	if err != nil {
+		log.Fatal(err)
+	}
 
 	lis, err := net.Listen("tcp", fmt.Sprintf("%s:%d", cnf.Server.IP, cnf.Server.Port))
 	if err != nil {
 		log.Fatal(err)
 	}
 	s := grpc.NewServer(grpc.UnaryInterceptor(interceptor.UnaryAuthInterceptor), grpc.StreamInterceptor(metrics_app.NewStreamMiddleware(registry).WrapHandler()))
-	service := server.NewChatService(recordsData, vector_data.NewChatRecordsData(cnf, vector.GetVdb()), cnf, logger, busMetrics)
+	service := server.NewChatService(recordsData, vectorRecordsData, cnf, logger, busMetrics)
 	proto.RegisterChatServer(s, service)
 
 	healthCheckSrv := health.NewServer()

ai-chat-service/chat-server/server/app.go

@@ -59,18 +59,10 @@ func (s *chatService) newApp(in *proto.ChatCompletionRequest, contextCache chat_
 		if in.ChatParam.Model != "" {
 			conf.Model = in.ChatParam.Model
 		}
-		if in.ChatParam.TopP != 0 {
-			conf.TopP = in.ChatParam.TopP
-		}
-		if in.ChatParam.FrequencyPenalty != 0 {
-			conf.FrequencyPenalty = in.ChatParam.FrequencyPenalty
-		}
-		if in.ChatParam.PresencePenalty != 0 {
-			conf.PresencePenalty = in.ChatParam.PresencePenalty
-		}
-		if in.ChatParam.Temperature != 0 {
-			conf.Temperature = in.ChatParam.Temperature
-		}
+		conf.TopP = in.ChatParam.TopP
+		conf.FrequencyPenalty = in.ChatParam.FrequencyPenalty
+		conf.PresencePenalty = in.ChatParam.PresencePenalty
+		conf.Temperature = in.ChatParam.Temperature
 		if in.ChatParam.BotDesc != "" {
 			conf.BotDesc = in.ChatParam.BotDesc
 		}

ai-chat-service/chat-server/server/server.go

@@ -62,7 +62,7 @@ func (s *chatService) ChatCompletion(ctx context.Context, in *proto.ChatCompleti
 		idStr, score, err := s.vectorData.QueryData(context.Background(), map[string][]string{"keywords": {strings.Join(keywords, ",")}})
 		if err != nil {
 			s.log.Error(err)
-		} else if score > 0.99 {
+		} else if score > s.config.Vector.Threshold {
 			id, err := strconv.ParseInt(idStr, 10, 64)
 			if err != nil {
 				s.log.Error(err)
@@ -200,7 +200,7 @@ func (s *chatService) ChatCompletionStream(in *proto.ChatCompletionRequest, stre
 		idStr, score, err := s.vectorData.QueryData(context.Background(), map[string][]string{"keywords": {strings.Join(keywords, ",")}})
 		if err != nil {
 			s.log.Error(err)
-		} else if score > 0.99 {
+		} else if score > s.config.Vector.Threshold {
 			id, err := strconv.ParseInt(idStr, 10, 64)
 			if err != nil {
 				s.log.Error(err)

ai-chat-service/chat-server/vector-data/chat_records.go

@@ -3,7 +3,7 @@ package vector_data
 import (
 	"ai-chat-service/pkg/config"
 	"context"
-	"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
+	"fmt"
 )
 
 const CHAT_RECORDS = "chat_records"
@@ -17,53 +17,13 @@ type IChatRecordsData interface {
 	QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error)
 }
 
-type chatRecordsData struct {
-	config   *config.Config
-	vectorDB *tcvectordb.Client
-}
-
-func NewChatRecordsData(config *config.Config, vectorDB *tcvectordb.Client) IChatRecordsData {
-	return &chatRecordsData{
-		config:   config,
-		vectorDB: vectorDB,
+func NewChatRecordsData(config *config.Config) (IChatRecordsData, error) {
+	switch config.Vector.Provider {
+	case "tencent", "":
+		return newTencentChatRecordsData(config)
+	case "pgvector":
+		return newPgvectorChatRecordsData(config)
+	default:
+		return nil, fmt.Errorf("unsupported vector provider: %s", config.Vector.Provider)
 	}
 }
-func (data *chatRecordsData) UpsertData(ctx context.Context, list []*ChatRecord) error {
-	database := data.config.VectorDB.Database
-	collection := CHAT_RECORDS
-	coll := data.vectorDB.Database(database).Collection(collection)
-	documentList := make([]tcvectordb.Document, 0, len(list))
-	for _, l := range list {
-		doc := tcvectordb.Document{
-			Id: l.ID,
-		}
-		doc.Fields = make(map[string]tcvectordb.Field, len(l.KVs))
-		for k, v := range l.KVs {
-			doc.Fields[k] = tcvectordb.Field{Val: v}
-		}
-		documentList = append(documentList, doc)
-	}
-	_, err := coll.Upsert(ctx, documentList)
-	if err != nil {
-		return err
-	}
-	return nil
-}
-func (data *chatRecordsData) QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error) {
-	database := data.config.VectorDB.Database
-	collection := CHAT_RECORDS
-	coll := data.vectorDB.Database(database).Collection(collection)
-	result, err := coll.SearchByText(ctx, text, &tcvectordb.SearchDocumentParams{
-		Params: &tcvectordb.SearchDocParams{Ef: 100},
-		Limit:  1,
-	})
-	if err != nil {
-		return "", 0, err
-	}
-	if len(result.Documents) > 0 && len(result.Documents[0]) > 0 {
-		doc := result.Documents[0][0]
-		return doc.Id, doc.Score, nil
-
-	}
-	return "", 0, nil
-}

ai-chat-service/chat-server/vector-data/pgvector.go

@@ -0,0 +1,121 @@
+package vector_data
+
+import (
+	"ai-chat-service/pkg/config"
+	"ai-chat-service/services/embedding"
+	"context"
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx"
+)
+
+type pgvectorChatRecordsData struct {
+	config   *config.Config
+	pool     *pgx.ConnPool
+	embedder embedding.Embedder
+}
+
+func newPgvectorChatRecordsData(config *config.Config) (IChatRecordsData, error) {
+	connConfig, err := pgx.ParseConnectionString(config.Vector.Pgvector.DSN)
+	if err != nil {
+		return nil, err
+	}
+	pool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
+		ConnConfig:     connConfig,
+		MaxConnections: config.Vector.Pgvector.MaxOpenConn,
+	})
+	if err != nil {
+		return nil, err
+	}
+	embedder, err := embedding.NewEmbedder(config)
+	if err != nil {
+		pool.Close()
+		return nil, err
+	}
+	return &pgvectorChatRecordsData{
+		config:   config,
+		pool:     pool,
+		embedder: embedder,
+	}, nil
+}
+
+func (data *pgvectorChatRecordsData) UpsertData(ctx context.Context, list []*ChatRecord) error {
+	table := data.config.Vector.Pgvector.Table
+	if table == "" {
+		table = "chat_record_vectors"
+	}
+	for _, item := range list {
+		recordID, err := strconv.ParseInt(item.ID, 10, 64)
+		if err != nil {
+			return err
+		}
+		keywordsText := embedding.BuildText(item.KVs["keywords"])
+		if keywordsText == "" {
+			continue
+		}
+		vector, err := data.embedder.Embed(ctx, keywordsText)
+		if err != nil {
+			return err
+		}
+		commandTag, err := data.pool.Exec(
+			fmt.Sprintf(
+				"INSERT INTO %s (record_id, keywords_text, embedding, created_at) VALUES ($1, $2, $3::vector, $4) ON CONFLICT (record_id) DO UPDATE SET keywords_text = EXCLUDED.keywords_text, embedding = EXCLUDED.embedding, created_at = EXCLUDED.created_at",
+				table,
+			),
+			recordID,
+			keywordsText,
+			vectorLiteral(vector),
+			time.Now().Unix(),
+		)
+		if err != nil {
+			return err
+		}
+		if commandTag.RowsAffected() == 0 {
+			return fmt.Errorf("pgvector upsert affected 0 rows for record_id=%d", recordID)
+		}
+	}
+	return nil
+}
+
+func (data *pgvectorChatRecordsData) QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error) {
+	keywordsText := embedding.BuildText(text["keywords"]...)
+	if keywordsText == "" {
+		return "", 0, nil
+	}
+	vector, err := data.embedder.Embed(ctx, keywordsText)
+	if err != nil {
+		return "", 0, err
+	}
+	table := data.config.Vector.Pgvector.Table
+	if table == "" {
+		table = "chat_record_vectors"
+	}
+	var recordID int64
+	err = data.pool.QueryRowEx(
+		ctx,
+		fmt.Sprintf(
+			"SELECT record_id, CAST(1 - (embedding <=> $1::vector) AS real) AS score FROM %s ORDER BY embedding <=> $1::vector LIMIT 1",
+			table,
+		),
+		nil,
+		vectorLiteral(vector),
+	).Scan(&recordID, &score)
+	if err != nil {
+		if err == pgx.ErrNoRows {
+			return "", 0, nil
+		}
+		return "", 0, err
+	}
+	return strconv.FormatInt(recordID, 10), score, nil
+}
+
+func vectorLiteral(values []float32) string {
+	parts := make([]string, 0, len(values))
+	for _, value := range values {
+		parts = append(parts, strconv.FormatFloat(float64(value), 'f', -1, 32))
+	}
+	return "[" + strings.Join(parts, ",") + "]"
+}

ai-chat-service/chat-server/vector-data/tencent.go

@@ -0,0 +1,66 @@
+package vector_data
+
+import (
+	"ai-chat-service/pkg/config"
+	"context"
+	"time"
+
+	"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
+)
+
+type tencentChatRecordsData struct {
+	config   *config.Config
+	vectorDB *tcvectordb.Client
+}
+
+func newTencentChatRecordsData(config *config.Config) (IChatRecordsData, error) {
+	option := &tcvectordb.ClientOption{
+		Timeout:            time.Second * time.Duration(config.Vector.Tencent.Timeout),
+		MaxIdldConnPerHost: config.Vector.Tencent.MaxIdleConnPerHost,
+		IdleConnTimeout:    time.Second * time.Duration(config.Vector.Tencent.IdleConnTimeout),
+		ReadConsistency:    tcvectordb.ReadConsistency(config.Vector.Tencent.ReadConsistency),
+	}
+	client, err := tcvectordb.NewClient(config.Vector.Tencent.Url, config.Vector.Tencent.Username, config.Vector.Tencent.Pwd, option)
+	if err != nil {
+		return nil, err
+	}
+	return &tencentChatRecordsData{
+		config:   config,
+		vectorDB: client,
+	}, nil
+}
+
+func (data *tencentChatRecordsData) UpsertData(ctx context.Context, list []*ChatRecord) error {
+	database := data.config.Vector.Tencent.Database
+	collection := CHAT_RECORDS
+	coll := data.vectorDB.Database(database).Collection(collection)
+	documentList := make([]tcvectordb.Document, 0, len(list))
+	for _, l := range list {
+		doc := tcvectordb.Document{Id: l.ID}
+		doc.Fields = make(map[string]tcvectordb.Field, len(l.KVs))
+		for k, v := range l.KVs {
+			doc.Fields[k] = tcvectordb.Field{Val: v}
+		}
+		documentList = append(documentList, doc)
+	}
+	_, err := coll.Upsert(ctx, documentList)
+	return err
+}
+
+func (data *tencentChatRecordsData) QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error) {
+	database := data.config.Vector.Tencent.Database
+	collection := CHAT_RECORDS
+	coll := data.vectorDB.Database(database).Collection(collection)
+	result, err := coll.SearchByText(ctx, text, &tcvectordb.SearchDocumentParams{
+		Params: &tcvectordb.SearchDocParams{Ef: 100},
+		Limit:  1,
+	})
+	if err != nil {
+		return "", 0, err
+	}
+	if len(result.Documents) > 0 && len(result.Documents[0]) > 0 {
+		doc := result.Documents[0][0]
+		return doc.Id, doc.Score, nil
+	}
+	return "", 0, nil
+}