Merge pull request #44 from bricks-cloud/v1.9.0

[V1.9.0]
bricks-cloud · Feb 6, 2024 · f1f6384 · f1f6384
2 parents a84aae2 + 7001d14
commit f1f6384
Show file tree

Hide file tree

Showing 21 changed files with 912 additions and 554 deletions.
diff --git a/cmd/bricksllm/main.go b/cmd/bricksllm/main.go
@@ -14,6 +14,7 @@ import (
 	"github.com/bricks-cloud/bricksllm/internal/config"
 	"github.com/bricks-cloud/bricksllm/internal/logger/zap"
 	"github.com/bricks-cloud/bricksllm/internal/manager"
+	"github.com/bricks-cloud/bricksllm/internal/message"
 	"github.com/bricks-cloud/bricksllm/internal/provider/anthropic"
 	"github.com/bricks-cloud/bricksllm/internal/provider/azure"
 	"github.com/bricks-cloud/bricksllm/internal/provider/custom"
@@ -171,10 +172,23 @@ func main() {
 		log.Sugar().Fatalf("error connecting to api redis cache: %v", err)
 	}
 
+	accessRedisCache := redis.NewClient(&redis.Options{
+		Addr:     fmt.Sprintf("%s:%s", cfg.RedisHosts, cfg.RedisPort),
+		Password: cfg.RedisPassword,
+		DB:       4,
+	})
+
+	ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	if err := apiRedisCache.Ping(ctx).Err(); err != nil {
+		log.Sugar().Fatalf("error connecting to api redis cache: %v", err)
+	}
+
 	rateLimitCache := redisStorage.NewCache(rateLimitRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout)
 	costLimitCache := redisStorage.NewCache(costLimitRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout)
 	costStorage := redisStorage.NewStore(costRedisStorage, cfg.RedisWriteTimeout, cfg.RedisReadTimeout)
 	apiCache := redisStorage.NewCache(apiRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout)
+	accessCache := redisStorage.NewAccessCache(accessRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout)
 
 	m := manager.NewManager(store)
 	krm := manager.NewReportingManager(costStorage, store, store)
@@ -209,7 +223,16 @@ func main() {
 
 	c := cache.NewCache(apiCache)
 
-	ps, err := proxy.NewProxyServer(log, *modePtr, *privacyPtr, c, m, rm, a, psm, cpm, store, memStore, ce, ace, aoe, v, rec, rlm, cfg.ProxyTimeout)
+	messageBus := message.NewMessageBus()
+	eventMessageChan := make(chan message.Message)
+	messageBus.Subscribe("event", eventMessageChan)
+
+	handler := message.NewHandler(rec, log, ace, ce, aoe, v, m, rlm, accessCache)
+
+	eventConsumer := message.NewConsumer(eventMessageChan, log, 4, handler.HandleEventWithRequestAndResponse)
+	eventConsumer.StartEventMessageConsumers()
+
+	ps, err := proxy.NewProxyServer(log, *modePtr, *privacyPtr, c, m, rm, a, psm, cpm, store, memStore, ce, ace, aoe, v, rec, messageBus, rlm, cfg.ProxyTimeout, accessCache)
 	if err != nil {
 		log.Sugar().Fatalf("error creating proxy http server: %v", err)
 	}
@@ -220,6 +243,7 @@ func main() {
 	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
 	<-quit
 
+	eventConsumer.Stop()
 	memStore.Stop()
 	psMemStore.Stop()
 	cpMemStore.Stop()

diff --git a/cmd/tool/main.go b/cmd/tool/main.go
diff --git a/go.mod b/go.mod
@@ -12,7 +12,7 @@ require (
 	github.com/mattn/go-colorable v0.1.13
 	github.com/pkoukk/tiktoken-go-loader v0.0.1
 	github.com/redis/go-redis/v9 v9.0.5
-	github.com/sashabaranov/go-openai v1.17.7
+	github.com/sashabaranov/go-openai v1.19.2
 	github.com/stretchr/testify v1.8.4
 	go.uber.org/zap v1.24.0
 )

diff --git a/go.sum b/go.sum
@@ -85,6 +85,8 @@ github.com/sashabaranov/go-openai v1.17.1 h1:tapFKbKE8ep0/qGkKp5Q3TtxWUD7m9VIFe9
 github.com/sashabaranov/go-openai v1.17.1/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/sashabaranov/go-openai v1.17.7 h1:MPcAwlwbeo7ZmhQczoOgZBHtIBY1TfZqsdx6+/ndloM=
 github.com/sashabaranov/go-openai v1.17.7/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
+github.com/sashabaranov/go-openai v1.19.2 h1:+dkuCADSnwXV02YVJkdphY8XD9AyHLUWwk6V7LB6EL8=
+github.com/sashabaranov/go-openai v1.19.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=

diff --git a/internal/config/config.go b/internal/config/config.go
@@ -7,25 +7,26 @@ import (
 )
 
 type Config struct {
-	PostgresqlHosts          string        `env:"POSTGRESQL_HOSTS" envSeparator:":" envDefault:"localhost"`
-	PostgresqlDbName         string        `env:"POSTGRESQL_DB_NAME"`
-	PostgresqlUsername       string        `env:"POSTGRESQL_USERNAME"`
-	PostgresqlPassword       string        `env:"POSTGRESQL_PASSWORD"`
-	PostgresqlSslMode        string        `env:"POSTGRESQL_SSL_MODE" envDefault:"disable"`
-	PostgresqlPort           string        `env:"POSTGRESQL_PORT" envDefault:"5432"`
-	RedisHosts               string        `env:"REDIS_HOSTS" envSeparator:":" envDefault:"localhost"`
-	RedisPort                string        `env:"REDIS_PORT" envDefault:"6379"`
-	RedisUsername            string        `env:"REDIS_USERNAME"`
-	RedisPassword            string        `env:"REDIS_PASSWORD"`
-	RedisReadTimeout         time.Duration `env:"REDIS_READ_TIME_OUT" envDefault:"1s"`
-	RedisWriteTimeout        time.Duration `env:"REDIS_WRITE_TIME_OUT" envDefault:"500ms"`
-	PostgresqlReadTimeout    time.Duration `env:"POSTGRESQL_READ_TIME_OUT" envDefault:"2s"`
-	PostgresqlWriteTimeout   time.Duration `env:"POSTGRESQL_WRITE_TIME_OUT" envDefault:"1s"`
-	InMemoryDbUpdateInterval time.Duration `env:"IN_MEMORY_DB_UPDATE_INTERVAL" envDefault:"5s"`
-	OpenAiKey                string        `env:"OPENAI_API_KEY"`
-	StatsProvider            string        `env:"STATS_PROVIDER"`
-	AdminPass                string        `env:"ADMIN_PASS"`
-	ProxyTimeout             time.Duration `env:"PROXY_TIMEOUT" envDefault:"600s"`
+	PostgresqlHosts               string        `env:"POSTGRESQL_HOSTS" envSeparator:":" envDefault:"localhost"`
+	PostgresqlDbName              string        `env:"POSTGRESQL_DB_NAME"`
+	PostgresqlUsername            string        `env:"POSTGRESQL_USERNAME"`
+	PostgresqlPassword            string        `env:"POSTGRESQL_PASSWORD"`
+	PostgresqlSslMode             string        `env:"POSTGRESQL_SSL_MODE" envDefault:"disable"`
+	PostgresqlPort                string        `env:"POSTGRESQL_PORT" envDefault:"5432"`
+	RedisHosts                    string        `env:"REDIS_HOSTS" envSeparator:":" envDefault:"localhost"`
+	RedisPort                     string        `env:"REDIS_PORT" envDefault:"6379"`
+	RedisUsername                 string        `env:"REDIS_USERNAME"`
+	RedisPassword                 string        `env:"REDIS_PASSWORD"`
+	RedisReadTimeout              time.Duration `env:"REDIS_READ_TIME_OUT" envDefault:"1s"`
+	RedisWriteTimeout             time.Duration `env:"REDIS_WRITE_TIME_OUT" envDefault:"500ms"`
+	PostgresqlReadTimeout         time.Duration `env:"POSTGRESQL_READ_TIME_OUT" envDefault:"2s"`
+	PostgresqlWriteTimeout        time.Duration `env:"POSTGRESQL_WRITE_TIME_OUT" envDefault:"1s"`
+	InMemoryDbUpdateInterval      time.Duration `env:"IN_MEMORY_DB_UPDATE_INTERVAL" envDefault:"5s"`
+	OpenAiKey                     string        `env:"OPENAI_API_KEY"`
+	StatsProvider                 string        `env:"STATS_PROVIDER"`
+	AdminPass                     string        `env:"ADMIN_PASS"`
+	ProxyTimeout                  time.Duration `env:"PROXY_TIMEOUT" envDefault:"600s"`
+	NumberOfEventMessageConsumers int           `env:"NUMBER_OF_EVENT_MESSAGE_CONSUMERS" envDefault:"3"`
 }
 
 func ParseEnvVariables() (*Config, error) {