67 lines
2.1 KiB
Go
67 lines
2.1 KiB
Go
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
|
|
}
|