429 lines
10 KiB
Go
429 lines
10 KiB
Go
package retrieval
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"net/url"
|
|
"strconv"
|
|
"strings"
|
|
|
|
sdk "github.com/qdrant/go-client/qdrant"
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/geo-platform/tenant-api/internal/shared/config"
|
|
"github.com/geo-platform/tenant-api/internal/shared/middleware"
|
|
)
|
|
|
|
const (
|
|
defaultQdrantHost = "localhost"
|
|
defaultQdrantGRPCPort = 6334
|
|
defaultQdrantCollection = "geo_knowledge_chunks"
|
|
)
|
|
|
|
type VectorStore interface {
|
|
Validate() error
|
|
EnsureCollection(ctx context.Context, vectorSize int) error
|
|
Upsert(ctx context.Context, points []VectorPoint) error
|
|
Search(ctx context.Context, req SearchRequest) ([]SearchPoint, error)
|
|
DeletePoints(ctx context.Context, pointIDs []string) error
|
|
}
|
|
|
|
type VectorPoint struct {
|
|
ID string
|
|
Vector []float64
|
|
Payload map[string]any
|
|
}
|
|
|
|
type SearchRequest struct {
|
|
Vector []float64
|
|
TenantID int64
|
|
BrandID int64
|
|
GroupIDs []int64
|
|
Limit int
|
|
}
|
|
|
|
type SearchPoint struct {
|
|
ID string
|
|
Score float64
|
|
Payload map[string]any
|
|
}
|
|
|
|
type qdrantStore struct {
|
|
client *sdk.Client
|
|
collection string
|
|
logger *zap.Logger
|
|
}
|
|
|
|
type disabledVectorStore struct {
|
|
reason string
|
|
}
|
|
|
|
func NewQdrantStore(cfg config.QdrantConfig, logger *zap.Logger) VectorStore {
|
|
host, port, useTLS, err := parseQdrantEndpoint(cfg.URL)
|
|
if err != nil {
|
|
return disabledVectorStore{reason: fmt.Sprintf("invalid qdrant.url: %v", err)}
|
|
}
|
|
|
|
collection := strings.TrimSpace(cfg.Collection)
|
|
if collection == "" {
|
|
collection = defaultQdrantCollection
|
|
}
|
|
|
|
client, err := sdk.NewClient(&sdk.Config{
|
|
Host: host,
|
|
Port: port,
|
|
APIKey: strings.TrimSpace(cfg.APIKey),
|
|
UseTLS: useTLS,
|
|
SkipCompatibilityCheck: true,
|
|
})
|
|
if err != nil {
|
|
return disabledVectorStore{reason: fmt.Sprintf("init qdrant client failed: %v", err)}
|
|
}
|
|
|
|
return &qdrantStore{
|
|
client: client,
|
|
collection: collection,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
func (s *qdrantStore) Validate() error {
|
|
if s == nil || s.client == nil {
|
|
return fmt.Errorf("%w: qdrant client is not initialized", ErrNotConfigured)
|
|
}
|
|
if strings.TrimSpace(s.collection) == "" {
|
|
return fmt.Errorf("%w: qdrant collection is empty", ErrNotConfigured)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *qdrantStore) EnsureCollection(ctx context.Context, vectorSize int) error {
|
|
if err := s.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if vectorSize <= 0 {
|
|
return fmt.Errorf("invalid qdrant vector size %d", vectorSize)
|
|
}
|
|
|
|
exists, err := s.client.CollectionExists(ctx, s.collection)
|
|
if err != nil {
|
|
s.logError(ctx, "qdrant collection exists check failed",
|
|
zap.String("collection", s.collection),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("query qdrant collection: %w", err)
|
|
}
|
|
if exists {
|
|
return nil
|
|
}
|
|
|
|
err = s.client.CreateCollection(ctx, &sdk.CreateCollection{
|
|
CollectionName: s.collection,
|
|
VectorsConfig: sdk.NewVectorsConfig(&sdk.VectorParams{
|
|
Size: uint64(vectorSize),
|
|
Distance: sdk.Distance_Cosine,
|
|
}),
|
|
})
|
|
if err != nil {
|
|
s.logError(ctx, "qdrant create collection failed",
|
|
zap.String("collection", s.collection),
|
|
zap.Int("vector_size", vectorSize),
|
|
zap.Error(err),
|
|
)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (s *qdrantStore) Upsert(ctx context.Context, points []VectorPoint) error {
|
|
if err := s.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if len(points) == 0 {
|
|
return nil
|
|
}
|
|
|
|
converted := make([]*sdk.PointStruct, 0, len(points))
|
|
for _, point := range points {
|
|
if strings.TrimSpace(point.ID) == "" || len(point.Vector) == 0 {
|
|
continue
|
|
}
|
|
converted = append(converted, &sdk.PointStruct{
|
|
Id: sdk.NewID(point.ID),
|
|
Vectors: sdk.NewVectors(float64SliceToFloat32(point.Vector)...),
|
|
Payload: sdk.NewValueMap(point.Payload),
|
|
})
|
|
}
|
|
if len(converted) == 0 {
|
|
return nil
|
|
}
|
|
|
|
wait := true
|
|
_, err := s.client.Upsert(ctx, &sdk.UpsertPoints{
|
|
CollectionName: s.collection,
|
|
Wait: &wait,
|
|
Points: converted,
|
|
})
|
|
if err != nil {
|
|
s.logError(ctx, "qdrant upsert failed",
|
|
zap.String("collection", s.collection),
|
|
zap.Int("point_count", len(converted)),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("upsert qdrant points: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *qdrantStore) Search(ctx context.Context, req SearchRequest) ([]SearchPoint, error) {
|
|
if err := s.Validate(); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(req.Vector) == 0 || req.TenantID <= 0 {
|
|
return []SearchPoint{}, nil
|
|
}
|
|
limit := req.Limit
|
|
if limit <= 0 {
|
|
limit = 10
|
|
}
|
|
limitValue := uint64(limit)
|
|
|
|
filter := &sdk.Filter{
|
|
Must: []*sdk.Condition{
|
|
sdk.NewMatchInt("tenant_id", req.TenantID),
|
|
sdk.NewMatchKeyword("status", "active"),
|
|
},
|
|
}
|
|
if req.BrandID > 0 {
|
|
filter.Must = append(filter.Must, sdk.NewMatchInt("brand_id", req.BrandID))
|
|
}
|
|
groupIDs := normalizeGroupIDs(req.GroupIDs)
|
|
if len(groupIDs) > 0 {
|
|
filter.Must = append(filter.Must, sdk.NewMatchInts("group_id", groupIDs...))
|
|
}
|
|
|
|
points, err := s.client.Query(ctx, &sdk.QueryPoints{
|
|
CollectionName: s.collection,
|
|
Query: sdk.NewQueryDense(float64SliceToFloat32(req.Vector)),
|
|
Filter: filter,
|
|
Limit: &limitValue,
|
|
WithPayload: sdk.NewWithPayload(true),
|
|
})
|
|
if err != nil {
|
|
s.logError(ctx, "qdrant search failed",
|
|
zap.String("collection", s.collection),
|
|
zap.Int64("tenant_id", req.TenantID),
|
|
zap.Int("group_count", len(groupIDs)),
|
|
zap.Int("limit", limit),
|
|
zap.Error(err),
|
|
)
|
|
return nil, fmt.Errorf("query qdrant points: %w", err)
|
|
}
|
|
|
|
results := make([]SearchPoint, 0, len(points))
|
|
for _, point := range points {
|
|
if point == nil {
|
|
continue
|
|
}
|
|
results = append(results, SearchPoint{
|
|
ID: pointIDToString(point.GetId()),
|
|
Score: float64(point.GetScore()),
|
|
Payload: valueMapToAny(point.GetPayload()),
|
|
})
|
|
}
|
|
return results, nil
|
|
}
|
|
|
|
func (s *qdrantStore) DeletePoints(ctx context.Context, pointIDs []string) error {
|
|
if err := s.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if len(pointIDs) == 0 {
|
|
return nil
|
|
}
|
|
|
|
ids := make([]*sdk.PointId, 0, len(pointIDs))
|
|
for _, pointID := range pointIDs {
|
|
pointID = strings.TrimSpace(pointID)
|
|
if pointID == "" {
|
|
continue
|
|
}
|
|
ids = append(ids, sdk.NewID(pointID))
|
|
}
|
|
if len(ids) == 0 {
|
|
return nil
|
|
}
|
|
|
|
wait := true
|
|
_, err := s.client.Delete(ctx, &sdk.DeletePoints{
|
|
CollectionName: s.collection,
|
|
Wait: &wait,
|
|
Points: sdk.NewPointsSelector(ids...),
|
|
})
|
|
if err != nil {
|
|
s.logError(ctx, "qdrant delete points failed",
|
|
zap.String("collection", s.collection),
|
|
zap.Int("point_count", len(ids)),
|
|
zap.Error(err),
|
|
)
|
|
return fmt.Errorf("delete qdrant points: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *qdrantStore) logError(ctx context.Context, msg string, fields ...zap.Field) {
|
|
if s == nil || s.logger == nil {
|
|
return
|
|
}
|
|
if requestID := middleware.RequestIDFromContext(ctx); requestID != "" {
|
|
fields = append(fields, zap.String("request_id", requestID))
|
|
}
|
|
s.logger.Error(msg, fields...)
|
|
}
|
|
|
|
func parseQdrantEndpoint(raw string) (string, int, bool, error) {
|
|
value := strings.TrimSpace(raw)
|
|
if value == "" {
|
|
return defaultQdrantHost, defaultQdrantGRPCPort, false, nil
|
|
}
|
|
|
|
if strings.Contains(value, "://") {
|
|
parsed, err := url.Parse(value)
|
|
if err != nil {
|
|
return "", 0, false, err
|
|
}
|
|
port := defaultQdrantGRPCPort
|
|
if parsed.Port() != "" {
|
|
port, err = strconv.Atoi(parsed.Port())
|
|
if err != nil {
|
|
return "", 0, false, err
|
|
}
|
|
}
|
|
host := parsed.Hostname()
|
|
if host == "" {
|
|
host = defaultQdrantHost
|
|
}
|
|
return host, port, parsed.Scheme == "https", nil
|
|
}
|
|
|
|
host, portText, err := net.SplitHostPort(value)
|
|
if err == nil {
|
|
port, convErr := strconv.Atoi(portText)
|
|
if convErr != nil {
|
|
return "", 0, false, convErr
|
|
}
|
|
return host, port, false, nil
|
|
}
|
|
|
|
if strings.Contains(err.Error(), "missing port in address") {
|
|
return value, defaultQdrantGRPCPort, false, nil
|
|
}
|
|
|
|
return "", 0, false, err
|
|
}
|
|
|
|
func float64SliceToFloat32(values []float64) []float32 {
|
|
result := make([]float32, 0, len(values))
|
|
for _, value := range values {
|
|
result = append(result, float32(value))
|
|
}
|
|
return result
|
|
}
|
|
|
|
func pointIDToString(id *sdk.PointId) string {
|
|
if id == nil {
|
|
return ""
|
|
}
|
|
if uuid := strings.TrimSpace(id.GetUuid()); uuid != "" {
|
|
return uuid
|
|
}
|
|
if num := id.GetNum(); num > 0 {
|
|
return strconv.FormatUint(num, 10)
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func valueMapToAny(values map[string]*sdk.Value) map[string]any {
|
|
if len(values) == 0 {
|
|
return map[string]any{}
|
|
}
|
|
result := make(map[string]any, len(values))
|
|
for key, value := range values {
|
|
result[key] = valueToAny(value)
|
|
}
|
|
return result
|
|
}
|
|
|
|
func valueToAny(value *sdk.Value) any {
|
|
if value == nil {
|
|
return nil
|
|
}
|
|
switch value.GetKind().(type) {
|
|
case *sdk.Value_StringValue:
|
|
return value.GetStringValue()
|
|
case *sdk.Value_IntegerValue:
|
|
return value.GetIntegerValue()
|
|
case *sdk.Value_DoubleValue:
|
|
return value.GetDoubleValue()
|
|
case *sdk.Value_BoolValue:
|
|
return value.GetBoolValue()
|
|
case *sdk.Value_ListValue:
|
|
list := value.GetListValue()
|
|
if list == nil {
|
|
return []any{}
|
|
}
|
|
items := make([]any, 0, len(list.GetValues()))
|
|
for _, item := range list.GetValues() {
|
|
items = append(items, valueToAny(item))
|
|
}
|
|
return items
|
|
case *sdk.Value_StructValue:
|
|
structValue := value.GetStructValue()
|
|
if structValue == nil {
|
|
return map[string]any{}
|
|
}
|
|
return valueMapToAny(structValue.GetFields())
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func normalizeGroupIDs(ids []int64) []int64 {
|
|
result := make([]int64, 0, len(ids))
|
|
seen := make(map[int64]struct{}, len(ids))
|
|
for _, id := range ids {
|
|
if id <= 0 {
|
|
continue
|
|
}
|
|
if _, ok := seen[id]; ok {
|
|
continue
|
|
}
|
|
seen[id] = struct{}{}
|
|
result = append(result, id)
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (s disabledVectorStore) Validate() error {
|
|
if s.reason == "" {
|
|
s.reason = "missing qdrant configuration"
|
|
}
|
|
return fmt.Errorf("%w: %s", ErrNotConfigured, s.reason)
|
|
}
|
|
|
|
func (s disabledVectorStore) EnsureCollection(context.Context, int) error {
|
|
return s.Validate()
|
|
}
|
|
|
|
func (s disabledVectorStore) Upsert(context.Context, []VectorPoint) error {
|
|
return s.Validate()
|
|
}
|
|
|
|
func (s disabledVectorStore) Search(context.Context, SearchRequest) ([]SearchPoint, error) {
|
|
return nil, s.Validate()
|
|
}
|
|
|
|
func (s disabledVectorStore) DeletePoints(context.Context, []string) error {
|
|
return s.Validate()
|
|
}
|