Skip to content

Commit 1cb0123

Browse files
Pratham-Mishra04Pratham-Mishra04
committed
feat: send back raw request support
1 parent feb7ebe commit 1cb0123

File tree

45 files changed

+687
-107
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

+687
-107
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,
@@ -635,6 +653,10 @@ func HandleAnthropicChatCompletionStreaming(
635653
return
636654
}
637655
}
656+
// Set raw request if enabled
657+
if sendBackRawRequest {
658+
providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonBody)
659+
}
638660
response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
639661
ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
640662
providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, response, nil, nil, nil), responseChan)
@@ -672,7 +694,7 @@ func (provider *AnthropicProvider) Responses(ctx context.Context, key schemas.Ke
672694
response := AcquireAnthropicMessageResponse()
673695
defer ReleaseAnthropicMessageResponse(response)
674696

675-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse))
697+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
676698
if bifrostErr != nil {
677699
return nil, bifrostErr
678700
}
@@ -686,6 +708,11 @@ func (provider *AnthropicProvider) Responses(ctx context.Context, key schemas.Ke
686708
bifrostResponse.ExtraFields.RequestType = schemas.ResponsesRequest
687709
bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
688710

711+
// Set raw request if enabled
712+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
713+
bifrostResponse.ExtraFields.RawRequest = rawRequest
714+
}
715+
689716
// Set raw response if enabled
690717
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
691718
bifrostResponse.ExtraFields.RawResponse = rawResponse
@@ -735,6 +762,7 @@ func (provider *AnthropicProvider) ResponsesStream(ctx context.Context, postHook
735762
jsonBody,
736763
headers,
737764
provider.networkConfig.ExtraHeaders,
765+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
738766
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
739767
provider.GetProviderKey(),
740768
postHookRunner,
@@ -752,6 +780,7 @@ func HandleAnthropicResponsesStream(
752780
jsonBody []byte,
753781
headers map[string]string,
754782
extraHeaders map[string]string,
783+
sendBackRawRequest bool,
755784
sendBackRawResponse bool,
756785
providerName schemas.ModelProvider,
757786
postHookRunner schemas.PostHookRunner,
@@ -959,6 +988,10 @@ func HandleAnthropicResponsesStream(
959988
response.Response = &schemas.BifrostResponsesResponse{}
960989
}
961990
response.Response.Usage = usage
991+
// Set raw request if enabled
992+
if sendBackRawRequest {
993+
providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonBody)
994+
}
962995
response.ExtraFields.Latency = time.Since(startTime).Milliseconds()
963996
ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
964997
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
@@ -472,6 +492,7 @@ func (provider *AzureProvider) ChatCompletionStream(ctx context.Context, postHoo
472492
jsonData,
473493
authHeader,
474494
provider.networkConfig.ExtraHeaders,
495+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
475496
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
476497
provider.GetProviderKey(),
477498
postHookRunner,
@@ -499,6 +520,7 @@ func (provider *AzureProvider) ChatCompletionStream(ctx context.Context, postHoo
499520
request,
500521
authHeader,
501522
provider.networkConfig.ExtraHeaders,
523+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
502524
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
503525
provider.GetProviderKey(),
504526
postHookRunner,
@@ -568,18 +590,19 @@ func (provider *AzureProvider) Responses(ctx context.Context, key schemas.Key, r
568590
}
569591

570592
response := &schemas.BifrostResponsesResponse{}
593+
var rawRequest interface{}
571594
var rawResponse interface{}
572595

573596
if schemas.IsAnthropicModel(deployment) {
574597
anthropicResponse := anthropic.AcquireAnthropicMessageResponse()
575598
defer anthropic.ReleaseAnthropicMessageResponse(anthropicResponse)
576-
rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, provider.sendBackRawResponse)
599+
rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
577600
if bifrostErr != nil {
578601
return nil, bifrostErr
579602
}
580603
response = anthropicResponse.ToBifrostResponsesResponse()
581604
} else {
582-
rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, response, provider.sendBackRawResponse)
605+
rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
583606
if bifrostErr != nil {
584607
return nil, bifrostErr
585608
}
@@ -591,6 +614,11 @@ func (provider *AzureProvider) Responses(ctx context.Context, key schemas.Key, r
591614
response.ExtraFields.Latency = latency.Milliseconds()
592615
response.ExtraFields.RequestType = schemas.ResponsesRequest
593616

617+
// Set raw request if enabled
618+
if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) {
619+
response.ExtraFields.RawRequest = rawRequest
620+
}
621+
594622
// Set raw response if enabled
595623
if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) {
596624
response.ExtraFields.RawResponse = rawResponse
@@ -647,6 +675,7 @@ func (provider *AzureProvider) ResponsesStream(ctx context.Context, postHookRunn
647675
jsonData,
648676
authHeader,
649677
provider.networkConfig.ExtraHeaders,
678+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
650679
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
651680
provider.GetProviderKey(),
652681
postHookRunner,
@@ -675,6 +704,7 @@ func (provider *AzureProvider) ResponsesStream(ctx context.Context, postHookRunn
675704
request,
676705
authHeader,
677706
provider.networkConfig.ExtraHeaders,
707+
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
678708
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
679709
provider.GetProviderKey(),
680710
postHookRunner,
@@ -724,7 +754,7 @@ func (provider *AzureProvider) Embedding(ctx context.Context, key schemas.Key, r
724754
response := &schemas.BifrostEmbeddingResponse{}
725755

726756
// Use enhanced response handler with pre-allocated response
727-
rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, provider.sendBackRawResponse)
757+
rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, response, jsonData, provider.sendBackRawRequest, provider.sendBackRawResponse)
728758
if bifrostErr != nil {
729759
return nil, bifrostErr
730760
}
@@ -735,6 +765,12 @@ func (provider *AzureProvider) Embedding(ctx context.Context, key schemas.Key, r
735765
response.ExtraFields.ModelDeployment = deployment
736766
response.ExtraFields.RequestType = schemas.EmbeddingRequest
737767

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

0 commit comments

Comments
 (0)