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
74 changes: 41 additions & 33 deletions transports/bifrost-http/integrations/genai/router.go
Original file line number Diff line number Diff line change
@@ -1,63 +1,71 @@
package genai

import (
"encoding/json"
"fmt"
"strings"

"github.com/fasthttp/router"
bifrost "github.com/maximhq/bifrost/core"
"github.com/maximhq/bifrost/transports/bifrost-http/lib"
"github.com/maximhq/bifrost/core/schemas"
"github.com/maximhq/bifrost/transports/bifrost-http/integrations"
"github.com/valyala/fasthttp"
)

// GenAIRouter holds route registrations for genai endpoints.
type GenAIRouter struct {
client *bifrost.Bifrost
*integrations.GenericRouter
}

// NewGenAIRouter creates a new GenAIRouter with the given bifrost client.
func NewGenAIRouter(client *bifrost.Bifrost) *GenAIRouter {
return &GenAIRouter{client: client}
}
routes := []integrations.RouteConfig{
{
Path: "/genai/v1beta/models/{model}",
Method: "POST",
RequestType: &GeminiChatRequest{},
RequestConverter: func(req interface{}) *schemas.BifrostRequest {
if geminiReq, ok := req.(*GeminiChatRequest); ok {
return geminiReq.ConvertToBifrostRequest()
}
return nil
},
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
ResponseFunc: func(resp *schemas.BifrostResponse) interface{} {
return DeriveGenAIFromBifrostResponse(resp)
},
PreCallback: extractAndSetModelFromURL,
},
}

// RegisterRoutes registers all genai routes on the given router.
func (g *GenAIRouter) RegisterRoutes(r *router.Router) {
r.POST("/genai/v1beta/models/{model}", g.handleChatCompletion)
return &GenAIRouter{
GenericRouter: integrations.NewGenericRouter(client, routes),
}
}

// handleChatCompletion handles POST /genai/v1beta/models/{model}
func (g *GenAIRouter) handleChatCompletion(ctx *fasthttp.RequestCtx) {
// extractAndSetModelFromURL extracts model from URL and sets it in the request
func extractAndSetModelFromURL(ctx *fasthttp.RequestCtx, req interface{}) error {
model := ctx.UserValue("model")
if model == nil {
ctx.SetStatusCode(fasthttp.StatusBadRequest)
ctx.SetBodyString("Model parameter is required")
return
return fmt.Errorf("model parameter is required")
}

modelStr := model.(string)
modelStr = modelStr[:len(modelStr)-len(":generateContent")]
// Remove :generateContent suffix if present
modelStr = strings.TrimSuffix(modelStr, ":generateContent")
// Remove trailing colon if present
if len(modelStr) > 0 && modelStr[len(modelStr)-1] == ':' {
modelStr = modelStr[:len(modelStr)-1]
}

var req GeminiChatRequest
if err := json.Unmarshal(ctx.PostBody(), &req); err != nil {
ctx.SetStatusCode(fasthttp.StatusBadRequest)
json.NewEncoder(ctx).Encode(err)
return
// Add google/ prefix for Bifrost if not already present
processedModel := modelStr
if !strings.HasPrefix(modelStr, "google/") {
processedModel = "google/" + modelStr
}

bifrostReq := req.ConvertToBifrostRequest("google/" + modelStr)

bifrostCtx := lib.ConvertToBifrostContext(ctx)

result, err := g.client.ChatCompletionRequest(*bifrostCtx, bifrostReq)
if err != nil {
ctx.SetStatusCode(fasthttp.StatusInternalServerError)
json.NewEncoder(ctx).Encode(err)
return
// Set the model in the request
if geminiReq, ok := req.(*GeminiChatRequest); ok {
geminiReq.Model = processedModel
return nil
}

genAIResponse := DeriveGenAIFromBifrostResponse(result)
ctx.SetStatusCode(fasthttp.StatusOK)
ctx.SetContentType("application/json")
json.NewEncoder(ctx).Encode(genAIResponse)
return fmt.Errorf("invalid request type for GenAI")
}
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
5 changes: 3 additions & 2 deletions transports/bifrost-http/integrations/genai/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
var fnTypePtr = bifrost.Ptr(string(schemas.ToolChoiceTypeFunction))

type GeminiChatRequest struct {
Model string `json:"model,omitempty"` // Model field for explicit model specification
Contents []genai_sdk.Content `json:"contents"`
GenerationConfig genai_sdk.GenerationConfig `json:"generationConfig,omitempty"`
SafetySettings []genai_sdk.SafetySetting `json:"safetySettings,omitempty"`
Expand All @@ -20,10 +21,10 @@ type GeminiChatRequest struct {
Labels map[string]string `json:"labels,omitempty"`
}

func (r *GeminiChatRequest) ConvertToBifrostRequest(modelStr string) *schemas.BifrostRequest {
func (r *GeminiChatRequest) ConvertToBifrostRequest() *schemas.BifrostRequest {
bifrostReq := &schemas.BifrostRequest{
Provider: schemas.Vertex,
Model: modelStr,
Model: r.Model,
Input: schemas.RequestInput{
ChatCompletionInput: &[]schemas.BifrostMessage{},
},
Expand Down
37 changes: 37 additions & 0 deletions transports/bifrost-http/integrations/openai/router.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
package openai

import (
bifrost "github.com/maximhq/bifrost/core"
"github.com/maximhq/bifrost/core/schemas"
"github.com/maximhq/bifrost/transports/bifrost-http/integrations"
)

// OpenAIRouter holds route registrations for OpenAI endpoints.
// It supports standard chat completions and image-enabled vision capabilities.
type OpenAIRouter struct {
*integrations.GenericRouter
}

// NewOpenAIRouter creates a new OpenAIRouter with the given bifrost client.
func NewOpenAIRouter(client *bifrost.Bifrost) *OpenAIRouter {
routes := []integrations.RouteConfig{
{
Path: "/openai/v1/chat/completions",
Method: "POST",
RequestType: &OpenAIChatRequest{},
RequestConverter: func(req interface{}) *schemas.BifrostRequest {
if openaiReq, ok := req.(*OpenAIChatRequest); ok {
return openaiReq.ConvertToBifrostRequest()
}
return nil
},
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
ResponseFunc: func(resp *schemas.BifrostResponse) interface{} {
return DeriveOpenAIFromBifrostResponse(resp)
},
},
}

return &OpenAIRouter{
GenericRouter: integrations.NewGenericRouter(client, routes),
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Loading