internal/gateway/request_model_resolution.go

Rama main · réplica de lectura del repositorio original.

package gateway

import (
	"context"
	"errors"
	"net/http"
	"strings"

	"github.com/enterpilot/gomodel/internal/core"
)

type modelCountProvider interface {
	ModelCount() int
}

type providerModelRefresher interface {
	RefreshProviderModels(ctx context.Context, providerSelector string) (int, error)
}

type modelRefreshTargetResolver interface {
	ResolveRefreshTarget(requested core.RequestedModelSelector) (core.ModelSelector, bool, error)
}

// ResolvedProviderName returns the configured provider instance name for a selector.
func ResolvedProviderName(provider core.RoutableProvider, selector core.ModelSelector, fallback string) string {
	fallback = strings.TrimSpace(fallback)
	if provider == nil {
		return fallback
	}
	if named, ok := provider.(core.ProviderNameResolver); ok {
		if providerName := strings.TrimSpace(named.GetProviderName(selector.QualifiedModel())); providerName != "" {
			return providerName
		}
	}
	// A model the catalog does not list (a pinned jev version) still belongs to
	// the provider its selector names, not to the fallback's.
	if providerName, providerType := configuredSelectorProvider(provider, selector); providerType != "" {
		return providerName
	}
	return fallback
}

// configuredSelectorProvider returns the provider a selector names explicitly
// and that provider's type, or empty strings when the selector names no
// configured provider.
func configuredSelectorProvider(provider core.RoutableProvider, selector core.ModelSelector) (string, string) {
	providerName := strings.TrimSpace(selector.Provider)
	named, ok := provider.(core.ProviderNameTypeResolver)
	if providerName == "" || !ok {
		return "", ""
	}
	providerType := strings.TrimSpace(named.GetProviderTypeForName(providerName))
	if providerType == "" {
		return "", ""
	}
	return providerName, providerType
}

// ResolvedWorkflowProviderName returns the configured provider name recorded in a resolution.
func ResolvedWorkflowProviderName(resolution *core.RequestModelResolution) string {
	if resolution == nil {
		return ""
	}
	if providerName := strings.TrimSpace(resolution.ProviderName); providerName != "" {
		return providerName
	}
	return strings.TrimSpace(resolution.ResolvedSelector.Provider)
}

// WorkflowProviderNameForType maps a provider type to its configured provider name when available.
func WorkflowProviderNameForType(provider core.RoutableProvider, providerType string) string {
	providerType = strings.TrimSpace(providerType)
	if providerType == "" || provider == nil {
		return ""
	}
	if named, ok := provider.(core.ProviderTypeNameResolver); ok {
		return strings.TrimSpace(named.GetProviderNameForType(providerType))
	}
	return ""
}

// ResolveRequestModelWithAuthorizer resolves and validates a requested selector.
func ResolveRequestModelWithAuthorizer(
	ctx context.Context,
	provider core.RoutableProvider,
	resolver ModelResolver,
	authorizer ModelAuthorizer,
	requested core.RequestedModelSelector,
) (*core.RequestModelResolution, error) {
	requested = core.NewRequestedModelSelector(requested.Model, requested.ProviderHint)

	resolvedSelector, route, aliasApplied, err := resolveExecutionRoute(ctx, provider, resolver, requested)
	refreshed := false
	if err != nil {
		var refreshErr error
		refreshed, refreshErr = refreshProviderModelsForResolution(ctx, provider, resolver, requested, resolvedSelector)
		if refreshErr != nil {
			return nil, refreshErr
		}
		if !refreshed {
			return nil, resolutionError(err)
		}
		resolvedSelector, route, aliasApplied, err = resolveExecutionRoute(ctx, provider, resolver, requested)
		if err != nil {
			return nil, resolutionError(err)
		}
	}
	if resolvedSelector == (core.ModelSelector{}) {
		resolvedSelector, err = requested.Normalize()
		if err != nil {
			return nil, core.NewInvalidRequestError(err.Error(), err)
		}
	}

	resolvedModel := resolvedSelector.QualifiedModel()
	if counted, ok := provider.(modelCountProvider); ok && counted.ModelCount() == 0 {
		if !refreshed {
			var refreshErr error
			refreshed, refreshErr = refreshProviderModelsForResolution(ctx, provider, resolver, requested, resolvedSelector)
			if refreshErr != nil {
				return nil, refreshErr
			}
			if refreshed {
				resolvedSelector, route, aliasApplied, err = resolveExecutionRoute(ctx, provider, resolver, requested)
				if err != nil {
					return nil, resolutionError(err)
				}
				resolvedModel = resolvedSelector.QualifiedModel()
			}
		}
	}
	if counted, ok := provider.(modelCountProvider); ok && counted.ModelCount() == 0 {
		return nil, core.NewProviderError("", http.StatusServiceUnavailable, "model registry not initialized: no models have been loaded yet, retry shortly", nil)
	}
	if !provider.Supports(resolvedModel) {
		if !refreshed {
			var refreshErr error
			refreshed, refreshErr = refreshProviderModelsForResolution(ctx, provider, resolver, requested, resolvedSelector)
			if refreshErr != nil {
				return nil, refreshErr
			}
			if refreshed {
				resolvedSelector, route, aliasApplied, err = resolveExecutionRoute(ctx, provider, resolver, requested)
				if err != nil {
					return nil, resolutionError(err)
				}
				resolvedModel = resolvedSelector.QualifiedModel()
			}
		}
	}
	if !provider.Supports(resolvedModel) {
		return nil, core.NewModelNotFoundError(resolvedModel)
	}
	if authorizer != nil {
		if err := authorizer.ValidateModelAccess(ctx, resolvedSelector); err != nil {
			return nil, err
		}
	}

	resolution := &core.RequestModelResolution{
		Requested:        requested,
		ResolvedSelector: resolvedSelector,
		ProviderType:     strings.TrimSpace(provider.GetProviderType(resolvedModel)),
		ProviderName:     ResolvedProviderName(provider, resolvedSelector, ""),
		AliasApplied:     aliasApplied,
		VirtualModelPath: route,
	}
	if slowdownResolver, ok := resolver.(ModelSlowdownResolver); ok {
		resolution.Slowdown = slowdownResolver.ResolveSlowdown(ctx, requested, resolvedSelector)
	}
	resolution.CacheDerivedSelectors()
	return resolution, nil
}

// resolutionError reports a failed model resolution. A typed gateway error
// keeps its status (the router answers 503 while the model registry is still
// empty, so clients retry); anything else is a selector the caller got wrong.
func resolutionError(err error) error {
	if gatewayErr, ok := errors.AsType[*core.GatewayError](err); ok {
		return gatewayErr
	}
	return core.NewInvalidRequestError(err.Error(), err)
}

func refreshProviderModelsForResolution(
	ctx context.Context,
	provider core.RoutableProvider,
	resolver ModelResolver,
	requested core.RequestedModelSelector,
	resolvedSelector core.ModelSelector,
) (bool, error) {
	refresher, ok := provider.(providerModelRefresher)
	if !ok {
		return false, nil
	}

	providerSelector := strings.TrimSpace(resolvedSelector.Provider)
	if providerSelector == "" {
		if targetResolver, ok := resolver.(modelRefreshTargetResolver); ok {
			selector, ok, err := targetResolver.ResolveRefreshTarget(requested)
			if err != nil {
				return false, err
			}
			if ok {
				providerSelector = strings.TrimSpace(selector.Provider)
			}
		}
	}
	if providerSelector == "" {
		selector, err := requested.Normalize()
		if err != nil {
			return false, nil
		}
		providerSelector = strings.TrimSpace(selector.Provider)
	}
	if providerSelector == "" {
		return false, nil
	}

	_, err := refresher.RefreshProviderModels(ctx, providerSelector)
	return true, err
}

// ResolveExecutionSelector applies explicit and provider-owned selector
// resolution. ctx carries the effective request user path so a resolver that
// implements UserPathModelResolver can apply user_path-scoped redirects.
func ResolveExecutionSelector(
	ctx context.Context,
	provider core.RoutableProvider,
	resolver ModelResolver,
	requested core.RequestedModelSelector,
) (core.ModelSelector, bool, error) {
	selector, _, aliasApplied, err := resolveExecutionRoute(ctx, provider, resolver, requested)
	return selector, aliasApplied, err
}

// resolveExecutionRoute is ResolveExecutionSelector that also returns the
// virtual-model route a RoutedModelResolver reports.
func resolveExecutionRoute(
	ctx context.Context,
	provider core.RoutableProvider,
	resolver ModelResolver,
	requested core.RequestedModelSelector,
) (core.ModelSelector, []string, bool, error) {
	requested = core.NewRequestedModelSelector(requested.Model, requested.ProviderHint)

	var (
		resolvedSelector core.ModelSelector
		route            []string
		aliasApplied     bool
		err              error
	)

	if resolver != nil {
		resolvedSelector, route, aliasApplied, err = resolveRequestedModel(ctx, resolver, requested)
		if err != nil {
			return core.ModelSelector{}, nil, false, err
		}
		requested = core.NewRequestedModelSelector(resolvedSelector.QualifiedModel(), "")
	}

	if providerResolver, ok := provider.(ModelResolver); ok {
		providerSelector, providerChanged, err := providerResolver.ResolveModel(requested)
		if err != nil {
			if resolvedSelector != (core.ModelSelector{}) {
				// Preserve alias targets so callers can refresh the concrete provider before retrying.
				return resolvedSelector, route, aliasApplied, err
			}
			return core.ModelSelector{}, nil, false, err
		}
		return providerSelector, route, aliasApplied || providerChanged, nil
	}

	if resolvedSelector != (core.ModelSelector{}) {
		return resolvedSelector, route, aliasApplied, nil
	}

	resolvedSelector, err = requested.Normalize()
	return resolvedSelector, route, aliasApplied, err
}

// resolveRequestedModel resolves through a RoutedModelResolver or
// UserPathModelResolver when the resolver implements one (applying
// user_path-scoped redirects), falling back to the unscoped ModelResolver
// otherwise. Only a RoutedModelResolver reports a route.
func resolveRequestedModel(ctx context.Context, resolver ModelResolver, requested core.RequestedModelSelector) (core.ModelSelector, []string, bool, error) {
	if routed, ok := resolver.(RoutedModelResolver); ok {
		return routed.ResolveModelRouteForUserPath(ctx, requested)
	}
	if scoped, ok := resolver.(UserPathModelResolver); ok {
		selector, changed, err := scoped.ResolveModelForUserPath(ctx, requested)
		return selector, nil, changed, err
	}
	selector, changed, err := resolver.ResolveModel(requested)
	return selector, nil, changed, err
}