init: initial commit

internal/vector/store.go

@@ -0,0 +1,146 @@
+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
+}