Skip to content

enhancement: serialisation optimisations #194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
akshaydeo merged 1 commit into from
Jul 29, 2025
Merged

enhancement: serialisation optimisations #194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion core/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ go 1.24.1
require (
github.com/aws/aws-sdk-go-v2 v1.36.3
github.com/aws/aws-sdk-go-v2/config v1.29.14
github.com/goccy/go-json v0.10.5
github.com/bytedance/sonic v1.14.0
github.com/mark3labs/mcp-go v0.32.0
github.com/valyala/fasthttp v1.60.0
golang.org/x/oauth2 v0.30.0
Expand All @@ -25,11 +25,16 @@ require (
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.1 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.33.19 // indirect
github.com/aws/smithy-go v1.22.3 // indirect
github.com/bytedance/sonic/loader v0.3.0 // indirect
github.com/cloudwego/base64x v0.1.5 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/klauspost/compress v1.18.0 // indirect
github.com/klauspost/cpuid/v2 v2.0.9 // indirect
github.com/spf13/cast v1.7.1 // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/valyala/bytebufferpool v1.0.0 // indirect
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
golang.org/x/arch v0.0.0-20210923205945-b76863e36670 // indirect
golang.org/x/net v0.39.0 // indirect
golang.org/x/text v0.24.0 // indirect
)
28 changes: 26 additions & 2 deletions core/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,18 +28,28 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.33.19 h1:1XuUZ8mYJw9B6lzAkXhqHlJd/Xv
github.com/aws/aws-sdk-go-v2/service/sts v1.33.19/go.mod h1:cQnB8CUnxbMU82JvlqjKR2HBOm3fe9pWorWBza6MBJ4=
github.com/aws/smithy-go v1.22.3 h1:Z//5NuZCSW6R4PhQ93hShNbyBbn8BWCmCVCt+Q8Io5k=
github.com/aws/smithy-go v1.22.3/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=
github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ=
github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA=
github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA=
github.com/bytedance/sonic/loader v0.3.0/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
Expand All @@ -52,8 +62,17 @@ github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZV
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
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=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
github.com/valyala/fasthttp v1.60.0 h1:kBRYS0lOhVJ6V+bYN8PqAHELKHtXqwq9zNMLKx1MBsw=
Expand All @@ -62,11 +81,16 @@ github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZ
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
golang.org/x/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VAiXCnxFY6NyDX0bHDmkU=
golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
golang.org/x/net v0.39.0 h1:ZCu7HMWDxpXpaiKdhzIfaltL9Lp31x/3fCP11bc6/fY=
golang.org/x/net v0.39.0/go.mod h1:X7NRbYVEA+ewNkCNyJ513WmMdQ3BineSwVtN2zD/d+E=
golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU=
golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0=
golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50=
23 changes: 11 additions & 12 deletions core/providers/anthropic.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,7 @@ import (
"sync"
"time"

"github.com/goccy/go-json"

"github.com/bytedance/sonic"
schemas "github.com/maximhq/bifrost/core/schemas"
"github.com/valyala/fasthttp"
)
Expand Down Expand Up @@ -263,7 +262,7 @@ func (provider *AnthropicProvider) prepareTextCompletionParams(params map[string
// Returns the response body or an error if the request fails.
func (provider *AnthropicProvider) completeRequest(ctx context.Context, requestBody map[string]interface{}, url string, key string) ([]byte, *schemas.BifrostError) {
// Marshal the request body
jsonData, err := json.Marshal(requestBody)
jsonData, err := sonic.Marshal(requestBody)
if err != nil {
return nil, newBifrostOperationError(schemas.ErrProviderJSONMarshaling, err, schemas.Anthropic)
}
Expand Down Expand Up @@ -548,7 +547,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
if toolCall.Function.Name != nil {
var input map[string]interface{}
if toolCall.Function.Arguments != "" {
if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &input); err != nil {
if err := sonic.Unmarshal([]byte(toolCall.Function.Arguments), &input); err != nil {
// If unmarshaling fails, use a simple string representation
input = map[string]interface{}{"arguments": toolCall.Function.Arguments}
}
Expand Down Expand Up @@ -712,7 +711,7 @@ func parseAnthropicResponse(response *AnthropicChatResponse, bifrostResponse *sc
Name: &c.Name,
}

args, err := json.Marshal(c.Input)
args, err := sonic.Marshal(c.Input)
if err != nil {
function.Arguments = fmt.Sprintf("%v", c.Input)
} else {
Expand Down Expand Up @@ -826,7 +825,7 @@ func handleAnthropicStreaming(
logger schemas.Logger,
) (chan *schemas.BifrostStream, *schemas.BifrostError) {

jsonBody, err := json.Marshal(requestBody)
jsonBody, err := sonic.Marshal(requestBody)
if err != nil {
return nil, newBifrostOperationError(schemas.ErrProviderJSONMarshaling, err, providerType)
}
Expand Down Expand Up @@ -903,7 +902,7 @@ func handleAnthropicStreaming(
switch eventType {
case "message_start":
var event AnthropicStreamEvent
if err := json.Unmarshal([]byte(eventData), &event); err != nil {
if err := sonic.Unmarshal([]byte(eventData), &event); err != nil {
logger.Warn(fmt.Sprintf("Failed to parse message_start event: %v", err))
continue
}
Expand All @@ -914,7 +913,7 @@ func handleAnthropicStreaming(

case "content_block_start":
var event AnthropicStreamEvent
if err := json.Unmarshal([]byte(eventData), &event); err != nil {
if err := sonic.Unmarshal([]byte(eventData), &event); err != nil {
logger.Warn(fmt.Sprintf("Failed to parse content_block_start event: %v", err))
continue
}
Expand Down Expand Up @@ -1002,7 +1001,7 @@ func handleAnthropicStreaming(

case "content_block_delta":
var event AnthropicStreamEvent
if err := json.Unmarshal([]byte(eventData), &event); err != nil {
if err := sonic.Unmarshal([]byte(eventData), &event); err != nil {
logger.Warn(fmt.Sprintf("Failed to parse content_block_delta event: %v", err))
continue
}
Expand Down Expand Up @@ -1122,7 +1121,7 @@ func handleAnthropicStreaming(

case "message_delta":
var event AnthropicStreamEvent
if err := json.Unmarshal([]byte(eventData), &event); err != nil {
if err := sonic.Unmarshal([]byte(eventData), &event); err != nil {
logger.Warn(fmt.Sprintf("Failed to parse message_delta event: %v", err))
continue
}
Expand Down Expand Up @@ -1160,7 +1159,7 @@ func handleAnthropicStreaming(

case "message_stop":
var event AnthropicStreamEvent
if err := json.Unmarshal([]byte(eventData), &event); err != nil {
if err := sonic.Unmarshal([]byte(eventData), &event); err != nil {
logger.Warn(fmt.Sprintf("Failed to parse message_stop event: %v", err))
continue
}
Expand Down Expand Up @@ -1203,7 +1202,7 @@ func handleAnthropicStreaming(

case "error":
var event AnthropicStreamEvent
if err := json.Unmarshal([]byte(eventData), &event); err != nil {
if err := sonic.Unmarshal([]byte(eventData), &event); err != nil {
logger.Warn(fmt.Sprintf("Failed to parse error event: %v", err))
continue
}
Expand Down
7 changes: 3 additions & 4 deletions core/providers/azure.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,7 @@ import (
"sync"
"time"

"github.com/goccy/go-json"

"github.com/bytedance/sonic"
schemas "github.com/maximhq/bifrost/core/schemas"
"github.com/valyala/fasthttp"
)
Expand Down Expand Up @@ -169,7 +168,7 @@ func (provider *AzureProvider) completeRequest(ctx context.Context, requestBody
}

// Marshal the request body
jsonData, err := json.Marshal(requestBody)
jsonData, err := sonic.Marshal(requestBody)
if err != nil {
return nil, newBifrostOperationError(schemas.ErrProviderJSONMarshaling, err, schemas.Azure)
}
Expand Down Expand Up @@ -392,7 +391,7 @@ func (provider *AzureProvider) Embedding(ctx context.Context, model string, key

// Parse response
var response AzureEmbeddingResponse
if err := json.Unmarshal(responseBody, &response); err != nil {
if err := sonic.Unmarshal(responseBody, &response); err != nil {
return nil, newBifrostOperationError(schemas.ErrProviderResponseUnmarshal, err, schemas.Azure)
}

Expand Down
30 changes: 15 additions & 15 deletions core/providers/bedrock.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
Expand All @@ -17,13 +18,12 @@ import (
"sync"
"time"

"github.com/goccy/go-json"

"bufio"

"github.com/aws/aws-sdk-go-v2/aws"
v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/bytedance/sonic"
schemas "github.com/maximhq/bifrost/core/schemas"
)

Expand Down Expand Up @@ -269,7 +269,7 @@ func (provider *BedrockProvider) completeRequest(ctx context.Context, requestBod
region = *provider.meta.GetRegion()
}

jsonBody, err := json.Marshal(requestBody)
jsonBody, err := sonic.Marshal(requestBody)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return nil, &schemas.BifrostError{
Expand Down Expand Up @@ -346,7 +346,7 @@ func (provider *BedrockProvider) completeRequest(ctx context.Context, requestBod
if resp.StatusCode != http.StatusOK {
var errorResp BedrockError

if err := json.Unmarshal(body, &errorResp); err != nil {
if err := sonic.Unmarshal(body, &errorResp); err != nil {
return nil, &schemas.BifrostError{
IsBifrostError: true,
StatusCode: &resp.StatusCode,
Expand Down Expand Up @@ -379,7 +379,7 @@ func (provider *BedrockProvider) getTextCompletionResult(result []byte, model st
fallthrough
case "anthropic.claude-v2:1":
var response BedrockAnthropicTextResponse
if err := json.Unmarshal(result, &response); err != nil {
if err := sonic.Unmarshal(result, &response); err != nil {
return nil, &schemas.BifrostError{
IsBifrostError: true,
Error: schemas.ErrorField{
Expand Down Expand Up @@ -421,7 +421,7 @@ func (provider *BedrockProvider) getTextCompletionResult(result []byte, model st
fallthrough
case "mistral.mistral-small-2402-v1:0":
var response BedrockMistralTextResponse
if err := json.Unmarshal(result, &response); err != nil {
if err := sonic.Unmarshal(result, &response); err != nil {
return nil, &schemas.BifrostError{
IsBifrostError: true,
Error: schemas.ErrorField{
Expand Down Expand Up @@ -465,7 +465,7 @@ func (provider *BedrockProvider) getTextCompletionResult(result []byte, model st
func parseBedrockAnthropicMessageToolCallContent(content string) map[string]interface{} {
toolResultContentBlock := map[string]interface{}{}
var parsedJSON interface{}
err := json.Unmarshal([]byte(content), &parsedJSON)
err := sonic.Unmarshal([]byte(content), &parsedJSON)
if err == nil {
if arr, ok := parsedJSON.([]interface{}); ok {
toolResultContentBlock["json"] = map[string]interface{}{"content": arr}
Expand Down Expand Up @@ -549,7 +549,7 @@ func (provider *BedrockProvider) prepareChatCompletionMessages(messages []schema
for _, toolCall := range *msg.AssistantMessage.ToolCalls {
var input map[string]interface{}
if toolCall.Function.Arguments != "" {
if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &input); err != nil {
if err := sonic.Unmarshal([]byte(toolCall.Function.Arguments), &input); err != nil {
input = map[string]interface{}{"arguments": toolCall.Function.Arguments}
}
}
Expand Down Expand Up @@ -838,7 +838,7 @@ func (provider *BedrockProvider) TextCompletion(ctx context.Context, model strin

// Parse raw response
var rawResponse interface{}
if err := json.Unmarshal(body, &rawResponse); err != nil {
if err := sonic.Unmarshal(body, &rawResponse); err != nil {
return nil, newBifrostOperationError("error parsing raw response", err, schemas.Bedrock)
}

Expand Down Expand Up @@ -978,7 +978,7 @@ func (provider *BedrockProvider) ChatCompletion(ctx context.Context, model strin
if input == nil {
input = map[string]any{}
}
arguments, err := json.Marshal(input)
arguments, err := sonic.Marshal(input)
if err != nil {
arguments = []byte("{}")
}
Expand Down Expand Up @@ -1159,7 +1159,7 @@ func (provider *BedrockProvider) handleTitanEmbedding(ctx context.Context, model
Embedding []float32 `json:"embedding"`
InputTextTokenCount int `json:"inputTextTokenCount"`
}
if err := json.Unmarshal(rawResponse, &titanResp); err != nil {
if err := sonic.Unmarshal(rawResponse, &titanResp); err != nil {
return nil, newBifrostOperationError("error parsing Titan embedding response", err, schemas.Bedrock)
}

Expand Down Expand Up @@ -1210,7 +1210,7 @@ func (provider *BedrockProvider) handleCohereEmbedding(ctx context.Context, mode
ID string `json:"id"`
Texts []string `json:"texts"`
}
if err := json.Unmarshal(rawResponse, &cohereResp); err != nil {
if err := sonic.Unmarshal(rawResponse, &cohereResp); err != nil {
return nil, newBifrostOperationError("error parsing Cohere embedding response", err, schemas.Bedrock)
}

Expand Down Expand Up @@ -1292,7 +1292,7 @@ func (provider *BedrockProvider) ChatCompletionStream(ctx context.Context, postH
}

// Create the streaming request
jsonBody, jsonErr := json.Marshal(requestBody)
jsonBody, jsonErr := sonic.Marshal(requestBody)
if jsonErr != nil {
return nil, newBifrostOperationError(schemas.ErrProviderJSONMarshaling, jsonErr, schemas.Bedrock)
}
Expand Down Expand Up @@ -1386,7 +1386,7 @@ func (provider *BedrockProvider) ChatCompletionStream(ctx context.Context, postH

// Parse the JSON event
var event map[string]interface{}
if err := json.Unmarshal([]byte(jsonStr), &event); err != nil {
if err := sonic.Unmarshal([]byte(jsonStr), &event); err != nil {
provider.logger.Debug(fmt.Sprintf("Failed to parse JSON from stream: %v, data: %s", err, jsonStr))
continue
}
Expand Down Expand Up @@ -1456,7 +1456,7 @@ func (provider *BedrockProvider) ChatCompletionStream(ctx context.Context, postH

// Extract and marshal input as arguments
if input, hasInput := toolUse["input"].(map[string]interface{}); hasInput {
inputBytes, err := json.Marshal(input)
inputBytes, err := sonic.Marshal(input)
if err != nil {
toolCall.Function.Arguments = "{}"
} else {
Expand Down
Loading