geo/server/internal/shared/retrieval/qdrant.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
	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"),
		},
	}
	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()
}