Skip to content

Commit f15544e

Browse files
fbzhongfbzhong
committed
feat: enhance Azure OpenAI compatibility and streaming reliability
- Add Azure OpenAI content filtering support (PromptFilterResults, ContentFilterResults) - Improve OpenAI streaming flow with consistent response handling - Add LLMUsage.Clone() method for safe usage object duplication - Remove unused SetModel method from AnthropicStreamResponse - Fix streaming termination issues with usage-only chunks - Enhance response structure compatibility across providers
1 parent e59957a commit f15544e

File tree

5 files changed

+112
-49
lines changed

5 files changed

+112
-49
lines changed

‎core/providers/anthropic.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1171,7 +1171,6 @@ func handleAnthropicStreaming(
11711171
// Handle delta changes to the top-level message
11721172
if event.Usage != nil && usage != nil {
11731173
usage.OutputTokens = event.Usage.OutputTokens
1174-
usage.CacheCreationInputTokens = event.Usage.CacheCreationInputTokens
11751174
}
11761175

11771176
// Send usage information immediately if present

‎core/providers/openai.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -594,6 +594,9 @@ func handleOpenAIStreaming(
594594

595595
// Handle usage-only chunks (when stream_options include_usage is true)
596596
if len(response.Choices) == 0 && response.Usage != nil {
597+
// Empty choices array.
598+
response.Choices = []schemas.BifrostResponseChoice{}
599+
597600
// This is a usage information chunk at the end of stream
598601
if params != nil {
599602
response.ExtraFields.Params = *params
@@ -619,9 +622,7 @@ func handleOpenAIStreaming(
619622
response.ExtraFields.Provider = providerType
620623

621624
processAndSendResponse(ctx, postHookRunner, &response, responseChan)
622-
623-
// End stream processing after finish reason
624-
break
625+
continue
625626
}
626627

627628
// Handle regular content chunks
@@ -632,6 +633,7 @@ func handleOpenAIStreaming(
632633
response.ExtraFields.Provider = providerType
633634

634635
processAndSendResponse(ctx, postHookRunner, &response, responseChan)
636+
continue
635637
}
636638
}
637639

‎core/schemas/bifrost.go‎

Lines changed: 68 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -378,18 +378,41 @@ type ImageURLStruct struct {
378378

379379
// BifrostResponse represents the complete result from any bifrost request.
380380
type BifrostResponse struct {
381-
ID string `json:"id,omitempty"`
382-
Object string `json:"object,omitempty"` // text.completion, chat.completion, or embedding
383-
Choices []BifrostResponseChoice `json:"choices,omitempty"`
384-
Embedding [][]float32 `json:"data,omitempty"` // Maps to "data" field in provider responses (e.g., OpenAI embedding format)
385-
Speech *BifrostSpeech `json:"speech,omitempty"` // Maps to "speech" field in provider responses (e.g., OpenAI speech format)
386-
Transcribe *BifrostTranscribe `json:"transcribe,omitempty"` // Maps to "transcribe" field in provider responses (e.g., OpenAI transcription format)
387-
Model string `json:"model,omitempty"`
388-
Created int `json:"created,omitempty"` // The Unix timestamp (in seconds).
389-
ServiceTier *string `json:"service_tier,omitempty"`
390-
SystemFingerprint *string `json:"system_fingerprint,omitempty"`
391-
Usage *LLMUsage `json:"usage,omitempty"`
392-
ExtraFields BifrostResponseExtraFields `json:"extra_fields"`
381+
ID string `json:"id,omitempty"`
382+
Object string `json:"object,omitempty"` // text.completion, chat.completion, or embedding
383+
Choices []BifrostResponseChoice `json:"choices,omitempty"`
384+
Embedding [][]float32 `json:"data,omitempty"` // Maps to "data" field in provider responses (e.g., OpenAI embedding format)
385+
Speech *BifrostSpeech `json:"speech,omitempty"` // Maps to "speech" field in provider responses (e.g., OpenAI speech format)
386+
Transcribe *BifrostTranscribe `json:"transcribe,omitempty"` // Maps to "transcribe" field in provider responses (e.g., OpenAI transcription format)
387+
Model string `json:"model,omitempty"`
388+
Created int `json:"created,omitempty"` // The Unix timestamp (in seconds).
389+
ServiceTier *string `json:"service_tier,omitempty"`
390+
SystemFingerprint *string `json:"system_fingerprint,omitempty"`
391+
Usage *LLMUsage `json:"usage,omitempty"`
392+
PromptFilterResults *[]PromptFilterResult `json:"prompt_filter_results,omitempty"` // Azure OpenAI Service
393+
ExtraFields BifrostResponseExtraFields `json:"extra_fields"`
394+
}
395+
396+
// FilterResult represents the result of a content filter.
397+
type FilterResult struct {
398+
Filtered bool `json:"filtered"`
399+
Severity bool `json:"severity"`
400+
}
401+
402+
// ContentFilterResult represents the result of a content filter.
403+
type ContentFilterResult struct {
404+
HateSpeech FilterResult `json:"hate_speech,omitempty"`
405+
SelfHarm FilterResult `json:"self_harm,omitempty"`
406+
Sexual FilterResult `json:"sexual,omitempty"`
407+
Violence FilterResult `json:"violence,omitempty"`
408+
Jailbreak FilterResult `json:"jailbreak,omitempty"`
409+
Profanity FilterResult `json:"profanity,omitempty"`
410+
}
411+
412+
// PromptFilterResult represents the result of a prompt filter.
413+
type PromptFilterResult struct {
414+
PromptIndex int `json:"prompt_index"`
415+
ContentFilterResults *ContentFilterResult `json:"content_filter_results"`
393416
}
394417

395418
// LLMUsage represents token usage information
@@ -401,6 +424,36 @@ type LLMUsage struct {
401424
CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
402425
}
403426

427+
func (u *LLMUsage) Clone() *LLMUsage {
428+
if u == nil {
429+
return nil
430+
}
431+
432+
ret := &LLMUsage{
433+
PromptTokens: u.PromptTokens,
434+
CompletionTokens: u.CompletionTokens,
435+
TotalTokens: u.TotalTokens,
436+
}
437+
438+
if u.TokenDetails != nil {
439+
ret.TokenDetails = &TokenDetails{
440+
CachedTokens: u.TokenDetails.CachedTokens,
441+
AudioTokens: u.TokenDetails.AudioTokens,
442+
}
443+
}
444+
445+
if u.CompletionTokensDetails != nil {
446+
ret.CompletionTokensDetails = &CompletionTokensDetails{
447+
ReasoningTokens: u.CompletionTokensDetails.ReasoningTokens,
448+
AudioTokens: u.CompletionTokensDetails.AudioTokens,
449+
AcceptedPredictionTokens: u.CompletionTokensDetails.AcceptedPredictionTokens,
450+
RejectedPredictionTokens: u.CompletionTokensDetails.RejectedPredictionTokens,
451+
}
452+
}
453+
454+
return ret
455+
}
456+
404457
type AudioLLMUsage struct {
405458
InputTokens int `json:"input_tokens"`
406459
InputTokensDetails *AudioTokenDetails `json:"input_tokens_details,omitempty"`
@@ -501,8 +554,9 @@ type Annotation struct {
501554
// IMPORTANT: Only one of BifrostNonStreamResponseChoice or BifrostStreamResponseChoice
502555
// should be non-nil at a time.
503556
type BifrostResponseChoice struct {
504-
Index int `json:"index"`
505-
FinishReason *string `json:"finish_reason,omitempty"`
557+
Index int `json:"index"`
558+
FinishReason *string `json:"finish_reason,omitempty"`
559+
ContentFilterResults *ContentFilterResult `json:"content_filter_results,omitempty"` // Azure OpenAI Service or DeepSeek
506560

507561
*BifrostNonStreamResponseChoice
508562
*BifrostStreamResponseChoice

‎transports/bifrost-http/integrations/anthropic/types.go‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -135,15 +135,6 @@ func (s *AnthropicStreamResponse) ToSSE() string {
135135
return fmt.Sprintf("event: %s\ndata: %s\n\n", s.Type, string(jsonData))
136136
}
137137

138-
func (s *AnthropicStreamResponse) SetModel(model string) {
139-
if s.Model != nil {
140-
*s.Model = model
141-
}
142-
if s.Message != nil && s.Message.Model != "" {
143-
s.Message.Model = model
144-
}
145-
}
146-
147138
// AnthropicStreamMessage represents the message structure in streaming events
148139
type AnthropicStreamMessage struct {
149140
ID string `json:"id"`

‎transports/bifrost-http/integrations/openai/types.go‎

Lines changed: 39 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package openai
22

33
import (
4+
"encoding/json"
5+
"fmt"
6+
47
"github.com/maximhq/bifrost/core/schemas"
58
"github.com/maximhq/bifrost/transports/bifrost-http/integrations"
69
)
@@ -70,14 +73,15 @@ func (r *OpenAITranscriptionRequest) IsStreamingRequested() bool {
7073

7174
// OpenAIChatResponse represents an OpenAI chat completion response
7275
type OpenAIChatResponse struct {
73-
ID string `json:"id"`
74-
Object string `json:"object"`
75-
Created int `json:"created"`
76-
Model string `json:"model"`
77-
Choices []schemas.BifrostResponseChoice `json:"choices"`
78-
Usage *schemas.LLMUsage `json:"usage,omitempty"` // Reuse schema type
79-
ServiceTier *string `json:"service_tier,omitempty"`
80-
SystemFingerprint *string `json:"system_fingerprint,omitempty"`
76+
ID string `json:"id"`
77+
Object string `json:"object"`
78+
Created int `json:"created"`
79+
Model string `json:"model"`
80+
Choices []schemas.BifrostResponseChoice `json:"choices"`
81+
Usage *schemas.LLMUsage `json:"usage,omitempty"` // Reuse schema type
82+
ServiceTier *string `json:"service_tier,omitempty"`
83+
SystemFingerprint *string `json:"system_fingerprint,omitempty"`
84+
PromptFilterResults *[]schemas.PromptFilterResult `json:"prompt_filter_results,omitempty"`
8185
}
8286

8387
// OpenAIChatError represents an OpenAI chat completion error response
@@ -93,6 +97,11 @@ type OpenAIChatError struct {
9397
} `json:"error"`
9498
}
9599

100+
func (e *OpenAIChatError) ToSSE() string {
101+
data, _ := json.Marshal(e)
102+
return fmt.Sprintf("data: %s\n\n", data)
103+
}
104+
96105
// OpenAIChatErrorStruct represents the error structure of an OpenAI chat completion error response
97106
type OpenAIChatErrorStruct struct {
98107
Type string `json:"type"` // Error type
@@ -104,10 +113,11 @@ type OpenAIChatErrorStruct struct {
104113

105114
// OpenAIStreamChoice represents a choice in a streaming response chunk
106115
type OpenAIStreamChoice struct {
107-
Index int `json:"index"`
108-
Delta *OpenAIStreamDelta `json:"delta,omitempty"`
109-
FinishReason *string `json:"finish_reason,omitempty"`
110-
LogProbs *schemas.LogProbs `json:"logprobs,omitempty"`
116+
Index int `json:"index"`
117+
Delta *OpenAIStreamDelta `json:"delta,omitempty"`
118+
FinishReason *string `json:"finish_reason,omitempty"`
119+
LogProbs *schemas.LogProbs `json:"logprobs,omitempty"`
120+
ContentFilterResults *schemas.ContentFilterResult `json:"content_filter_results,omitempty"`
111121
}
112122

113123
// OpenAIStreamDelta represents the incremental content in a streaming chunk
@@ -128,6 +138,11 @@ type OpenAIStreamResponse struct {
128138
Usage *schemas.LLMUsage `json:"usage,omitempty"`
129139
}
130140

141+
func (r *OpenAIStreamResponse) ToSSE() string {
142+
data, _ := json.Marshal(r)
143+
return fmt.Sprintf("data: %s\n\n", data)
144+
}
145+
131146
// ConvertToBifrostRequest converts an OpenAI chat request to Bifrost format
132147
func (r *OpenAIChatRequest) ConvertToBifrostRequest() *schemas.BifrostRequest {
133148
provider, model := integrations.ParseModelString(r.Model, schemas.OpenAI)
@@ -314,14 +329,15 @@ func DeriveOpenAIFromBifrostResponse(bifrostResp *schemas.BifrostResponse) *Open
314329
}
315330

316331
openaiResp := &OpenAIChatResponse{
317-
ID: bifrostResp.ID,
318-
Object: bifrostResp.Object,
319-
Created: bifrostResp.Created,
320-
Model: bifrostResp.Model,
321-
Choices: bifrostResp.Choices,
322-
Usage: bifrostResp.Usage,
323-
ServiceTier: bifrostResp.ServiceTier,
324-
SystemFingerprint: bifrostResp.SystemFingerprint,
332+
ID: bifrostResp.ID,
333+
Object: bifrostResp.Object,
334+
Created: bifrostResp.Created,
335+
Model: bifrostResp.Model,
336+
Choices: bifrostResp.Choices,
337+
Usage: bifrostResp.Usage,
338+
ServiceTier: bifrostResp.ServiceTier,
339+
SystemFingerprint: bifrostResp.SystemFingerprint,
340+
PromptFilterResults: bifrostResp.PromptFilterResults,
325341
}
326342

327343
return openaiResp
@@ -413,8 +429,9 @@ func DeriveOpenAIStreamFromBifrostResponse(bifrostResp *schemas.BifrostResponse)
413429
// Convert choices to streaming format
414430
for _, choice := range bifrostResp.Choices {
415431
streamChoice := OpenAIStreamChoice{
416-
Index: choice.Index,
417-
FinishReason: choice.FinishReason,
432+
Index: choice.Index,
433+
FinishReason: choice.FinishReason,
434+
ContentFilterResults: choice.ContentFilterResults,
418435
}
419436

420437
var delta *OpenAIStreamDelta

0 commit comments

Comments
 (0)