audita/internal/framework/validators/llm_validators.go

package validators

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"

	"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/outcomes"
	"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/stagename"
	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/internal/prompts"
)

type LLMBackedValidator struct {
	name          string
	validatorType LLMValidatorType
	estimator     chunking.TokenEstimator
}

func (v *LLMBackedValidator) Name() string {
	return v.name
}

// SetTokenEstimator allows deterministic test control over batching behavior.
func (v *LLMBackedValidator) SetTokenEstimator(estimator chunking.TokenEstimator) {
	if v == nil || estimator == nil {
		return
	}
	v.estimator = estimator
}

func NewLLMBackedValidator(name string, validatorType LLMValidatorType) (*LLMBackedValidator, error) {
	if validatorPromptMetadata(validatorType).PromptID == "" {
		return nil, fmt.Errorf("unsupported LLM validator type %q", validatorType)
	}
	if strings.TrimSpace(name) == "" {
		return nil, fmt.Errorf("validator name must not be empty")
	}
	return &LLMBackedValidator{
		name:          name,
		validatorType: validatorType,
		estimator:     chunking.NewSimpleTokenEstimator(),
	}, nil
}

func (v *LLMBackedValidator) Validate(ctx context.Context, req Request) (Result, error) {
	if v == nil {
		return Result{}, fmt.Errorf("validator is nil")
	}
	if len(req.CandidateProposal) == 0 {
		return Result{ValidatorName: v.name, Decisions: nil}, nil
	}
	if req.LLMClient == nil {
		return Result{}, fmt.Errorf("LLM-backed validator %q requires a structured LLM client", v.name)
	}
	validationReq, immediate := BuildLLMValidationRequest(v.name, v.validatorType, req)
	if len(validationReq.Items) == 0 {
		all := append([]Decision(nil), immediate...)
		sort.SliceStable(all, func(i, j int) bool { return all[i].ProposalIndex < all[j].ProposalIndex })
		if err := EnforceDecisionCardinality(req.CandidateProposal, all); err != nil {
			return Result{}, err
		}
		return Result{ValidatorName: v.name, Decisions: all}, nil
	}

	maxTokens := config.DefaultValidationMaxPromptTokens
	if req.Config != nil && req.Config.Validation.MaxPromptTokens > 0 {
		maxTokens = req.Config.Validation.MaxPromptTokens
	}

	oversized := oversizedValidationDecisions(validationReq.Items, maxTokens, v.estimator)
	operations := make([]outcomes.Operation, 0)
	sectionIndex := 0
	if req.SectionIndex != nil {
		sectionIndex = *req.SectionIndex
	}
	operationRef := func(kind outcomes.OperationKind, batchIndex *int) outcomes.Reference {
		return outcomes.Reference{
			ModuleKey: req.ModuleKey, ModuleInstance: req.ModuleInstance,
			SectionIndex: sectionIndex, Kind: kind, ValidatorName: v.name,
			ValidatorOrder: req.ValidatorOrder, BatchIndex: batchIndex,
		}
	}
	if len(oversized) > 0 {
		indices := make([]int, len(oversized))
		for i, decision := range oversized {
			indices[i] = decision.ProposalIndex
		}
		operations = append(operations, outcomes.Operation{
			Reference:       operationRef(outcomes.OperationValidationInput, nil),
			FailureReason:   outcomes.ReasonValidationInputTooLarge,
			ProposalIndices: indices,
		})
	}
	itemsForBatching := filterItemsByDecision(validationReq.Items, oversized)
	if len(itemsForBatching) == 0 {
		all := append([]Decision(nil), immediate...)
		all = append(all, oversized...)
		if err := EnforceDecisionCardinality(req.CandidateProposal, all); err != nil {
			return Result{ValidatorName: v.name, Decisions: all, Operations: operations}, err
		}
		sort.SliceStable(all, func(i, j int) bool { return all[i].ProposalIndex < all[j].ProposalIndex })
		return Result{ValidatorName: v.name, Decisions: all, Operations: operations}, nil
	}

	batches, err := ChunkLLMValidationItems(itemsForBatching, maxTokens, v.estimator)
	if err != nil {
		return Result{ValidatorName: v.name, Decisions: append(append([]Decision(nil), immediate...), oversized...), Operations: operations}, err
	}

	llmDecisions := make([]Decision, 0)
	for _, batch := range batches {
		if err := ctx.Err(); err != nil {
			return Result{ValidatorName: v.name, Decisions: append(append(append([]Decision(nil), immediate...), oversized...), llmDecisions...), Operations: operations}, err
		}
		transcriptDescription := ""
		if req.Config != nil {
			transcriptDescription = req.Config.Context.Description
		}
		variables, err := BuildValidationVariables(batch.Items, transcriptDescription)
		if err != nil {
			return Result{ValidatorName: v.name, Decisions: append(append(append([]Decision(nil), immediate...), oversized...), llmDecisions...), Operations: operations}, err
		}

		var response LLMValidationResponse
		var completion StructuredCompletionResponse
		responseSchema := responseschema.MustLookup(responseschema.ValidatorDecisionSetKey)
		promptMetadata := validatorPromptMetadata(v.validatorType)
		stageInstance := req.ModuleInstance
		if req.SectionIndex != nil {
			stageInstance = fmt.Sprintf("%s:section-%04d", stageInstance, *req.SectionIndex)
		}
		stage := stagename.ValidatorBatch(stageInstance, v.name, batch.BatchIndex)
		completion, err = req.LLMClient.CompleteStructured(ctx, StructuredCompletionRequest{
			PromptID: promptMetadata.PromptID, PromptVersion: promptMetadata.PromptVersion,
			Variables: variables, Schema: responseschema.ValidatorDecisionSetKey,
			Context: structuredoutput.DiagnosticContext{Role: "validation", ModuleKey: req.ModuleKey, ModuleInstance: req.ModuleInstance, SectionIndex: req.SectionIndex, BatchIndex: &batch.BatchIndex, ValidatorName: v.name},
		}, &response)
		batchIndex := batch.BatchIndex
		operation := outcomes.Operation{
			Reference: operationRef(outcomes.OperationValidation, &batchIndex),
			Counts:    completion.Execution,
		}

		artifacts := InteractionArtifacts{}
		if req.DiagnosticsWriter != nil {
			var diagnosticErr error
			artifacts, diagnosticErr = req.DiagnosticsWriter.WriteInteraction(
				stage,
				map[string]any{
					"validator_name":  v.name,
					"validator_type":  v.validatorType,
					"batch_index":     batch.BatchIndex,
					"prompt":          completion.Prompt,
					"response_schema": responseSchema.DiagnosticsMap(),
				},
				completion.Prepared,
				map[string]any{"decoded": response, "final_content": completion.FinalContent, "model": completion.Model},
				errPayload(err),
			)
			if diagnosticErr != nil && req.ReportDiagnosticError != nil {
				req.ReportDiagnosticError(fmt.Errorf("validation interaction %s: %w", stage, diagnosticErr))
			}
		}
		if err != nil {
			var failure *structuredoutput.Failure
			if errors.As(err, &failure) {
				operation.FailureReason = failure.Reason
			}
			if operation.FailureReason == "" {
				operation.FailureReason = outcomes.ReasonInternalFailure
			}
			operation.ProposalIndices = batchProposalIndices(batch.Items)
			operations = append(operations, operation)
			if structuredoutput.IsRecoverable(err) && failure != nil && failure.Reason != "" {
				llmDecisions = append(llmDecisions, rejectBatch(batch.Items, string(failure.Reason), "model validation was unavailable")...)
				continue
			}
			return Result{ValidatorName: v.name, Decisions: append(append(append([]Decision(nil), immediate...), oversized...), llmDecisions...), Operations: operations}, fmt.Errorf("LLM validator %q completion failed: %w", v.name, err)
		}

		batchDecisions, err := mapLLMResponseToDecisions(batch.Items, response)
		if err != nil {
			operation.FailureReason = outcomes.ReasonValidationDecisionsInvalid
			operation.ProposalIndices = batchProposalIndices(batch.Items)
			operations = append(operations, operation)
			llmDecisions = append(llmDecisions, rejectBatch(batch.Items, string(outcomes.ReasonValidationDecisionsInvalid), "model validation decisions were unusable")...)
			continue
		}
		operation.Succeeded = true
		if completion.Execution.StructuralRepairs > 0 {
			operation.RecoveryReasons = append(operation.RecoveryReasons, outcomes.RecoveryOutputRepaired)
		}
		if completion.Execution.TransientRetries > 0 {
			operation.RecoveryReasons = append(operation.RecoveryReasons, outcomes.RecoveryServiceRecovered)
		}
		operations = append(operations, operation)
		for i := range batchDecisions {
			batchDecisions[i].DiagnosticArtifactPath = artifacts.ResponsePayloadPath
		}
		llmDecisions = append(llmDecisions, batchDecisions...)
	}

	all := append([]Decision(nil), immediate...)
	all = append(all, oversized...)
	all = append(all, llmDecisions...)
	if err := EnforceDecisionCardinality(req.CandidateProposal, all); err != nil {
		return Result{ValidatorName: v.name, Decisions: all, Operations: operations}, err
	}
	sort.SliceStable(all, func(i, j int) bool { return all[i].ProposalIndex < all[j].ProposalIndex })
	return Result{ValidatorName: v.name, Decisions: all, Operations: operations}, nil
}

func validatorPromptMetadata(validatorType LLMValidatorType) prompts.Metadata {
	switch validatorType {
	case LLMValidatorTypeSpokenFormPlausibility:
		return prompts.MustLookupMetadata(prompts.PromptIDValidatorSpokenFormPlausibility)
	case LLMValidatorTypeMeaningReversal:
		return prompts.MustLookupMetadata(prompts.PromptIDValidatorMeaningReversalReview)
	case LLMValidatorTypeEditorialReview:
		return prompts.MustLookupMetadata(prompts.PromptIDValidatorEditorialReview)
	case LLMValidatorTypeGrammarReview:
		return prompts.MustLookupMetadata(prompts.PromptIDValidatorGrammarReview)
	case LLMValidatorTypeSpokenWordReview:
		return prompts.MustLookupMetadata(prompts.PromptIDValidatorSpokenWordReview)
	default:
		return prompts.Metadata{}
	}
}

func errPayload(err error) any {
	if err == nil {
		return nil
	}
	return map[string]any{"error": err.Error()}
}

func BuildLLMValidationRequest(validatorName string, validatorType LLMValidatorType, req Request) (LLMValidationRequest, []Decision) {
	items := make([]LLMValidationItem, 0, len(req.CandidateProposal))
	immediate := make([]Decision, 0)
	segments := make(map[int]schema.Segment)
	if req.WorkingTranscript != nil {
		segments = make(map[int]schema.Segment, len(req.WorkingTranscript.Segments))
		for _, seg := range req.WorkingTranscript.Segments {
			segments[seg.ID] = seg
		}
	}

	for _, p := range req.CandidateProposal {
		seg, ok := segments[p.TargetSegmentID]
		if !ok {
			immediate = append(immediate, rejection(p.ProposalIndex, ReasonMissingTargetSegment, "target segment was not found"))
			continue
		}

		preview := proposals.PreviewProposalForSegment(&seg, p.CorrectionProposal, req.ReplacementPolicy)
		if !preview.Applicable {
			immediate = append(immediate, rejection(p.ProposalIndex, string(preview.SkipReason), "proposal is not previewable for LLM validation"))
			continue
		}

		items = append(items, LLMValidationItem{
			CorrectionIndex:      p.ProposalIndex,
			SegmentID:            p.TargetSegmentID,
			OriginalText:         p.OriginalText,
			CorrectedText:        p.CorrectedText,
			OriginalSegmentText:  seg.Text,
			CorrectedSegmentText: preview.CorrectedSegmentText,
			Categories:           append([]string(nil), seg.Categories...),
		})
	}

	return LLMValidationRequest{
		ValidatorName:     validatorName,
		ValidatorType:     validatorType,
		ModuleKey:         req.ModuleKey,
		ModuleInstance:    req.ModuleInstance,
		ReplacementPolicy: string(req.ReplacementPolicy),
		Glossary:          req.Glossary,
		Items:             items,
	}, immediate
}

func mapLLMResponseToDecisions(items []LLMValidationItem, response LLMValidationResponse) ([]Decision, error) {
	expected := make(map[int]LLMValidationItem, len(items))
	for _, item := range items {
		expected[item.CorrectionIndex] = item
	}
	if len(response.Validations) == 0 {
		return nil, fmt.Errorf("missing validations in structured response")
	}

	seen := make(map[int]LLMValidationDecision, len(response.Validations))
	for _, d := range response.Validations {
		if d.Confidence < 0.0 || d.Confidence > 1.0 {
			return nil, fmt.Errorf("confidence for correction_index %d must be between 0.0 and 1.0", d.CorrectionIndex)
		}
		if _, ok := expected[d.CorrectionIndex]; !ok {
			return nil, fmt.Errorf("unknown correction_index %d", d.CorrectionIndex)
		}
		if _, exists := seen[d.CorrectionIndex]; exists {
			return nil, fmt.Errorf("duplicate correction_index %d", d.CorrectionIndex)
		}
		seen[d.CorrectionIndex] = d
	}

	decisions := make([]Decision, 0, len(items))
	for _, item := range items {
		d, ok := seen[item.CorrectionIndex]
		if !ok {
			return nil, fmt.Errorf("missing correction_index %d", item.CorrectionIndex)
		}
		reasonCode := ReasonApproved
		if !d.Approved {
			reasonCode = "llm_rejected"
		}
		decisions = append(decisions, Decision{
			ProposalIndex: item.CorrectionIndex,
			Approved:      d.Approved,
			Disposition:   decisionDisposition(d.Approved),
			ReasonCode:    reasonCode,
			Message:       strings.TrimSpace(d.Reason),
		})
	}
	return decisions, nil
}

func oversizedValidationDecisions(items []LLMValidationItem, maxPromptTokens int, estimator chunking.TokenEstimator) []Decision {
	out := make([]Decision, 0)
	for _, item := range items {
		singleTokens, err := estimateBatchTokens(estimator, []LLMValidationItem{item})
		if err != nil || singleTokens <= maxPromptTokens {
			continue
		}
		out = append(out, unavailable(item.CorrectionIndex, ReasonValidatorInputTooLarge, "validation input exceeds max prompt tokens"))
	}
	return out
}

func filterItemsByDecision(items []LLMValidationItem, decisions []Decision) []LLMValidationItem {
	if len(decisions) == 0 {
		return append([]LLMValidationItem(nil), items...)
	}
	rejected := make(map[int]struct{}, len(decisions))
	for _, decision := range decisions {
		rejected[decision.ProposalIndex] = struct{}{}
	}
	out := make([]LLMValidationItem, 0, len(items))
	for _, item := range items {
		if _, ok := rejected[item.CorrectionIndex]; ok {
			continue
		}
		out = append(out, item)
	}
	return out
}

func rejectBatch(items []LLMValidationItem, reasonCode string, message string) []Decision {
	out := make([]Decision, 0, len(items))
	for _, item := range items {
		out = append(out, unavailable(item.CorrectionIndex, reasonCode, message))
	}
	return out
}

func batchProposalIndices(items []LLMValidationItem) []int {
	indices := make([]int, len(items))
	for i, item := range items {
		indices[i] = item.CorrectionIndex
	}
	return indices
}

func decisionDisposition(approved bool) DecisionDisposition {
	if approved {
		return DispositionApproved
	}
	return DispositionRejected
}