70 lines
1.9 KiB
Go
70 lines
1.9 KiB
Go
package vector_data
|
|
|
|
import (
|
|
"ai-chat-service/pkg/config"
|
|
"context"
|
|
"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
|
|
)
|
|
|
|
const CHAT_RECORDS = "chat_records"
|
|
|
|
type ChatRecord struct {
|
|
ID string
|
|
KVs map[string]string
|
|
}
|
|
type IChatRecordsData interface {
|
|
UpsertData(ctx context.Context, list []*ChatRecord) error
|
|
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 (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
|
|
}
|