Add intra-module pipeline for LLM validation

This commit is contained in:
2026-05-12 18:39:26 -05:00
parent 1afd753fad
commit 3d45571bb0
6 changed files with 835 additions and 216 deletions

View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"path/filepath"
"sort"
"sync"
"time"
@@ -134,125 +135,38 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q chunking failed: %w", spec.InstanceName, err)
}
enriched, proposeErr := collectSectionProposals(ctx, collectSectionProposalsInput{
Module: module,
Spec: spec,
Policy: policy,
Config: input.Config,
Glossary: input.Glossary,
Sections: sections,
ProposalClient: input.ProposalLLMClient,
ProposalScheduler: input.ProposalLLMScheduler,
DiagnosticsDir: input.ProposalDiagnosticsDir,
pipelineResult, pipelineErr := runModulePipeline(ctx, collectSectionProposalsInput{
Module: module,
Spec: spec,
Policy: policy,
Config: input.Config,
Glossary: input.Glossary,
Sections: sections,
ProposalClient: input.ProposalLLMClient,
ProposalScheduler: input.ProposalLLMScheduler,
DiagnosticsDir: input.ProposalDiagnosticsDir,
ValidationClient: input.ValidationLLMClient,
ValidationScheduler: input.ValidationLLMScheduler,
ValidationDiagnosticsDir: input.ValidationDiagnosticsDir,
})
if proposeErr != nil {
if pipelineErr != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ErrorMessage: proposeErr.Error(),
StartedAt: startedAt,
CompletedAt: time.Now().UTC(),
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ProposalCount: pipelineResult.ProposalCount,
ValidatorDecisions: pipelineResult.ValidatorDecisions,
ValidatorRejected: pipelineResult.ValidatorRejected,
ErrorMessage: pipelineErr.Error(),
StartedAt: startedAt,
CompletedAt: time.Now().UTC(),
}
results = append(results, failed)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q failed: %w", spec.InstanceName, proposeErr)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q failed: %w", spec.InstanceName, pipelineErr)
}
validatorDecisions := make([]ValidatorDecisionRecord, 0)
validatorRejected := make([]ValidatorRejectedChange, 0)
eligible := enriched
for _, validator := range module.Validators() {
var diagnosticsWriter validators.InteractionDiagnosticsWriter
if input.ValidationDiagnosticsDir != "" {
diagnosticsWriter = &llmDiagnosticsWriterAdapter{
writer: llm.NewDiagnosticsWriter(
filepath.Join(input.ValidationDiagnosticsDir, spec.InstanceName),
validatorSecrets(input.Config),
),
}
}
vResult, vErr := validator.Validate(ctx, contracts.ValidationRequest{
WorkingTranscript: working,
CandidateProposal: eligible,
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Glossary: input.Glossary,
Config: input.Config,
LLMClient: validationLLMClientAdapter{client: input.ValidationLLMClient},
Scheduler: input.ValidationLLMScheduler,
DiagnosticsWriter: diagnosticsWriter,
})
if vErr != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ProposalCount: len(enriched),
ValidatorDecisions: validatorDecisions,
ValidatorRejected: validatorRejected,
ErrorMessage: vErr.Error(),
StartedAt: startedAt,
CompletedAt: time.Now().UTC(),
}
results = append(results, failed)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q validator %q failed: %w", spec.InstanceName, validator.Name(), vErr)
}
if err := validators.EnforceDecisionCardinality(eligible, vResult.Decisions); err != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ProposalCount: len(enriched),
ValidatorDecisions: validatorDecisions,
ValidatorRejected: validatorRejected,
ErrorMessage: err.Error(),
StartedAt: startedAt,
CompletedAt: time.Now().UTC(),
}
results = append(results, failed)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q validator %q cardinality failed: %w", spec.InstanceName, validator.Name(), err)
}
nextEligible := make([]proposals.EnrichedCorrectionProposal, 0, len(eligible))
byIndex := make(map[int]proposals.EnrichedCorrectionProposal, len(eligible))
for _, p := range eligible {
byIndex[p.ProposalIndex] = p
}
for _, d := range vResult.Decisions {
validatorDecisions = append(validatorDecisions, ValidatorDecisionRecord{
ValidatorName: validator.Name(),
ProposalIndex: d.ProposalIndex,
Approved: d.Approved,
ReasonCode: d.ReasonCode,
Message: d.Message,
DiagnosticArtifactPath: d.DiagnosticArtifactPath,
})
if d.Approved {
nextEligible = append(nextEligible, byIndex[d.ProposalIndex])
continue
}
p := byIndex[d.ProposalIndex]
validatorRejected = append(validatorRejected, ValidatorRejectedChange{
ValidatorName: validator.Name(),
ProposalIndex: p.ProposalIndex,
ModuleKey: p.ModuleKey,
ModuleInstance: p.ModuleInstance,
TargetSegmentID: p.TargetSegmentID,
OriginalText: p.OriginalText,
CorrectedText: p.CorrectedText,
ReasonCode: d.ReasonCode,
Message: d.Message,
})
}
eligible = nextEligible
}
applyResult := proposals.ApplyProposals(working, eligible, policy)
applyResult := proposals.ApplyProposals(working, pipelineResult.Approved, policy)
working = applyResult.Transcript
results = append(results, ModuleResult{
@@ -260,9 +174,9 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusSuccess,
ProposalCount: len(enriched),
ValidatorDecisions: validatorDecisions,
ValidatorRejected: validatorRejected,
ProposalCount: pipelineResult.ProposalCount,
ValidatorDecisions: pipelineResult.ValidatorDecisions,
ValidatorRejected: pipelineResult.ValidatorRejected,
AppliedChanges: applyResult.Applied,
SkippedChanges: applyResult.Skipped,
StartedAt: startedAt,
@@ -274,15 +188,18 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
}
type collectSectionProposalsInput struct {
Module contracts.TranscriptModule
Spec contracts.ModuleRunSpec
Policy proposals.ReplacementPolicy
Config *config.Config
Glossary *schema.Glossary
Sections []chunking.Section
ProposalClient contracts.StructuredLLMClient
ProposalScheduler contracts.LLMScheduler
DiagnosticsDir string
Module contracts.TranscriptModule
Spec contracts.ModuleRunSpec
Policy proposals.ReplacementPolicy
Config *config.Config
Glossary *schema.Glossary
Sections []chunking.Section
ProposalClient contracts.StructuredLLMClient
ProposalScheduler contracts.LLMScheduler
DiagnosticsDir string
ValidationClient contracts.StructuredLLMClient
ValidationScheduler ValidationScheduler
ValidationDiagnosticsDir string
}
type sectionProposals struct {
@@ -290,98 +207,341 @@ type sectionProposals struct {
corrected []proposals.CorrectionProposal
}
func collectSectionProposals(ctx context.Context, input collectSectionProposalsInput) ([]proposals.EnrichedCorrectionProposal, error) {
if len(input.Sections) == 0 {
return []proposals.EnrichedCorrectionProposal{}, nil
}
type sectionProposalResult struct {
sectionPos int
section chunking.Section
corrected []proposals.CorrectionProposal
err error
}
maxWorkers := 1
if input.Config != nil && input.Config.EffectiveProposalLLMConcurrency() > 1 {
maxWorkers = input.Config.EffectiveProposalLLMConcurrency()
}
if maxWorkers > len(input.Sections) {
maxWorkers = len(input.Sections)
}
type sectionValidationResult struct {
sectionPos int
enriched []proposals.EnrichedCorrectionProposal
approved []proposals.EnrichedCorrectionProposal
decisions []ValidatorDecisionRecord
rejected []ValidatorRejectedChange
err error
}
sectionResults := make([]sectionProposals, len(input.Sections))
type modulePipelineResult struct {
ProposalCount int
Approved []proposals.EnrichedCorrectionProposal
ValidatorDecisions []ValidatorDecisionRecord
ValidatorRejected []ValidatorRejectedChange
}
func collectSectionProposals(ctx context.Context, input collectSectionProposalsInput) (context.Context, <-chan sectionProposalResult, context.CancelFunc) {
results := make(chan sectionProposalResult, len(input.Sections))
runCtx, cancel := context.WithCancel(ctx)
go func() {
defer close(results)
var wg sync.WaitGroup
for sectionPos, section := range input.Sections {
sectionPos := sectionPos
section := section
wg.Add(1)
go func() {
defer wg.Done()
meta := contracts.SectionMetadataFromSection(section)
corrected, err := input.Module.Propose(runCtx, contracts.ProposalRequest{
ExecutionContext: contracts.ExecutionContext{
Config: input.Config,
WorkingTranscript: transcriptFromSection(section),
Glossary: input.Glossary,
Section: &meta,
DiagnosticsDir: input.DiagnosticsDir,
},
RunSpec: contracts.ModuleRunSpec{
ModuleKey: input.Spec.ModuleKey,
InstanceName: input.Spec.InstanceName,
ReplacementPolicy: input.Policy,
},
LLMClient: input.ProposalClient,
LLMScheduler: input.ProposalScheduler,
})
select {
case results <- sectionProposalResult{
sectionPos: sectionPos,
section: section,
corrected: corrected,
err: err,
}:
case <-runCtx.Done():
return
}
}()
}
wg.Wait()
}()
return runCtx, results, cancel
}
func runModulePipeline(ctx context.Context, input collectSectionProposalsInput) (modulePipelineResult, error) {
out := modulePipelineResult{
Approved: make([]proposals.EnrichedCorrectionProposal, 0),
ValidatorDecisions: make([]ValidatorDecisionRecord, 0),
ValidatorRejected: make([]ValidatorRejectedChange, 0),
}
if len(input.Sections) == 0 {
return out, nil
}
runCtx, proposalResults, cancel := collectSectionProposals(ctx, input)
defer cancel()
sem := make(chan struct{}, maxWorkers)
pending := make(map[int]sectionProposalResult, len(input.Sections))
validationResults := make(chan sectionValidationResult, len(input.Sections))
validationBySection := make(map[int]sectionValidationResult, len(input.Sections))
validatorsOrdered, validatorOrder := reorderValidatorsForPipeline(input.Module.Validators())
nextSectionToProcess := 0
nextProposalIndex := 0
validationLaunches := 0
var (
wg sync.WaitGroup
errOnce sync.Once
firstErr error
errOnce sync.Once
vwg sync.WaitGroup
)
for sectionPos, section := range input.Sections {
sectionPos := sectionPos
section := section
wg.Add(1)
go func() {
defer wg.Done()
select {
case sem <- struct{}{}:
case <-runCtx.Done():
return
}
defer func() { <-sem }()
setErr := func(err error) {
if err == nil {
return
}
errOnce.Do(func() {
firstErr = err
cancel()
})
}
meta := contracts.SectionMetadataFromSection(section)
corrected, err := input.Module.Propose(runCtx, contracts.ProposalRequest{
ExecutionContext: contracts.ExecutionContext{
Config: input.Config,
WorkingTranscript: transcriptFromSection(section),
Glossary: input.Glossary,
Section: &meta,
DiagnosticsDir: input.DiagnosticsDir,
},
RunSpec: contracts.ModuleRunSpec{
ModuleKey: input.Spec.ModuleKey,
InstanceName: input.Spec.InstanceName,
ReplacementPolicy: input.Policy,
},
LLMClient: input.ProposalClient,
LLMScheduler: input.ProposalScheduler,
})
if err != nil {
errOnce.Do(func() {
firstErr = err
cancel()
for result := range proposalResults {
if result.err != nil {
setErr(result.err)
continue
}
if firstErr != nil {
continue
}
pending[result.sectionPos] = result
for {
sectionResult, ok := pending[nextSectionToProcess]
if !ok {
break
}
delete(pending, nextSectionToProcess)
sectionMeta := contracts.SectionMetadataFromSection(sectionResult.section)
sectionEnriched := make([]proposals.EnrichedCorrectionProposal, 0, len(sectionResult.corrected))
for i, corrected := range sectionResult.corrected {
sectionIndex := sectionMeta.Index
sectionEnriched = append(sectionEnriched, proposals.EnrichedCorrectionProposal{
CorrectionProposal: corrected,
ProposalMetadata: proposals.ProposalMetadata{
ProposalIndex: nextProposalIndex + i,
ModuleKey: input.Spec.ModuleKey,
ModuleInstance: input.Spec.InstanceName,
SectionIndex: &sectionIndex,
},
})
return
}
nextProposalIndex += len(sectionEnriched)
out.ProposalCount += len(sectionEnriched)
sectionResults[sectionPos] = sectionProposals{
meta: meta,
corrected: corrected,
}
}()
}
wg.Wait()
if firstErr != nil {
return nil, firstErr
}
enriched := make([]proposals.EnrichedCorrectionProposal, 0)
nextProposalIndex := 0
for _, sectionResult := range sectionResults {
for _, corrected := range sectionResult.corrected {
sectionIndex := sectionResult.meta.Index
enriched = append(enriched, proposals.EnrichedCorrectionProposal{
CorrectionProposal: corrected,
ProposalMetadata: proposals.ProposalMetadata{
ProposalIndex: nextProposalIndex,
ModuleKey: input.Spec.ModuleKey,
ModuleInstance: input.Spec.InstanceName,
SectionIndex: &sectionIndex,
},
})
nextProposalIndex++
validationLaunches++
vwg.Add(1)
go func(sectionPos int, enriched []proposals.EnrichedCorrectionProposal, sectionMetadata contracts.SectionMetadata) {
defer vwg.Done()
validated, err := validateSectionCandidates(runCtx, validateSectionCandidatesInput{
Spec: input.Spec,
Policy: input.Policy,
Glossary: input.Glossary,
Config: input.Config,
WorkingTranscript: transcriptFromSection(sectionResult.section),
ModuleInstanceForStages: fmt.Sprintf("%s:section-%04d", input.Spec.InstanceName, sectionMetadata.Index),
Validators: validatorsOrdered,
SectionEnriched: enriched,
ValidationLLMClient: input.ValidationClient,
ValidationScheduler: input.ValidationScheduler,
DiagnosticsDir: input.ValidationDiagnosticsDir,
})
validationResults <- sectionValidationResult{
sectionPos: sectionPos,
enriched: enriched,
approved: validated.approved,
decisions: validated.decisions,
rejected: validated.rejected,
err: err,
}
}(nextSectionToProcess, sectionEnriched, sectionMeta)
nextSectionToProcess++
}
}
return enriched, nil
vwg.Wait()
close(validationResults)
for v := range validationResults {
validationBySection[v.sectionPos] = v
if v.err != nil {
setErr(v.err)
}
}
for i := 0; i < validationLaunches; i++ {
res, ok := validationBySection[i]
if !ok {
break
}
out.Approved = append(out.Approved, res.approved...)
out.ValidatorDecisions = append(out.ValidatorDecisions, res.decisions...)
out.ValidatorRejected = append(out.ValidatorRejected, res.rejected...)
}
sort.SliceStable(out.ValidatorDecisions, func(i, j int) bool {
if out.ValidatorDecisions[i].ProposalIndex != out.ValidatorDecisions[j].ProposalIndex {
return out.ValidatorDecisions[i].ProposalIndex < out.ValidatorDecisions[j].ProposalIndex
}
return validatorOrder[out.ValidatorDecisions[i].ValidatorName] < validatorOrder[out.ValidatorDecisions[j].ValidatorName]
})
sort.SliceStable(out.ValidatorRejected, func(i, j int) bool {
if out.ValidatorRejected[i].ProposalIndex != out.ValidatorRejected[j].ProposalIndex {
return out.ValidatorRejected[i].ProposalIndex < out.ValidatorRejected[j].ProposalIndex
}
return validatorOrder[out.ValidatorRejected[i].ValidatorName] < validatorOrder[out.ValidatorRejected[j].ValidatorName]
})
if firstErr != nil {
return out, firstErr
}
return out, nil
}
type validateSectionCandidatesInput struct {
Spec contracts.ModuleRunSpec
Policy proposals.ReplacementPolicy
Glossary *schema.Glossary
Config *config.Config
WorkingTranscript *schema.Transcript
ModuleInstanceForStages string
Validators []contracts.Validator
SectionEnriched []proposals.EnrichedCorrectionProposal
ValidationLLMClient contracts.StructuredLLMClient
ValidationScheduler ValidationScheduler
DiagnosticsDir string
}
type validateSectionCandidatesResult struct {
approved []proposals.EnrichedCorrectionProposal
decisions []ValidatorDecisionRecord
rejected []ValidatorRejectedChange
}
func validateSectionCandidates(ctx context.Context, input validateSectionCandidatesInput) (validateSectionCandidatesResult, error) {
decisions := make([]ValidatorDecisionRecord, 0)
rejected := make([]ValidatorRejectedChange, 0)
eligible := append([]proposals.EnrichedCorrectionProposal(nil), input.SectionEnriched...)
for _, validator := range input.Validators {
var diagnosticsWriter validators.InteractionDiagnosticsWriter
if input.DiagnosticsDir != "" {
diagnosticsWriter = &llmDiagnosticsWriterAdapter{
writer: llm.NewDiagnosticsWriter(
filepath.Join(input.DiagnosticsDir, input.Spec.InstanceName),
validatorSecrets(input.Config),
),
}
}
vResult, err := validator.Validate(ctx, contracts.ValidationRequest{
WorkingTranscript: input.WorkingTranscript,
CandidateProposal: eligible,
ModuleKey: input.Spec.ModuleKey,
ModuleInstance: input.ModuleInstanceForStages,
ReplacementPolicy: input.Policy,
Glossary: input.Glossary,
Config: input.Config,
LLMClient: validationLLMClientAdapter{client: input.ValidationLLMClient},
Scheduler: input.ValidationScheduler,
DiagnosticsWriter: diagnosticsWriter,
})
if err != nil {
return validateSectionCandidatesResult{
approved: eligible,
decisions: decisions,
rejected: rejected,
}, fmt.Errorf("validator %q failed: %w", validator.Name(), err)
}
if err := validators.EnforceDecisionCardinality(eligible, vResult.Decisions); err != nil {
return validateSectionCandidatesResult{
approved: eligible,
decisions: decisions,
rejected: rejected,
}, fmt.Errorf("validator %q cardinality failed: %w", validator.Name(), err)
}
nextEligible := make([]proposals.EnrichedCorrectionProposal, 0, len(eligible))
byIndex := make(map[int]proposals.EnrichedCorrectionProposal, len(eligible))
for _, p := range eligible {
byIndex[p.ProposalIndex] = p
}
for _, d := range vResult.Decisions {
decisions = append(decisions, ValidatorDecisionRecord{
ValidatorName: validator.Name(),
ProposalIndex: d.ProposalIndex,
Approved: d.Approved,
ReasonCode: d.ReasonCode,
Message: d.Message,
DiagnosticArtifactPath: d.DiagnosticArtifactPath,
})
if d.Approved {
nextEligible = append(nextEligible, byIndex[d.ProposalIndex])
continue
}
p := byIndex[d.ProposalIndex]
rejected = append(rejected, ValidatorRejectedChange{
ValidatorName: validator.Name(),
ProposalIndex: p.ProposalIndex,
ModuleKey: p.ModuleKey,
ModuleInstance: p.ModuleInstance,
TargetSegmentID: p.TargetSegmentID,
OriginalText: p.OriginalText,
CorrectedText: p.CorrectedText,
ReasonCode: d.ReasonCode,
Message: d.Message,
})
}
eligible = nextEligible
}
return validateSectionCandidatesResult{
approved: eligible,
decisions: decisions,
rejected: rejected,
}, nil
}
func reorderValidatorsForPipeline(in []contracts.Validator) ([]contracts.Validator, map[string]int) {
deterministic := make([]contracts.Validator, 0, len(in))
llmBacked := make([]contracts.Validator, 0, len(in))
for _, validator := range in {
if _, ok := validator.(*validators.LLMBackedValidator); ok {
llmBacked = append(llmBacked, validator)
continue
}
deterministic = append(deterministic, validator)
}
ordered := append(deterministic, llmBacked...)
order := make(map[string]int, len(ordered))
for idx, validator := range ordered {
order[validator.Name()] = idx
}
return ordered, order
}
func chunkWorkingTranscript(cfg *config.Config, transcript *schema.Transcript) ([]chunking.Section, error) {

View File

@@ -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{