eric/audita
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

Implemented an --llm-concurrency flag in the Go application that enforces a global LLM concurrency cap

Browse Source
This commit is contained in:
Eric Rakestraw
2026-05-12 12:59:46 -05:00
parent cad172a758
commit af84249da0
11 changed files with 520 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
internal/cli/run.go
View File

@@ -194,21 +194,28 @@ var processRunner = func(inv processInvocation, stdout io.Writer) (*normalizatio
validationLLMClient = client
}
}
if proposalScheduler == nil {
globalScheduler := proposalScheduler
if globalScheduler == nil {
primaryCfg := llm.ResolvePrimaryConfig(inv.Config)
s, sErr := llm.NewScheduler(primaryCfg.Concurrency)
if sErr != nil {
return fail("runner_setup", sErr, nil)
}
proposalScheduler = s
globalScheduler = s
}
if proposalScheduler == nil {
proposalScheduler = globalScheduler
}
if validationScheduler == nil {
validationCfg := llm.ResolveValidationConfig(inv.Config)
s, sErr := llm.NewScheduler(validationCfg.Concurrency)
if sErr != nil {
return fail("runner_setup", sErr, nil)
validationScheduler = globalScheduler
if inv.Config.ValidationLLM.Concurrency != nil && inv.Config.PrimaryLLM.Concurrency > inv.Config.EffectiveValidationLLMConfig().Concurrency {
validationCfg := llm.ResolveValidationConfig(inv.Config)
s, sErr := llm.NewScheduler(validationCfg.Concurrency)
if sErr != nil {
return fail("runner_setup", sErr, nil)
}
validationScheduler = composeSchedulers(globalScheduler, s)
}
validationScheduler = s
}
}
@@ -276,6 +283,42 @@ func sourceToCanonicalTranscript(source *schema.SourceTranscript) *schema.Transc
return &schema.Transcript{Segments: segments}
}
type chainedScheduler struct {
schedulers []runner.ValidationScheduler
}
func (s chainedScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
if len(s.schedulers) == 0 {
return fn(ctx)
}
run := fn
for i := len(s.schedulers) - 1; i >= 0; i-- {
scheduler := s.schedulers[i]
next := run
run = func(callCtx context.Context) error {
return scheduler.Run(callCtx, next)
}
}
return run(ctx)
}
func composeSchedulers(schedulers ...runner.ValidationScheduler) runner.ValidationScheduler {
filtered := make([]runner.ValidationScheduler, 0, len(schedulers))
for _, scheduler := range schedulers {
if scheduler != nil {
filtered = append(filtered, scheduler)
}
}
switch len(filtered) {
case 0:
return nil
case 1:
return filtered[0]
default:
return chainedScheduler{schedulers: filtered}
}
}
// Run executes the Audita CLI with the provided arguments and streams.
func Run(args []string, stdout, stderr io.Writer) int {
if len(args) == 0 {
@@ -349,6 +392,8 @@ func runProcess(args []string, stdout, stderr io.Writer) int {
overrides.ValidationBaseURL = pFlags.validationBaseURL
case "llm-timeout-seconds":
overrides.PrimaryLLMTimeoutSeconds = pFlags.llmTimeoutSeconds
case "llm-concurrency":
overrides.PrimaryLLMConcurrency = pFlags.llmConcurrency
case "validation-llm-timeout-seconds":
overrides.ValidationLLMTimeoutSeconds = pFlags.validationLLMTimeoutSeconds
case "max-retries":
@@ -637,6 +682,7 @@ type processFlags struct {
baseURL *string
validationBaseURL *string
llmTimeoutSeconds *int
llmConcurrency *int
validationLLMTimeoutSeconds *int
validationMaxPromptTokens *int
targetSections *int
@@ -693,6 +739,7 @@ func newProcessFlagSet(cfg config.Config, stderr io.Writer) (*flag.FlagSet, proc
baseURL: fs.String("base-url", cfg.PrimaryLLM.BaseURL, "Primary OpenAI-compatible base URL"),
validationBaseURL: fs.String("validation-base-url", cfg.ValidationLLM.BaseURL, "Validation OpenAI-compatible base URL"),
llmTimeoutSeconds: fs.Int("llm-timeout-seconds", cfg.PrimaryLLM.TimeoutSeconds, "Primary LLM timeout in seconds"),
llmConcurrency: fs.Int("llm-concurrency", cfg.PrimaryLLM.Concurrency, "Primary LLM concurrency"),
validationLLMTimeoutSeconds: fs.Int("validation-llm-timeout-seconds", validationTimeoutSecondsDefault, "Validation LLM timeout in seconds"),
validationMaxPromptTokens: fs.Int("validation-max-prompt-tokens", cfg.ValidationMaxPromptTokens, "Validation max prompt tokens"),
targetSections: fs.Int("target-sections", targetSectionsDefault, "Target number of transcript sections"),

132
internal/cli/run_test.go
View File

@@ -9,12 +9,16 @@ import (
"path/filepath"
"reflect"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"gitea.maximumdirect.net/eric/audita/internal/core/normalization"
"gitea.maximumdirect.net/eric/audita/internal/core/reporting"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
"gitea.maximumdirect.net/eric/audita/internal/framework/llm"
"gitea.maximumdirect.net/eric/audita/internal/framework/modules"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposal_generation"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
@@ -59,6 +63,7 @@ func TestRunProcessHelpListsExpectedFlags(t *testing.T) {
"--base-url",
"--validation-base-url",
"--llm-timeout-seconds",
"--llm-concurrency",
"--validation-llm-timeout-seconds",
"--validation-max-prompt-tokens",
"--target-sections",
@@ -241,6 +246,133 @@ func TestRunProcessCLIOverridesEnvironment(t *testing.T) {
}
}
func TestRunProcessRejectsValidationConcurrencyAbovePrimaryConcurrency(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
transcriptPath := writeFile(t, "transcript.json", `[
{"id":1,"speaker":"Alice","start":0.0,"end":1.0,"text":"Hello"}
]`)
exitCode := Run([]string{
"process",
transcriptPath,
"--glossary",
fixturePath("tiny_glossary.yaml"),
"--llm-concurrency",
"1",
"--validation-llm-concurrency",
"2",
}, &stdout, &stderr)
if exitCode == 0 {
t.Fatalf("expected nonzero exit code for invalid llm concurrency combination")
}
if stdout.Len() != 0 {
t.Fatalf("expected empty stdout, got %q", stdout.String())
}
if !strings.Contains(stderr.String(), "invalid CLI configuration") {
t.Fatalf("expected invalid CLI configuration error, got %q", stderr.String())
}
if !strings.Contains(stderr.String(), "validation llm concurrency must be less than or equal to primary llm concurrency") {
t.Fatalf("expected validation/primary concurrency error details, got %q", stderr.String())
}
}
func TestRunProcessLLMConcurrencyDrivesEffectiveValidationConcurrencyWhenUnset(t *testing.T) {
processModuleFactory = fakeModuleFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{
key: "m",
policy: proposals.ReplacementPolicyRequireUnique,
validators: []contracts.Validator{
fakeValidator{name: "capture-config", validateF: func(req contracts.ValidationRequest) (validators.Result, error) {
if req.Config == nil {
t.Fatal("expected config in validation request")
}
if req.Config.PrimaryLLM.Concurrency != 4 {
t.Fatalf("expected primary llm concurrency 4, got %d", req.Config.PrimaryLLM.Concurrency)
}
if req.Config.ValidationLLM.Concurrency != nil {
t.Fatalf("expected validation concurrency unset, got %#v", req.Config.ValidationLLM.Concurrency)
}
if req.Config.EffectiveValidationLLMConfig().Concurrency != 4 {
t.Fatalf("expected inherited effective validation concurrency 4, got %d", req.Config.EffectiveValidationLLMConfig().Concurrency)
}
return validators.Result{ValidatorName: "capture-config", Decisions: nil}, nil
}},
},
proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
return nil, nil
},
},
}}
t.Cleanup(func() { processModuleFactory = nil })
var stdout bytes.Buffer
var stderr bytes.Buffer
transcriptPath := writeFile(t, "transcript.json", `[
{"id":1,"speaker":"Alice","start":0.0,"end":1.0,"text":"Hello"}
]`)
exitCode := Run([]string{
"process",
transcriptPath,
"--glossary",
fixturePath("tiny_glossary.yaml"),
"--modules",
"m",
"--llm-concurrency",
"4",
}, &stdout, &stderr)
if exitCode != 0 {
t.Fatalf("expected success, got %d stderr=%q", exitCode, stderr.String())
}
}
func TestComposeSchedulersEnforcesStricterSubcap(t *testing.T) {
global, err := llm.NewScheduler(3)
if err != nil {
t.Fatalf("NewScheduler(global): %v", err)
}
subcap, err := llm.NewScheduler(1)
if err != nil {
t.Fatalf("NewScheduler(subcap): %v", err)
}
composed := composeSchedulers(global, subcap)
if composed == nil {
t.Fatal("expected composed scheduler")
}
var inFlight int32
var maxInFlight int32
var wg sync.WaitGroup
for i := 0; i < 6; i++ {
wg.Add(1)
go func() {
defer wg.Done()
runErr := composed.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
}
}
time.Sleep(20 * time.Millisecond)
atomic.AddInt32(&inFlight, -1)
return nil
})
if runErr != nil {
t.Errorf("scheduler run error: %v", runErr)
}
}()
}
wg.Wait()
if maxInFlight > 1 {
t.Fatalf("expected stricter subcap to govern composed scheduler, got max in-flight %d", maxInFlight)
}
}
func TestRunProcessReportJSONSuccessIncludesNormalizationSummary(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer

45
internal/core/config/config_test.go
View File

@@ -127,10 +127,12 @@ func TestApplyCLIOverridesPrecedence(t *testing.T) {
model := "cli-model"
workDir := "/cli/work"
modules := "grammar"
llmConcurrency := 5
overrides := CLIOverrides{
PrimaryModel: &model,
WorkDir: &workDir,
ModulesCSV: &modules,
PrimaryModel: &model,
WorkDir: &workDir,
ModulesCSV: &modules,
PrimaryLLMConcurrency: &llmConcurrency,
}
if err := cfg.ApplyCLIOverrides(overrides); err != nil {
@@ -146,12 +148,17 @@ func TestApplyCLIOverridesPrecedence(t *testing.T) {
if !reflect.DeepEqual(cfg.Modules, []string{"grammar"}) {
t.Fatalf("unexpected modules: %#v", cfg.Modules)
}
if cfg.PrimaryLLM.Concurrency != 5 {
t.Fatalf("expected CLI llm concurrency override, got %d", cfg.PrimaryLLM.Concurrency)
}
}
func TestValidationFailures(t *testing.T) {
cfg := Default()
cfg.PrimaryLLM.TimeoutSeconds = -1
cfg.PrimaryLLM.Concurrency = 0
validationConcurrency := 5
cfg.ValidationLLM.Concurrency = &validationConcurrency
cfg.ValidationMaxPromptTokens = 0
cfg.MaxSectionTokens = 100
cfg.MinSectionTokens = 200
@@ -167,6 +174,7 @@ func TestValidationFailures(t *testing.T) {
for _, expected := range []string{
"primary llm timeout seconds",
"primary llm concurrency",
"validation llm concurrency must be less than or equal to primary llm concurrency",
"validation max prompt tokens",
"min section tokens",
"grammar confidence threshold",
@@ -208,6 +216,37 @@ func TestEffectiveValidationLLMInheritance(t *testing.T) {
}
}
func TestValidationLLMConcurrencyCannotExceedPrimary(t *testing.T) {
cfg := Default()
cfg.PrimaryLLM.Concurrency = 2
validationConcurrency := 3
cfg.ValidationLLM.Concurrency = &validationConcurrency
if err := cfg.Validate(); err == nil {
t.Fatal("expected validation error when validation llm concurrency exceeds primary")
}
validationConcurrency = 2
cfg.ValidationLLM.Concurrency = &validationConcurrency
if err := cfg.Validate(); err != nil {
t.Fatalf("expected equal concurrency to validate, got %v", err)
}
}
func TestCLIPrimaryLLMConcurrencyOverrideDrivesEffectiveValidationConcurrencyWhenValidationUnset(t *testing.T) {
cfg := Default()
llmConcurrency := 6
if err := cfg.ApplyCLIOverrides(CLIOverrides{PrimaryLLMConcurrency: &llmConcurrency}); err != nil {
t.Fatalf("ApplyCLIOverrides failed: %v", err)
}
if cfg.ValidationLLM.Concurrency != nil {
t.Fatalf("expected validation concurrency to remain unset, got %#v", cfg.ValidationLLM.Concurrency)
}
if cfg.EffectiveValidationLLMConfig().Concurrency != 6 {
t.Fatalf("expected inherited validation concurrency 6, got %d", cfg.EffectiveValidationLLMConfig().Concurrency)
}
}
func TestRedactedConfig(t *testing.T) {
cfg := Default()
cfg.PrimaryLLM.APIKey = "secret-primary"

4
internal/core/config/flags.go
View File

@@ -11,6 +11,7 @@ type CLIOverrides struct {
PrimaryBaseURL *string
ValidationBaseURL *string
PrimaryLLMTimeoutSeconds *int
PrimaryLLMConcurrency *int
ValidationLLMTimeoutSeconds *int
MaxRetries *int
ValidationMaxRetries *int
@@ -61,6 +62,9 @@ func (c *Config) ApplyCLIOverrides(overrides CLIOverrides) error {
if overrides.PrimaryLLMTimeoutSeconds != nil {
c.PrimaryLLM.TimeoutSeconds = *overrides.PrimaryLLMTimeoutSeconds
}
if overrides.PrimaryLLMConcurrency != nil {
c.PrimaryLLM.Concurrency = *overrides.PrimaryLLMConcurrency
}
if overrides.ValidationLLMTimeoutSeconds != nil {
value := *overrides.ValidationLLMTimeoutSeconds
c.ValidationLLM.TimeoutSeconds = &value

3
internal/core/config/validation.go
View File

@@ -37,6 +37,9 @@ func (c Config) Validate() error {
if c.ValidationLLM.Concurrency != nil && *c.ValidationLLM.Concurrency <= 0 {
issues = append(issues, "validation llm concurrency must be greater than zero")
}
if c.ValidationLLM.Concurrency != nil && *c.ValidationLLM.Concurrency > c.PrimaryLLM.Concurrency {
issues = append(issues, "validation llm concurrency must be less than or equal to primary llm concurrency")
}
if c.ValidationMaxPromptTokens <= 0 {
issues = append(issues, "validation max prompt tokens must be greater than zero")

180
internal/framework/runner/runner.go
View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"path/filepath"
"sync"
"time"
"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
@@ -133,53 +134,29 @@ 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 := make([]proposals.EnrichedCorrectionProposal, 0)
nextProposalIndex := 0
for _, section := range sections {
sectionMeta := contracts.SectionMetadataFromSection(section)
sectionTranscript := transcriptFromSection(section)
proposed, proposeErr := module.Propose(ctx, contracts.ProposalRequest{
ExecutionContext: contracts.ExecutionContext{
Config: input.Config,
WorkingTranscript: sectionTranscript,
Glossary: input.Glossary,
Section: &sectionMeta,
DiagnosticsDir: input.ProposalDiagnosticsDir,
},
RunSpec: contracts.ModuleRunSpec{
ModuleKey: spec.ModuleKey,
InstanceName: spec.InstanceName,
ReplacementPolicy: policy,
},
LLMClient: input.ProposalLLMClient,
LLMScheduler: input.ProposalLLMScheduler,
})
if proposeErr != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ErrorMessage: proposeErr.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)
}
for _, p := range proposed {
sectionIndex := sectionMeta.Index
enriched = append(enriched, proposals.EnrichedCorrectionProposal{
CorrectionProposal: p,
ProposalMetadata: proposals.ProposalMetadata{
ProposalIndex: nextProposalIndex,
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
SectionIndex: &sectionIndex,
},
})
nextProposalIndex++
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,
})
if proposeErr != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ErrorMessage: proposeErr.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)
}
validatorDecisions := make([]ValidatorDecisionRecord, 0)
@@ -296,6 +273,117 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
return RunOutput{FinalTranscript: working, ModuleResults: results}, nil
}
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
}
type sectionProposals struct {
meta contracts.SectionMetadata
corrected []proposals.CorrectionProposal
}
func collectSectionProposals(ctx context.Context, input collectSectionProposalsInput) ([]proposals.EnrichedCorrectionProposal, error) {
if len(input.Sections) == 0 {
return []proposals.EnrichedCorrectionProposal{}, nil
}
maxWorkers := 1
if input.Config != nil && input.Config.PrimaryLLM.Concurrency > 1 {
maxWorkers = input.Config.PrimaryLLM.Concurrency
}
if maxWorkers > len(input.Sections) {
maxWorkers = len(input.Sections)
}
sectionResults := make([]sectionProposals, len(input.Sections))
runCtx, cancel := context.WithCancel(ctx)
defer cancel()
sem := make(chan struct{}, maxWorkers)
var (
wg sync.WaitGroup
errOnce sync.Once
firstErr error
)
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 }()
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()
})
return
}
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++
}
}
return enriched, nil
}
func chunkWorkingTranscript(cfg *config.Config, transcript *schema.Transcript) ([]chunking.Section, error) {
if cfg == nil {
if transcript == nil || len(transcript.Segments) == 0 {

128
internal/framework/runner/runner_test.go
View File

@@ -6,7 +6,10 @@ import (
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"gitea.maximumdirect.net/eric/audita/internal/core/config"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
@@ -112,15 +115,20 @@ func TestRunnerProposalsExecutePerChunkSection(t *testing.T) {
{ID: 3, Text: "five six"},
}}
seenSections := make([]contracts.SectionMetadata, 0)
seenSegmentCounts := make([]int, 0)
var (
mu sync.Mutex
seenBySection = make(map[int]contracts.SectionMetadata)
segmentCountBySec = make(map[int]int)
)
r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
if req.Section == nil {
t.Fatalf("expected section metadata on proposal request")
}
seenSections = append(seenSections, *req.Section)
seenSegmentCounts = append(seenSegmentCounts, len(req.WorkingTranscript.Segments))
mu.Lock()
seenBySection[req.Section.Index] = *req.Section
segmentCountBySec[req.Section.Index] = len(req.WorkingTranscript.Segments)
mu.Unlock()
return nil, nil
}},
}})
@@ -138,15 +146,117 @@ func TestRunnerProposalsExecutePerChunkSection(t *testing.T) {
if len(out.ModuleResults) != 1 {
t.Fatalf("expected one module result, got %+v", out.ModuleResults)
}
if len(seenSections) != 3 {
t.Fatalf("expected 3 chunked proposal calls, got %d", len(seenSections))
if len(seenBySection) != 3 {
t.Fatalf("expected 3 chunked proposal calls, got %d", len(seenBySection))
}
for i, section := range seenSections {
for i := 0; i < 3; i++ {
section, ok := seenBySection[i]
if !ok {
t.Fatalf("expected section %d to be observed", i)
}
if section.Index != i {
t.Fatalf("expected section index %d, got %+v", i, section)
}
if seenSegmentCounts[i] != 1 {
t.Fatalf("expected one segment per section call, got %d at index %d", seenSegmentCounts[i], i)
if segmentCountBySec[i] != 1 {
t.Fatalf("expected one segment per section call, got %d at index %d", segmentCountBySec[i], i)
}
}
}
func TestRunnerProposalSectionConcurrencyBoundedByPrimaryLLMConcurrency(t *testing.T) {
cfg := config.Default()
cfg.MaxSectionTokens = 3
cfg.MinSectionTokens = 0
cfg.PrimaryLLM.Concurrency = 2
transcript := &schema.Transcript{Segments: []schema.Segment{
{ID: 1, Text: "one two"},
{ID: 2, Text: "three four"},
{ID: 3, Text: "five six"},
{ID: 4, Text: "seven eight"},
}}
var inFlight int32
var maxInFlight int32
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
}
}
time.Sleep(25 * time.Millisecond)
atomic.AddInt32(&inFlight, -1)
return nil, nil
}},
}})
_, err := r.Run(context.Background(), RunInput{
Config: &cfg,
Transcript: transcript,
ModuleSpecs: []contracts.ModuleRunSpec{
{ModuleKey: "m", InstanceName: "m"},
},
})
if err != nil {
t.Fatalf("Run error: %v", err)
}
if maxInFlight < 2 {
t.Fatalf("expected proposal execution to run concurrently, got max in-flight %d", maxInFlight)
}
if maxInFlight > 2 {
t.Fatalf("expected proposal concurrency <= 2, got %d", maxInFlight)
}
}
func TestRunnerProposalIndexOrderingIsDeterministicAcrossParallelSections(t *testing.T) {
cfg := config.Default()
cfg.MaxSectionTokens = 3
cfg.MinSectionTokens = 0
cfg.PrimaryLLM.Concurrency = 3
transcript := &schema.Transcript{Segments: []schema.Segment{
{ID: 1, Text: "alpha one"},
{ID: 2, Text: "bravo two"},
{ID: 3, Text: "charlie three"},
}}
r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
if len(req.WorkingTranscript.Segments) != 1 {
t.Fatalf("expected one segment per section, got %d", len(req.WorkingTranscript.Segments))
}
seg := req.WorkingTranscript.Segments[0]
// Sleep inversely by ID to force completion order to differ from section order.
time.Sleep(time.Duration(4-seg.ID) * 10 * time.Millisecond)
return []proposals.CorrectionProposal{
{TargetSegmentID: seg.ID, OriginalText: seg.Text, CorrectedText: strings.ToUpper(seg.Text), Confidence: 1},
}, nil
}},
}})
out, err := r.Run(context.Background(), RunInput{
Config: &cfg,
Transcript: transcript,
ModuleSpecs: []contracts.ModuleRunSpec{
{ModuleKey: "m", InstanceName: "m"},
},
})
if err != nil {
t.Fatalf("Run error: %v", err)
}
if len(out.ModuleResults) != 1 {
t.Fatalf("expected one module result, got %+v", out.ModuleResults)
}
if got := len(out.ModuleResults[0].AppliedChanges); got != 3 {
t.Fatalf("expected three applied changes, got %d", got)
}
for i, change := range out.ModuleResults[0].AppliedChanges {
wantID := i + 1
if change.TargetSegmentID != wantID {
t.Fatalf("expected deterministic applied-change order by section/segment, got %+v", out.ModuleResults[0].AppliedChanges)
}
}
}

9
internal/modules/glossary/module.go
View File

@@ -69,7 +69,7 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
Glossary: req.Glossary,
Config: req.Config,
Messages: messages,
StageName: fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName),
StageName: proposalStageName(req),
StartIndex: 0,
LLMClient: req.LLMClient,
Scheduler: req.LLMScheduler,
@@ -93,3 +93,10 @@ func transcriptForSection(transcript *schema.Transcript, section *contracts.Sect
}
return &schema.Transcript{Segments: segments}
}
func proposalStageName(req contracts.ProposalRequest) string {
if req.Section == nil || req.Section.Index == 0 {
return fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName)
}
return fmt.Sprintf("%s:proposal:section-%04d", req.RunSpec.InstanceName, req.Section.Index)
}

9
internal/modules/grammar/module.go
View File

@@ -69,7 +69,7 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
Glossary: req.Glossary,
Config: req.Config,
Messages: messages,
StageName: fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName),
StageName: proposalStageName(req),
StartIndex: 0,
LLMClient: req.LLMClient,
Scheduler: req.LLMScheduler,
@@ -93,3 +93,10 @@ func transcriptForSection(transcript *schema.Transcript, section *contracts.Sect
}
return &schema.Transcript{Segments: segments}
}
func proposalStageName(req contracts.ProposalRequest) string {
if req.Section == nil || req.Section.Index == 0 {
return fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName)
}
return fmt.Sprintf("%s:proposal:section-%04d", req.RunSpec.InstanceName, req.Section.Index)
}

9
internal/modules/homophones/module.go
View File

@@ -69,7 +69,7 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
Glossary: req.Glossary,
Config: req.Config,
Messages: messages,
StageName: fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName),
StageName: proposalStageName(req),
StartIndex: 0,
LLMClient: req.LLMClient,
Scheduler: req.LLMScheduler,
@@ -93,3 +93,10 @@ func transcriptForSection(transcript *schema.Transcript, section *contracts.Sect
}
return &schema.Transcript{Segments: segments}
}
func proposalStageName(req contracts.ProposalRequest) string {
if req.Section == nil || req.Section.Index == 0 {
return fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName)
}
return fmt.Sprintf("%s:proposal:section-%04d", req.RunSpec.InstanceName, req.Section.Index)
}

9
internal/modules/spoken_word/module.go
View File

@@ -69,7 +69,7 @@ func (m *Module) Propose(ctx context.Context, req contracts.ProposalRequest) ([]
Glossary: req.Glossary,
Config: req.Config,
Messages: messages,
StageName: fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName),
StageName: proposalStageName(req),
StartIndex: 0,
LLMClient: req.LLMClient,
Scheduler: req.LLMScheduler,
@@ -93,3 +93,10 @@ func transcriptForSection(transcript *schema.Transcript, section *contracts.Sect
}
return &schema.Transcript{Segments: segments}
}
func proposalStageName(req contracts.ProposalRequest) string {
if req.Section == nil || req.Section.Index == 0 {
return fmt.Sprintf("%s:proposal", req.RunSpec.InstanceName)
}
return fmt.Sprintf("%s:proposal:section-%04d", req.RunSpec.InstanceName, req.Section.Index)
}