Skip to content

Commit a825810

Browse files
Update middlewares.go
1 parent 2e6cb90 commit a825810

File tree

119 files changed

+4494
-2376
lines changed

Some content is hidden

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

119 files changed

+4494
-2376
lines changed

Makefile

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@ help: ## Show this help message
3737
cleanup-enterprise: ## Clean up enterprise directories if present
3838
@echo "$(GREEN)Cleaning up enterprise...$(NC)"
3939
@if [ -d "ui/app/enterprise" ]; then rm -rf ui/app/enterprise; fi
40-
@if [ -d "ui/app/enterprise" ]; then rm -rf ui/app/enterprise; fi
4140
@echo "$(GREEN)Enterprise cleaned up$(NC)"
4241

4342
install-ui: cleanup-enterprise

core/bifrost.go

Lines changed: 76 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -294,36 +294,87 @@ func (bifrost *Bifrost) TranscriptionStreamRequest(ctx context.Context, req *sch
294294
return bifrost.handleStreamRequest(ctx, req, schemas.TranscriptionStreamRequest)
295295
}
296296

297+
// RemovePlugin removes a plugin from the server.
298+
func (bifrost *Bifrost) RemovePlugin(name string) error {
299+
300+
for {
301+
oldPlugins := bifrost.plugins.Load()
302+
if oldPlugins == nil {
303+
return nil
304+
}
305+
var pluginToCleanup schemas.Plugin
306+
found := false
307+
// Create new slice with replaced plugin
308+
newPlugins := make([]schemas.Plugin, len(*oldPlugins))
309+
copy(newPlugins, *oldPlugins)
310+
for i, p := range newPlugins {
311+
if p.GetName() == name {
312+
pluginToCleanup = p
313+
bifrost.logger.Debug("removing plugin %s", name)
314+
newPlugins = append(newPlugins[:i], newPlugins[i+1:]...)
315+
found = true
316+
break
317+
}
318+
}
319+
if !found {
320+
return nil
321+
}
322+
if pluginToCleanup != nil {
323+
// Atomic compare-and-swap
324+
if bifrost.plugins.CompareAndSwap(oldPlugins, &newPlugins) {
325+
// Cleanup the old plugin
326+
err := pluginToCleanup.Cleanup()
327+
if err != nil {
328+
bifrost.logger.Warn("failed to cleanup old plugin %s: %v", pluginToCleanup.GetName(), err)
329+
}
330+
return nil
331+
}
332+
}
333+
// Retrying as swapping did not work
334+
}
335+
}
336+
297337
// ReloadPlugin reloads a plugin with new instance
298338
// During the reload - it's stop the world phase where we take a global lock on the plugin mutex
299-
func (bifrost *Bifrost) ReloadPlugin(plugin schemas.Plugin) error {
339+
func (bifrost *Bifrost) ReloadPlugin(plugin schemas.Plugin) error {
300340
for {
341+
var pluginToCleanup schemas.Plugin
342+
found := false
301343
oldPlugins := bifrost.plugins.Load()
302344
if oldPlugins == nil {
303345
return nil
304346
}
305347
// Create new slice with replaced plugin
306348
newPlugins := make([]schemas.Plugin, len(*oldPlugins))
307349
copy(newPlugins, *oldPlugins)
308-
found := false
309350
for i, p := range newPlugins {
310351
if p.GetName() == plugin.GetName() {
352+
// Cleaning up old plugin before replacing it
353+
pluginToCleanup = p
354+
bifrost.logger.Debug("replacing plugin %s with new instance", plugin.GetName())
311355
newPlugins[i] = plugin
312356
found = true
313357
break
314358
}
315359
}
316-
if !found{
360+
if !found {
317361
// This means that user is adding a new plugin
362+
bifrost.logger.Debug("adding new plugin %s", plugin.GetName())
318363
newPlugins = append(newPlugins, plugin)
319364
}
320365
// Atomic compare-and-swap
321366
if bifrost.plugins.CompareAndSwap(oldPlugins, &newPlugins) {
367+
// Cleanup the old plugin
368+
if found && pluginToCleanup != nil {
369+
err := pluginToCleanup.Cleanup()
370+
if err != nil {
371+
bifrost.logger.Warn("failed to cleanup old plugin %s: %v", pluginToCleanup.GetName(), err)
372+
}
373+
}
322374
return nil
323375
}
324376
// Retrying as swapping did not work
325377
}
326-
327378
}
328379

329380
// UpdateProviderConcurrency dynamically updates the queue size and concurrency for an existing provider.
@@ -1023,7 +1074,7 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
10231074

10241075
msg := bifrost.getChannelMessage(*preReq, requestType)
10251076
msg.Context = ctx
1026-
1077+
startTime := time.Now()
10271078
select {
10281079
case queue <- *msg:
10291080
// Message was sent successfully
@@ -1047,9 +1098,14 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
10471098

10481099
var result *schemas.BifrostResponse
10491100
var resp *schemas.BifrostResponse
1101+
pluginCount := len(*bifrost.plugins.Load())
10501102
select {
10511103
case result = <-msg.Response:
1052-
resp, bifrostErr := pipeline.RunPostHooks(&ctx, result, nil, len(*bifrost.plugins.Load()))
1104+
latency := time.Since(startTime).Milliseconds()
1105+
if result.ExtraFields.Latency == nil {
1106+
result.ExtraFields.Latency = &latency
1107+
}
1108+
resp, bifrostErr := pipeline.RunPostHooks(&ctx, result, nil, pluginCount)
10531109
if bifrostErr != nil {
10541110
bifrost.releaseChannelMessage(msg)
10551111
return nil, bifrostErr
@@ -1058,7 +1114,7 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
10581114
return resp, nil
10591115
case bifrostErrVal := <-msg.Err:
10601116
bifrostErrPtr := &bifrostErrVal
1061-
resp, bifrostErrPtr = pipeline.RunPostHooks(&ctx, nil, bifrostErrPtr, len(*bifrost.plugins.Load()))
1117+
resp, bifrostErrPtr = pipeline.RunPostHooks(&ctx, nil, bifrostErrPtr, pluginCount)
10621118
bifrost.releaseChannelMessage(msg)
10631119
if bifrostErrPtr != nil {
10641120
return nil, bifrostErrPtr
@@ -1172,7 +1228,7 @@ func (bifrost *Bifrost) tryStreamRequest(req *schemas.BifrostRequest, ctx contex
11721228
// Marking final chunk
11731229
ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
11741230
// On error we will complete post-hooks
1175-
recoveredResp, recoveredErr := pipeline.RunPostHooks(&ctx, nil, &bifrostErrVal, len(bifrost.plugins))
1231+
recoveredResp, recoveredErr := pipeline.RunPostHooks(&ctx, nil, &bifrostErrVal, len(*bifrost.plugins.Load()))
11761232
bifrost.releaseChannelMessage(msg)
11771233
if recoveredErr != nil {
11781234
return nil, recoveredErr
@@ -1332,7 +1388,7 @@ func handleProviderRequest(provider schemas.Provider, req *ChannelMessage, key s
13321388
case schemas.TextCompletionRequest:
13331389
return provider.TextCompletion(req.Context, req.Model, key, *req.Input.TextCompletionInput, req.Params)
13341390
case schemas.ChatCompletionRequest:
1335-
return provider.ChatCompletion(req.Context, req.Model, key, *req.Input.ChatCompletionInput, req.Params)
1391+
return provider.ChatCompletion(req.Context, req.Model, key, req.Input.ChatCompletionInput, req.Params)
13361392
case schemas.EmbeddingRequest:
13371393
return provider.Embedding(req.Context, req.Model, key, req.Input.EmbeddingInput, req.Params)
13381394
case schemas.SpeechRequest:
@@ -1353,7 +1409,7 @@ func handleProviderRequest(provider schemas.Provider, req *ChannelMessage, key s
13531409
func handleProviderStreamRequest(provider schemas.Provider, req *ChannelMessage, key schemas.Key, postHookRunner schemas.PostHookRunner, reqType schemas.RequestType) (chan *schemas.BifrostStream, *schemas.BifrostError) {
13541410
switch reqType {
13551411
case schemas.ChatCompletionStreamRequest:
1356-
return provider.ChatCompletionStream(req.Context, postHookRunner, req.Model, key, *req.Input.ChatCompletionInput, req.Params)
1412+
return provider.ChatCompletionStream(req.Context, postHookRunner, req.Model, key, req.Input.ChatCompletionInput, req.Params)
13571413
case schemas.SpeechStreamRequest:
13581414
return provider.SpeechStream(req.Context, postHookRunner, req.Model, key, req.Input.SpeechInput, req.Params)
13591415
case schemas.TranscriptionStreamRequest:
@@ -1375,6 +1431,7 @@ func (p *PluginPipeline) RunPreHooks(ctx *context.Context, req *schemas.BifrostR
13751431
var shortCircuit *schemas.PluginShortCircuit
13761432
var err error
13771433
for i, plugin := range p.plugins {
1434+
p.logger.Debug("running pre-hook for plugin %s", plugin.GetName())
13781435
req, shortCircuit, err = plugin.PreHook(ctx, req)
13791436
if err != nil {
13801437
p.preHookErrors = append(p.preHookErrors, err)
@@ -1391,17 +1448,19 @@ func (p *PluginPipeline) RunPreHooks(ctx *context.Context, req *schemas.BifrostR
13911448
// RunPostHooks executes PostHooks in reverse order for the plugins whose PreHook ran.
13921449
// Accepts the response and error, and allows plugins to transform either (e.g., recover from error, or invalidate a response).
13931450
// Returns the final response and error after all hooks. If both are set, error takes precedence unless error is nil.
1394-
func (p *PluginPipeline) RunPostHooks(ctx *context.Context, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError, count int) (*schemas.BifrostResponse, *schemas.BifrostError) {
1451+
// runFrom is the count of plugins whose PreHooks ran; PostHooks will run in reverse from index (runFrom - 1) down to 0
1452+
func (p *PluginPipeline) RunPostHooks(ctx *context.Context, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError, runFrom int) (*schemas.BifrostResponse, *schemas.BifrostError) {
13951453
// Defensive: ensure count is within valid bounds
1396-
if count < 0 {
1397-
count = 0
1454+
if runFrom < 0 {
1455+
runFrom = 0
13981456
}
1399-
if count > len(p.plugins) {
1400-
count = len(p.plugins)
1457+
if runFrom > len(p.plugins) {
1458+
runFrom = len(p.plugins)
14011459
}
14021460
var err error
1403-
for i := count - 1; i >= 0; i-- {
1461+
for i := runFrom - 1; i >= 0; i-- {
14041462
plugin := p.plugins[i]
1463+
p.logger.Debug("running post-hook for plugin %s", plugin.GetName())
14051464
resp, bifrostErr, err = plugin.PostHook(ctx, resp, bifrostErr)
14061465
if err != nil {
14071466
p.postHookErrors = append(p.postHookErrors, err)
@@ -1618,4 +1677,5 @@ func (bifrost *Bifrost) Shutdown() {
16181677
bifrost.logger.Warn(fmt.Sprintf("Error cleaning up plugin: %s", err.Error()))
16191678
}
16201679
}
1680+
bifrost.logger.Info("all request channels closed")
16211681
}

core/changelog.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
<!-- The pattern we follow here is to keep the changelog for the latest version -->
22
<!-- Old changelogs are automatically attached to the GitHub releases -->
33

4-
- Fix: Updates token calculation for streaming responses. #520
4+
- Feature: Adds dynamic reloads for plugins. This removes requirement of restarts on updating plugins.

core/mcp.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -872,9 +872,9 @@ func (m *MCPManager) addMCPToolsToBifrostRequest(ctx context.Context, req *schem
872872
req.Params = &schemas.ModelParameters{}
873873
}
874874
if req.Params.Tools == nil {
875-
req.Params.Tools = &[]schemas.Tool{}
875+
req.Params.Tools = []schemas.Tool{}
876876
}
877-
tools := *req.Params.Tools
877+
tools := req.Params.Tools
878878

879879
// Create a map of existing tool names for O(1) lookup
880880
existingToolsMap := make(map[string]bool)
@@ -890,7 +890,7 @@ func (m *MCPManager) addMCPToolsToBifrostRequest(ctx context.Context, req *schem
890890
existingToolsMap[mcpTool.Function.Name] = true
891891
}
892892
}
893-
req.Params.Tools = &tools
893+
req.Params.Tools = tools
894894

895895
}
896896
return req

core/providers/anthropic.go

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,7 @@ const (
210210
AnthropicDefaultMaxTokens = 4096
211211
)
212212

213-
// mapAnthropicFinishReasonToOpenAI maps Anthropic finish reasons to OpenAI-compatible ones
213+
// MapAnthropicFinishReason maps Anthropic finish reasons to OpenAI-compatible ones
214214
func MapAnthropicFinishReason(anthropicReason string) string {
215215
switch anthropicReason {
216216
case "end_turn":
@@ -516,7 +516,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
516516
Text: *msg.Content.ContentStr,
517517
})
518518
} else if msg.Content.ContentBlocks != nil {
519-
for _, block := range *msg.Content.ContentBlocks {
519+
for _, block := range msg.Content.ContentBlocks {
520520
if block.Text != nil {
521521
systemMessages = append(systemMessages, BedrockAnthropicSystemMessage{
522522
Text: *block.Text,
@@ -547,7 +547,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
547547
"text": *msg.Content.ContentStr,
548548
})
549549
} else if msg.Content.ContentBlocks != nil {
550-
for _, block := range *msg.Content.ContentBlocks {
550+
for _, block := range msg.Content.ContentBlocks {
551551
if block.Text != nil {
552552
toolCallResultContent = append(toolCallResultContent, map[string]interface{}{
553553
"type": "text",
@@ -567,7 +567,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
567567
"text": *msg.Content.ContentStr,
568568
})
569569
} else if msg.Content.ContentBlocks != nil {
570-
for _, block := range *msg.Content.ContentBlocks {
570+
for _, block := range msg.Content.ContentBlocks {
571571
if block.Text != nil && *block.Text != "" {
572572
content = append(content, map[string]interface{}{
573573
"type": "text",
@@ -596,7 +596,7 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
596596

597597
// Add tool calls as content if present
598598
if msg.AssistantMessage != nil && msg.AssistantMessage.ToolCalls != nil {
599-
for _, toolCall := range *msg.AssistantMessage.ToolCalls {
599+
for _, toolCall := range msg.AssistantMessage.ToolCalls {
600600
if toolCall.Function.Name != nil {
601601
var input map[string]interface{}
602602
if toolCall.Function.Arguments != "" {
@@ -639,9 +639,9 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
639639
}
640640

641641
// Transform tools if present
642-
if params != nil && params.Tools != nil && len(*params.Tools) > 0 {
642+
if params != nil && params.Tools != nil && len(params.Tools) > 0 {
643643
var tools []map[string]interface{}
644-
for _, tool := range *params.Tools {
644+
for _, tool := range params.Tools {
645645
tools = append(tools, map[string]interface{}{
646646
"name": tool.Function.Name,
647647
"description": tool.Function.Description,
@@ -791,7 +791,7 @@ func parseAnthropicResponse(response *AnthropicChatResponse, bifrostResponse *sc
791791
if len(toolCalls) > 0 || thinking != "" {
792792
assistantMessage = &schemas.AssistantMessage{}
793793
if len(toolCalls) > 0 {
794-
assistantMessage.ToolCalls = &toolCalls
794+
assistantMessage.ToolCalls = toolCalls
795795
}
796796
if thinking != "" {
797797
assistantMessage.Thought = &thinking
@@ -807,7 +807,7 @@ func parseAnthropicResponse(response *AnthropicChatResponse, bifrostResponse *sc
807807
Message: schemas.BifrostMessage{
808808
Role: schemas.ModelChatMessageRoleAssistant,
809809
Content: schemas.MessageContent{
810-
ContentBlocks: &contentBlocks,
810+
ContentBlocks: contentBlocks,
811811
},
812812
AssistantMessage: assistantMessage,
813813
},

0 commit comments

Comments
 (0)