package main
import (
"fmt"
"strings"
"github.com/maximhq/bifrost/core"
"github.com/maximhq/bifrost/core/schemas"
)
// Define typed arguments for your tool
type CalculatorArgs struct {
Operation string `json:"operation"` // add, subtract, multiply, divide
A float64 `json:"a"`
B float64 `json:"b"`
}
// Define typed tool handler
func calculatorHandler(args CalculatorArgs) (string, error) {
switch strings.ToLower(args.Operation) {
case "add":
return fmt.Sprintf("%.2f", args.A + args.B), nil
case "subtract":
return fmt.Sprintf("%.2f", args.A - args.B), nil
case "multiply":
return fmt.Sprintf("%.2f", args.A * args.B), nil
case "divide":
if args.B == 0 {
return "", fmt.Errorf("cannot divide by zero")
}
return fmt.Sprintf("%.2f", args.A / args.B), nil
default:
return "", fmt.Errorf("unsupported operation: %s", args.Operation)
}
}
func main() {
// Initialize Bifrost (tool registry creates in-process MCP automatically)
client, err := bifrost.Init(schemas.BifrostConfig{
Account: account,
Logger: bifrost.NewDefaultLogger(schemas.LogLevelInfo),
})
// Define tool schema
calculatorSchema := schemas.Tool{
Type: "function",
Function: schemas.Function{
Name: "calculator",
Description: "Perform basic arithmetic operations",
Parameters: schemas.FunctionParameters{
Type: "object",
Properties: map[string]interface{}{
"operation": map[string]interface{}{
"type": "string",
"description": "The operation to perform",
"enum": []string{"add", "subtract", "multiply", "divide"},
},
"a": map[string]interface{}{
"type": "number",
"description": "First number",
},
"b": map[string]interface{}{
"type": "number",
"description": "Second number",
},
},
Required: []string{"operation", "a", "b"},
},
},
}
// Register the typed tool
err = client.RegisterMCPTool("calculator", "Perform arithmetic calculations",
func(args any) (string, error) {
// Convert args to typed struct
calculatorArgs := CalculatorArgs{}
if jsonBytes, err := json.Marshal(args); err == nil {
json.Unmarshal(jsonBytes, &calculatorArgs)
}
return calculatorHandler(calculatorArgs)
}, calculatorSchema)
if err != nil {
panic(fmt.Sprintf("Failed to register tool: %v", err))
}
// Now use the tool in requests
request := &schemas.BifrostRequest{
Provider: schemas.OpenAI,
Model: "gpt-4o-mini",
Input: schemas.RequestInput{
ChatCompletionInput: &[]schemas.BifrostMessage{
{
Role: schemas.ModelChatMessageRoleUser,
Content: schemas.MessageContent{
ContentStr: bifrost.Ptr("Calculate 15.5 + 24.3"),
},
},
},
},
Params: &schemas.ModelParameters{
Temperature: bifrost.Ptr(0.7),
},
}
response, err := client.ChatCompletionRequest(context.Background(), request)
// The model can now use the calculator tool automatically
}