performance_issues

cmd/bricksllm/main.go

@@ -37,6 +37,8 @@ import (
 	"github.com/redis/go-redis/v9"
 )
 
+const numEventConsumers = 16
+
 func main() {
 	modePtr := flag.String("m", "dev", "select the mode that bricksllm runs in")
 	privacyPtr := flag.String("p", "strict", "select the privacy mode that bricksllm runs in")
@@ -361,12 +363,12 @@ func main() {
 	c := cache.NewCache(apiCache)
 
 	messageBus := message.NewMessageBus()
-	eventMessageChan := make(chan message.Message)
+	eventMessageChan := make(chan message.Message, 1000)
 	messageBus.Subscribe("event", eventMessageChan)
 
 	handler := message.NewHandler(rec, log, ace, ce, vllme, aoe, v, uv, m, um, rlm, accessCache, userAccessCache)
 
-	eventConsumer := message.NewConsumer(eventMessageChan, log, 4, handler.HandleEventWithRequestAndResponse)
+	eventConsumer := message.NewConsumer(eventMessageChan, log, numEventConsumers, handler.HandleEventWithRequestAndResponse)
 	eventConsumer.StartEventMessageConsumers()
 
 	detector, err := amazon.NewClient(cfg.AmazonRequestTimeout, cfg.AmazonConnectionTimeout, log, cfg.AmazonRegion)

internal/message/consumer.go

@@ -1,6 +1,8 @@
 package message
 
 import (
+	"sync"
+
 	"github.com/bricks-cloud/bricksllm/internal/event"
 	"github.com/bricks-cloud/bricksllm/internal/key"
 	"go.uber.org/zap"
@@ -12,6 +14,7 @@ type Consumer struct {
 	log                 *zap.Logger
 	numOfEventConsumers int
 	handle              func(Message) error
+	wg                  sync.WaitGroup
 }
 
 type recorder interface {
@@ -33,12 +36,25 @@ func NewConsumer(mc <-chan Message, log *zap.Logger, num int, handle func(Messag
 
 func (c *Consumer) StartEventMessageConsumers() {
 	for i := 0; i < c.numOfEventConsumers; i++ {
+		c.wg.Add(1)
 		go func() {
+			defer c.wg.Done()
+
 			for {
 				select {
 				case <-c.done:
-					c.log.Info("event message consumer stoped...")
-					return
+					for {
+						select {
+						case m := <-c.messageChan:
+							err := c.handle(m)
+							if err != nil {
+								continue
+							}
+						default:
+							c.log.Info("event message consumer stoped...")
+							return
+						}
+					}
 
 				case m := <-c.messageChan:
 					err := c.handle(m)
@@ -56,5 +72,6 @@ func (c *Consumer) StartEventMessageConsumers() {
 func (c *Consumer) Stop() {
 	c.log.Info("shutting down consumer...")
 
-	c.done <- true
+	close(c.done)
+	c.wg.Wait()
 }

internal/provider/openai/cost.go

@@ -359,6 +359,10 @@ func init() {
 	for model := range OpenAiPerThousandTokenCost["images-tokens-input"] {
 		imageModelsWithTokensCost[model] = struct{}{}
 	}
+
+	for _, tool := range AllowedTools {
+		AllowedToolsSet[tool] = struct{}{}
+	}
 }
 
 var OpenAiPerThousandCallsToolCost = map[string]float64{
@@ -401,6 +405,8 @@ var AllowedTools = []string{
 	"skills",
 }
 
+var AllowedToolsSet = map[string]struct{}{}
+
 type tokenCounter interface {
 	Count(model string, input string) (int, error)
 }

internal/server/web/proxy/anthropic.go

@@ -355,7 +355,7 @@ func getMessagesHandler(prod, private bool, client http.Client, e anthropicEstim
 		model := c.GetString("model")
 
 		if !isStreaming && res.StatusCode == http.StatusOK {
-			dur := time.Now().Sub(start)
+			dur := time.Since(start)
 			telemetry.Timing("bricksllm.proxy.get_messages_handler.latency", dur, nil, 1)
 
 			bytes, err := io.ReadAll(res.Body)