audita/internal/framework/validators/llm_validators.go

package validators

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

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/chunking"
	"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/v3/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/outcomes"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/stagename"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/v3/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
	}

	for i := range validationReq.Items {
		validationReq.Items[i] = fitValidationContext(validationReq.Items[i], maxTokens, v.estimator)
	}
	oversized := oversizedValidationDecisions(validationReq.Items, maxTokens, v.estimator)
	operations := make([]outcomes.Operation, 0)
	llmDecisions := make([]Decision, 0)
	partialResult := func() Result {
		decisions := append([]Decision(nil), immediate...)
		decisions = append(decisions, oversized...)
		decisions = append(decisions, llmDecisions...)
		return Result{ValidatorName: v.name, Decisions: decisions, Operations: operations}
	}
	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 {
		result := partialResult()
		all := result.Decisions
		if err := EnforceDecisionCardinality(req.CandidateProposal, all); err != nil {
			return result, err
		}
		sort.SliceStable(all, func(i, j int) bool { return all[i].ProposalIndex < all[j].ProposalIndex })
		return result, nil
	}

	batches, err := ChunkLLMValidationItems(itemsForBatching, maxTokens, v.estimator)
	if err != nil {
		return partialResult(), err
	}

	for _, batch := range batches {
		if err := ctx.Err(); err != nil {
			return partialResult(), err
		}
		transcriptDescription := ""
		if req.Config != nil {
			transcriptDescription = req.Config.Context.Description
		}
		variables, err := BuildValidationVariables(batch.Items, transcriptDescription)
		if err != nil {
			return partialResult(), 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,
		}

		writeInteraction := func(responsePayload any, callErr error) InteractionArtifacts {
			if req.DiagnosticsWriter == nil {
				return InteractionArtifacts{}
			}
			var diagnosticErr error
			artifacts, diagnosticErr := req.DiagnosticsWriter.WriteInteraction(
				stage,
				map[string]any{
					"validator_name":   v.name,
					"validator_type":   v.validatorType,
					"batch_index":      batch.BatchIndex,
					"proposal_indices": batchProposalIndices(batch.Items),
					"prompt":           completion.Prompt,
					"response_schema":  responseSchema.DiagnosticsMap(),
				},
				completion.Prepared,
				responsePayload,
				errPayload(callErr),
			)
			if diagnosticErr != nil && req.ReportDiagnosticError != nil {
				req.ReportDiagnosticError(fmt.Errorf("validation interaction %s: %w", stage, diagnosticErr))
			}
			return artifacts
		}
		if err != nil {
			writeInteraction(map[string]any{"decoded": response, "final_content": completion.FinalContent, "model": completion.Model}, err)
			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 partialResult(), fmt.Errorf("LLM validator %q completion failed: %w", v.name, err)
		}

		reconciled := reconcileLLMResponse(batch.Items, response)
		artifacts := writeInteraction(map[string]any{"decoded": response, "final_content": completion.FinalContent, "model": completion.Model, "reconciliation": reconciled}, nil)
		if len(reconciled.UnavailableIndices) > 0 {
			operation.FailureReason = outcomes.ReasonValidationDecisionsInvalid
			operation.ProposalIndices = reconciled.UnavailableIndices
		} else {
			operation.Succeeded = true
			operation.RecoveryReasons = outcomes.ExecutionRecoveries(completion.Execution)
			if len(reconciled.IgnoredIndices) > 0 {
				operation.RecoveryReasons = append(operation.RecoveryReasons, outcomes.RecoveryValidationDecisionsFiltered)
			}
		}
		operations = append(operations, operation)
		for i := range reconciled.Decisions {
			reconciled.Decisions[i].DiagnosticArtifactPath = artifacts.ResponsePayloadPath
		}
		llmDecisions = append(llmDecisions, reconciled.Decisions...)
	}

	result := partialResult()
	all := result.Decisions
	if err := EnforceDecisionCardinality(req.CandidateProposal, all); err != nil {
		return result, err
	}
	sort.SliceStable(all, func(i, j int) bool { return all[i].ProposalIndex < all[j].ProposalIndex })
	return result, 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)
	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)
	var segments []schema.Segment
	positions := make(map[int]int)
	vocabulary := NewProtectedVocabulary(req.Glossary)
	if req.WorkingTranscript != nil {
		segments = req.WorkingTranscript.Segments
		for i, seg := range segments {
			positions[seg.ID] = i
		}
	}

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

		seg := segments[position]
		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
		}

		neighbors, glossaryEntries := validationContext(segments, position, req.Glossary, vocabulary, preview.CorrectedSegmentText)
		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...),
			NeighboringSegments:  neighbors,
			GlossaryEntries:      glossaryEntries,
		})
	}

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

type validationReconciliation struct {
	Decisions          []Decision `json:"decisions"`
	UnavailableIndices []int      `json:"unavailable_indices,omitempty"`
	IgnoredIndices     []int      `json:"ignored_indices,omitempty"`
}

func reconcileLLMResponse(items []LLMValidationItem, response LLMValidationResponse) validationReconciliation {
	result := validationReconciliation{Decisions: make([]Decision, 0, len(items))}
	seen := make(map[int][]LLMValidationDecision, len(items))
	for _, d := range response.Validations {
		if d.CorrectionIndex < 1 || d.CorrectionIndex > len(items) {
			result.IgnoredIndices = append(result.IgnoredIndices, d.CorrectionIndex)
			continue
		}
		seen[d.CorrectionIndex] = append(seen[d.CorrectionIndex], d)
	}

	for i, item := range items {
		entries := seen[i+1]
		if len(entries) == 0 {
			result.Decisions = append(result.Decisions, unavailable(item.CorrectionIndex, ReasonMissingValidationDecision, "required model validation decision was missing"))
			result.UnavailableIndices = append(result.UnavailableIndices, item.CorrectionIndex)
			continue
		}
		if len(entries) > 1 {
			result.Decisions = append(result.Decisions, unavailable(item.CorrectionIndex, ReasonDuplicateValidationDecision, "required model validation decision appeared more than once"))
			result.UnavailableIndices = append(result.UnavailableIndices, item.CorrectionIndex)
			continue
		}
		d := entries[0]
		if d.Confidence < 0.0 || d.Confidence > 1.0 {
			result.Decisions = append(result.Decisions, unavailable(item.CorrectionIndex, ReasonInvalidValidationDecision, "required model validation decision had invalid confidence"))
			result.UnavailableIndices = append(result.UnavailableIndices, item.CorrectionIndex)
			continue
		}
		reasonCode := ReasonApproved
		if !d.Approved {
			reasonCode = "llm_rejected"
		}
		result.Decisions = append(result.Decisions, Decision{
			ProposalIndex: item.CorrectionIndex,
			Approved:      d.Approved,
			Disposition:   decisionDisposition(d.Approved),
			ReasonCode:    reasonCode,
			Message:       strings.TrimSpace(d.Reason),
		})
	}
	return result
}

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
}