scriptorium/internal/usecase/runner.go

package usecase

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"time"

	"gitea.maximumdirect.net/eric/scriptorium/internal/artifact"
	"gitea.maximumdirect.net/eric/scriptorium/internal/domain"
	"gitea.maximumdirect.net/eric/scriptorium/internal/llm"
	"gitea.maximumdirect.net/eric/scriptorium/internal/profile"
	"gitea.maximumdirect.net/eric/scriptorium/internal/prompt"
	"gitea.maximumdirect.net/eric/scriptorium/internal/validate"
)

var (
	ErrInvalidRequest = errors.New("invalid run request")
	ErrProfileLoad    = errors.New("failed to load profile")
	ErrArtifactLoad   = errors.New("failed to load artifact")
	ErrPromptRender   = errors.New("failed to render prompt")
	ErrLLMGenerate    = errors.New("failed to generate output")
	ErrValidation     = errors.New("failed to validate output")
)

// Runner executes the Scriptorium core use case.
type Runner struct {
	profiles  profile.Repository
	artifacts artifact.Reader
	renderer  prompt.Renderer
	llm       llm.Client
	validator validate.Validator
	repairer  OutputRepairer
}

func NewRunner(
	profiles profile.Repository,
	artifacts artifact.Reader,
	renderer prompt.Renderer,
	llmClient llm.Client,
	validator validate.Validator,
) *Runner {
	return NewRunnerWithRepairer(profiles, artifacts, renderer, llmClient, validator, nil)
}

func NewRunnerWithRepairer(
	profiles profile.Repository,
	artifacts artifact.Reader,
	renderer prompt.Renderer,
	llmClient llm.Client,
	validator validate.Validator,
	repairer OutputRepairer,
) *Runner {
	return &Runner{
		profiles:  profiles,
		artifacts: artifacts,
		renderer:  renderer,
		llm:       llmClient,
		validator: validator,
		repairer:  repairer,
	}
}

func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunResult, error) {
	if strings.TrimSpace(req.ProfileID) == "" {
		return nil, fmt.Errorf("%w: profile id is required", ErrInvalidRequest)
	}

	start := time.Now().UTC()

	prof, err := r.profiles.GetProfile(ctx, req.ProfileID, req.ProfileVersion)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrProfileLoad, err)
	}

	effectiveModel := mergeModelTarget(prof.ModelDefaults, req.Model)
	effectiveContract := resolveOutputContract(prof, req.Validation)

	resolvedInputs := make(map[string]*domain.Artifact, len(req.Inputs))
	inputHashes := make(map[string]string, len(req.Inputs))
	for name, ref := range req.Inputs {
		art, readErr := r.artifacts.Read(ctx, ref)
		if readErr != nil {
			return nil, fmt.Errorf("%w: input %q: %w", ErrArtifactLoad, name, readErr)
		}
		if art.Name == "" {
			art.Name = name
		}
		resolvedInputs[name] = art
		inputHashes[name] = art.Hash
	}

	renderedPrompt, err := r.renderer.Render(ctx, prof, resolvedInputs, req.Vars)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrPromptRender, err)
	}

	promptHash := hashRenderedPrompt(*renderedPrompt)

	genResp, err := r.llm.Generate(ctx, domain.GenerateRequest{
		Prompt: *renderedPrompt,
		Target: effectiveModel,
	})
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrLLMGenerate, err)
	}

	outputArtifact := buildOutputArtifact(genResp.Content, effectiveContract.Format)
	validationResult, err := r.validateOutput(ctx, &outputArtifact, effectiveContract, 0)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrValidation, err)
	}

	if r.shouldAttemptRepair(effectiveContract, validationResult) {
		attemptsUsed := 0
		for attemptsUsed < effectiveContract.RepairAttempts && validationResult.Status == domain.ValidationFailed {
			attemptsUsed++

			repairResp, repairErr := r.repairer.Repair(ctx, RepairRequest{
				PreviousOutput:   genResp.Content,
				ValidationErrors: validationResult.Errors,
				Target:           effectiveModel,
				Attempt:          attemptsUsed,
				MaxAttempts:      effectiveContract.RepairAttempts,
				Mode:             effectiveContract.ValidationMode,
			})
			if repairErr != nil {
				return nil, fmt.Errorf("%w: %w", ErrValidation, repairErr)
			}
			if repairResp == nil {
				return nil, fmt.Errorf("%w: repairer returned nil response", ErrValidation)
			}

			genResp = repairResp
			outputArtifact = buildOutputArtifact(genResp.Content, effectiveContract.Format)

			validationResult, err = r.validateOutput(ctx, &outputArtifact, effectiveContract, attemptsUsed)
			if err != nil {
				return nil, fmt.Errorf("%w: %w", ErrValidation, err)
			}
		}
	}

	end := time.Now().UTC()

	return &domain.RunResult{
		Artifact:       outputArtifact,
		RawOutput:      genResp.Content,
		Validation:     validationResult,
		ProfileID:      prof.ID,
		ProfileVersion: prof.Version,
		ModelName:      effectiveModel.Model,
		Endpoint:       effectiveModel.Endpoint,
		InputHashes:    inputHashes,
		PromptHash:     promptHash,
		Usage:          genResp.Usage,
		StartTime:      start,
		EndTime:        end,
	}, nil
}

func (r *Runner) validateOutput(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract, attemptsUsed int) (domain.ValidationResult, error) {
	if r.validator == nil || contract.ValidationMode == domain.ValidationNone {
		return domain.ValidationResult{
			Status:         domain.ValidationSkipped,
			Mode:           contract.ValidationMode,
			SchemaPath:     contract.SchemaPath,
			RepairAttempts: attemptsUsed,
			IsValid:        true,
		}, nil
	}

	res, err := r.validator.Validate(ctx, artifact, contract)
	if err != nil {
		return domain.ValidationResult{}, err
	}
	res.RepairAttempts = attemptsUsed
	return res, nil
}

func (r *Runner) shouldAttemptRepair(contract domain.OutputContract, validationResult domain.ValidationResult) bool {
	if r.repairer == nil {
		return false
	}
	if contract.RepairAttempts <= 0 {
		return false
	}
	if validationResult.Status != domain.ValidationFailed {
		return false
	}
	return contract.ValidationMode == domain.ValidationJSON || contract.ValidationMode == domain.ValidationJSONSchema
}

func mergeModelTarget(base domain.ModelTarget, override *domain.ModelTarget) domain.ModelTarget {
	if override == nil {
		return base
	}

	out := base
	if override.Endpoint != "" {
		out.Endpoint = override.Endpoint
	}
	if override.Model != "" {
		out.Model = override.Model
	}
	if override.Temperature != 0 {
		out.Temperature = override.Temperature
	}
	if override.MaxTokens != 0 {
		out.MaxTokens = override.MaxTokens
	}
	if override.TopP != 0 {
		out.TopP = override.TopP
	}
	if override.TimeoutSeconds != 0 {
		out.TimeoutSeconds = override.TimeoutSeconds
	}
	return out
}

func resolveOutputContract(prof *domain.PromptProfile, override *domain.OutputContract) domain.OutputContract {
	contract := prof.Validation
	if contract.Format == "" {
		contract.Format = prof.OutputFormat
	}
	if override != nil {
		contract = *override
	}
	if contract.Format == "" {
		contract.Format = domain.FormatText
	}
	return contract
}

func hashRenderedPrompt(p domain.RenderedPrompt) string {
	var b strings.Builder
	for _, msg := range p.Messages {
		b.WriteString(msg.Role)
		b.WriteByte('\n')
		b.WriteString(msg.Content)
		b.WriteString("\n---\n")
	}
	h := sha256.Sum256([]byte(b.String()))
	return hex.EncodeToString(h[:])
}

func buildOutputArtifact(content string, format domain.OutputFormat) domain.Artifact {
	body := []byte(content)
	hash := sha256.Sum256(body)

	contentType := "text/plain"
	switch format {
	case domain.FormatMarkdown:
		contentType = "text/markdown"
	case domain.FormatJSON:
		contentType = "application/json"
	}

	return domain.Artifact{
		Name:        "output",
		ContentType: contentType,
		Body:        body,
		Size:        int64(len(body)),
		Hash:        hex.EncodeToString(hash[:]),
	}
}