recommender/cmd/create-collection.go

package cmd

import (
	"context"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"github.com/spf13/cobra"
	"leafdev.top/Ecosystem/recommender/internal/base"
)

func init() {
	RootCmd.AddCommand(createCollectionCmd)
	createCollectionCmd.Flags().String("dim", "768", "模型的维度")
	// 将 dim 参数标记为必填
	err := createCollectionCmd.MarkFlagRequired("dim")
	if err != nil {
		panic(err)
	}
}

var createCollectionCmd = &cobra.Command{
	Use: "create-collection",
	Run: func(cmd *cobra.Command, args []string) {
		app, err := CreateApp()
		if err != nil {
			panic(err)
			return
		}
		// 获取 flag
		dim, err := cmd.Flags().GetString("dim")
		if err != nil {
			panic(err)
		}

		createMilvusTagCollection(app, dim)
		createMilvusUserSummaryCollection(app, dim)
		createMilvusPostCollection(app, dim)
	},
}

func createMilvusTagCollection(app *base.Application, dim string) {
	var ctx = context.Background()
	var field = []*entity.Field{
		{
			Name:       "tag_id",
			PrimaryKey: true,
			AutoID:     false,
			DataType:   entity.FieldTypeInt64,
		},
		{
			Name:       "vector",
			PrimaryKey: false,
			DataType:   entity.FieldTypeFloatVector,
			TypeParams: map[string]string{
				"dim": dim,
			},
		},
	}

	var schema = &entity.Schema{
		CollectionName:     app.Config.Milvus.TagCollection,
		Description:        "",
		AutoID:             true,
		Fields:             field,
		EnableDynamicField: true,
	}

	err := app.Milvus.CreateCollection(ctx, schema, 2)
	if err != nil {
		panic(err)
	}

	index := entity.NewGenericIndex("idx_tag_id", entity.Inverted, map[string]string{})
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.TagCollection, "tag_id", index, false)
	if err != nil {
		panic(err)
	}

	index, err = entity.NewIndexAUTOINDEX(entity.L2)
	if err != nil {
		panic(err)
	}
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.TagCollection, "vector", index, false)
	if err != nil {
		panic(err)
	}
}

func createMilvusUserSummaryCollection(app *base.Application, dim string) {
	var ctx = context.Background()
	var field = []*entity.Field{
		{
			Name:       "external_user_id",
			PrimaryKey: true,
			AutoID:     false,
			DataType:   entity.FieldTypeInt64,
		},
		{
			Name:       "application_id",
			PrimaryKey: false,
			AutoID:     false,
			DataType:   entity.FieldTypeInt64,
		},
		{
			Name:       "vector",
			PrimaryKey: false,
			DataType:   entity.FieldTypeFloatVector,
			TypeParams: map[string]string{
				"dim": dim,
			},
		},
	}

	var schema = &entity.Schema{
		CollectionName:     app.Config.Milvus.UserSummaryCollection,
		Description:        "",
		AutoID:             true,
		Fields:             field,
		EnableDynamicField: true,
	}

	err := app.Milvus.CreateCollection(ctx, schema, 2)
	if err != nil {
		panic(err)
	}

	index := entity.NewGenericIndex("idx_external_user_id", entity.Inverted, map[string]string{})
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.UserSummaryCollection, "external_user_id", index, false)
	if err != nil {
		panic(err)
	}

	index = entity.NewGenericIndex("idx_application_id", entity.Inverted, map[string]string{})
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.UserSummaryCollection, "application_id", index, false)
	if err != nil {
		panic(err)
	}

	index, err = entity.NewIndexAUTOINDEX(entity.L2)
	if err != nil {
		panic(err)
	}
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.UserSummaryCollection, "vector", index, false)
	if err != nil {
		panic(err)
	}
}

func createMilvusPostCollection(app *base.Application, dim string) {
	var ctx = context.Background()
	var field = []*entity.Field{
		{
			Name:       "post_id",
			PrimaryKey: true,
			AutoID:     false,
			DataType:   entity.FieldTypeInt64,
		},
		{
			Name:       "category_id",
			PrimaryKey: false,
			AutoID:     false,
			DataType:   entity.FieldTypeInt64,
		},
		{
			Name:       "application_id",
			PrimaryKey: false,
			AutoID:     false,
			DataType:   entity.FieldTypeInt64,
		},
		{
			Name:       "vector",
			PrimaryKey: false,
			DataType:   entity.FieldTypeFloatVector,
			TypeParams: map[string]string{
				"dim": dim,
			},
		},
	}

	var schema = &entity.Schema{
		CollectionName:     app.Config.Milvus.PostCollection,
		Description:        "",
		AutoID:             true,
		Fields:             field,
		EnableDynamicField: true,
	}

	err := app.Milvus.CreateCollection(ctx, schema, 2)
	if err != nil {
		panic(err)
	}

	index := entity.NewGenericIndex("idx_post_id", entity.Inverted, map[string]string{})
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.PostCollection, "post_id", index, false)
	if err != nil {
		panic(err)
	}

	index = entity.NewGenericIndex("idx_category_id", entity.Inverted, map[string]string{})
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.PostCollection, "category_id", index, false)
	if err != nil {
		panic(err)
	}

	index = entity.NewGenericIndex("idx_application_id", entity.Inverted, map[string]string{})
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.PostCollection, "application_id", index, false)
	if err != nil {
		panic(err)
	}

	index, err = entity.NewIndexAUTOINDEX(entity.L2)
	if err != nil {
		panic(err)
	}
	err = app.Milvus.CreateIndex(ctx, app.Config.Milvus.PostCollection, "vector", index, false)
	if err != nil {
		panic(err)
	}
}