Skip to content

Commit f5b6ab1

Browse files
sergei-bronnikovsergei-bronnikov
authored
unmarshal anthropic msg content. abort if err unmarshal (#27)
1 parent 0ddb013 commit f5b6ab1

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
@@ -393,6 +393,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
393393
body, err := io.ReadAll(c.Request.Body)
394394
if err != nil {
395395
logError(logWithCid, "error when reading request body", prod, err)
396+
JSON(c, http.StatusInternalServerError, "[BricksLLM] error reading request body")
397+
c.Abort()
396398
return
397399
}
398400

@@ -415,6 +417,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
415417
err = json.Unmarshal(body, cr)
416418
if err != nil {
417419
logError(logWithCid, "error when unmarshalling anthropic completion request", prod, err)
420+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling anthropic completion request")
421+
c.Abort()
418422
return
419423
}
420424

@@ -440,6 +444,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
440444
err = json.Unmarshal(body, cr)
441445
if err != nil {
442446
logError(logWithCid, "error when unmarshalling bedrock anthropic completion request", prod, err)
447+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling bedrock anthropic completion request")
448+
c.Abort()
443449
return
444450
}
445451

@@ -465,6 +471,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
465471
err = json.Unmarshal(body, mr)
466472
if err != nil {
467473
logError(logWithCid, "error when unmarshalling anthropic messages request", prod, err)
474+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling anthropic messages request")
475+
c.Abort()
468476
return
469477
}
470478

@@ -488,6 +496,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
488496
err = json.Unmarshal(body, mr)
489497
if err != nil {
490498
logError(logWithCid, "error when unmarshalling anthropic messages request", prod, err)
499+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling anthropic messages request")
500+
c.Abort()
491501
return
492502
}
493503

@@ -624,6 +634,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
624634
err = json.Unmarshal(body, ccr)
625635
if err != nil {
626636
logError(logWithCid, "error when unmarshalling vllm chat completions request", prod, err)
637+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling vllm chat completions request")
638+
c.Abort()
627639
return
628640
}
629641

@@ -645,6 +657,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
645657
err = json.Unmarshal(body, cr)
646658
if err != nil {
647659
logError(logWithCid, "error when unmarshalling vllm completions request", prod, err)
660+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling vllm completions request")
661+
c.Abort()
648662
return
649663
}
650664

@@ -666,6 +680,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
666680
err = json.Unmarshal(body, ccr)
667681
if err != nil {
668682
logError(logWithCid, "error when unmarshalling deepinfra chat completions request", prod, err)
683+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling deepinfra chat completions request")
684+
c.Abort()
669685
return
670686
}
671687

@@ -686,6 +702,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
686702
err = json.Unmarshal(body, cr)
687703
if err != nil {
688704
logError(logWithCid, "error when unmarshalling deepinfra completions request", prod, err)
705+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling deepinfra completions request")
706+
c.Abort()
689707
return
690708
}
691709

@@ -706,6 +724,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
706724
err = json.Unmarshal(body, er)
707725
if err != nil {
708726
logError(logWithCid, "error when unmarshalling deepinfra embeddings request", prod, err)
727+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling deepinfra embeddings request")
728+
c.Abort()
709729
return
710730
}
711731

@@ -723,6 +743,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
723743
err = json.Unmarshal(body, ccr)
724744
if err != nil {
725745
logError(logWithCid, "error when unmarshalling azure openai chat completion request", prod, err)
746+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling azure openai chat completion request")
747+
c.Abort()
726748
return
727749
}
728750

@@ -744,6 +766,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
744766
err = json.Unmarshal(body, cr)
745767
if err != nil {
746768
logError(logWithCid, "error when unmarshalling azure openai completions request", prod, err)
769+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling azure openai completions request")
770+
c.Abort()
747771
return
748772
}
749773

@@ -765,6 +789,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
765789
err = json.Unmarshal(body, er)
766790
if err != nil {
767791
logError(logWithCid, "error when unmarshalling azure openai embedding request", prod, err)
792+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling azure openai embedding request")
793+
c.Abort()
768794
return
769795
}
770796

@@ -789,6 +815,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
789815
err = json.Unmarshal([]byte(cleaned), ccr)
790816
if err != nil {
791817
logError(logWithCid, "error when unmarshalling chat completion request", prod, err)
818+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling chat completion request")
819+
c.Abort()
792820
return
793821
}
794822

@@ -812,6 +840,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
812840
err = json.Unmarshal(body, responsesReq)
813841
if err != nil {
814842
logError(logWithCid, "error when unmarshalling openai responses request", prod, err)
843+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling openai responses request")
844+
c.Abort()
815845
return
816846
}
817847

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

@@ -894,6 +926,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
894926
err := json.Unmarshal(body, ir)
895927
if err != nil {
896928
logError(logWithCid, "error when unmarshalling create image request", prod, err)
929+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling create image request")
930+
c.Abort()
897931
return
898932
}
899933
enrichedEvent.Request = ir
@@ -913,6 +947,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
913947
err := json.Unmarshal(body, ier)
914948
if err != nil {
915949
logError(logWithCid, "error when unmarshalling edit image request", prod, err)
950+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling edit image request")
951+
c.Abort()
916952
return
917953
}
918954
enrichedEvent.Request = ier
@@ -941,6 +977,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
941977
err := json.Unmarshal(body, ir)
942978
if err != nil {
943979
logError(logWithCid, "error when unmarshalling image variations request", prod, err)
980+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling image variations request")
981+
c.Abort()
944982
return
945983
}
946984
enrichedEvent.Request = ir
@@ -968,6 +1006,8 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
9681006
err := json.Unmarshal(body, sr)
9691007
if err != nil {
9701008
logError(logWithCid, "error when unmarshalling create speech request", prod, err)
1009+
JSON(c, http.StatusBadRequest, "[BricksLLM] error when unmarshalling create speech request")
1010+
c.Abort()
9711011
return
9721012
}
9731013

0 commit comments

Comments
 (0)