ingest.go

package recon

import (
	"context"
	"errors"

	milvus "github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"github.com/neo4j/neo4j-go-driver/v5/neo4j"
	"github.com/pulpchat/pulpspot/db/models"
	"go.uber.org/zap"
)

type Recommendation struct{}

func RecommendationIngest(ctx context.Context, logger *zap.SugaredLogger, neoSess neo4j.SessionWithContext, milvusClient milvus.Client, user *models.User) ([]*models.Post, error) {
	tx, _ := user.GetUserWithTracks(ctx)

	res, err := neoSess.ExecuteRead(ctx, tx)
	if err != nil {
		logger.Errorf("error getting user with tracks: %v", err)

		return nil, err
	}

	user = res.(*models.User)

	recommendedPosts := make([]*models.Post, 0)

	for _, track := range user.TopTracks {
		err := milvusClient.LoadCollection(
			ctx,
			"lyric_collection",
			false,
		)

		if err != nil {
			logger.Errorf("error loading collection: %v", err)

			return nil, err
		}

		sp, _ := entity.NewIndexFlatSearchParam()

		sr, err := milvusClient.Search(ctx, "lyric_collection", []string{}, "", []string{"ID, Name, PostID"}, []entity.Vector{entity.FloatVector(track.LyricsEmbeddings)}, "Vector",
			entity.L2, 45, sp)

		if err != nil {
			logger.Errorf("error searching: %v", err)

			return nil, err
		}

		for _, result := range sr {
			var idColumn *entity.ColumnVarChar
			var nameColumn *entity.ColumnVarChar
			var postIDColumn *entity.ColumnVarChar

			for _, field := range result.Fields {
				if field.Name() == "ID" {
					c, ok := field.(*entity.ColumnVarChar)
					if ok {
						idColumn = c
					} else {
						return nil, errors.New("id column not found")
					}
				} else if field.Name() == "Name" {
					c, ok := field.(*entity.ColumnVarChar)
					if ok {
						nameColumn = c
					} else {
						return nil, errors.New("name column not found")
					}
				} else if field.Name() == "PostID" {
					c, ok := field.(*entity.ColumnVarChar)
					if ok {
						postIDColumn = c
					} else {
						return nil, errors.New("post id column not found")
					}
				}
			}

			if idColumn == nil {
				return nil, errors.New("id column not found")
			}

			if nameColumn == nil {
				return nil, errors.New("name column not found")
			}

			if postIDColumn == nil {
				return nil, errors.New("post id column not found")
			}

			post := models.DefaultPost()
			post.ID = postIDColumn.Data()[0]

			tx, _ := post.GetPostWithTrack(ctx)

			res, err := neoSess.ExecuteRead(ctx, tx)
			if err != nil {
				logger.Errorf("error getting post with track: %v", err)

				return nil, err
			}

			post = res.(*models.Post)

			recommendedPosts = append(recommendedPosts, post)

			logger.Infof("Found post %s with name %s", postIDColumn.Data()[0], nameColumn.Data()[0])
		}
	}

	return recommendedPosts, nil
}