AI-Writie-Assistant/internal/vector/store.go

package vector

import (
	"context"
	"fmt"
	"log/slog"

	pb "github.com/qdrant/go-client/qdrant"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

// Chunk is a text segment with its vector.
type Chunk struct {
	ID       string
	SourceID string
	Text     string
	Vector   []float32
}

// QdrantStore implements vector storage via remote Qdrant gRPC.
type QdrantStore struct {
	conn       *grpc.ClientConn
	points     pb.PointsClient
	collection pb.CollectionsClient
	dimension  uint64
}

// NewQdrantStore connects to a Qdrant instance.
func NewQdrantStore(endpoint string, dimension uint64) (*QdrantStore, error) {
	conn, err := grpc.NewClient(endpoint, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, fmt.Errorf("connect qdrant: %w", err)
	}
	return &QdrantStore{
		conn:       conn,
		points:     pb.NewPointsClient(conn),
		collection: pb.NewCollectionsClient(conn),
		dimension:  dimension,
	}, nil
}

// EnsureCollection creates a collection if it doesn't exist.
func (s *QdrantStore) EnsureCollection(ctx context.Context, name string) error {
	_, err := s.collection.Get(ctx, &pb.GetCollectionInfoRequest{CollectionName: name})
	if err == nil {
		return nil // already exists
	}

	_, err = s.collection.Create(ctx, &pb.CreateCollection{
		CollectionName: name,
		VectorsConfig: &pb.VectorsConfig{
			Config: &pb.VectorsConfig_Params{
				Params: &pb.VectorParams{
					Size:     s.dimension,
					Distance: pb.Distance_Cosine,
				},
			},
		},
	})
	if err != nil {
		return fmt.Errorf("create collection %s: %w", name, err)
	}
	slog.Info("created qdrant collection", "name", name, "dim", s.dimension)
	return nil
}

// Insert upserts chunks into the specified collection.
func (s *QdrantStore) Insert(ctx context.Context, collectionName string, chunks []Chunk) error {
	points := make([]*pb.PointStruct, len(chunks))
	for i, c := range chunks {
		points[i] = &pb.PointStruct{
			Id: &pb.PointId{
				PointIdOptions: &pb.PointId_Uuid{Uuid: c.ID},
			},
			Vectors: &pb.Vectors{
				VectorsOptions: &pb.Vectors_Vector{
					Vector: &pb.Vector{Data: c.Vector},
				},
			},
			Payload: map[string]*pb.Value{
				"text":      {Kind: &pb.Value_StringValue{StringValue: c.Text}},
				"source_id": {Kind: &pb.Value_StringValue{StringValue: c.SourceID}},
			},
		}
	}

	_, err := s.points.Upsert(ctx, &pb.UpsertPoints{
		CollectionName: collectionName,
		Points:         points,
	})
	return err
}

// Search performs KNN search and returns top-k results.
func (s *QdrantStore) Search(ctx context.Context, collectionName string, queryVec []float32, topK uint64) ([]Chunk, error) {
	resp, err := s.points.Search(ctx, &pb.SearchPoints{
		CollectionName: collectionName,
		Vector:         queryVec,
		Limit:          topK,
		WithPayload:    &pb.WithPayloadSelector{SelectorOptions: &pb.WithPayloadSelector_Enable{Enable: true}},
	})
	if err != nil {
		return nil, fmt.Errorf("qdrant search: %w", err)
	}

	results := make([]Chunk, 0, len(resp.Result))
	for _, hit := range resp.Result {
		text := ""
		sourceID := ""
		if v, ok := hit.Payload["text"]; ok {
			text = v.GetStringValue()
		}
		if v, ok := hit.Payload["source_id"]; ok {
			sourceID = v.GetStringValue()
		}
		results = append(results, Chunk{
			ID:       hit.Id.GetUuid(),
			SourceID: sourceID,
			Text:     text,
		})
	}
	return results, nil
}

// DeleteCollection removes a collection.
func (s *QdrantStore) DeleteCollection(ctx context.Context, name string) error {
	_, err := s.collection.Delete(ctx, &pb.DeleteCollection{CollectionName: name})
	return err
}

// Close closes the gRPC connection.
func (s *QdrantStore) Close() {
	if s.conn != nil {
		s.conn.Close()
	}
}

// TestConnection verifies the Qdrant server is reachable.
func (s *QdrantStore) TestConnection(ctx context.Context) (bool, error) {
	_, err := s.collection.List(ctx, &pb.ListCollectionsRequest{})
	if err != nil {
		return false, err
	}
	return true, nil
}