redis缓存替换+pgvector向量替换

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