Support using memory/redis as index storage

cmd/httpd/httpd.go

@@ -15,7 +15,6 @@ import (
 	"github.com/pandodao/botastic/handler/hc"
 	"github.com/pandodao/botastic/internal/chanhub"
 	"github.com/pandodao/botastic/internal/gpt"
-	"github.com/pandodao/botastic/internal/milvus"
 	"github.com/pandodao/botastic/internal/tiktoken"
 	appServ "github.com/pandodao/botastic/service/app"
 	botServ "github.com/pandodao/botastic/service/bot"
@@ -79,13 +78,15 @@ func NewCmdHttpd() *cobra.Command {
 			bots := bot.New(h)
 			orders := order.New(h)
 
-			milvusClient, err := milvus.Init(ctx, cfg.Milvus.Address)
+			indexes, err := index.Init(ctx, cfg.IndexStore)
 			if err != nil {
 				return err
 			}
-			indexes := index.New(ctx, milvusClient)
-			models := model.New(h)
+			if err := indexes.Init(ctx); err != nil {
+				return err
+			}
 
+			models := model.New(h)
 			tiktokenHandler, err := tiktoken.Init()
 			if err != nil {
 				return err

cmd/migrate/migrate.go

@@ -5,8 +5,6 @@ import (
 	"strconv"
 
 	"github.com/pandodao/botastic/config"
-	"github.com/pandodao/botastic/core"
-	"github.com/pandodao/botastic/internal/milvus"
 	"github.com/pandodao/botastic/store"
 	"github.com/spf13/cobra"
 )
@@ -92,24 +90,5 @@ func NewCmdMigrate() *cobra.Command {
 		},
 	})
 
-	cmd.AddCommand(&cobra.Command{
-		Use:   "milvus",
-		Short: "Update milvus collection",
-		RunE: func(cmd *cobra.Command, args []string) error {
-			ctx := cmd.Context()
-			cfg := config.C()
-			client, err := milvus.Init(ctx, cfg.Milvus.Address)
-			if err != nil {
-				return err
-			}
-
-			index := core.Index{}
-			if err := client.CreateCollectionIfNotExist(ctx, index.Schema(), 2); err != nil {
-				return err
-			}
-			return client.BuildIndex(ctx, index.CollectionName(), "vectors")
-		},
-	})
-
 	return cmd
 }

config.example.yaml

@@ -11,8 +11,17 @@ sys:
   extra_rate: 0.1
   init_user_credits: 0.1
 
-milvus:
-  address: "localhost:19530"
+index_store:
+  driver: memory # redis, milvus
+  redis:
+    address: "localhost:6379"
+    password: ""
+    db: 0
+    key_prefix: app
+  milvus:
+    address: "localhost:19530"
+    collection: indexes
+    collection_shards_num: 2
 
 mixpay:
   payee_id: "USER_UUID"

config/config.go

@@ -1,23 +1,39 @@
 package config
 
 import (
+	"fmt"
 	"io/ioutil"
 	"log"
 	"time"
 
 	"gopkg.in/yaml.v3"
 )
 
+type IndexStoreDriver string
+
+const (
+	IndexStoreMemory IndexStoreDriver = "memory"
+	IndexStoreMilvus IndexStoreDriver = "milvus"
+	IndexStoreRedis  IndexStoreDriver = "redis"
+)
+
 type Config struct {
 	DB          DBConfig          `yaml:"db"`
-	Milvus      Milvus            `yaml:"milvus"`
+	IndexStore  IndexStoreConfig  `yaml:"index_store"`
 	Sys         System            `yaml:"sys"`
 	OpenAI      OpenAIConfig      `yaml:"openai"`
 	Auth        Auth              `yaml:"auth"`
 	Mixpay      Mixpay            `yaml:"mixpay"`
 	OrderSyncer OrderSyncerConfig `yaml:"order_syncer"`
 }
 
+type IndexStoreConfig struct {
+	Driver    IndexStoreDriver        `yaml:"driver"`
+	Dimension int                     `yaml:"dimension"`
+	Milvus    *IndexStoreMilvusConfig `yaml:"milvus"`
+	Redis     *IndexStoreRedisConfig  `yaml:"redis"`
+}
+
 type OrderSyncerConfig struct {
 	Interval       time.Duration `yaml:"interval"`
 	CheckInterval  time.Duration `yaml:"check_interval"`
@@ -33,8 +49,17 @@ type Mixpay struct {
 	FailedReturnTo    string `yaml:"failed_return_to"`
 }
 
-type Milvus struct {
-	Address string `yaml:"address"`
+type IndexStoreMilvusConfig struct {
+	Collection          string `yaml:"collection"`
+	CollectionShardsNum int32  `yaml:"collection_shards_num"`
+	Address             string `yaml:"address"`
+}
+
+type IndexStoreRedisConfig struct {
+	Address   string `yaml:"address"`
+	Password  string `yaml:"password"`
+	DB        int    `yaml:"db"`
+	KeyPrefix string `yaml:"key_prefix"`
 }
 
 type OpenAIConfig struct {
@@ -61,6 +86,19 @@ type Auth struct {
 }
 
 func (c Config) validate() error {
+	switch c.IndexStore.Driver {
+	case IndexStoreMemory:
+	case IndexStoreMilvus:
+		if c.IndexStore.Milvus == nil {
+			return fmt.Errorf("index_store.milvus is required")
+		}
+	case IndexStoreRedis:
+		if c.IndexStore.Redis == nil {
+			return fmt.Errorf("index_store.redis is required")
+		}
+	default:
+		return fmt.Errorf("index_store.driver is invalid: %s", c.IndexStore.Driver)
+	}
 	return nil
 }
 
@@ -77,6 +115,9 @@ func defaultConfig() *Config {
 			CheckInterval:  10 * time.Second,
 			CancelInterval: 2 * time.Hour,
 		},
+		IndexStore: IndexStoreConfig{
+			Driver: IndexStoreMemory,
+		},
 	}
 }
 

core/index.go

@@ -2,100 +2,38 @@ package core
 
 import (
 	"context"
-	"strings"
-
-	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 )
 
 type (
 	Index struct {
-		ID         string    `json:"-"`
-		AppID      string    `json:"-"`
+		AppID      string    `json:"app_id,omitempty"`
 		Data       string    `json:"data"`
-		DataToken  int64     `json:"-"`
-		Vectors    []float32 `json:"-"`
+		DataToken  int64     `json:"data_token,omitempty"`
+		Vectors    []float32 `json:"vectors,omitempty"`
 		ObjectID   string    `json:"object_id"`
 		Category   string    `json:"category"`
 		Properties string    `json:"properties"`
-		CreatedAt  int64     `db:"created_at" json:"created_at"`
-		Score      float32   `json:"score"`
+		CreatedAt  int64     `json:"created_at"`
+		Score      float64   `json:"score"`
 	}
 
 	IndexStore interface {
-		CreateIndexes(ctx context.Context, idx []*Index) error
+		Init(ctx context.Context) error
+		Upsert(ctx context.Context, appID string, idx []*Index) error
 		Search(ctx context.Context, appId string, vectors []float32, n int) ([]*Index, error)
 		Reset(ctx context.Context, appId string) error
-		DeleteByPks(ctx context.Context, items []*Index) error
+		Delete(ctx context.Context, appID string, items []*Index) error
 	}
 
 	IndexService interface {
-		CreateIndexes(ctx context.Context, userID uint64, items []*Index) error
+		CreateIndexes(ctx context.Context, userID uint64, appId string, items []*Index) error
 		SearchIndex(ctx context.Context, userID uint64, data string, limit int) ([]*Index, error)
 		ResetIndexes(ctx context.Context, appID string) error
 	}
 )
 
-func (i Index) CollectionName() string {
-	return "indices"
-}
-
-func (i Index) PartitionName() string {
-	return strings.ReplaceAll(i.AppID, "-", "_")
-}
-
-// func (i Index) PartitionName() string {
-// 	return fmt.Sprintf("%d_%s", i.AppID, i.IndexName)
-// }
-
-func (i Index) Schema() *entity.Schema {
-	return &entity.Schema{
-		CollectionName: i.CollectionName(),
-		AutoID:         true,
-		Fields: []*entity.Field{
-			{
-				Name:       "id",
-				DataType:   entity.FieldTypeVarChar,
-				PrimaryKey: true,
-				TypeParams: map[string]string{"max_length": "64"},
-			},
-			{
-				Name:       "app_id",
-				DataType:   entity.FieldTypeVarChar,
-				TypeParams: map[string]string{"max_length": "32"},
-			},
-			{
-				Name:       "object_id",
-				DataType:   entity.FieldTypeVarChar,
-				TypeParams: map[string]string{"max_length": "32"},
-			},
-			{
-				Name:       "data",
-				DataType:   entity.FieldTypeVarChar,
-				TypeParams: map[string]string{"max_length": "2048"},
-			},
-			{
-				Name:     "data_token",
-				DataType: entity.FieldTypeInt64,
-			},
-			{
-				Name:       "vectors",
-				DataType:   entity.FieldTypeFloatVector,
-				TypeParams: map[string]string{"dim": "1536"},
-			},
-			{
-				Name:       "category",
-				DataType:   entity.FieldTypeVarChar,
-				TypeParams: map[string]string{"max_length": "32"},
-			},
-			{
-				Name:       "properties",
-				DataType:   entity.FieldTypeVarChar,
-				TypeParams: map[string]string{"max_length": "1024"},
-			},
-			{
-				Name:     "created_at",
-				DataType: entity.FieldTypeInt64,
-			},
-		},
-	}
+func (ix *Index) Mask() {
+	ix.AppID = ""
+	ix.DataToken = 0
+	ix.Vectors = nil
 }

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/pandodao/mixpay-go v0.0.0-20230405103808-9f656bd75b03
 	github.com/pkoukk/tiktoken-go v0.1.1-0.20230410080739-eb45b82397d3
 	github.com/pressly/goose/v3 v3.9.0
+	github.com/redis/go-redis/v9 v9.0.3
 	github.com/rs/cors v1.8.3
 	github.com/sashabaranov/go-openai v1.5.7
 	github.com/sirupsen/logrus v1.9.0
@@ -39,10 +40,12 @@ require (
 	github.com/btcsuite/btcd/btcec/v2 v2.3.2 // indirect
 	github.com/btcsuite/btcd/chaincfg/chainhash v1.0.2 // indirect
 	github.com/btcsuite/btcutil v1.0.3-0.20201208143702-a53e38424cce // indirect
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/cockroachdb/errors v1.9.1 // indirect
 	github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b // indirect
 	github.com/cockroachdb/redact v1.1.3 // indirect
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.1.0 // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/dlclark/regexp2 v1.9.0 // indirect
 	github.com/ethereum/go-ethereum v1.11.5 // indirect
 	github.com/fkgi/abnf v1.0.0 // indirect

go.sum

@@ -60,6 +60,8 @@ github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3d
 github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
 github.com/aymerick/raymond v2.0.3-0.20180322193309-b565731e1464+incompatible/go.mod h1:osfaiScAUVup+UC9Nfq76eWqDhXlp+4UYaA8uhTBO6g=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
+github.com/bsm/ginkgo/v2 v2.7.0 h1:ItPMPH90RbmZJt5GtkcNvIRuGEdwlBItdNVoyzaNQao=
+github.com/bsm/gomega v1.26.0 h1:LhQm+AFcgV2M0WyKroMASzAzCAJVpAxQXv4SaI9a69Y=
 github.com/btcsuite/btcd v0.20.1-beta/go.mod h1:wVuoA8VJLEcwgqHBwHmzLRazpKxTv13Px/pDuV7OomQ=
 github.com/btcsuite/btcd/btcec/v2 v2.3.2 h1:5n0X6hX0Zk+6omWcihdYvdAlGf2DfasC0GMf7DClJ3U=
 github.com/btcsuite/btcd/btcec/v2 v2.3.2/go.mod h1:zYzJ8etWJQIv1Ogk7OzpWjowwOdXY1W/17j2MW85J04=
@@ -77,6 +79,7 @@ github.com/btcsuite/websocket v0.0.0-20150119174127-31079b680792/go.mod h1:ghJtE
 github.com/btcsuite/winsvc v1.0.0/go.mod h1:jsenWakMcC0zFBFurPLEAyrnc/teJEM1O46fmI40EZs=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
+github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
 github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
 github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
@@ -114,6 +117,8 @@ github.com/decred/dcrd/dcrec/secp256k1/v4 v4.1.0 h1:HbphB4TFFXpv7MNrT52FGrrgVXF1
 github.com/decred/dcrd/dcrec/secp256k1/v4 v4.1.0/go.mod h1:DZGJHZMqrU4JJqFAWUS2UO1+lbSKsdiOoYi9Zzey7Fc=
 github.com/dgraph-io/badger v1.6.0/go.mod h1:zwt7syl517jmP8s94KqSxTlM6IMsdhYy6psNgSztDR4=
 github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/dlclark/regexp2 v1.9.0 h1:pTK/l/3qYIKaRXuHnEnIf7Y5NxfRPfpb7dis6/gdlVI=
 github.com/dlclark/regexp2 v1.9.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/drone/signal v1.0.0 h1:NrnM2M/4yAuU/tXs6RP1a1ZfxnaHwYkd0kJurA1p6uI=
@@ -488,6 +493,8 @@ github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:
 github.com/prometheus/client_model v0.3.0 h1:UBgGFHqYdG/TPFD1B1ogZywDqEkwp3fBMvqdiQ7Xew4=
 github.com/prometheus/common v0.39.0 h1:oOyhkDq05hPZKItWVBkJ6g6AtGxi+fy7F4JvUV8uhsI=
 github.com/prometheus/procfs v0.9.0 h1:wzCHvIvM5SxWqYvwgVL7yJY8Lz3PKn49KQtpgMYJfhI=
+github.com/redis/go-redis/v9 v9.0.3 h1:+7mmR26M0IvyLxGZUHxu4GiBkJkVDid0Un+j4ScYu4k=
+github.com/redis/go-redis/v9 v9.0.3/go.mod h1:WqMKv5vnQbRuZstUwxQI195wHy+t4PuXDOjzMvcuQHk=
 github.com/relvacode/iso8601 v1.3.0 h1:HguUjsGpIMh/zsTczGN3DVJFxTU/GX+MMmzcKoMO7ko=
 github.com/relvacode/iso8601 v1.3.0/go.mod h1:FlNp+jz+TXpyRqgmM7tnzHHzBnz776kmAH2h3sZCn0I=
 github.com/remyoudompheng/bigfft v0.0.0-20220927061507-ef77025ab5aa h1:tEkEyxYeZ43TR55QU/hsIt9aRGBxbgGuz9CGykjvogY=

handler/index/index.go

@@ -63,7 +63,7 @@ func CreateIndex(indexes core.IndexService) http.HandlerFunc {
 			return
 		}
 
-		err := indexes.CreateIndexes(r.Context(), app.UserID, is)
+		err := indexes.CreateIndexes(r.Context(), app.UserID, app.AppID, is)
 		if err != nil {
 			render.Error(w, http.StatusInternalServerError, err)
 			return
@@ -84,7 +84,7 @@ func Delete(apps core.AppStore, indexes core.IndexStore) http.HandlerFunc {
 			return
 		}
 
-		if err := indexes.DeleteByPks(ctx, []*core.Index{
+		if err := indexes.Delete(ctx, app.AppID, []*core.Index{
 			{
 				AppID:    app.AppID,
 				ObjectID: objectID,
@@ -125,6 +125,9 @@ func Search(apps core.AppStore, indexes core.IndexService) http.HandlerFunc {
 			render.Error(w, http.StatusInternalServerError, err)
 			return
 		}
+		for _, r := range result {
+			r.Mask()
+		}
 
 		render.JSON(w, result)
 	}