71 lines
2.0 KiB
Go
71 lines
2.0 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.
|
|
type Router struct {
|
|
defaultModel string
|
|
byRole map[Role]string
|
|
visionModel 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 }
|
|
}
|
|
|
|
// 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{}}
|
|
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
|
|
}
|
|
|
|
// 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
|
|
}
|