modelRT/model/redis_recommend.go

// Package model define model struct of model runtime service
package model

import (
	"context"
	"fmt"
	"math"
	"strings"

	"modelRT/constants"
	"modelRT/diagram"
	"modelRT/logger"

	"github.com/RediSearch/redisearch-go/v2/redisearch"
	redigo "github.com/gomodule/redigo/redis"
)

var ac *redisearch.Autocompleter

// InitAutocompleterWithPool define func of initialize the Autocompleter with redigo pool
func InitAutocompleterWithPool(pool *redigo.Pool) {
	ac = redisearch.NewAutocompleterFromPool(pool, constants.RedisSearchDictName)
}

// RedisSearchRecommend define func of redis search by input string and return recommend results
func RedisSearchRecommend(ctx context.Context, input string) ([]string, error) {
	rdb := diagram.GetRedisClientInstance()

	if input == "" {
		// 返回所有 grid 名
		return getAllGridKeys(ctx, constants.RedisAllGridSetKey)
	}

	inputs := strings.Split(input, ".")
	inputLen := len(inputs)

	switch inputLen {
	case 1:
		gridExist, err := rdb.SIsMember(ctx, constants.RedisAllGridSetKey, input).Result()
		if err != nil {
			logger.Error(ctx, "check grid key exist failed ", "grid_key", input, "error", err)
			return []string{}, err
		}
		// TODO delete debug info
		fmt.Println("gridExist", gridExist)
		if !gridExist {
			results, err := ac.SuggestOpts(input, redisearch.SuggestOptions{
				// TODO 测试如何返回全部的可能模糊结果
				Num:          math.MaxInt16,
				Fuzzy:        true,
				WithScores:   false,
				WithPayloads: false,
			})
			// TODO delete debug info
			fmt.Printf("results:%+v\n", results)
			fmt.Println("err", err)
			if err != nil {
				logger.Error(ctx, "query info by fuzzy failed", "query_key", input, "error", err)
				return []string{}, err
			}

			if len(results) == 0 {
				// TODO 构建 for 循环返回所有可能的补全
				// continue
			}

			var grids []string
			for _, result := range results {
				grids = append(grids, result.Term)
			}
			// 返回模糊查询结果
			return grids, nil
		}

		// 处理 input 不为空、不含有.并且 input 是一个完整的 grid key 的情况
		if strings.HasSuffix(input, ".") == false {
			return []string{"."}, nil
		}

		// 处理 input 不为空并且以.结尾的情况
		if strings.HasSuffix(input, ".") == true {
			setKey := fmt.Sprintf(constants.RedisSpecGridZoneSetKey, input)
			return getSpecificZoneKeys(ctx, setKey)
		}
	default:
		lastToken := inputs[inputLen-1]
		// 判断 queryKey 是否是空值，空值则返回上一级别下的所有key
		if lastToken == "" {
			setKey := getConstantsKeyByLength(inputLen - 1)
			targetSet := diagram.NewRedisSet(ctx, setKey, 10, true)

			results, err := targetSet.SMembers(setKey)
			if err != nil {
				logger.Error(ctx, "get all recommend key by setKey failed", "set_key", setKey, "error", err)
				return []string{}, fmt.Errorf("get all recommend key by setKey failed,%w", err)
			}
			return results, nil
		}

		querykey := lastToken
		setKey := getConstantsKeyByLength(inputLen)
		targetSet := diagram.NewRedisSet(ctx, setKey, 10, true)
		exist, err := targetSet.SIsMember(setKey, querykey)
		if err != nil {
			logger.Error(ctx, "check keys exist failed", "set_key", setKey, "query_key", querykey, "error", err)
			return []string{}, fmt.Errorf("check keys failed,%w", err)
		}

		if !exist {
			logger.Info(ctx, "use fuzzy query", "input", input)
			results, err := ac.SuggestOpts(input, redisearch.SuggestOptions{
				// TODO 测试如何返回全部的可能模糊结果
				Num:          math.MaxInt16,
				Fuzzy:        true,
				WithScores:   true,
				WithPayloads: true,
			})
			if err != nil {
				logger.Error(ctx, "query info by fuzzy failed", "query_key", input, "error", err)
				return []string{}, err
			}

			var terms []string
			for _, result := range results {
				terms = append(terms, result.Term)
			}
			// 返回模糊查询结果
			return terms, nil
		}
		return []string{}, nil
	}
	return []string{}, nil
}

func getAllGridKeys(ctx context.Context, setKey string) ([]string, error) {
	// 从redis set 中获取所有的 grid key
	gridSets := diagram.NewRedisSet(ctx, setKey, 10, true)
	keys, err := gridSets.SMembers("grid_keys")
	if err != nil {
		return []string{}, fmt.Errorf("get all root keys failed, error: %v", err)
	}
	return keys, nil
}

func getSpecificZoneKeys(ctx context.Context, setKey string) ([]string, error) {
	// TODO 从redis set 中获取指定 grid 下的 zone key
	zoneSets := diagram.NewRedisSet(ctx, setKey, 10, true)
	keys, err := zoneSets.SMembers("grid_keys")
	if err != nil {
		return []string{}, fmt.Errorf("get all root keys failed, error: %v", err)
	}
	return keys, nil
}

func getConstantsKeyByLength(inputLen int) string {
	switch inputLen {
	case 1:
		return constants.RedisAllGridSetKey
	case 2:
		return constants.RedisAllZoneSetKey
	case 3:
		return constants.RedisAllStationSetKey
	case 4:
		return constants.RedisAllComponentSetKey
	default:
		return constants.RedisAllGridSetKey
	}
}

// GetLongestCommonPrefixLength define func of get longest common prefix length between two strings
func GetLongestCommonPrefixLength(input string, recommendResult string) int {
	if input == "" {
		return 0
	}

	minLen := min(len(input), len(recommendResult))

	for i := range minLen {
		if input[i] != recommendResult[i] {
			return i
		}
	}
	return minLen
}