Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
125 changes: 91 additions & 34 deletions core/bifrost.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"time"

"github.com/google/uuid"
"github.com/maximhq/bifrost/core/mcp"
"github.com/maximhq/bifrost/core/providers/anthropic"
"github.com/maximhq/bifrost/core/providers/azure"
"github.com/maximhq/bifrost/core/providers/bedrock"
Expand Down Expand Up @@ -63,7 +64,8 @@ type Bifrost struct {
pluginPipelinePool sync.Pool // Pool for PluginPipeline objects
bifrostRequestPool sync.Pool // Pool for BifrostRequest objects
logger schemas.Logger // logger instance, default logger is used if not provided
mcpManager *MCPManager // MCP integration manager (nil if MCP not configured)
mcpManager *mcp.MCPManager // MCP integration manager (nil if MCP not configured)
mcpInitOnce sync.Once // Ensures MCP manager is initialized only once
dropExcessRequests atomic.Bool // If true, in cases where the queue is full, requests will not wait for the queue to be empty and will be dropped instead.
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
keySelector schemas.KeySelector // Custom key selector function
}
Expand Down Expand Up @@ -176,13 +178,10 @@ func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) {

// Initialize MCP manager if configured
if config.MCPConfig != nil {
mcpManager, err := newMCPManager(bifrostCtx, *config.MCPConfig, bifrost.logger)
if err != nil {
bifrost.logger.Warn(fmt.Sprintf("failed to initialize MCP manager: %v", err))
} else {
bifrost.mcpManager = mcpManager
bifrost.mcpInitOnce.Do(func() {
bifrost.mcpManager = mcp.NewMCPManager(bifrostCtx, *config.MCPConfig, bifrost.logger)
bifrost.logger.Info("MCP integration initialized successfully")
}
})
}

// Create buffered channels for each provider and start workers
Expand Down Expand Up @@ -492,8 +491,7 @@ func (bifrost *Bifrost) TextCompletionStreamRequest(ctx context.Context, req *sc
return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// ChatCompletionRequest sends a chat completion request to the specified provider.
func (bifrost *Bifrost) ChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
func (bifrost *Bifrost) makeChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
if req == nil {
return nil, &schemas.BifrostError{
IsBifrostError: false,
Expand All @@ -519,10 +517,35 @@ func (bifrost *Bifrost) ChatCompletionRequest(ctx context.Context, req *schemas.
if err != nil {
return nil, err
}
//TODO: Release the response

return response.ChatResponse, nil
}

// ChatCompletionRequest sends a chat completion request to the specified provider.
func (bifrost *Bifrost) ChatCompletionRequest(ctx context.Context, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
// If ctx is nil, use the bifrost context (defensive check for mcp agent mode)
if ctx == nil {
ctx = bifrost.ctx
}

response, err := bifrost.makeChatCompletionRequest(ctx, req)
if err != nil {
return nil, err
}

// Check if we should enter agent mode
if bifrost.mcpManager != nil {
return bifrost.mcpManager.CheckAndExecuteAgentForChatRequest(
&ctx,
req,
response,
bifrost.makeChatCompletionRequest,
)
}

return response, nil
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

// ChatCompletionStreamRequest sends a chat completion stream request to the specified provider.
func (bifrost *Bifrost) ChatCompletionStreamRequest(ctx context.Context, req *schemas.BifrostChatRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
if req == nil {
Expand All @@ -549,8 +572,7 @@ func (bifrost *Bifrost) ChatCompletionStreamRequest(ctx context.Context, req *sc
return bifrost.handleStreamRequest(ctx, bifrostReq)
}

// ResponsesRequest sends a responses request to the specified provider.
func (bifrost *Bifrost) ResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
func (bifrost *Bifrost) makeResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
if req == nil {
return nil, &schemas.BifrostError{
IsBifrostError: false,
Expand All @@ -576,10 +598,34 @@ func (bifrost *Bifrost) ResponsesRequest(ctx context.Context, req *schemas.Bifro
if err != nil {
return nil, err
}
//TODO: Release the response
return response.ResponsesResponse, nil
}

// ResponsesRequest sends a responses request to the specified provider.
func (bifrost *Bifrost) ResponsesRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
// If ctx is nil, use the bifrost context (defensive check for mcp agent mode)
if ctx == nil {
ctx = bifrost.ctx
}

response, err := bifrost.makeResponsesRequest(ctx, req)
if err != nil {
return nil, err
}

// Check if we should enter agent mode
if bifrost.mcpManager != nil {
return bifrost.mcpManager.CheckAndExecuteAgentForResponsesRequest(
&ctx,
req,
response,
bifrost.makeResponsesRequest,
)
}

return response, nil
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

// ResponsesStreamRequest sends a responses stream request to the specified provider.
func (bifrost *Bifrost) ResponsesStreamRequest(ctx context.Context, req *schemas.BifrostResponsesRequest) (chan *schemas.BifrostStream, *schemas.BifrostError) {
if req == nil {
Expand Down Expand Up @@ -1089,7 +1135,7 @@ func (bifrost *Bifrost) RegisterMCPTool(name, description string, handler func(a
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

return bifrost.mcpManager.registerTool(name, description, handler, toolSchema)
return bifrost.mcpManager.RegisterTool(name, description, handler, toolSchema)
}

// ExecuteMCPTool executes an MCP tool call and returns the result as a tool message.
Expand All @@ -1112,13 +1158,12 @@ func (bifrost *Bifrost) ExecuteMCPTool(ctx context.Context, toolCall schemas.Cha
}
}

result, err := bifrost.mcpManager.executeTool(ctx, toolCall)
result, err := bifrost.mcpManager.ExecuteTool(ctx, toolCall)
if err != nil {
return nil, &schemas.BifrostError{
IsBifrostError: false,
Error: &schemas.ErrorField{
Message: err.Error(),
Error: err,
},
}
}
Expand All @@ -1141,12 +1186,9 @@ func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
return nil, fmt.Errorf("MCP is not configured in this Bifrost instance")
}

clients, err := bifrost.mcpManager.GetClients()
if err != nil {
return nil, err
}

clients := bifrost.mcpManager.GetClients()
clientsInConfig := make([]schemas.MCPClient, 0, len(clients))

for _, client := range clients {
tools := make([]schemas.ChatToolFunction, 0, len(client.ToolMap))
for _, tool := range client.ToolMap {
Expand Down Expand Up @@ -1192,13 +1234,17 @@ func (bifrost *Bifrost) GetMCPClients() ([]schemas.MCPClient, error) {
// })
func (bifrost *Bifrost) AddMCPClient(config schemas.MCPClientConfig) error {
if bifrost.mcpManager == nil {
manager := &MCPManager{
ctx: bifrost.ctx,
clientMap: make(map[string]*MCPClient),
logger: bifrost.logger,
}
// Use sync.Once to ensure thread-safe initialization
bifrost.mcpInitOnce.Do(func() {
bifrost.mcpManager = mcp.NewMCPManager(bifrost.ctx, schemas.MCPConfig{
ClientConfigs: []schemas.MCPClientConfig{config},
}, bifrost.logger)
})
}

bifrost.mcpManager = manager
// Handle case where initialization succeeded elsewhere but manager is still nil
if bifrost.mcpManager == nil {
return fmt.Errorf("MCP manager is not initialized")
}

return bifrost.mcpManager.AddClient(config)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Expand Down Expand Up @@ -1266,6 +1312,20 @@ func (bifrost *Bifrost) ReconnectMCPClient(id string) error {
return bifrost.mcpManager.ReconnectClient(id)
}

// UpdateToolManagerConfig updates the tool manager config for the MCP manager.
// This allows for hot-reloading of the tool manager config at runtime.
func (bifrost *Bifrost) UpdateToolManagerConfig(maxAgentDepth int, toolExecutionTimeoutInSeconds int) error {
if bifrost.mcpManager == nil {
return fmt.Errorf("MCP is not configured in this Bifrost instance")
}

bifrost.mcpManager.UpdateToolManagerConfig(&schemas.MCPToolManagerConfig{
MaxAgentDepth: maxAgentDepth,
ToolExecutionTimeout: time.Duration(toolExecutionTimeoutInSeconds) * time.Second,
})
return nil
}

// PROVIDER MANAGEMENT

// createBaseProvider creates a provider based on the base provider type
Expand Down Expand Up @@ -1764,11 +1824,8 @@ func (bifrost *Bifrost) tryRequest(ctx context.Context, req *schemas.BifrostRequ
}

// Add MCP tools to request if MCP is configured and requested
if req.RequestType != schemas.EmbeddingRequest &&
req.RequestType != schemas.SpeechRequest &&
req.RequestType != schemas.TranscriptionRequest &&
bifrost.mcpManager != nil {
req = bifrost.mcpManager.addMCPToolsToBifrostRequest(ctx, req)
if bifrost.mcpManager != nil {
req = bifrost.mcpManager.AddToolsToRequest(ctx, req)
}

pipeline := bifrost.getPluginPipeline()
Expand Down Expand Up @@ -1854,7 +1911,7 @@ func (bifrost *Bifrost) tryStreamRequest(ctx context.Context, req *schemas.Bifro

// Add MCP tools to request if MCP is configured and requested
if req.RequestType != schemas.SpeechStreamRequest && req.RequestType != schemas.TranscriptionStreamRequest && bifrost.mcpManager != nil {
req = bifrost.mcpManager.addMCPToolsToBifrostRequest(ctx, req)
req = bifrost.mcpManager.AddToolsToRequest(ctx, req)
}

pipeline := bifrost.getPluginPipeline()
Expand Down Expand Up @@ -2596,7 +2653,7 @@ func (bifrost *Bifrost) Shutdown() {

// Cleanup MCP manager
if bifrost.mcpManager != nil {
err := bifrost.mcpManager.cleanup()
err := bifrost.mcpManager.Cleanup()
if err != nil {
bifrost.logger.Warn(fmt.Sprintf("Error cleaning up MCP manager: %s", err.Error()))
}
Expand Down
6 changes: 3 additions & 3 deletions core/chatbot_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -515,15 +515,15 @@ func (s *ChatSession) SendMessage(message string) (string, error) {
s.history = append(s.history, *assistantMessage)

// Check if assistant wants to use tools
if assistantMessage.ToolCalls != nil && len(assistantMessage.ToolCalls) > 0 {
if len(assistantMessage.ToolCalls) > 0 {
return s.handleToolCalls(*assistantMessage)
}

// Extract text content for regular responses
var responseText string
if assistantMessage.Content.ContentStr != nil {
responseText = *assistantMessage.Content.ContentStr
} else if assistantMessage.Content.ContentBlocks != nil && len(assistantMessage.Content.ContentBlocks) > 0 {
} else if len(assistantMessage.Content.ContentBlocks) > 0 {
var textParts []string
for _, block := range assistantMessage.Content.ContentBlocks {
if block.Text != nil {
Expand Down Expand Up @@ -633,7 +633,7 @@ func (s *ChatSession) synthesizeToolResults() (string, error) {
synthesisRequest := &schemas.BifrostChatRequest{
Provider: s.config.Provider,
Model: s.config.Model,
Input: conversationWithSynthesis,
Input: conversationWithSynthesis,
Params: &schemas.ChatParameters{
Temperature: s.config.Temperature,
MaxCompletionTokens: s.config.MaxTokens,
Expand Down
5 changes: 5 additions & 0 deletions core/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ require (
github.com/aws/aws-sdk-go-v2/config v1.31.13
github.com/aws/smithy-go v1.23.1
github.com/bytedance/sonic v1.14.1
github.com/clarkmcc/go-typescript v0.7.0
github.com/dop251/goja v0.0.0-20251103141225-af2ceb9156d7
github.com/google/uuid v1.6.0
github.com/hajimehoshi/go-mp3 v0.3.4
github.com/mark3labs/mcp-go v0.41.1
Expand Down Expand Up @@ -37,6 +39,9 @@ require (
github.com/bytedance/sonic/loader v0.3.0 // indirect
github.com/cloudwego/base64x v0.1.6 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/dlclark/regexp2 v1.11.4 // indirect
github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect
github.com/google/pprof v0.0.0-20240625030939-27f56978b8b0 // indirect
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
github.com/invopop/jsonschema v0.13.0 // indirect
github.com/klauspost/compress v1.18.1 // indirect
github.com/klauspost/cpuid/v2 v2.3.0 // indirect
Expand Down
14 changes: 14 additions & 0 deletions core/go.sum
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs=
cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10=
github.com/Masterminds/semver/v3 v3.2.1 h1:RN9w6+7QoMeJVGyfmbcgs28Br8cvmnucEXnY0rYXWg0=
github.com/Masterminds/semver/v3 v3.2.1/go.mod h1:qvl/7zhW3nngYb5+80sSMF+FG2BjYrf8m9wsX0PNOMQ=
github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ=
github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
github.com/aws/aws-sdk-go-v2 v1.39.5 h1:e/SXuia3rkFtapghJROrydtQpfQaaUgd1cUvyO1mp2w=
Expand Down Expand Up @@ -40,18 +42,28 @@ github.com/bytedance/sonic v1.14.1 h1:FBMC0zVz5XUmE4z9wF4Jey0An5FueFvOsTKKKtwIl7
github.com/bytedance/sonic v1.14.1/go.mod h1:gi6uhQLMbTdeP0muCnrjHLeCUPyb70ujhnNlhOylAFc=
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/clarkmcc/go-typescript v0.7.0 h1:3nVeaPYyTCWjX6Lf8GoEOTxME2bM5tLuWmwhSZ86uxg=
github.com/clarkmcc/go-typescript v0.7.0/go.mod h1:IZ/nzoVeydAmyfX7l6Jmp8lJDOEnae3jffoXwP4UyYg=
github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M=
github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dlclark/regexp2 v1.11.4 h1:rPYF9/LECdNymJufQKmri9gV604RvvABwgOA8un7yAo=
github.com/dlclark/regexp2 v1.11.4/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/dop251/goja v0.0.0-20251103141225-af2ceb9156d7 h1:jxmXU5V9tXxJnydU5v/m9SG8TRUa/Z7IXODBpMs/P+U=
github.com/dop251/goja v0.0.0-20251103141225-af2ceb9156d7/go.mod h1:MxLav0peU43GgvwVgNbLAj1s/bSGboKkhuULvq/7hx4=
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/go-sourcemap/sourcemap v2.1.3+incompatible h1:W1iEw64niKVGogNgBN3ePyLFfuisuzeidWPMPWmECqU=
github.com/go-sourcemap/sourcemap v2.1.3+incompatible/go.mod h1:F8jJfvm2KbVjc5NqelyYJmf/v5J0dwNLS2mL4sNA1Jg=
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/pprof v0.0.0-20240625030939-27f56978b8b0 h1:e+8XbKB6IMn8A4OAyZccO4pYfB3s7bt6azNIPE7AnPg=
github.com/google/pprof v0.0.0-20240625030939-27f56978b8b0/go.mod h1:K1liHPHnj73Fdn/EKuT8nrFqBihUSKXoLYU0BuatOYo=
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/hajimehoshi/go-mp3 v0.3.4 h1:NUP7pBYH8OguP4diaTZ9wJbUbk3tC0KlfzsEpWmYj68=
Expand Down Expand Up @@ -129,6 +141,8 @@ golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
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=
Loading