notarius/internal/framework/pipeline/runner.go

package pipeline

import (
	"context"
	"fmt"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/validate"
)

type ExtractorFactory interface {
	Build(key string) (contracts.Extractor, error)
}

type Runner struct {
	extractors ExtractorFactory
}

func New(extractors ExtractorFactory) *Runner {
	return &Runner{extractors: extractors}
}

type RunInput struct {
	Source        *source.SourceDocument
	ExtractorKeys []string
	LLMClient     contracts.StructuredLLMClient
	Metadata      map[string]any
}

type RunOutput struct {
	Approved []artifacts.Artifact         `json:"approved,omitempty"`
	Rejected []artifacts.RejectedArtifact `json:"rejected,omitempty"`
	Warnings []contracts.Warning          `json:"warnings,omitempty"`
}

func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
	var output RunOutput
	if r == nil {
		return output, fmt.Errorf("runner must not be nil")
	}
	if r.extractors == nil {
		return output, fmt.Errorf("runner extractor factory must not be nil")
	}
	if err := source.ValidateDocument(input.Source); err != nil {
		return output, fmt.Errorf("validate source document: %w", err)
	}
	if len(input.ExtractorKeys) == 0 {
		return output, fmt.Errorf("extractor keys must not be empty")
	}

	nextCandidateIndex := 0
	for _, extractorKey := range input.ExtractorKeys {
		extractor, err := r.extractors.Build(extractorKey)
		if err != nil {
			return output, fmt.Errorf("build extractor %q: %w", extractorKey, err)
		}
		if extractor == nil {
			return output, fmt.Errorf("build extractor %q: returned nil extractor", extractorKey)
		}

		result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
			Source:    input.Source,
			LLMClient: input.LLMClient,
			Metadata:  input.Metadata,
		})
		output.Warnings = append(output.Warnings, result.Warnings...)
		if err != nil {
			return output, fmt.Errorf("extract with extractor %q: %w", extractor.Key(), err)
		}

		candidates, err := normalizeCandidates(extractor, result.Candidates, &nextCandidateIndex)
		if err != nil {
			return output, err
		}

		approved, rejected, warnings, err := runValidators(ctx, extractor, input.Source, candidates, input.Metadata)
		output.Warnings = append(output.Warnings, warnings...)
		output.Rejected = append(output.Rejected, rejected...)
		if err != nil {
			return output, err
		}

		for _, candidate := range approved {
			output.Approved = append(output.Approved, artifacts.ArtifactFromCandidate(candidate))
		}
	}

	return output, nil
}

func normalizeCandidates(extractor contracts.Extractor, candidates []artifacts.ArtifactCandidate, nextIndex *int) ([]artifacts.ArtifactCandidate, error) {
	normalized := make([]artifacts.ArtifactCandidate, 0, len(candidates))
	for _, candidate := range candidates {
		candidate.Index = *nextIndex
		*nextIndex = *nextIndex + 1

		if candidate.ExtractorKey == "" {
			candidate.ExtractorKey = extractor.Key()
		} else if candidate.ExtractorKey != extractor.Key() {
			return nil, fmt.Errorf("candidate extractor_key %q does not match extractor %q", candidate.ExtractorKey, extractor.Key())
		}

		if candidate.ArtifactType == "" {
			candidate.ArtifactType = extractor.ArtifactType()
		} else if candidate.ArtifactType != extractor.ArtifactType() {
			return nil, fmt.Errorf("candidate artifact_type %q does not match extractor %q artifact type %q", candidate.ArtifactType, extractor.Key(), extractor.ArtifactType())
		}

		if candidate.SchemaVersion == "" {
			candidate.SchemaVersion = extractor.SchemaVersion()
		} else if candidate.SchemaVersion != extractor.SchemaVersion() {
			return nil, fmt.Errorf("candidate schema_version %q does not match extractor %q schema version %q", candidate.SchemaVersion, extractor.Key(), extractor.SchemaVersion())
		}

		normalized = append(normalized, candidate)
	}
	return normalized, nil
}

func runValidators(ctx context.Context, extractor contracts.Extractor, doc *source.SourceDocument, candidates []artifacts.ArtifactCandidate, metadata map[string]any) ([]artifacts.ArtifactCandidate, []artifacts.RejectedArtifact, []contracts.Warning, error) {
	eligible := candidates
	var rejected []artifacts.RejectedArtifact
	var warnings []contracts.Warning

	for validatorIndex, validator := range extractor.Validators() {
		if validator == nil {
			return nil, rejected, warnings, fmt.Errorf("extractor %q validator[%d] must not be nil", extractor.Key(), validatorIndex)
		}
		result, err := validator.Validate(ctx, contracts.ValidationRequest{
			Source:     doc,
			Candidates: eligible,
			Metadata:   metadata,
		})
		warnings = append(warnings, result.Warnings...)
		if err != nil {
			return nil, rejected, warnings, fmt.Errorf("validate extractor %q with validator %q: %w", extractor.Key(), validator.Name(), err)
		}
		if result.ValidatorName != validator.Name() {
			return nil, rejected, warnings, fmt.Errorf("validator %q returned result for %q", validator.Name(), result.ValidatorName)
		}
		if err := validate.EnforceDecisionCardinality(eligible, result.Decisions); err != nil {
			return nil, rejected, warnings, fmt.Errorf("validate extractor %q with validator %q: %w", extractor.Key(), validator.Name(), err)
		}

		decisions := make(map[int]contracts.ValidationDecision, len(result.Decisions))
		for _, decision := range result.Decisions {
			decisions[decision.CandidateIndex] = decision
		}

		nextEligible := make([]artifacts.ArtifactCandidate, 0, len(eligible))
		for _, candidate := range eligible {
			decision := decisions[candidate.Index]
			if decision.Approved {
				nextEligible = append(nextEligible, candidate)
				continue
			}
			rejected = append(rejected, artifacts.RejectedArtifact{
				Candidate:     candidate,
				ValidatorName: result.ValidatorName,
				ReasonCode:    decision.ReasonCode,
				Message:       decision.Message,
			})
		}
		eligible = nextEligible
	}

	return eligible, rejected, warnings, nil
}