promptkit/internal/usecase/smoke_selection.go

package usecase

import (
	"context"
	"fmt"
	"net/url"

	"gitea.maximumdirect.net/eric/promptkit/internal/defaults"
	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

const (
	smokeSelectionSourceBackendDefault = "backend_default"
	smokeSelectionSourceProfile        = "profile"
	smokeSelectionSourcePrompt         = "prompt"
	smokeSelectionSourceBackend        = "backend"
)

type smokeBackendLookup interface {
	GetBackend(string) (domain.Backend, error)
}

type smokeTestSelection struct {
	target           domain.SmokeTestTarget
	backend          domain.Backend
	transientRetries uint64
}

func (r *Runner) resolveSmokeTestTarget(ctx context.Context, req domain.SmokeTestRequest) (*smokeTestSelection, error) {
	if !validSmokeTestKind(req.Kind) {
		return nil, fmt.Errorf("%w: unknown smoke-test kind %q", ErrInvalidRequest, req.Kind)
	}
	transientRetries, err := domain.ResolveTransientRetries(0, req.TransientRetries)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
	}

	if req.PromptID == "" && req.PromptVersion != "" {
		return nil, fmt.Errorf("%w: prompt version requires a prompt id", ErrInvalidRequest)
	}

	var promptInspection *domain.PromptInspection
	if req.PromptID != "" {
		promptInspection, err = r.resolvePromptInspection(ctx, req.PromptID, req.PromptVersion)
		if err != nil {
			return nil, err
		}
	}

	selection := &smokeTestSelection{transientRetries: transientRetries}
	var selectedProfileID string
	switch {
	case req.ProfileID != "":
		selectedProfileID = req.ProfileID
		selection.target.SelectionSource = smokeSelectionSourceProfile
	case promptInspection != nil && (promptInspection.DefaultProfileID != "" || isInferenceSmokeTest(req.Kind)):
		selectedProfileID = promptInspection.DefaultProfileID
		selection.target.SelectionSource = smokeSelectionSourcePrompt
	case req.BackendID != "" && isInferenceSmokeTest(req.Kind):
		backend, err := r.resolveSmokeBackend(req.BackendID)
		if err != nil {
			return nil, err
		}
		if backend.SmokeTests == nil || backend.SmokeTests.DefaultProfile == nil {
			return nil, fmt.Errorf("%w: backend %q has no default smoke-test profile", ErrProfileRequired, req.BackendID)
		}
		selectedProfileID = *backend.SmokeTests.DefaultProfile
		selection.backend = backend
		selection.target.DefaultSourceBackendID = backend.ID
		selection.target.SelectionSource = smokeSelectionSourceBackendDefault
	case req.BackendID != "":
		backend, err := r.resolveSmokeBackend(req.BackendID)
		if err != nil {
			return nil, err
		}
		selection.backend = backend
		selection.target.SelectionSource = smokeSelectionSourceBackend
	default:
		if isInferenceSmokeTest(req.Kind) {
			return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, ErrProfileRequired)
		}
		return nil, fmt.Errorf("%w: backend id, profile id, or prompt id is required", ErrInvalidRequest)
	}
	if selection.target.SelectionSource != smokeSelectionSourceBackend && selectedProfileID == "" {
		return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, ErrProfileRequired)
	}

	if selectedProfileID != "" {
		selection.backend = domain.Backend{}
		profileSelection, err := r.resolveProfileSelection(ctx, selectedProfileID)
		if err != nil {
			return nil, err
		}
		settings, presence := resolveExecutionTarget(profileSelection.profile, req.Execution)
		settings.APIKey = req.APIKey
		if isInferenceSmokeTest(req.Kind) {
			settings, err = normalizeResolvedExecutionTarget(settings)
		} else {
			settings, err = normalizeResolvedExecutionSettings(settings)
		}
		if err != nil {
			return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
		}
		selection.target.ProfileID = profileSelection.id
		selection.target.BackendID = settings.BackendID
		selection.target.EffectiveSettings = settings
		selection.target.SettingsPresence = presence
		if settings.BackendID != "" {
			selection.backend, err = r.resolveSmokeBackend(settings.BackendID)
			if err != nil {
				return nil, err
			}
		}
	} else {
		settings := defaults.ExecutionTargetDefault()
		settings.BackendID = selection.backend.ID
		settings.Endpoint = selection.backend.Endpoint
		settings.APIKeyEnv = selection.backend.APIKeyEnv
		settings.ExtraParams = copyExtraParams(selection.backend.ExtraParams)
		presence := domain.ExecutionTargetPresence{}
		if req.Execution != nil {
			settings, presence = mergeExecutionTargetOverride(settings, presence, *req.Execution)
		}
		settings.APIKey = req.APIKey
		settings, err = normalizeResolvedExecutionSettings(settings)
		if err != nil {
			return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
		}
		selection.target.BackendID = settings.BackendID
		selection.target.EffectiveSettings = settings
		selection.target.SettingsPresence = presence
	}

	if promptInspection != nil {
		selection.target.PromptID = promptInspection.PromptID
		selection.target.PromptVersion = promptInspection.PromptVersion
	}
	selection.target.Endpoint, err = resolveSmokeTestEndpoint(req.Kind, selection.target.EffectiveSettings.Endpoint, selection.backend.SmokeTests)
	if err != nil {
		return nil, fmt.Errorf("%w: diagnostic endpoint: %w", ErrInvalidRequest, err)
	}
	return selection, nil
}

func (r *Runner) resolveSmokeBackend(id string) (domain.Backend, error) {
	lookup, ok := r.profiles.(smokeBackendLookup)
	if !ok {
		return domain.Backend{}, fmt.Errorf("%w: backend lookup is not configured", ErrProfileLoad)
	}
	backend, err := lookup.GetBackend(id)
	if err != nil {
		return domain.Backend{}, fmt.Errorf("%w: backend %q: %v", ErrProfileLoad, id, err)
	}
	return backend, nil
}

func validSmokeTestKind(kind domain.SmokeTestKind) bool {
	switch kind {
	case domain.SmokeTestReachability, domain.SmokeTestAuthentication,
		domain.SmokeTestTextInference, domain.SmokeTestStructuredOutput:
		return true
	default:
		return false
	}
}

func isInferenceSmokeTest(kind domain.SmokeTestKind) bool {
	return kind == domain.SmokeTestTextInference || kind == domain.SmokeTestStructuredOutput
}

func resolveSmokeTestEndpoint(kind domain.SmokeTestKind, baseEndpoint string, smokeTests *domain.BackendSmokeTestConfig) (string, error) {
	switch kind {
	case domain.SmokeTestReachability:
		if smokeTests != nil && smokeTests.Reachability != nil && smokeTests.Reachability.Endpoint != nil {
			return *smokeTests.Reachability.Endpoint, nil
		}
	case domain.SmokeTestAuthentication:
		if smokeTests != nil && smokeTests.Authentication != nil {
			authentication := smokeTests.Authentication
			if authentication.Endpoint != nil {
				return *authentication.Endpoint, nil
			}
			if authentication.Kind != nil && *authentication.Kind == "openrouter_key" {
				return url.JoinPath(baseEndpoint, "key")
			}
		}
	}
	return baseEndpoint, nil
}