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

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 {