redis缓存替换+pgvector向量替换

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