From 4dca7906482e38691feb8aca47bd65c06da3d7e5 Mon Sep 17 00:00:00 2001 From: Sergei Bronnikov Date: Thu, 11 Jun 2026 12:33:44 +0100 Subject: [PATCH 1/4] refactor --- cmd/bricksllm/main.go | 6 ++++-- internal/server/web/proxy/middleware.go | 8 +++++--- internal/server/web/proxy/proxy.go | 9 ++++++++- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/cmd/bricksllm/main.go b/cmd/bricksllm/main.go index 035e08d..40913bc 100644 --- a/cmd/bricksllm/main.go +++ b/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) diff --git a/internal/server/web/proxy/middleware.go b/internal/server/web/proxy/middleware.go index 4c304ef..d031211 100644 --- a/internal/server/web/proxy/middleware.go +++ b/internal/server/web/proxy/middleware.go @@ -197,9 +197,6 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag c.Set("removeUserAgent", removeUserAgent) } - blw := &responseWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer} - c.Writer = blw - cid := util.NewUuid() c.Set(util.STRING_CORRELATION_ID, cid) logWithCid := log.With(zap.String(util.STRING_CORRELATION_ID, cid)) @@ -321,6 +318,11 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } kc, settings, err := a.AuthenticateHttpRequest(c.Request, c.Param(xcustom.XProviderIdParam)) + if kc.ShouldLogResponse { + blw := &responseWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer} + c.Writer = blw + } + enrichedEvent.Key = kc _, ok := err.(notAuthorizedError) if ok { diff --git a/internal/server/web/proxy/proxy.go b/internal/server/web/proxy/proxy.go index 4740867..0a62ba8 100644 --- a/internal/server/web/proxy/proxy.go +++ b/internal/server/web/proxy/proxy.go @@ -91,7 +91,14 @@ func NewProxyServer(log *zap.Logger, mode, privacyMode string, c cache, m KeyMan router.Use(getTimeoutMiddleware(timeout)) router.Use(getMiddleware(cpm, rm, pm, a, prod, private, log, pub, "proxy", ac, uac, http.Client{}, scanner, cd, um, removeAgentHeaders)) - client := http.Client{} + client := http.Client{ + Transport: &http.Transport{ + MaxIdleConns: 200, + MaxIdleConnsPerHost: 100, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + }, + } // health check router.POST("/api/health", getGetHealthCheckHandler()) From 7957c3d88029d3b069206ccb4ab7de45182f7a02 Mon Sep 17 00:00:00 2001 From: Sergei Bronnikov Date: Thu, 11 Jun 2026 14:04:09 +0100 Subject: [PATCH 2/4] refactor --- internal/server/web/proxy/middleware.go | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/internal/server/web/proxy/middleware.go b/internal/server/web/proxy/middleware.go index d031211..701657e 100644 --- a/internal/server/web/proxy/middleware.go +++ b/internal/server/web/proxy/middleware.go @@ -197,6 +197,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag c.Set("removeUserAgent", removeUserAgent) } + var blw *responseWriter + cid := util.NewUuid() c.Set(util.STRING_CORRELATION_ID, cid) logWithCid := log.With(zap.String(util.STRING_CORRELATION_ID, cid)) @@ -318,10 +320,6 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } kc, settings, err := a.AuthenticateHttpRequest(c.Request, c.Param(xcustom.XProviderIdParam)) - if kc.ShouldLogResponse { - blw := &responseWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer} - c.Writer = blw - } enrichedEvent.Key = kc _, ok := err.(notAuthorizedError) @@ -353,6 +351,11 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag c.Set("key", kc) c.Set("settings", settings) + if kc.ShouldLogResponse { + blw = &responseWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer} + c.Writer = blw + } + if len(settings) >= 1 { selected := settings[0] @@ -1384,7 +1387,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if !c.GetBool("stream") { + if !c.GetBool("stream") && blw != nil { responseData := blw.body.Bytes() if len(responseData) != 0 { From 45ff3dad95b4c168b0da72c730ac8abb949e1e24 Mon Sep 17 00:00:00 2001 From: Sergei Bronnikov Date: Wed, 17 Jun 2026 09:52:52 +0100 Subject: [PATCH 3/4] refactor --- internal/server/web/proxy/anthropic.go | 2 +- internal/server/web/proxy/middleware.go | 185 ++++++++++++------------ internal/server/web/proxy/video.go | 1 - 3 files changed, 94 insertions(+), 94 deletions(-) diff --git a/internal/server/web/proxy/anthropic.go b/internal/server/web/proxy/anthropic.go index 5853ed3..79b479c 100644 --- a/internal/server/web/proxy/anthropic.go +++ b/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) diff --git a/internal/server/web/proxy/middleware.go b/internal/server/web/proxy/middleware.go index 701657e..d5b576b 100644 --- a/internal/server/web/proxy/middleware.go +++ b/internal/server/web/proxy/middleware.go @@ -133,13 +133,13 @@ type publisher interface { Publish(message.Message) } -func getProvider(c *gin.Context) string { +func getProvider(c *gin.Context, fullPath string) string { existing := c.GetString("provider") if len(existing) != 0 { return existing } - parts := strings.Split(c.FullPath(), "/") + parts := strings.Split(fullPath, "/") spaceRemoved := []string{} @@ -149,17 +149,17 @@ func getProvider(c *gin.Context) string { } } - if strings.HasPrefix(c.FullPath(), "/api/providers/") { + if strings.HasPrefix(fullPath, "/api/providers/") { if len(spaceRemoved) >= 3 { return spaceRemoved[2] } } - if strings.HasPrefix(c.FullPath(), "/api/custom/providers/") { + if strings.HasPrefix(fullPath, "/api/custom/providers/") { return c.Param("provider") } - if strings.HasPrefix(c.FullPath(), "/api/routes/") { + if strings.HasPrefix(fullPath, "/api/routes/") { return c.Param("provider") } @@ -188,7 +188,10 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag return } - if c.FullPath() == "/api/health" { + fullPath := c.FullPath() + method := c.Request.Method + + if fullPath == "/api/health" { c.Abort() return } @@ -224,7 +227,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag latency := int(dur.Milliseconds()) if !prod { - logWithCid.Sugar().Infof("%s | %d | %s | %s | %dms", prefix, c.Writer.Status(), c.Request.Method, c.FullPath(), latency) + logWithCid.Sugar().Infof("%s | %d | %s | %s | %dms", prefix, c.Writer.Status(), method, fullPath, latency) } keyId := "" @@ -248,15 +251,15 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag telemetry.Timing("bricksllm.proxy.get_middleware.proxy_latency_in_ms", dur, nil, 1) - selectedProvider := getProvider(c) + selectedProvider := getProvider(c, fullPath) if prod { logWithCid.Info("response to proxy", zap.String("provider", selectedProvider), zap.String("keyId", keyId), zap.Int("code", c.Writer.Status()), - zap.String("method", c.Request.Method), - zap.String("path", c.FullPath()), + zap.String("method", method), + zap.String("path", fullPath), zap.Int("lantecyInMs", latency), ) } @@ -278,7 +281,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag CompletionTokenCount: c.GetInt("completionTokenCount"), LatencyInMs: latency, Path: c.Request.URL.Path, - Method: c.Request.Method, + Method: method, CustomId: customId, Request: requestBytes, Response: responseBytes, @@ -312,7 +315,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag }) }() - if len(c.FullPath()) == 0 { + if len(fullPath) == 0 { telemetry.Incr("bricksllm.proxy.get_middleware.route_does_not_exist", nil, 1) JSON(c, http.StatusNotFound, "[BricksLLM] route not supported") c.Abort() @@ -364,13 +367,13 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag c.Set("cost_map", selected.CostMap) } - if strings.HasPrefix(c.FullPath(), "/api/providers/azure/openai") { + if strings.HasPrefix(fullPath, "/api/providers/azure/openai") { if selected != nil && len(selected.Setting["resourceName"]) != 0 { c.Set("resourceName", selected.Setting["resourceName"]) } } - if strings.HasPrefix(c.FullPath(), "/api/providers/bedrock/anthropic") { + if strings.HasPrefix(fullPath, "/api/providers/bedrock/anthropic") { if selected != nil && len(selected.Setting["awsAccessKeyId"]) != 0 { c.Set("awsAccessKeyId", selected.Setting["awsAccessKeyId"]) } @@ -384,7 +387,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if strings.HasPrefix(c.FullPath(), "/api/providers/vllm") { + if strings.HasPrefix(fullPath, "/api/providers/vllm") { if selected != nil && len(selected.Setting["url"]) != 0 { c.Set("vllmUrl", selected.Setting["url"]) } @@ -409,11 +412,11 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag c.Set("requestBytes", requestBytes) } - if c.Request.Method != http.MethodGet { + if method != http.MethodGet { c.Request.Body = io.NopCloser(bytes.NewReader(body)) } - if c.FullPath() == "/api/providers/anthropic/v1/complete" { + if fullPath == "/api/providers/anthropic/v1/complete" { logCompletionRequest(logWithCid, body, prod, private) cr := &anthropic.CompletionRequest{} @@ -438,7 +441,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = cr } - if c.FullPath() == "/api/providers/bedrock/anthropic/v1/complete" { + if fullPath == "/api/providers/bedrock/anthropic/v1/complete" { logCompletionRequest(logWithCid, body, prod, private) cr := &anthropic.CompletionRequest{} @@ -463,7 +466,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = cr } - if c.FullPath() == "/api/providers/bedrock/anthropic/v1/messages" { + if fullPath == "/api/providers/bedrock/anthropic/v1/messages" { logCreateMessageRequest(logWithCid, body, prod, private) mr := &anthropic.MessagesRequest{} @@ -486,7 +489,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = mr } - if c.FullPath() == "/api/providers/anthropic/v1/messages" { + if fullPath == "/api/providers/anthropic/v1/messages" { logCreateMessageRequest(logWithCid, body, prod, private) mr := &anthropic.MessagesRequest{} @@ -509,7 +512,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = mr } - if strings.HasPrefix(c.FullPath(), "/api/custom/providers/:provider") { + if strings.HasPrefix(fullPath, "/api/custom/providers/:provider") { providerName := c.Param("provider") rc := cpm.GetRouteConfigFromMem(providerName, c.Param("wildcard")) @@ -547,19 +550,20 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag // c.Set("promptTokenCount", tks) - result := gjson.Get(string(body), rc.StreamLocation) + bodyStr := string(body) + result := gjson.Get(bodyStr, rc.StreamLocation) if result.IsBool() { c.Set("stream", result.Bool()) } - result = gjson.Get(string(body), rc.ModelLocation) + result = gjson.Get(bodyStr, rc.ModelLocation) if len(result.Str) != 0 { c.Set("model", result.Str) } } - if strings.HasPrefix(c.FullPath(), "/api/routes") { + if strings.HasPrefix(fullPath, "/api/routes") { r := c.Param("route") rc := rm.GetRouteFromMemDb(r) @@ -624,7 +628,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if c.FullPath() == "/api/providers/vllm/v1/chat/completions" { + if fullPath == "/api/providers/vllm/v1/chat/completions" { ccr := &vllm.ChatRequest{} err = json.Unmarshal(body, ccr) if err != nil { @@ -645,7 +649,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = ccr } - if c.FullPath() == "/api/providers/vllm/v1/completions" { + if fullPath == "/api/providers/vllm/v1/completions" { cr := &vllm.CompletionRequest{} err = json.Unmarshal(body, cr) if err != nil { @@ -666,7 +670,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = cr } - if c.FullPath() == "/api/providers/deepinfra/v1/chat/completions" { + if fullPath == "/api/providers/deepinfra/v1/chat/completions" { ccr := &vllm.ChatRequest{} err = json.Unmarshal(body, ccr) if err != nil { @@ -686,7 +690,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = ccr } - if c.FullPath() == "/api/providers/deepinfra/v1/completions" { + if fullPath == "/api/providers/deepinfra/v1/completions" { cr := &vllm.CompletionRequest{} err = json.Unmarshal(body, cr) if err != nil { @@ -706,7 +710,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = cr } - if c.FullPath() == "/api/providers/deepinfra/v1/embeddings" { + if fullPath == "/api/providers/deepinfra/v1/embeddings" { er := &goopenai.EmbeddingRequest{} err = json.Unmarshal(body, er) if err != nil { @@ -723,7 +727,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = er } - if c.FullPath() == "/api/providers/azure/openai/deployments/:deployment_id/chat/completions" { + if fullPath == "/api/providers/azure/openai/deployments/:deployment_id/chat/completions" { ccr := &goopenai.ChatCompletionRequest{} err = json.Unmarshal(body, ccr) if err != nil { @@ -744,7 +748,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = ccr } - if c.FullPath() == "/api/providers/azure/openai/deployments/:deployment_id/completions" { + if fullPath == "/api/providers/azure/openai/deployments/:deployment_id/completions" { cr := &goopenai.CompletionRequest{} err = json.Unmarshal(body, cr) if err != nil { @@ -765,7 +769,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = cr } - if c.FullPath() == "/api/providers/azure/openai/deployments/:deployment_id/embeddings" { + if fullPath == "/api/providers/azure/openai/deployments/:deployment_id/embeddings" { er := &goopenai.EmbeddingRequest{} err = json.Unmarshal(body, er) if err != nil { @@ -783,7 +787,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = er } - if c.FullPath() == "/api/providers/openai/v1/chat/completions" { + if fullPath == "/api/providers/openai/v1/chat/completions" { ccr := &goopenai.ChatCompletionRequest{} // this is a hack around an open issue in go-openai. // https://github.com/sashabaranov/go-openai/issues/884 @@ -812,7 +816,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = ccr } - if strings.HasPrefix(c.FullPath(), "/api/providers/openai/v1/responses") { + if strings.HasPrefix(fullPath, "/api/providers/openai/v1/responses") { responsesReq := &openai.ResponseRequest{} err = json.Unmarshal(body, responsesReq) if err != nil { @@ -830,10 +834,16 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } hasNotAllowedTools := false + isCreateContainerTool := false + var containerMemLimit string for _, tool := range responsesReq.Tools { if !slices.Contains(openai.AllowedTools, tool.Type) { hasNotAllowedTools = true - break + } + + if tool.GetContainerAsResponseRequestToolContainer() != nil { + isCreateContainerTool = true + containerMemLimit = tool.GetContainerAsResponseRequestToolContainer().GetMemoryLimit() } } @@ -844,15 +854,6 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag return } - isCreateContainerTool := false - var containerMemLimit string - for _, tool := range responsesReq.Tools { - if tool.GetContainerAsResponseRequestToolContainer() != nil { - isCreateContainerTool = true - containerMemLimit = tool.GetContainerAsResponseRequestToolContainer().GetMemoryLimit() - break - } - } if isCreateContainerTool { _, ok := openai.OpenAiCodeInterpreterContainerCost[containerMemLimit] if !ok { @@ -876,7 +877,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = responsesReq } - if c.FullPath() == "/api/providers/openai/v1/embeddings" { + if fullPath == "/api/providers/openai/v1/embeddings" { er := &goopenai.EmbeddingRequest{} err = json.Unmarshal(body, er) if err != nil { @@ -894,7 +895,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag policyInput = er } - if c.FullPath() == "/api/providers/openai/v1/images/generations" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/images/generations" && method == http.MethodPost { ir := &goopenai.ImageRequest{} err := json.Unmarshal(body, ir) if err != nil { @@ -913,7 +914,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag logCreateImageRequest(logWithCid, ir, prod, private) } - if c.FullPath() == "/api/providers/openai/v1/images/edits" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/images/edits" && method == http.MethodPost { ier := &goopenai.ImageEditRequest{} err := json.Unmarshal(body, ier) if err != nil { @@ -941,7 +942,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag logEditImageRequest(logWithCid, prompt, model, n, size, responseFormat, user, prod, private) } - if c.FullPath() == "/api/providers/openai/v1/images/variations" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/images/variations" && method == http.MethodPost { ir := &goopenai.ImageVariRequest{} err := json.Unmarshal(body, ir) if err != nil { @@ -968,7 +969,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag logImageVariationsRequest(logWithCid, model, n, size, responseFormat, user, prod) } - if c.FullPath() == "/api/providers/openai/v1/audio/speech" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/audio/speech" && method == http.MethodPost { sr := &goopenai.CreateSpeechRequest{} err := json.Unmarshal(body, sr) if err != nil { @@ -983,7 +984,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag logCreateSpeechRequest(logWithCid, sr, prod, private) } - if c.FullPath() == "/api/providers/openai/v1/audio/transcriptions" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/audio/transcriptions" && method == http.MethodPost { model := c.PostForm("model") language := c.PostForm("language") prompt := c.PostForm("prompt") @@ -996,7 +997,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag logCreateTranscriptionRequest(logWithCid, model, language, prompt, responseFormat, converted, prod, private) } - if c.FullPath() == "/api/providers/openai/v1/audio/translations" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/audio/translations" && method == http.MethodPost { model := c.PostForm("model") prompt := c.PostForm("prompt") responseFormat := c.PostForm("response_format") @@ -1008,7 +1009,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag logCreateTranslationRequest(logWithCid, model, prompt, responseFormat, converted, prod, private) } - if strings.HasPrefix(c.FullPath(), "/api/providers/openai/v1/videos") && c.Request.Method == http.MethodPost { + if strings.HasPrefix(fullPath, "/api/providers/openai/v1/videos") && method == http.MethodPost { model := c.PostForm("model") if model == "" { vr := &openai.VideoRequest{} @@ -1022,7 +1023,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag c.Set("model", model) } - if len(kc.AllowedPaths) != 0 && !containsPath(kc.AllowedPaths, c.FullPath(), c.Request.Method) { + if len(kc.AllowedPaths) != 0 && !containsPath(kc.AllowedPaths, fullPath, method) { telemetry.Incr("bricksllm.proxy.get_middleware.path_not_allowed", nil, 1) JSON(c, http.StatusForbidden, "[BricksLLM] path is not allowed") c.Abort() @@ -1036,7 +1037,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag c.Abort() return } - if !isModelSupported(c.FullPath(), model) { + if !isModelSupported(fullPath, model) { telemetry.Incr("bricksllm.proxy.get_middleware.model_not_supported", nil, 1) JSON(c, http.StatusBadRequest, "[BricksLLM] model is not supported") c.Abort() @@ -1068,7 +1069,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag qm["before"] = val } - if c.FullPath() == "/api/providers/openai/v1/assistants" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/assistants" && method == http.MethodPost { logCreateAssistantRequest(logWithCid, body, prod, private) ar := &goopenai.AssistantRequest{} @@ -1085,39 +1086,39 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if c.FullPath() == "/api/providers/openai/v1/assistants/:assistant_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/assistants/:assistant_id" && method == http.MethodGet { logRetrieveAssistantRequest(logWithCid, prod, aid) } - if c.FullPath() == "/api/providers/openai/v1/assistants/:assistant_id" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/assistants/:assistant_id" && method == http.MethodPost { logModifyAssistantRequest(logWithCid, body, prod, private, aid) } - if c.FullPath() == "/api/providers/openai/v1/assistants/:assistant_id" && c.Request.Method == http.MethodDelete { + if fullPath == "/api/providers/openai/v1/assistants/:assistant_id" && method == http.MethodDelete { logDeleteAssistantRequest(logWithCid, prod, aid) } - if c.FullPath() == "/api/providers/openai/v1/assistants" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/assistants" && method == http.MethodGet { logListAssistantsRequest(logWithCid, prod) } - if c.FullPath() == "/api/providers/openai/v1/assistants/:assistant_id/files" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/assistants/:assistant_id/files" && method == http.MethodPost { logCreateAssistantFileRequest(logWithCid, body, prod, aid) } - if c.FullPath() == "/api/providers/openai/v1/assistants/:assistant_id/files/:file_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/assistants/:assistant_id/files/:file_id" && method == http.MethodGet { logRetrieveAssistantFileRequest(logWithCid, prod, fid, aid) } - if c.FullPath() == "/api/providers/openai/v1/assistants/:assistant_id/files/:file_id" && c.Request.Method == http.MethodDelete { + if fullPath == "/api/providers/openai/v1/assistants/:assistant_id/files/:file_id" && method == http.MethodDelete { logDeleteAssistantFileRequest(logWithCid, prod, fid, aid) } - if c.FullPath() == "/api/providers/openai/v1/assistants/:assistant_id/files" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/assistants/:assistant_id/files" && method == http.MethodGet { logListAssistantFilesRequest(logWithCid, prod, aid, qm) } - if c.FullPath() == "/api/providers/openai/v1/threads" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads" && method == http.MethodPost { logCreateThreadRequest(logWithCid, body, prod, private) tr := &openai.ThreadRequest{} @@ -1132,19 +1133,19 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id" && method == http.MethodGet { logRetrieveThreadRequest(logWithCid, prod, tid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/:thread_id" && method == http.MethodPost { logModifyThreadRequest(logWithCid, body, prod, tid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id" && c.Request.Method == http.MethodDelete { + if fullPath == "/api/providers/openai/v1/threads/:thread_id" && method == http.MethodDelete { logDeleteThreadRequest(logWithCid, prod, tid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/messages" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/messages" && method == http.MethodPost { logCreateMessageRequest(logWithCid, body, prod, private) mr := &openai.MessageRequest{} @@ -1158,27 +1159,27 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id" && method == http.MethodGet { logRetrieveMessageRequest(logWithCid, prod, mid, tid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id" && method == http.MethodPost { logModifyMessageRequest(logWithCid, body, prod, private, tid, mid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/messages" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/messages" && method == http.MethodGet { logListMessagesRequest(logWithCid, prod, aid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id/files/:file_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id/files/:file_id" && method == http.MethodGet { logRetrieveMessageFileRequest(logWithCid, prod, mid, tid, fid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id/files" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/messages/:message_id/files" && method == http.MethodGet { logListMessageFilesRequest(logWithCid, prod, tid, mid, qm) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs" && method == http.MethodPost { logCreateRunRequest(logWithCid, body, prod, private) rr := &goopenai.RunRequest{} @@ -1193,27 +1194,27 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id" && method == http.MethodGet { logRetrieveRunRequest(logWithCid, prod, tid, rid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id" && method == http.MethodPost { logModifyRunRequest(logWithCid, body, prod, tid, rid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs" && method == http.MethodGet { logListRunsRequest(logWithCid, prod, tid, qm) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/submit_tool_outputs" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/submit_tool_outputs" && method == http.MethodPost { logSubmitToolOutputsRequest(logWithCid, body, prod, tid, rid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/cancel" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/cancel" && method == http.MethodPost { logCancelARunRequest(logWithCid, prod, tid, rid) } - if c.FullPath() == "/api/providers/openai/v1/threads/runs" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/threads/runs" && method == http.MethodPost { logCreateThreadAndRunRequest(logWithCid, body, prod, private) r := &openai.CreateThreadAndRunRequest{} @@ -1228,44 +1229,44 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/steps/:step_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/steps/:step_id" && method == http.MethodGet { logRetrieveRunStepRequest(logWithCid, prod, tid, rid, sid) } - if c.FullPath() == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/steps" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/threads/:thread_id/runs/:run_id/steps" && method == http.MethodGet { logListRunStepsRequest(logWithCid, prod, tid, rid, qm) } - if c.FullPath() == "/api/providers/openai/v1/moderations" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/moderations" && method == http.MethodPost { logCreateModerationRequest(logWithCid, body, prod, private) } - if c.FullPath() == "/api/providers/openai/v1/models/:model" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/models/:model" && method == http.MethodGet { logRetrieveModelRequest(logWithCid, prod, md) } - if c.FullPath() == "/api/providers/openai/v1/models/:model" && c.Request.Method == http.MethodDelete { + if fullPath == "/api/providers/openai/v1/models/:model" && method == http.MethodDelete { logDeleteModelRequest(logWithCid, prod, md) } - if c.FullPath() == "/api/providers/openai/v1/files" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/files" && method == http.MethodGet { logListFilesRequest(logWithCid, prod, qm) } - if c.FullPath() == "/api/providers/openai/v1/files" && c.Request.Method == http.MethodPost { + if fullPath == "/api/providers/openai/v1/files" && method == http.MethodPost { purpose := c.PostForm("purpose") logUploadFileRequest(logWithCid, prod, purpose) } - if c.FullPath() == "/api/providers/openai/v1/files/:file_id" && c.Request.Method == http.MethodDelete { + if fullPath == "/api/providers/openai/v1/files/:file_id" && method == http.MethodDelete { logDeleteFileRequest(logWithCid, prod, fid) } - if c.FullPath() == "/api/providers/openai/v1/files/:file_id" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/files/:file_id" && method == http.MethodGet { logRetrieveFileRequest(logWithCid, prod, fid) } - if c.FullPath() == "/api/providers/openai/v1/files/:file_id/content" && c.Request.Method == http.MethodGet { + if fullPath == "/api/providers/openai/v1/files/:file_id/content" && method == http.MethodGet { logRetrieveFileContentRequest(logWithCid, prod, fid) } @@ -1300,9 +1301,9 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag return } - if len(us[0].AllowedPaths) != 0 && !containsPath(us[0].AllowedPaths, c.FullPath(), c.Request.Method) { + if len(us[0].AllowedPaths) != 0 && !containsPath(us[0].AllowedPaths, fullPath, method) { telemetry.Incr("bricksllm.proxy.get_middleware.user_requested_path_not_allowed", nil, 1) - JSON(c, http.StatusForbidden, fmt.Sprintf("[BricksLLM] path: %s forbidden for user: %s", c.FullPath(), userId)) + JSON(c, http.StatusForbidden, fmt.Sprintf("[BricksLLM] path: %s forbidden for user: %s", fullPath, userId)) c.Abort() return } diff --git a/internal/server/web/proxy/video.go b/internal/server/web/proxy/video.go index 7e35130..f9270e9 100644 --- a/internal/server/web/proxy/video.go +++ b/internal/server/web/proxy/video.go @@ -116,7 +116,6 @@ func getVideoHandler(prod bool, client http.Client, e estimator) gin.HandlerFunc } ginCtx.Set("costInUsd", cost) ginCtx.Data(res.StatusCode, res.Header.Get("Content-Type"), bytes) - return } } From c1a8050e580e36abc66257dce978e3478d981a29 Mon Sep 17 00:00:00 2001 From: Sergei Bronnikov Date: Wed, 17 Jun 2026 10:28:45 +0100 Subject: [PATCH 4/4] refactor --- internal/provider/openai/cost.go | 6 ++++++ internal/server/web/proxy/middleware.go | 25 +++++++------------------ 2 files changed, 13 insertions(+), 18 deletions(-) diff --git a/internal/provider/openai/cost.go b/internal/provider/openai/cost.go index e83a100..93fd7ff 100644 --- a/internal/provider/openai/cost.go +++ b/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) } diff --git a/internal/server/web/proxy/middleware.go b/internal/server/web/proxy/middleware.go index d5b576b..72ced98 100644 --- a/internal/server/web/proxy/middleware.go +++ b/internal/server/web/proxy/middleware.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "slices" + "strconv" "strings" "time" @@ -239,14 +240,7 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag } if len(metadata) != 0 { - data, err := json.Marshal(metadata) - if err != nil { - telemetry.Incr("bricksllm.proxy.get_middleware.json_marshal_metadata_err", nil, 1) - } - - if err == nil { - metadataBytes = data - } + metadataBytes = []byte(metadata) } telemetry.Timing("bricksllm.proxy.get_middleware.proxy_latency_in_ms", dur, nil, 1) @@ -837,13 +831,14 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag isCreateContainerTool := false var containerMemLimit string for _, tool := range responsesReq.Tools { - if !slices.Contains(openai.AllowedTools, tool.Type) { + if _, ok := openai.AllowedToolsSet[tool.Type]; !ok { hasNotAllowedTools = true + break } - if tool.GetContainerAsResponseRequestToolContainer() != nil { + if container := tool.GetContainerAsResponseRequestToolContainer(); container != nil { isCreateContainerTool = true - containerMemLimit = tool.GetContainerAsResponseRequestToolContainer().GetMemoryLimit() + containerMemLimit = container.GetMemoryLimit() } } @@ -1404,13 +1399,7 @@ type StreamingData struct { } func contains(arr []string, target string) bool { - for _, str := range arr { - if str == target { - return true - } - } - - return false + return slices.Contains(arr, target) } func isModelAllowed(model string, settings []*provider.Setting) bool {