Skip to content

Commit 45a92d3

Browse files
sergei-bronnikovsergei-bronnikov
committed
Merge branch 'main' into 17720_performance_issues
2 parents 56d30cb + f5b6ab1 commit 45a92d3

4 files changed

Lines changed: 116 additions & 5 deletions

File tree

‎internal/policy/policy.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ func (p *Policy) Filter(client http.Client, input any, scanner Scanner, cd Custo
438438

439439
contents := []string{}
440440
for _, message := range converted.Messages {
441-
contents = append(contents, message.Content)
441+
contents = append(contents, message.Content.String())
442442
}
443443

444444
result, err := p.scan(contents, scanner, cd, log)
@@ -460,7 +460,7 @@ func (p *Policy) Filter(client http.Client, input any, scanner Scanner, cd Custo
460460

461461
for index, c := range result.Updated {
462462
newMessages = append(newMessages, anthropic.Message{
463-
Content: c,
463+
Content: anthropic.FlexContent{Text: c},
464464
Role: converted.Messages[index].Role,
465465
})
466466
}

‎internal/provider/anthropic/anthropic.go‎

Lines changed: 73 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
package anthropic
22

3+
import (
4+
"encoding/json"
5+
"strings"
6+
)
7+
38
type Metadata struct {
49
UserId string `json:"user_id"`
510
}
@@ -16,9 +21,75 @@ type CompletionRequest struct {
1621
Stream bool `json:"stream,omitempty"`
1722
}
1823

24+
type FlexContent struct {
25+
Text string
26+
Raw []interface{}
27+
}
28+
29+
func (f *FlexContent) UnmarshalJSON(data []byte) error {
30+
var s string
31+
if err := json.Unmarshal(data, &s); err == nil {
32+
f.Text = s
33+
return nil
34+
}
35+
36+
var arr []interface{}
37+
if err := json.Unmarshal(data, &arr); err != nil {
38+
return err
39+
}
40+
f.Raw = arr
41+
return nil
42+
}
43+
44+
func (f *FlexContent) String() string {
45+
if f == nil {
46+
return ""
47+
}
48+
if f.Text != "" {
49+
return f.Text
50+
}
51+
if len(f.Raw) > 0 {
52+
var builder strings.Builder
53+
for _, block := range f.Raw {
54+
builder.WriteString(extractAnthropicContentText(block))
55+
}
56+
return builder.String()
57+
}
58+
return ""
59+
}
60+
61+
func extractAnthropicContentText(value any) string {
62+
switch v := value.(type) {
63+
case string:
64+
return v
65+
case []any:
66+
var builder strings.Builder
67+
for _, item := range v {
68+
builder.WriteString(extractAnthropicContentText(item))
69+
}
70+
return builder.String()
71+
case map[string]any:
72+
if text, ok := v["text"].(string); ok {
73+
return text
74+
}
75+
if content, ok := v["content"].(string); ok {
76+
return content
77+
}
78+
if raw, ok := v["raw"].([]any); ok {
79+
var builder strings.Builder
80+
for _, item := range raw {
81+
builder.WriteString(extractAnthropicContentText(item))
82+
}
83+
return builder.String()
84+
}
85+
}
86+
87+
return ""
88+
}
89+
1990
type Message struct {
20-
Content string `json:"content"`
21-
Role string `json:"role"`
91+
Content FlexContent `json:"content"`
92+
Role string `json:"role"`
2293
}
2394

2495
type MessagesRequest struct {

‎internal/provider/anthropic/cost.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ func (ce *CostEstimator) CountMessagesTokens(messages []Message) int {
198198
count := 0
199199

200200
for _, message := range messages {
201-
count += ce.tc.Count(message.Content) + anthropicMessageOverhead
201+
count += ce.tc.Count(message.Content.String()) + anthropicMessageOverhead
202202
}
203203

204204
return count + anthropicMessageOverhead

‎internal/server/web/proxy/middleware.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
395395
body, err := io.ReadAll(c.Request.Body)
396396
if err != nil {
397397
logError(logWithCid, "error when reading request body", prod, err)
398+
JSON(c, http.StatusInternalServerError, "[BricksLLM] error reading request body")
399+
c.Abort()
398400
return
399401
}
400402

@@ -417,6 +419,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
417419
err = json.Unmarshal(body, cr)
418420
if err != nil {
419421
logError(logWithCid, "error when unmarshalling anthropic completion request", prod, err)
422+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling anthropic completion request")
423+
c.Abort()
420424
return
421425
}
422426

@@ -442,6 +446,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
442446
err = json.Unmarshal(body, cr)
443447
if err != nil {
444448
logError(logWithCid, "error when unmarshalling bedrock anthropic completion request", prod, err)
449+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling bedrock anthropic completion request")
450+
c.Abort()
445451
return
446452
}
447453

@@ -467,6 +473,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
467473
err = json.Unmarshal(body, mr)
468474
if err != nil {
469475
logError(logWithCid, "error when unmarshalling anthropic messages request", prod, err)
476+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling anthropic messages request")
477+
c.Abort()
470478
return
471479
}
472480

@@ -490,6 +498,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
490498
err = json.Unmarshal(body, mr)
491499
if err != nil {
492500
logError(logWithCid, "error when unmarshalling anthropic messages request", prod, err)
501+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling anthropic messages request")
502+
c.Abort()
493503
return
494504
}
495505

@@ -627,6 +637,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
627637
err = json.Unmarshal(body, ccr)
628638
if err != nil {
629639
logError(logWithCid, "error when unmarshalling vllm chat completions request", prod, err)
640+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling vllm chat completions request")
641+
c.Abort()
630642
return
631643
}
632644

@@ -648,6 +660,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
648660
err = json.Unmarshal(body, cr)
649661
if err != nil {
650662
logError(logWithCid, "error when unmarshalling vllm completions request", prod, err)
663+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling vllm completions request")
664+
c.Abort()
651665
return
652666
}
653667

@@ -669,6 +683,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
669683
err = json.Unmarshal(body, ccr)
670684
if err != nil {
671685
logError(logWithCid, "error when unmarshalling deepinfra chat completions request", prod, err)
686+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling deepinfra chat completions request")
687+
c.Abort()
672688
return
673689
}
674690

@@ -689,6 +705,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
689705
err = json.Unmarshal(body, cr)
690706
if err != nil {
691707
logError(logWithCid, "error when unmarshalling deepinfra completions request", prod, err)
708+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling deepinfra completions request")
709+
c.Abort()
692710
return
693711
}
694712

@@ -709,6 +727,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
709727
err = json.Unmarshal(body, er)
710728
if err != nil {
711729
logError(logWithCid, "error when unmarshalling deepinfra embeddings request", prod, err)
730+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling deepinfra embeddings request")
731+
c.Abort()
712732
return
713733
}
714734

@@ -726,6 +746,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
726746
err = json.Unmarshal(body, ccr)
727747
if err != nil {
728748
logError(logWithCid, "error when unmarshalling azure openai chat completion request", prod, err)
749+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling azure openai chat completion request")
750+
c.Abort()
729751
return
730752
}
731753

@@ -747,6 +769,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
747769
err = json.Unmarshal(body, cr)
748770
if err != nil {
749771
logError(logWithCid, "error when unmarshalling azure openai completions request", prod, err)
772+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling azure openai completions request")
773+
c.Abort()
750774
return
751775
}
752776

@@ -768,6 +792,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
768792
err = json.Unmarshal(body, er)
769793
if err != nil {
770794
logError(logWithCid, "error when unmarshalling azure openai embedding request", prod, err)
795+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling azure openai embedding request")
796+
c.Abort()
771797
return
772798
}
773799

@@ -792,6 +818,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
792818
err = json.Unmarshal([]byte(cleaned), ccr)
793819
if err != nil {
794820
logError(logWithCid, "error when unmarshalling chat completion request", prod, err)
821+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling chat completion request")
822+
c.Abort()
795823
return
796824
}
797825

@@ -815,6 +843,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
815843
err = json.Unmarshal(body, responsesReq)
816844
if err != nil {
817845
logError(logWithCid, "error when unmarshalling openai responses request", prod, err)
846+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling openai responses request")
847+
c.Abort()
818848
return
819849
}
820850

@@ -877,6 +907,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
877907
err = json.Unmarshal(body, er)
878908
if err != nil {
879909
logError(logWithCid, "error when unmarshalling embedding request", prod, err)
910+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling embedding request")
911+
c.Abort()
880912
return
881913
}
882914

@@ -895,6 +927,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
895927
err := json.Unmarshal(body, ir)
896928
if err != nil {
897929
logError(logWithCid, "error when unmarshalling create image request", prod, err)
930+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling create image request")
931+
c.Abort()
898932
return
899933
}
900934
enrichedEvent.Request = ir
@@ -914,6 +948,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
914948
err := json.Unmarshal(body, ier)
915949
if err != nil {
916950
logError(logWithCid, "error when unmarshalling edit image request", prod, err)
951+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling edit image request")
952+
c.Abort()
917953
return
918954
}
919955
enrichedEvent.Request = ier
@@ -942,6 +978,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
942978
err := json.Unmarshal(body, ir)
943979
if err != nil {
944980
logError(logWithCid, "error when unmarshalling image variations request", prod, err)
981+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling image variations request")
982+
c.Abort()
945983
return
946984
}
947985
enrichedEvent.Request = ir
@@ -969,6 +1007,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
9691007
err := json.Unmarshal(body, sr)
9701008
if err != nil {
9711009
logError(logWithCid, "error when unmarshalling create speech request", prod, err)
1010+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling create speech request")
1011+
c.Abort()
9721012
return
9731013
}
9741014

0 commit comments

Comments
 (0)