Cyrene_v2/internal/runtime/model/router.go

package model

import "strings"

// Router maps roles to model names. An unknown role resolves to the configured
// default 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.
//
// Per 裁决 03.4.1 the "role" of per-role routing is the **persona** a turn
// answers as, not the LLM message role: a persona binding wins over the
// message-role binding, and a persona without a binding keeps the previous
// behavior. The binding carrier is the active models.json group
// (runtimeconfig.Group.PersonaModels / PersonaProviders).
//
// Image routing is deliberately NOT this router's job any more (裁决
// 14.4.1/14.4.5): the single authority for layered image routing is
// runtime/multimodal/vision.Router, which the loop calls with the main model's
// real image-input capability. A "has image → switch to the vision model" rule
// here used to bypass that decision.
type Router struct {
	defaultModel    string
	byRole          map[Role]string
	visionModel     string
	defaultProvider string
	providerByRole  map[Role]string
	personaModel    map[string]string
	personaProvider map[string]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. It is kept for
// backwards compatibility (tests and callers that route a pure image request by
// model name) but production no longer sets it: image routing belongs to
// vision.Router (裁决 14.4.1/14.4.5).
func WithVisionModel(model string) RouterOption {
	return func(r *Router) { r.visionModel = model }
}

// WithPersonaModel binds a model to a persona (裁决 03.4.1). A blank persona or
// model is ignored so a partially configured group cannot create a binding that
// routes a persona to an empty model.
func WithPersonaModel(persona, model string) RouterOption {
	return func(r *Router) {
		if p := trimKey(persona); p != "" && model != "" {
			r.personaModel[p] = model
		}
	}
}

// WithPersonaProvider binds a vendor protocol to a persona, so one persona can
// answer through a different vendor than the group default (裁决 03.4.1). Pair it
// with WithPersonaModel.
func WithPersonaProvider(persona, provider string) RouterOption {
	return func(r *Router) {
		if p := trimKey(persona); p != "" && provider != "" {
			r.personaProvider[p] = provider
		}
	}
}

// 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{},
		personaModel:    map[string]string{},
		personaProvider: map[string]string{},
	}
	for _, o := range opts {
		o(r)
	}
	return r
}

// trimKey normalizes a persona id or provider name before it is stored, so
// " cyrene" and "cyrene" cannot become two different bindings.
func trimKey(k string) string { return strings.TrimSpace(k) }

// 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
// carries image content. 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).
//
// Deprecated in the production path: the loop routes by persona via
// RouteRequestForPersona, and image routing belongs to vision.Router (裁决
// 14.4.1/14.4.5). This method stays because it is the ModelRouter contract the
// agent package and the existing tests are built on.
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
}

// PersonaRouter is the optional persona-aware routing extension of ModelRouter
// (裁决 03.4.1). The loop detects it on the agent's router and prefers it over
// RouteRequest, so a per-persona model/provider binding actually reaches the
// request. It is a separate interface so a router that only knows message roles
// keeps working unchanged.
type PersonaRouter interface {
	RouteRequestForPersona(personaID string, req Request) (model string, needsVision bool)
}

// RouteForPersona returns the model bound to persona, falling back to the
// router's default model. Unlike Route (message role) this is the persona-level
// binding: 裁决 03.4.1 defines the "role" of per-role routing as the persona.
func (r *Router) RouteForPersona(persona string) string {
	if m, ok := r.personaModel[trimKey(persona)]; ok && m != "" {
		return m
	}
	return r.defaultModel
}

// RouteProviderForPersona returns the provider bound to persona, falling back to
// the router's default provider (itself the active group's protocol).
func (r *Router) RouteProviderForPersona(persona string) string {
	if p, ok := r.personaProvider[trimKey(persona)]; ok && p != "" {
		return p
	}
	return r.defaultProvider
}

// RouteRequestForPersona routes a request as the given persona: a persona
// binding wins over the message-role decision, and a persona without a binding
// falls back to RouteRequest so nothing changes for the unconfigured case. It
// reports whether the request carries an image, but never switches the model to
// the vision model — that decision belongs to vision.Router (裁决 14.4.1/14.4.5).
func (r *Router) RouteRequestForPersona(personaID string, req Request) (model string, needsVision bool) {
	needsVision = HasImage(req.Messages)
	if m, ok := r.personaModel[trimKey(personaID)]; ok && m != "" {
		return m, needsVision
	}
	role := RoleUser
	if n := len(req.Messages); n > 0 {
		role = req.Messages[n-1].Role
	}
	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
}