Add intra-module pipeline for LLM validation

2026-05-12 18:39:26 -05:00
parent 1afd753fad
commit 3d45571bb0
6 changed files with 835 additions and 216 deletions
--- a/internal/framework/runner/runner_test.go
+++ b/internal/framework/runner/runner_test.go
@@ -5,6 +5,8 @@ import (
 	"errors"
 	"os"
 	"path/filepath"
+	"regexp"
+	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -181,18 +183,28 @@ func TestRunnerProposalSectionConcurrencyBoundedByProposalLLMConcurrency(t *test
 	var maxInFlight int32
 	entered := make(chan struct{}, len(transcript.Segments))
 	release := make(chan struct{})
+	scheduler, err := llm.NewScheduler(2)
+	if err != nil {
+		t.Fatalf("NewScheduler: %v", err)
+	}
 	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
 		"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
-			current := atomic.AddInt32(&inFlight, 1)
-			for {
-				prior := atomic.LoadInt32(&maxInFlight)
-				if current <= prior || atomic.CompareAndSwapInt32(&maxInFlight, prior, current) {
-					break
+			err := req.LLMScheduler.Run(context.Background(), func(context.Context) error {
+				current := atomic.AddInt32(&inFlight, 1)
+				for {
+					prior := atomic.LoadInt32(&maxInFlight)
+					if current <= prior || atomic.CompareAndSwapInt32(&maxInFlight, prior, current) {
+						break
+					}
 				}
+				entered <- struct{}{}
+				<-release
+				atomic.AddInt32(&inFlight, -1)
+				return nil
+			})
+			if err != nil {
+				return nil, err
 			}
-			entered <- struct{}{}
-			<-release
-			atomic.AddInt32(&inFlight, -1)
 			return nil, nil
 		}},
 	}})
@@ -200,8 +212,9 @@ func TestRunnerProposalSectionConcurrencyBoundedByProposalLLMConcurrency(t *test
 	resultCh := make(chan error, 1)
 	go func() {
 		_, err := r.Run(context.Background(), RunInput{
-			Config:     &cfg,
-			Transcript: transcript,
+			Config:               &cfg,
+			Transcript:           transcript,
+			ProposalLLMScheduler: scheduler,
 			ModuleSpecs: []contracts.ModuleRunSpec{
 				{ModuleKey: "m", InstanceName: "m"},
 			},
@@ -212,7 +225,7 @@ func TestRunnerProposalSectionConcurrencyBoundedByProposalLLMConcurrency(t *test
 	waitForRunnerEntries(t, entered, 2, "proposal workers to enter")
 	close(release)

-	err := <-resultCh
+	err = <-resultCh
 	if err != nil {
 		t.Fatalf("Run error: %v", err)
 	}
@@ -342,6 +355,436 @@ func mapKeys(values map[int]struct{}) []int {
 	return keys
 }

+func TestRunnerValidationStartsBeforeAllSectionProposalsComplete(t *testing.T) {
+	cfg := config.Default()
+	cfg.MaxSectionTokens = 3
+	cfg.MinSectionTokens = 0
+	cfg.TotalLLMConcurrency = 2
+	cfg.ProposalLLMConcurrency = 2
+
+	transcript := &schema.Transcript{Segments: []schema.Segment{
+		{ID: 1, Text: "alpha one"},
+		{ID: 2, Text: "bravo two"},
+	}}
+
+	releaseSectionOne := make(chan struct{})
+	validatorStarted := make(chan struct{}, 1)
+
+	validator := fakeValidator{
+		name: "capture",
+		validateF: func(req contracts.ValidationRequest) (validators.Result, error) {
+			if len(req.CandidateProposal) == 1 && req.CandidateProposal[0].TargetSegmentID == 1 {
+				select {
+				case validatorStarted <- struct{}{}:
+				default:
+				}
+			}
+			decisions := make([]validators.Decision, 0, len(req.CandidateProposal))
+			for _, p := range req.CandidateProposal {
+				decisions = append(decisions, validators.Decision{
+					ProposalIndex: p.ProposalIndex,
+					Approved:      true,
+					ReasonCode:    validators.ReasonApproved,
+					Message:       "ok",
+				})
+			}
+			return validators.Result{ValidatorName: "capture", Decisions: decisions}, nil
+		},
+	}
+
+	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
+		"m": fakeModule{
+			key:        "m",
+			policy:     proposals.ReplacementPolicyRequireUnique,
+			validators: []contracts.Validator{validator},
+			proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
+				seg := req.WorkingTranscript.Segments[0]
+				if req.Section != nil && req.Section.Index == 1 {
+					<-releaseSectionOne
+				}
+				return []proposals.CorrectionProposal{
+					{TargetSegmentID: seg.ID, OriginalText: seg.Text, CorrectedText: strings.ToUpper(seg.Text), Confidence: 1},
+				}, nil
+			},
+		},
+	}})
+
+	resultCh := make(chan error, 1)
+	go func() {
+		_, err := r.Run(context.Background(), RunInput{
+			Config:     &cfg,
+			Transcript: transcript,
+			ModuleSpecs: []contracts.ModuleRunSpec{
+				{ModuleKey: "m", InstanceName: "m"},
+			},
+		})
+		resultCh <- err
+	}()
+
+	select {
+	case <-validatorStarted:
+	case <-time.After(350 * time.Millisecond):
+		t.Fatal("expected section-level validator work before all proposal sections complete")
+	}
+
+	close(releaseSectionOne)
+	if err := <-resultCh; err != nil {
+		t.Fatalf("Run error: %v", err)
+	}
+}
+
+func TestRunnerProposalAndValidationLLMCanOverlapWithinModule(t *testing.T) {
+	cfg := config.Default()
+	cfg.MaxSectionTokens = 3
+	cfg.MinSectionTokens = 0
+	cfg.TotalLLMConcurrency = 2
+	cfg.ProposalLLMConcurrency = 2
+	validationCap := 2
+	cfg.ValidationLLMConcurrency = &validationCap
+
+	transcript := &schema.Transcript{Segments: []schema.Segment{
+		{ID: 1, Text: "alpha one"},
+		{ID: 2, Text: "bravo two"},
+	}}
+
+	global, err := llm.NewScheduler(2)
+	if err != nil {
+		t.Fatalf("NewScheduler: %v", err)
+	}
+	scheduler := &trackingScheduler{inner: global}
+
+	proposalSectionOneStarted := make(chan struct{}, 1)
+	releaseProposalSectionOne := make(chan struct{})
+
+	releaseValidation := make(chan struct{})
+	client := &stageAwareStructuredClient{
+		startedSection:    make(chan int, 8),
+		releaseValidation: releaseValidation,
+	}
+	llmValidator, err := validators.NewLLMBackedValidator("spoken_form_plausibility_review", validators.LLMValidatorTypeSpokenFormPlausibility, "")
+	if err != nil {
+		t.Fatalf("NewLLMBackedValidator: %v", err)
+	}
+
+	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
+		"m": fakeModule{
+			key:        "m",
+			policy:     proposals.ReplacementPolicyRequireUnique,
+			validators: []contracts.Validator{llmValidator},
+			proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
+				seg := req.WorkingTranscript.Segments[0]
+				err := req.LLMScheduler.Run(context.Background(), func(context.Context) error {
+					if req.Section != nil && req.Section.Index == 1 {
+						select {
+						case proposalSectionOneStarted <- struct{}{}:
+						default:
+						}
+						<-releaseProposalSectionOne
+					}
+					return nil
+				})
+				if err != nil {
+					return nil, err
+				}
+				return []proposals.CorrectionProposal{
+					{TargetSegmentID: seg.ID, OriginalText: seg.Text, CorrectedText: strings.ToUpper(seg.Text), Confidence: 1},
+				}, nil
+			},
+		},
+	}})
+
+	resultCh := make(chan struct {
+		out RunOutput
+		err error
+	}, 1)
+	go func() {
+		out, err := r.Run(context.Background(), RunInput{
+			Config:                 &cfg,
+			Transcript:             transcript,
+			ModuleSpecs:            []contracts.ModuleRunSpec{{ModuleKey: "m", InstanceName: "m"}},
+			ProposalLLMScheduler:   scheduler,
+			ValidationLLMScheduler: scheduler,
+			ValidationLLMClient:    client,
+		})
+		resultCh <- struct {
+			out RunOutput
+			err error
+		}{out: out, err: err}
+	}()
+
+	select {
+	case <-proposalSectionOneStarted:
+	case <-time.After(350 * time.Millisecond):
+		t.Fatal("expected section-1 proposal to start")
+	}
+
+	select {
+	case section := <-client.startedSection:
+		if section != 0 {
+			t.Fatalf("expected first validation to start for section 0, got section %d", section)
+		}
+	case <-time.After(350 * time.Millisecond):
+		t.Fatal("expected validation call for section 0 while later section proposal still running")
+	}
+
+	close(releaseProposalSectionOne)
+	close(releaseValidation)
+
+	result := <-resultCh
+	if result.err != nil {
+		t.Fatalf("Run error: %v", result.err)
+	}
+	if got := atomic.LoadInt32(&scheduler.maxInFlight); got > 2 {
+		t.Fatalf("expected combined proposal+validation in-flight <= 2, got %d", got)
+	}
+	if got := atomic.LoadInt32(&scheduler.maxInFlight); got < 2 {
+		t.Fatalf("expected observed overlap/in-flight utilization of at least 2, got %d", got)
+	}
+}
+
+func TestRunnerValidationLLMConcurrencyRespected(t *testing.T) {
+	cfg := config.Default()
+	cfg.MaxSectionTokens = 3
+	cfg.MinSectionTokens = 0
+	cfg.TotalLLMConcurrency = 4
+	cfg.ProposalLLMConcurrency = 4
+	validationCap := 1
+	cfg.ValidationLLMConcurrency = &validationCap
+
+	transcript := &schema.Transcript{Segments: []schema.Segment{
+		{ID: 1, Text: "one one"},
+		{ID: 2, Text: "two two"},
+		{ID: 3, Text: "three three"},
+	}}
+
+	validationSchedulerInner, err := llm.NewScheduler(1)
+	if err != nil {
+		t.Fatalf("NewScheduler(validation): %v", err)
+	}
+	validationScheduler := &trackingScheduler{inner: validationSchedulerInner}
+
+	releaseValidation := make(chan struct{})
+	client := &stageAwareStructuredClient{
+		startedSection:    make(chan int, 16),
+		releaseValidation: releaseValidation,
+	}
+	llmValidator, err := validators.NewLLMBackedValidator("spoken_form_plausibility_review", validators.LLMValidatorTypeSpokenFormPlausibility, "")
+	if err != nil {
+		t.Fatalf("NewLLMBackedValidator: %v", err)
+	}
+
+	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
+		"m": fakeModule{
+			key:        "m",
+			policy:     proposals.ReplacementPolicyRequireUnique,
+			validators: []contracts.Validator{llmValidator},
+			proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
+				seg := req.WorkingTranscript.Segments[0]
+				return []proposals.CorrectionProposal{
+					{TargetSegmentID: seg.ID, OriginalText: seg.Text, CorrectedText: strings.ToUpper(seg.Text), Confidence: 1},
+				}, nil
+			},
+		},
+	}})
+
+	resultCh := make(chan error, 1)
+	go func() {
+		_, err := r.Run(context.Background(), RunInput{
+			Config:                 &cfg,
+			Transcript:             transcript,
+			ModuleSpecs:            []contracts.ModuleRunSpec{{ModuleKey: "m", InstanceName: "m"}},
+			ValidationLLMScheduler: validationScheduler,
+			ValidationLLMClient:    client,
+		})
+		resultCh <- err
+	}()
+
+	select {
+	case <-client.startedSection:
+	case <-time.After(350 * time.Millisecond):
+		t.Fatal("expected validation to start")
+	}
+	close(releaseValidation)
+
+	if err := <-resultCh; err != nil {
+		t.Fatalf("Run error: %v", err)
+	}
+	if got := atomic.LoadInt32(&validationScheduler.maxInFlight); got > 1 {
+		t.Fatalf("expected validation in-flight <= 1, got %d", got)
+	}
+}
+
+func TestRunnerMixedProposalValidationFIFOOrder(t *testing.T) {
+	cfg := config.Default()
+	cfg.MaxSectionTokens = 3
+	cfg.MinSectionTokens = 0
+	cfg.TotalLLMConcurrency = 1
+	cfg.ProposalLLMConcurrency = 1
+	validationCap := 1
+	cfg.ValidationLLMConcurrency = &validationCap
+
+	transcript := &schema.Transcript{Segments: []schema.Segment{
+		{ID: 1, Text: "alpha one"},
+		{ID: 2, Text: "bravo two"},
+	}}
+
+	global, err := llm.NewScheduler(1)
+	if err != nil {
+		t.Fatalf("NewScheduler: %v", err)
+	}
+	proposalScheduler := global
+	validationScheduler := global
+
+	events := make(chan string, 8)
+	releaseProposalSectionOne := make(chan struct{})
+	releaseValidation := make(chan struct{})
+	sectionZeroEntered := make(chan struct{})
+
+	client := &stageAwareStructuredClient{
+		startedSection:    make(chan int, 8),
+		releaseValidation: releaseValidation,
+		eventSink:         events,
+	}
+	llmValidator, err := validators.NewLLMBackedValidator("spoken_form_plausibility_review", validators.LLMValidatorTypeSpokenFormPlausibility, "")
+	if err != nil {
+		t.Fatalf("NewLLMBackedValidator: %v", err)
+	}
+
+	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
+		"m": fakeModule{
+			key:        "m",
+			policy:     proposals.ReplacementPolicyRequireUnique,
+			validators: []contracts.Validator{llmValidator},
+			proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
+				seg := req.WorkingTranscript.Segments[0]
+				if req.Section != nil && req.Section.Index == 1 {
+					<-sectionZeroEntered
+				}
+				err := req.LLMScheduler.Run(context.Background(), func(context.Context) error {
+					if req.Section != nil {
+						events <- "p" + strconv.Itoa(req.Section.Index)
+						if req.Section.Index == 0 {
+							select {
+							case sectionZeroEntered <- struct{}{}:
+							default:
+							}
+						}
+						if req.Section.Index == 1 {
+							<-releaseProposalSectionOne
+						}
+					}
+					return nil
+				})
+				if err != nil {
+					return nil, err
+				}
+				return []proposals.CorrectionProposal{
+					{TargetSegmentID: seg.ID, OriginalText: seg.Text, CorrectedText: strings.ToUpper(seg.Text), Confidence: 1},
+				}, nil
+			},
+		},
+	}})
+
+	resultCh := make(chan error, 1)
+	go func() {
+		_, err := r.Run(context.Background(), RunInput{
+			Config:                 &cfg,
+			Transcript:             transcript,
+			ModuleSpecs:            []contracts.ModuleRunSpec{{ModuleKey: "m", InstanceName: "m"}},
+			ProposalLLMScheduler:   proposalScheduler,
+			ValidationLLMScheduler: validationScheduler,
+			ValidationLLMClient:    client,
+		})
+		resultCh <- err
+	}()
+
+	if got := <-events; got != "p0" {
+		t.Fatalf("expected first event p0, got %q", got)
+	}
+	if got := <-events; got != "p1" {
+		t.Fatalf("expected second event p1 (FIFO queued proposal), got %q", got)
+	}
+
+	close(releaseProposalSectionOne)
+	close(releaseValidation)
+
+	if got := <-events; got != "v0" {
+		t.Fatalf("expected validator event v0 after queued p1, got %q", got)
+	}
+
+	if err := <-resultCh; err != nil {
+		t.Fatalf("Run error: %v", err)
+	}
+}
+
+var sectionStagePattern = regexp.MustCompile(`section-(\d+)`)
+
+type stageAwareStructuredClient struct {
+	startedSection    chan int
+	releaseValidation <-chan struct{}
+	eventSink         chan<- string
+}
+
+func (c *stageAwareStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
+	section := parseSectionFromStage(req.StageName)
+	if c.startedSection != nil {
+		select {
+		case c.startedSection <- section:
+		default:
+		}
+	}
+	if c.eventSink != nil {
+		c.eventSink <- "v" + strconv.Itoa(section)
+	}
+	if c.releaseValidation != nil {
+		select {
+		case <-c.releaseValidation:
+		case <-ctx.Done():
+			return contracts.StructuredCompletionResponse{}, ctx.Err()
+		}
+	}
+
+	target := out.(*validators.LLMValidationResponse)
+	*target = validators.LLMValidationResponse{
+		Validations: []validators.LLMValidationDecision{
+			{CorrectionIndex: section, Approved: true, Confidence: 0.99, Reason: "ok"},
+		},
+	}
+	return contracts.StructuredCompletionResponse{}, nil
+}
+
+func parseSectionFromStage(stage string) int {
+	match := sectionStagePattern.FindStringSubmatch(stage)
+	if len(match) != 2 {
+		return 0
+	}
+	n, err := strconv.Atoi(match[1])
+	if err != nil {
+		return 0
+	}
+	return n
+}
+
+type trackingScheduler struct {
+	inner       contracts.LLMScheduler
+	inFlight    int32
+	maxInFlight int32
+}
+
+func (s *trackingScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
+	return s.inner.Run(ctx, func(callCtx context.Context) error {
+		current := atomic.AddInt32(&s.inFlight, 1)
+		for {
+			prior := atomic.LoadInt32(&s.maxInFlight)
+			if current <= prior || atomic.CompareAndSwapInt32(&s.maxInFlight, prior, current) {
+				break
+			}
+		}
+		defer atomic.AddInt32(&s.inFlight, -1)
+		return fn(callCtx)
+	})
+}
+
 func TestRunnerSkippedRecorded(t *testing.T) {
 	transcript := &schema.Transcript{Segments: []schema.Segment{{ID: 1, Speaker: "A", Start: 0, End: 1, Text: "word word"}}}
 	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{