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, bool, error) {
	rdb := diagram.GetRedisClientInstance()

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

	inputs := strings.Split(input, ".")
	inputLen := len(inputs)
	fmt.Printf("inputLen:%d\n", inputLen)
	switch inputLen {
	case 1:
		// TODO 优化成NewSet的形式
		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{}, false, err
		}

		// TODO delete debug info
		fmt.Println("gridExist", gridExist)

		searchInput := input
		inputLen := len(searchInput)
		for inputLen != 0 && !gridExist {
			results, err := ac.SuggestOpts(searchInput, redisearch.SuggestOptions{
				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{}, false, err
			}

			if len(results) == 0 {
				// TODO 构建 for 循环返回所有可能的补全
				searchInput = input[:inputLen-1]
				inputLen = len(searchInput)
				fmt.Printf("next search input:%s\n", searchInput)
				fmt.Printf("next search input len:%d\n", inputLen)
				continue
			}

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

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

		// 处理 input 不为空并且以.结尾的情况
		if strings.HasSuffix(input, ".") == true {
			setKey := fmt.Sprintf(constants.RedisSpecGridZoneSetKey, input)
			return getSpecificZoneKeys(ctx, setKey)
		}
	default:
		lastInput := inputs[inputLen-1]
		// 判断 queryKey 是否是空值，空值则返回上一级别下的所有key
		if lastInput == "" {
			fmt.Println("last token is empty")
			// TODO 根据所在层级拼接相关setKey
			setKey := getCombinedConstantsKeyByLength(inputs[inputLen-2], inputLen)
			fmt.Printf("default case set key:%s\n", setKey)
			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{}, false, fmt.Errorf("get all recommend key by setKey failed,%w", err)
			}
			return results, false, nil
		}

		setKey := getCombinedConstantsKeyByLength(inputs[inputLen-2], inputLen)
		fmt.Printf("default case set key:%s\n", setKey)
		targetSet := diagram.NewRedisSet(ctx, setKey, 10, true)
		exist, err := targetSet.SIsMember(setKey, lastInput)
		if err != nil {
			logger.Error(ctx, "check keys exist failed", "set_key", setKey, "query_key", lastInput, "error", err)
			return []string{}, false, fmt.Errorf("check keys failed,%w", err)
		}

		searchInput := input
		inputLen := len(searchInput)
		for inputLen != 0 && !exist {
			logger.Info(ctx, "use fuzzy query", "input", input)
			results, err := ac.SuggestOpts(searchInput, redisearch.SuggestOptions{
				Num:          math.MaxInt16,
				Fuzzy:        true,
				WithScores:   false,
				WithPayloads: false,
			})
			if err != nil {
				logger.Error(ctx, "query info by fuzzy failed", "query_key", input, "error", err)
				return []string{}, false, err
			}
			if len(results) == 0 {
				// TODO 构建 for 循环返回所有可能的补全
				searchInput = input[:inputLen-1]
				inputLen = len(searchInput)
				fmt.Printf("next search input:%s\n", searchInput)
				fmt.Printf("next search input len:%d\n", inputLen)
				continue
			}

			fmt.Printf("default fuzzy search results:%v\n", results)
			var terms []string
			for _, result := range results {
				terms = append(terms, result.Term)
			}
			// 返回模糊查询结果
			return terms, true, nil
		}
		return []string{input}, false, nil
	}
	return []string{}, false, nil
}

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

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

func getCombinedConstantsKeyByLength(key string, inputLen int) string {
	switch inputLen {
	case 2:
		return fmt.Sprintf(constants.RedisSpecGridZoneSetKey, key)
	case 3:
		return fmt.Sprintf(constants.RedisSpecZoneStationSetKey, key)
	case 4:
		return fmt.Sprintf(constants.RedisSpecStationComponentSetKey, key)
	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
}