106 lines
3.5 KiB
Go
106 lines
3.5 KiB
Go
package model
|
|
|
|
// Router maps roles to model names. An unknown role resolves to the configured
|
|
// default model. When a request carries image content and a vision model is
|
|
// configured, the request routes to the vision model. A Router also carries a
|
|
// provider per role (defaulting to the global provider) so a multi-vendor
|
|
// deployment can pair a role's model with the correct protocol; the provider
|
|
// selection is surfaced via RouteProvider.
|
|
type Router struct {
|
|
defaultModel string
|
|
byRole map[Role]string
|
|
visionModel string
|
|
defaultProvider string
|
|
providerByRole map[Role]string
|
|
}
|
|
|
|
// RouterOption configures a Router.
|
|
type RouterOption func(*Router)
|
|
|
|
// WithRoleModel binds a model to a role.
|
|
func WithRoleModel(role Role, model string) RouterOption {
|
|
return func(r *Router) { r.byRole[role] = model }
|
|
}
|
|
|
|
// WithVisionModel configures the dedicated vision model.
|
|
func WithVisionModel(model string) RouterOption {
|
|
return func(r *Router) { r.visionModel = model }
|
|
}
|
|
|
|
// WithDefaultProvider sets the provider returned by RouteProvider for roles with
|
|
// no explicit binding. It defaults to openai.
|
|
func WithDefaultProvider(provider string) RouterOption {
|
|
return func(r *Router) { r.defaultProvider = provider }
|
|
}
|
|
|
|
// WithRoleProvider binds a provider to a role, so a multi-vendor deployment can
|
|
// route a role to a different vendor's protocol (e.g. one role on anthropic, one
|
|
// on openai). Pair it with WithRoleModel for the role's model name.
|
|
func WithRoleProvider(role Role, provider string) RouterOption {
|
|
return func(r *Router) { r.providerByRole[role] = provider }
|
|
}
|
|
|
|
// NewRouter builds a Router with a default model applied to any role that has
|
|
// no explicit binding.
|
|
func NewRouter(defaultModel string, opts ...RouterOption) *Router {
|
|
r := &Router{
|
|
defaultModel: defaultModel,
|
|
byRole: map[Role]string{},
|
|
defaultProvider: ProviderOpenAI,
|
|
providerByRole: map[Role]string{},
|
|
}
|
|
for _, o := range opts {
|
|
o(r)
|
|
}
|
|
return r
|
|
}
|
|
|
|
// Route returns the model for a role, falling back to the default model.
|
|
func (r *Router) Route(role Role) string {
|
|
if m, ok := r.byRole[role]; ok && m != "" {
|
|
return m
|
|
}
|
|
return r.defaultModel
|
|
}
|
|
|
|
// RouteProvider returns the provider for a role, falling back to the router's
|
|
// default provider. It is the routing hook for a multi-vendor deployment: each
|
|
// role can carry its own provider+model pairing, and roles without a binding fall
|
|
// back to the global provider. The runtime/loop is expected to build a per-role
|
|
// adapter or chain from this decision (the provider-aware factory is the builder).
|
|
func (r *Router) RouteProvider(role Role) string {
|
|
if p, ok := r.providerByRole[role]; ok && p != "" {
|
|
return p
|
|
}
|
|
return r.defaultProvider
|
|
}
|
|
|
|
// RouteRequest picks the model for a request and reports whether the request
|
|
// needs a vision model. When the request carries image content and a vision
|
|
// model is configured, the vision model is returned with needsVision set to
|
|
// true. Otherwise the role of the last message drives routing (user by
|
|
// default).
|
|
func (r *Router) RouteRequest(req Request) (model string, needsVision bool) {
|
|
role := RoleUser
|
|
if n := len(req.Messages); n > 0 {
|
|
role = req.Messages[n-1].Role
|
|
}
|
|
needsVision = HasImage(req.Messages)
|
|
if needsVision && r.visionModel != "" {
|
|
return r.visionModel, true
|
|
}
|
|
return r.Route(role), needsVision
|
|
}
|
|
|
|
// HasImage reports whether any message carries an image content part.
|
|
func HasImage(msgs []Message) bool {
|
|
for _, m := range msgs {
|
|
for _, part := range m.Content {
|
|
if part.Type == ContentTypeImage {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
return false
|
|
}
|