77 lines
2.0 KiB
Go
77 lines
2.0 KiB
Go
package router
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"gophergate/internal/db"
|
|
)
|
|
|
|
// Decision holds the result of a routing decision.
|
|
type Decision struct {
|
|
SelectedModel string `json:"selected_model"`
|
|
Strategy string `json:"strategy"` // "heuristic" or "classifier"
|
|
Reason string `json:"reason"`
|
|
}
|
|
|
|
// ClassifierFunc is the callback for classifier-based routing.
|
|
type ClassifierFunc func(ctx context.Context, selectorModel, systemPrompt, userMessage string) (string, error)
|
|
|
|
// Router resolves model groups to concrete models.
|
|
type Router struct {
|
|
groups map[string]db.ModelGroup
|
|
classify ClassifierFunc
|
|
}
|
|
|
|
// New creates a Router. classify may be nil if no classifier groups exist.
|
|
func New(groups []db.ModelGroup, classify ClassifierFunc) *Router {
|
|
r := &Router{
|
|
groups: make(map[string]db.ModelGroup),
|
|
classify: classify,
|
|
}
|
|
for _, g := range groups {
|
|
r.groups[g.ID] = g
|
|
}
|
|
return r
|
|
}
|
|
|
|
// IsGroup returns true if the model name is a group ID.
|
|
func (r *Router) IsGroup(modelID string) bool {
|
|
_, ok := r.groups[modelID]
|
|
return ok
|
|
}
|
|
|
|
// Route resolves a group to a concrete model.
|
|
func (r *Router) Route(ctx context.Context, groupID string, userMessage string) (*Decision, error) {
|
|
group, ok := r.groups[groupID]
|
|
if !ok {
|
|
return nil, fmt.Errorf("unknown model group: %s", groupID)
|
|
}
|
|
|
|
var targets []string
|
|
if err := json.Unmarshal([]byte(group.Targets), &targets); err != nil || len(targets) == 0 {
|
|
return nil, fmt.Errorf("invalid or empty targets for group %s", groupID)
|
|
}
|
|
|
|
switch group.Strategy {
|
|
case "heuristic":
|
|
return routeHeuristic(group, targets, userMessage)
|
|
case "classifier":
|
|
if r.classify == nil {
|
|
return routeHeuristic(group, targets, userMessage)
|
|
}
|
|
return routeClassifier(ctx, r.classify, group, targets, userMessage)
|
|
default:
|
|
return nil, fmt.Errorf("unknown strategy: %s", group.Strategy)
|
|
}
|
|
}
|
|
|
|
// Reload replaces the group definitions without recreating the router.
|
|
func (r *Router) Reload(groups []db.ModelGroup) {
|
|
r.groups = make(map[string]db.ModelGroup)
|
|
for _, g := range groups {
|
|
r.groups[g.ID] = g
|
|
}
|
|
}
|