Skip to content

Commit 10060d1

Browse files
Pratham-Mishra04Pratham-Mishra04
committed
feat: send back raw request support
1 parent 37ebe6b commit 10060d1

File tree

45 files changed

+686
-97
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

45 files changed

+686
-97
lines changed

‎core/changelog.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
feat: send back raw request in extra fields

‎core/providers/anthropic/anthropic.go‎

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ type AnthropicProvider struct {
2323
client *fasthttp.Client // HTTP client for API requests
2424
apiVersion string // API version for the provider
2525
networkConfig schemas.NetworkConfig // Network configuration including extra headers
26+
sendBackRawRequest bool // Whether to include raw request in BifrostResponse
2627
sendBackRawResponse bool // Whether to include raw response in BifrostResponse
2728
customProviderConfig *schemas.CustomProviderConfig // Custom provider config
2829
}
@@ -103,6 +104,7 @@ func NewAnthropicProvider(config *schemas.ProviderConfig, logger schemas.Logger)
103104
client: client,
104105
apiVersion: "2023-06-01",
105106
networkConfig: config.NetworkConfig,
107+
sendBackRawRequest: config.SendBackRawRequest,
106108
sendBackRawResponse: config.SendBackRawResponse,
107109
customProviderConfig: config.CustomProviderConfig,
108110
}
@@ -199,7 +201,7 @@ func (provider *AnthropicProvider) listModelsByKey(ctx context.Context, key sche
199201

200202
// Parse Anthropic's response
201203
var anthropicResponse AnthropicListModelsResponse
202-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(resp.Body(), &anthropicResponse, providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse))
204+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(resp.Body(), &anthropicResponse, nil, provider.sendBackRawRequest, provider.sendBackRawResponse)
203205
if bifrostErr != nil {
204206
return nil, bifrostErr
205207
}
@@ -208,6 +210,11 @@ func (provider *AnthropicProvider) listModelsByKey(ctx context.Context, key sche
208210
response := anthropicResponse.ToBifrostListModelsResponse(provider.GetProviderKey(), key.Models)
209211
response.ExtraFields.Latency = latency.Milliseconds()
210212

213+
// Set raw request if enabled
214+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
215+
response.ExtraFields.RawRequest = rawRequest
216+
}
217+
211218
// Set raw response if enabled
212219
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
213220
response.ExtraFields.RawResponse = rawResponse
@@ -264,7 +271,7 @@ func (provider *AnthropicProvider) TextCompletion(ctx context.Context, key schem
264271
response := acquireAnthropicTextResponse()
265272
defer releaseAnthropicTextResponse(response)
266273

267-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, provider.sendBackRawResponse)
274+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
268275
if bifrostErr != nil {
269276
return nil, bifrostErr
270277
}
@@ -277,6 +284,11 @@ func (provider *AnthropicProvider) TextCompletion(ctx context.Context, key schem
277284
bifrostResponse.ExtraFields.RequestType = schemas.TextCompletionRequest
278285
bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
279286

287+
// Set raw request if enabled
288+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
289+
bifrostResponse.ExtraFields.RawRequest = rawRequest
290+
}
291+
280292
// Set raw response if enabled
281293
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
282294
bifrostResponse.ExtraFields.RawResponse = rawResponse
@@ -320,11 +332,10 @@ func (provider *AnthropicProvider) ChatCompletion(ctx context.Context, key schem
320332
response := AcquireAnthropicMessageResponse()
321333
defer ReleaseAnthropicMessageResponse(response)
322334

323-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse))
335+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
324336
if bifrostErr != nil {
325337
return nil, bifrostErr
326338
}
327-
328339
// Create final response
329340
bifrostResponse := response.ToBifrostChatResponse()
330341

@@ -334,6 +345,11 @@ func (provider *AnthropicProvider) ChatCompletion(ctx context.Context, key schem
334345
bifrostResponse.ExtraFields.RequestType = schemas.ChatCompletionRequest
335346
bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
336347

348+
// Set raw request if enabled
349+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
350+
bifrostResponse.ExtraFields.RawRequest = rawRequest
351+
}
352+
337353
// Set raw response if enabled
338354
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
339355
bifrostResponse.ExtraFields.RawResponse = rawResponse
@@ -386,6 +402,7 @@ func (provider *AnthropicProvider) ChatCompletionStream(ctx context.Context, pos
386402
jsonData,
387403
headers,
388404
provider.networkConfig.ExtraHeaders,
405+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
389406
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
390407
provider.GetProviderKey(),
391408
postHookRunner,
@@ -403,6 +420,7 @@ func HandleAnthropicChatCompletionStreaming(
403420
jsonBody []byte,
404421
headers map[string]string,
405422
extraHeaders map[string]string,
423+
sendBackRawRequest bool,
406424
sendBackRawResponse bool,
407425
providerName schemas.ModelProvider,
408426
postHookRunner schemas.PostHookRunner,
@@ -627,6 +645,10 @@ func HandleAnthropicChatCompletionStreaming(
627645
return
628646
}
629647
}
648+
// Set raw request if enabled
649+
if sendBackRawRequest {
650+
providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonBody)
651+
}
630652
response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
631653
ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
632654
providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, response, nil, nil, nil), responseChan)
@@ -664,7 +686,7 @@ func (provider *AnthropicProvider) Responses(ctx context.Context, key schemas.Ke
664686
response := AcquireAnthropicMessageResponse()
665687
defer ReleaseAnthropicMessageResponse(response)
666688

667-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse))
689+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
668690
if bifrostErr != nil {
669691
return nil, bifrostErr
670692
}
@@ -678,6 +700,11 @@ func (provider *AnthropicProvider) Responses(ctx context.Context, key schemas.Ke
678700
bifrostResponse.ExtraFields.RequestType = schemas.ResponsesRequest
679701
bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
680702

703+
// Set raw request if enabled
704+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
705+
bifrostResponse.ExtraFields.RawRequest = rawRequest
706+
}
707+
681708
// Set raw response if enabled
682709
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
683710
bifrostResponse.ExtraFields.RawResponse = rawResponse
@@ -727,6 +754,7 @@ func (provider *AnthropicProvider) ResponsesStream(ctx context.Context, postHook
727754
jsonBody,
728755
headers,
729756
provider.networkConfig.ExtraHeaders,
757+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
730758
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
731759
provider.GetProviderKey(),
732760
postHookRunner,
@@ -744,6 +772,7 @@ func HandleAnthropicResponsesStream(
744772
jsonBody []byte,
745773
headers map[string]string,
746774
extraHeaders map[string]string,
775+
sendBackRawRequest bool,
747776
sendBackRawResponse bool,
748777
providerName schemas.ModelProvider,
749778
postHookRunner schemas.PostHookRunner,
@@ -943,6 +972,10 @@ func HandleAnthropicResponsesStream(
943972
response.Response = &schemas.BifrostResponsesResponse{}
944973
}
945974
response.Response.Usage = usage
975+
// Set raw request if enabled
976+
if sendBackRawRequest {
977+
providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonBody)
978+
}
946979
response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
947980
ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
948981
providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, nil, response, nil, nil), responseChan)

‎core/providers/azure/azure.go‎

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ type AzureProvider struct {
2323
logger schemas.Logger // Logger for provider operations
2424
client *fasthttp.Client // HTTP client for API requests
2525
networkConfig schemas.NetworkConfig // Network configuration including extra headers
26+
sendBackRawRequest bool // Whether to include raw request in BifrostResponse
2627
sendBackRawResponse bool // Whether to include raw response in BifrostResponse
2728
}
2829

@@ -47,6 +48,7 @@ func NewAzureProvider(config *schemas.ProviderConfig, logger schemas.Logger) (*A
4748
logger: logger,
4849
client: client,
4950
networkConfig: config.NetworkConfig,
51+
sendBackRawRequest: config.SendBackRawRequest,
5052
sendBackRawResponse: config.SendBackRawResponse,
5153
}, nil
5254
}
@@ -192,7 +194,7 @@ func (provider *AzureProvider) listModelsByKey(ctx context.Context, key schemas.
192194

193195
// Parse Azure-specific response
194196
azureResponse := &AzureListModelsResponse{}
195-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, azureResponse, provider.sendBackRawResponse)
197+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, azureResponse, nil, provider.sendBackRawRequest, provider.sendBackRawResponse)
196198
if bifrostErr != nil {
197199
return nil, bifrostErr
198200
}
@@ -205,6 +207,12 @@ func (provider *AzureProvider) listModelsByKey(ctx context.Context, key schemas.
205207

206208
response.ExtraFields.Latency = latency.Milliseconds()
207209

210+
// Set raw request if enabled
211+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
212+
response.ExtraFields.RawRequest = rawRequest
213+
}
214+
215+
// Set raw response if enabled
208216
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
209217
response.ExtraFields.RawResponse = rawResponse
210218
}
@@ -263,7 +271,7 @@ func (provider *AzureProvider) TextCompletion(ctx context.Context, key schemas.K
263271

264272
response := &schemas.BifrostTextCompletionResponse{}
265273

266-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, provider.sendBackRawResponse)
274+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
267275
if bifrostErr != nil {
268276
return nil, bifrostErr
269277
}
@@ -274,6 +282,11 @@ func (provider *AzureProvider) TextCompletion(ctx context.Context, key schemas.K
274282
response.ExtraFields.RequestType = schemas.TextCompletionRequest
275283
response.ExtraFields.Latency = latency.Milliseconds()
276284

285+
// Set raw request if enabled
286+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
287+
response.ExtraFields.RawRequest = rawRequest
288+
}
289+
277290
// Set raw response if enabled
278291
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
279292
response.ExtraFields.RawResponse = rawResponse
@@ -324,6 +337,7 @@ func (provider *AzureProvider) TextCompletionStream(ctx context.Context, postHoo
324337
request,
325338
authHeader,
326339
provider.networkConfig.ExtraHeaders,
340+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
327341
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
328342
provider.GetProviderKey(),
329343
postHookRunner,
@@ -388,18 +402,19 @@ func (provider *AzureProvider) ChatCompletion(ctx context.Context, key schemas.K
388402
}
389403

390404
response := &schemas.BifrostChatResponse{}
405+
var rawRequest interface{}
391406
var rawResponse interface{}
392407

393408
if schemas.IsAnthropicModel(deployment) {
394409
anthropicResponse := anthropic.AcquireAnthropicMessageResponse()
395410
defer anthropic.ReleaseAnthropicMessageResponse(anthropicResponse)
396-
rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, provider.sendBackRawResponse)
411+
rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
397412
if bifrostErr != nil {
398413
return nil, bifrostErr
399414
}
400415
response = anthropicResponse.ToBifrostChatResponse()
401416
} else {
402-
rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, response, provider.sendBackRawResponse)
417+
rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
403418
if bifrostErr != nil {
404419
return nil, bifrostErr
405420
}
@@ -411,6 +426,11 @@ func (provider *AzureProvider) ChatCompletion(ctx context.Context, key schemas.K
411426
response.ExtraFields.Latency = latency.Milliseconds()
412427
response.ExtraFields.RequestType = schemas.ChatCompletionRequest
413428

429+
// Set raw request if enabled
430+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
431+
response.ExtraFields.RawRequest = rawRequest
432+
}
433+
414434
// Set raw response if enabled
415435
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
416436
response.ExtraFields.RawResponse = rawResponse
@@ -470,6 +490,7 @@ func (provider *AzureProvider) ChatCompletionStream(ctx context.Context, postHoo
470490
jsonData,
471491
authHeader,
472492
provider.networkConfig.ExtraHeaders,
493+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
473494
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
474495
provider.GetProviderKey(),
475496
postHookRunner,
@@ -497,6 +518,7 @@ func (provider *AzureProvider) ChatCompletionStream(ctx context.Context, postHoo
497518
request,
498519
authHeader,
499520
provider.networkConfig.ExtraHeaders,
521+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
500522
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
501523
provider.GetProviderKey(),
502524
postHookRunner,
@@ -566,18 +588,19 @@ func (provider *AzureProvider) Responses(ctx context.Context, key schemas.Key, r
566588
}
567589

568590
response := &schemas.BifrostResponsesResponse{}
591+
var rawRequest interface{}
569592
var rawResponse interface{}
570593

571594
if schemas.IsAnthropicModel(deployment) {
572595
anthropicResponse := anthropic.AcquireAnthropicMessageResponse()
573596
defer anthropic.ReleaseAnthropicMessageResponse(anthropicResponse)
574-
rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, provider.sendBackRawResponse)
597+
rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
575598
if bifrostErr != nil {
576599
return nil, bifrostErr
577600
}
578601
response = anthropicResponse.ToBifrostResponsesResponse()
579602
} else {
580-
rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, response, provider.sendBackRawResponse)
603+
rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
581604
if bifrostErr != nil {
582605
return nil, bifrostErr
583606
}
@@ -589,6 +612,11 @@ func (provider *AzureProvider) Responses(ctx context.Context, key schemas.Key, r
589612
response.ExtraFields.Latency = latency.Milliseconds()
590613
response.ExtraFields.RequestType = schemas.ResponsesRequest
591614

615+
// Set raw request if enabled
616+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
617+
response.ExtraFields.RawRequest = rawRequest
618+
}
619+
592620
// Set raw response if enabled
593621
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
594622
response.ExtraFields.RawResponse = rawResponse
@@ -645,6 +673,7 @@ func (provider *AzureProvider) ResponsesStream(ctx context.Context, postHookRunn
645673
jsonData,
646674
authHeader,
647675
provider.networkConfig.ExtraHeaders,
676+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
648677
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
649678
provider.GetProviderKey(),
650679
postHookRunner,
@@ -673,6 +702,7 @@ func (provider *AzureProvider) ResponsesStream(ctx context.Context, postHookRunn
673702
request,
674703
authHeader,
675704
provider.networkConfig.ExtraHeaders,
705+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
676706
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
677707
provider.GetProviderKey(),
678708
postHookRunner,
@@ -722,7 +752,7 @@ func (provider *AzureProvider) Embedding(ctx context.Context, key schemas.Key, r
722752
response := &schemas.BifrostEmbeddingResponse{}
723753

724754
// Use enhanced response handler with pre-allocated response
725-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, provider.sendBackRawResponse)
755+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
726756
if bifrostErr != nil {
727757
return nil, bifrostErr
728758
}
@@ -733,6 +763,12 @@ func (provider *AzureProvider) Embedding(ctx context.Context, key schemas.Key, r
733763
response.ExtraFields.ModelDeployment = deployment
734764
response.ExtraFields.RequestType = schemas.EmbeddingRequest
735765

766+
// Set raw request if enabled
767+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
768+
response.ExtraFields.RawRequest = rawRequest
769+
}
770+
771+
// Set raw response if enabled
736772
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
737773
response.ExtraFields.RawResponse = rawResponse
738774
}

0 commit comments

Comments
 (0)