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

Add explicit LLM concurrency controls

Browse Source
This commit is contained in:
Eric Rakestraw
2026-05-12 21:17:35 +00:00
parent a48f6da1f4
commit 509436cc4a
19 changed files with 875 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
internal/core/config/config.go
View File

@@ -36,6 +36,9 @@ type Config struct {
Modules []string
PrimaryLLM LLMConfig
ValidationLLM ValidationLLMConfig
TotalLLMConcurrency int
ProposalLLMConcurrency int
ValidationLLMConcurrency *int
ValidationMaxPromptTokens int
MaxSectionTokens int
MinSectionTokens int
@@ -52,7 +55,9 @@ type LLMConfig struct {
BaseURL string
TimeoutSeconds int
MaxRetries int
Concurrency int
// Concurrency is retained as a backward-compatible alias for
// TotalLLMConcurrency.
Concurrency int
}
type ValidationLLMConfig struct {
@@ -61,7 +66,9 @@ type ValidationLLMConfig struct {
BaseURL string
TimeoutSeconds *int
MaxRetries *int
Concurrency *int
// Concurrency is retained as a backward-compatible alias for
// ValidationLLMConcurrency.
Concurrency *int
}
type ConfidenceThresholds struct {
@@ -91,6 +98,9 @@ func Default() Config {
Concurrency: DefaultLLMConcurrency,
},
ValidationLLM: ValidationLLMConfig{},
TotalLLMConcurrency: DefaultLLMConcurrency,
ProposalLLMConcurrency: DefaultLLMConcurrency,
ValidationLLMConcurrency: nil,
ValidationMaxPromptTokens: DefaultValidationMaxPromptTokens,
MaxSectionTokens: DefaultMaxSectionTokens,
MinSectionTokens: DefaultMinSectionTokens,
@@ -146,9 +156,37 @@ func (c Config) EffectiveValidationLLMConfig() LLMConfig {
if c.ValidationLLM.MaxRetries != nil {
effective.MaxRetries = *c.ValidationLLM.MaxRetries
}
if c.ValidationLLM.Concurrency != nil {
effective.Concurrency = *c.ValidationLLM.Concurrency
}
effective.Concurrency = c.EffectiveValidationLLMConcurrency()
return effective
}
func (c Config) EffectiveValidationLLMConcurrency() int {
if c.ValidationLLMConcurrency != nil {
return *c.ValidationLLMConcurrency
}
return c.TotalLLMConcurrency
}
func (c Config) EffectiveProposalLLMConcurrency() int {
if c.ProposalLLMConcurrency > 0 {
return c.ProposalLLMConcurrency
}
return c.TotalLLMConcurrency
}
func (c *Config) syncLegacyConcurrencyAliases() {
if c == nil {
return
}
c.PrimaryLLM.Concurrency = c.TotalLLMConcurrency
c.ValidationLLM.Concurrency = intPtr(c.ValidationLLMConcurrency)
}
func intPtr(v *int) *int {
if v == nil {
return nil
}
x := *v
return &x
}

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

@@ -24,8 +24,17 @@ func TestDefaultConfigValues(t *testing.T) {
if cfg.PrimaryLLM.MaxRetries != DefaultMaxRetries {
t.Fatalf("unexpected default max retries: %d", cfg.PrimaryLLM.MaxRetries)
}
if cfg.PrimaryLLM.Concurrency != DefaultLLMConcurrency {
t.Fatalf("unexpected default llm concurrency: %d", cfg.PrimaryLLM.Concurrency)
if cfg.TotalLLMConcurrency != DefaultLLMConcurrency {
t.Fatalf("unexpected default total llm concurrency: %d", cfg.TotalLLMConcurrency)
}
if cfg.ProposalLLMConcurrency != DefaultLLMConcurrency {
t.Fatalf("unexpected default proposal llm concurrency: %d", cfg.ProposalLLMConcurrency)
}
if cfg.ValidationLLMConcurrency != nil {
t.Fatalf("expected validation llm concurrency to be unset by default")
}
if cfg.PrimaryLLM.Concurrency != cfg.TotalLLMConcurrency {
t.Fatalf("expected primary llm concurrency alias to mirror total, got primary=%d total=%d", cfg.PrimaryLLM.Concurrency, cfg.TotalLLMConcurrency)
}
if cfg.ValidationLLM.TimeoutSeconds != nil {
t.Fatalf("expected validation timeout to be unset by default")
@@ -34,7 +43,7 @@ func TestDefaultConfigValues(t *testing.T) {
t.Fatalf("expected validation max retries to be unset by default")
}
if cfg.ValidationLLM.Concurrency != nil {
t.Fatalf("expected validation llm concurrency to be unset by default")
t.Fatalf("expected legacy validation llm concurrency alias to be unset by default")
}
if cfg.TargetSections != nil {
t.Fatalf("expected target sections to be unset by default")
@@ -56,7 +65,8 @@ func TestLoadFromEnvOverridesAndFallback(t *testing.T) {
"AUDITA_BASE_URL": "https://api.openai.com/v1",
"AUDITA_LLM_TIMEOUT_SECONDS": "120",
"AUDITA_MAX_RETRIES": "7",
"AUDITA_LLM_CONCURRENCY": "6",
"AUDITA_TOTAL_LLM_CONCURRENCY": "6",
"AUDITA_PROPOSAL_LLM_CONCURRENCY": "4",
"AUDITA_VALIDATION_LLM_CONCURRENCY": "2",
"AUDITA_VALIDATION_MAX_PROMPT_TOKENS": "4096",
"AUDITA_MAX_SECTION_TOKENS": "9000",
@@ -92,17 +102,63 @@ func TestLoadFromEnvOverridesAndFallback(t *testing.T) {
if cfg.TargetSections == nil || *cfg.TargetSections != 5 {
t.Fatalf("unexpected target sections: %#v", cfg.TargetSections)
}
if cfg.TotalLLMConcurrency != 6 {
t.Fatalf("unexpected total llm concurrency: %d", cfg.TotalLLMConcurrency)
}
if cfg.ProposalLLMConcurrency != 4 {
t.Fatalf("unexpected proposal llm concurrency: %d", cfg.ProposalLLMConcurrency)
}
if cfg.ValidationLLMConcurrency == nil || *cfg.ValidationLLMConcurrency != 2 {
t.Fatalf("unexpected validation llm concurrency: %#v", cfg.ValidationLLMConcurrency)
}
if cfg.PrimaryLLM.Concurrency != 6 {
t.Fatalf("unexpected primary llm concurrency: %d", cfg.PrimaryLLM.Concurrency)
t.Fatalf("expected primary alias concurrency 6, got %d", cfg.PrimaryLLM.Concurrency)
}
if cfg.ValidationLLM.Concurrency == nil || *cfg.ValidationLLM.Concurrency != 2 {
t.Fatalf("unexpected validation llm concurrency: %#v", cfg.ValidationLLM.Concurrency)
t.Fatalf("expected validation alias concurrency 2, got %#v", cfg.ValidationLLM.Concurrency)
}
if cfg.WorkDirRetention != WorkDirRetentionAlways {
t.Fatalf("unexpected work dir retention: %q", cfg.WorkDirRetention)
}
}
func TestLoadFromEnvLegacyLLMConcurrencyAliasForTotalAndProposal(t *testing.T) {
env := map[string]string{
"AUDITA_LLM_CONCURRENCY": "5",
}
cfg, err := loadFromLookup(mapLookup(env))
if err != nil {
t.Fatalf("loadFromLookup returned error: %v", err)
}
if cfg.TotalLLMConcurrency != 5 {
t.Fatalf("expected total concurrency from legacy alias, got %d", cfg.TotalLLMConcurrency)
}
if cfg.ProposalLLMConcurrency != 5 {
t.Fatalf("expected proposal concurrency to inherit legacy total, got %d", cfg.ProposalLLMConcurrency)
}
}
func TestLoadFromEnvCanonicalTotalWinsLegacyAlias(t *testing.T) {
env := map[string]string{
"AUDITA_TOTAL_LLM_CONCURRENCY": "4",
"AUDITA_LLM_CONCURRENCY": "9",
}
cfg, err := loadFromLookup(mapLookup(env))
if err != nil {
t.Fatalf("loadFromLookup returned error: %v", err)
}
if cfg.TotalLLMConcurrency != 4 {
t.Fatalf("expected canonical total to win over legacy alias, got %d", cfg.TotalLLMConcurrency)
}
if cfg.ProposalLLMConcurrency != 4 {
t.Fatalf("expected proposal to inherit canonical total when unset, got %d", cfg.ProposalLLMConcurrency)
}
}
func TestLoadFromEnvUsesAuditaLLMAPIKeyOverFallback(t *testing.T) {
env := map[string]string{
"AUDITA_LLM_API_KEY": "primary-key",
@@ -127,12 +183,14 @@ func TestApplyCLIOverridesPrecedence(t *testing.T) {
model := "cli-model"
workDir := "/cli/work"
modules := "grammar"
llmConcurrency := 5
totalLLMConcurrency := 5
proposalLLMConcurrency := 3
overrides := CLIOverrides{
PrimaryModel: &model,
WorkDir: &workDir,
ModulesCSV: &modules,
PrimaryLLMConcurrency: &llmConcurrency,
PrimaryModel: &model,
WorkDir: &workDir,
ModulesCSV: &modules,
TotalLLMConcurrency: &totalLLMConcurrency,
ProposalLLMConcurrency: &proposalLLMConcurrency,
}
if err := cfg.ApplyCLIOverrides(overrides); err != nil {
@@ -148,17 +206,54 @@ 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)
if cfg.TotalLLMConcurrency != 5 {
t.Fatalf("expected CLI total concurrency override, got %d", cfg.TotalLLMConcurrency)
}
if cfg.ProposalLLMConcurrency != 3 {
t.Fatalf("expected CLI proposal concurrency override, got %d", cfg.ProposalLLMConcurrency)
}
}
func TestApplyCLIOverridesLegacyLLMConcurrencyAlias(t *testing.T) {
cfg := Default()
aliasConcurrency := 6
if err := cfg.ApplyCLIOverrides(CLIOverrides{PrimaryLLMConcurrency: &aliasConcurrency}); err != nil {
t.Fatalf("ApplyCLIOverrides failed: %v", err)
}
if cfg.TotalLLMConcurrency != 6 {
t.Fatalf("expected legacy --llm-concurrency alias to set total, got %d", cfg.TotalLLMConcurrency)
}
if cfg.ProposalLLMConcurrency != 6 {
t.Fatalf("expected proposal to inherit aliased total when unset, got %d", cfg.ProposalLLMConcurrency)
}
}
func TestApplyCLIOverridesCanonicalTotalWinsLegacyAlias(t *testing.T) {
cfg := Default()
canonicalTotal := 4
legacyAlias := 9
if err := cfg.ApplyCLIOverrides(CLIOverrides{TotalLLMConcurrency: &canonicalTotal, PrimaryLLMConcurrency: &legacyAlias}); err != nil {
t.Fatalf("ApplyCLIOverrides failed: %v", err)
}
if cfg.TotalLLMConcurrency != 4 {
t.Fatalf("expected canonical total concurrency to win, got %d", cfg.TotalLLMConcurrency)
}
if cfg.ProposalLLMConcurrency != 4 {
t.Fatalf("expected proposal to inherit canonical total when proposal is unset, got %d", cfg.ProposalLLMConcurrency)
}
}
func TestValidationFailures(t *testing.T) {
cfg := Default()
cfg.PrimaryLLM.TimeoutSeconds = -1
cfg.PrimaryLLM.Concurrency = 0
cfg.TotalLLMConcurrency = 0
cfg.ProposalLLMConcurrency = 0
validationConcurrency := 5
cfg.ValidationLLM.Concurrency = &validationConcurrency
cfg.ValidationLLMConcurrency = &validationConcurrency
cfg.ValidationMaxPromptTokens = 0
cfg.MaxSectionTokens = 100
cfg.MinSectionTokens = 200
@@ -173,8 +268,9 @@ func TestValidationFailures(t *testing.T) {
message := err.Error()
for _, expected := range []string{
"primary llm timeout seconds",
"primary llm concurrency",
"validation llm concurrency must be less than or equal to primary llm concurrency",
"total llm concurrency",
"proposal llm concurrency",
"validation llm concurrency must be less than or equal to total llm concurrency",
"validation max prompt tokens",
"min section tokens",
"grammar confidence threshold",
@@ -193,7 +289,8 @@ func TestEffectiveValidationLLMInheritance(t *testing.T) {
cfg.PrimaryLLM.BaseURL = "https://primary.example/v1"
cfg.PrimaryLLM.TimeoutSeconds = 111
cfg.PrimaryLLM.MaxRetries = 2
cfg.PrimaryLLM.Concurrency = 7
cfg.TotalLLMConcurrency = 7
cfg.syncLegacyConcurrencyAliases()
effective := cfg.EffectiveValidationLLMConfig()
if effective.APIKey != "primary-key" || effective.Model != "primary-model" || effective.BaseURL != "https://primary.example/v1" || effective.TimeoutSeconds != 111 || effective.MaxRetries != 2 || effective.Concurrency != 7 {
@@ -208,7 +305,8 @@ func TestEffectiveValidationLLMInheritance(t *testing.T) {
cfg.ValidationLLM.TimeoutSeconds = &validationTimeout
cfg.ValidationLLM.MaxRetries = &validationRetries
validationConcurrency := 4
cfg.ValidationLLM.Concurrency = &validationConcurrency
cfg.ValidationLLMConcurrency = &validationConcurrency
cfg.syncLegacyConcurrencyAliases()
effective = cfg.EffectiveValidationLLMConfig()
if effective.APIKey != "validation-key" || effective.Model != "validation-model" || effective.BaseURL != "https://validation.example/v1" || effective.TimeoutSeconds != 222 || effective.MaxRetries != 9 || effective.Concurrency != 4 {
@@ -216,34 +314,50 @@ func TestEffectiveValidationLLMInheritance(t *testing.T) {
}
}
func TestValidationLLMConcurrencyCannotExceedPrimary(t *testing.T) {
func TestValidationLLMConcurrencyCannotExceedTotal(t *testing.T) {
cfg := Default()
cfg.PrimaryLLM.Concurrency = 2
cfg.TotalLLMConcurrency = 2
cfg.ProposalLLMConcurrency = 2
validationConcurrency := 3
cfg.ValidationLLM.Concurrency = &validationConcurrency
cfg.ValidationLLMConcurrency = &validationConcurrency
if err := cfg.Validate(); err == nil {
t.Fatal("expected validation error when validation llm concurrency exceeds primary")
t.Fatal("expected validation error when validation llm concurrency exceeds total")
}
validationConcurrency = 2
cfg.ValidationLLM.Concurrency = &validationConcurrency
cfg.ValidationLLMConcurrency = &validationConcurrency
if err := cfg.Validate(); err != nil {
t.Fatalf("expected equal concurrency to validate, got %v", err)
}
}
func TestCLIPrimaryLLMConcurrencyOverrideDrivesEffectiveValidationConcurrencyWhenValidationUnset(t *testing.T) {
func TestProposalLLMConcurrencyCannotExceedTotal(t *testing.T) {
cfg := Default()
llmConcurrency := 6
if err := cfg.ApplyCLIOverrides(CLIOverrides{PrimaryLLMConcurrency: &llmConcurrency}); err != nil {
cfg.TotalLLMConcurrency = 2
cfg.ProposalLLMConcurrency = 3
if err := cfg.Validate(); err == nil {
t.Fatal("expected validation error when proposal llm concurrency exceeds total")
}
cfg.ProposalLLMConcurrency = 2
if err := cfg.Validate(); err != nil {
t.Fatalf("expected equal concurrency to validate, got %v", err)
}
}
func TestCLITotalLLMConcurrencyOverrideDrivesEffectiveValidationConcurrencyWhenValidationUnset(t *testing.T) {
cfg := Default()
totalLLMConcurrency := 6
if err := cfg.ApplyCLIOverrides(CLIOverrides{TotalLLMConcurrency: &totalLLMConcurrency}); 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.ValidationLLMConcurrency != nil {
t.Fatalf("expected validation concurrency to remain unset, got %#v", cfg.ValidationLLMConcurrency)
}
if cfg.EffectiveValidationLLMConfig().Concurrency != 6 {
t.Fatalf("expected inherited validation concurrency 6, got %d", cfg.EffectiveValidationLLMConfig().Concurrency)
if cfg.EffectiveValidationLLMConcurrency() != 6 {
t.Fatalf("expected inherited validation concurrency 6, got %d", cfg.EffectiveValidationLLMConcurrency())
}
}

31
internal/core/config/env.go
View File

@@ -68,12 +68,37 @@ func loadFromLookup(lookup func(string) (string, bool)) (Config, error) {
}
cfg.PrimaryLLM.MaxRetries = value
}
totalConcurrencySet := false
if raw, ok := lookup("AUDITA_TOTAL_LLM_CONCURRENCY"); ok {
value, err := parseInt(raw)
if err != nil {
return Config{}, fmt.Errorf("AUDITA_TOTAL_LLM_CONCURRENCY: %w", err)
}
cfg.TotalLLMConcurrency = value
totalConcurrencySet = true
}
if raw, ok := lookup("AUDITA_LLM_CONCURRENCY"); ok {
value, err := parseInt(raw)
if err != nil {
return Config{}, fmt.Errorf("AUDITA_LLM_CONCURRENCY: %w", err)
}
cfg.PrimaryLLM.Concurrency = value
if !totalConcurrencySet {
cfg.TotalLLMConcurrency = value
totalConcurrencySet = true
}
}
proposalConcurrencySet := false
if raw, ok := lookup("AUDITA_PROPOSAL_LLM_CONCURRENCY"); ok {
value, err := parseInt(raw)
if err != nil {
return Config{}, fmt.Errorf("AUDITA_PROPOSAL_LLM_CONCURRENCY: %w", err)
}
cfg.ProposalLLMConcurrency = value
proposalConcurrencySet = true
}
if totalConcurrencySet && !proposalConcurrencySet {
cfg.ProposalLLMConcurrency = cfg.TotalLLMConcurrency
}
if raw, ok := lookup("AUDITA_VALIDATION_MAX_RETRIES"); ok {
@@ -88,7 +113,7 @@ func loadFromLookup(lookup func(string) (string, bool)) (Config, error) {
if err != nil {
return Config{}, fmt.Errorf("AUDITA_VALIDATION_LLM_CONCURRENCY: %w", err)
}
cfg.ValidationLLM.Concurrency = &value
cfg.ValidationLLMConcurrency = &value
}
if raw, ok := lookup("AUDITA_VALIDATION_MAX_PROMPT_TOKENS"); ok {
@@ -188,6 +213,8 @@ func loadFromLookup(lookup func(string) (string, bool)) (Config, error) {
cfg.WorkDirRetention = WorkDirRetention(raw)
}
cfg.syncLegacyConcurrencyAliases()
if err := cfg.Validate(); err != nil {
return Config{}, err
}

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

@@ -11,6 +11,8 @@ type CLIOverrides struct {
PrimaryBaseURL *string
ValidationBaseURL *string
PrimaryLLMTimeoutSeconds *int
TotalLLMConcurrency *int
ProposalLLMConcurrency *int
PrimaryLLMConcurrency *int
ValidationLLMTimeoutSeconds *int
MaxRetries *int
@@ -62,8 +64,24 @@ 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
totalConcurrencySet := false
if overrides.TotalLLMConcurrency != nil {
c.TotalLLMConcurrency = *overrides.TotalLLMConcurrency
totalConcurrencySet = true
}
// Backward-compatible alias: --llm-concurrency maps to total concurrency
// only when --total-llm-concurrency is not set in the same CLI invocation.
if overrides.PrimaryLLMConcurrency != nil && !totalConcurrencySet {
c.TotalLLMConcurrency = *overrides.PrimaryLLMConcurrency
totalConcurrencySet = true
}
proposalConcurrencySet := false
if overrides.ProposalLLMConcurrency != nil {
c.ProposalLLMConcurrency = *overrides.ProposalLLMConcurrency
proposalConcurrencySet = true
}
if totalConcurrencySet && !proposalConcurrencySet {
c.ProposalLLMConcurrency = c.TotalLLMConcurrency
}
if overrides.ValidationLLMTimeoutSeconds != nil {
value := *overrides.ValidationLLMTimeoutSeconds
@@ -78,7 +96,7 @@ func (c *Config) ApplyCLIOverrides(overrides CLIOverrides) error {
}
if overrides.ValidationLLMConcurrency != nil {
value := *overrides.ValidationLLMConcurrency
c.ValidationLLM.Concurrency = &value
c.ValidationLLMConcurrency = &value
}
if overrides.ValidationMaxPromptTokens != nil {
c.ValidationMaxPromptTokens = *overrides.ValidationMaxPromptTokens
@@ -124,5 +142,7 @@ func (c *Config) ApplyCLIOverrides(overrides CLIOverrides) error {
c.WorkDirRetention = WorkDirRetention(*overrides.WorkDirRetention)
}
c.syncLegacyConcurrencyAliases()
return c.Validate()
}

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

@@ -24,8 +24,14 @@ func (c Config) Validate() error {
if c.PrimaryLLM.MaxRetries < 0 {
issues = append(issues, "max retries must be zero or greater")
}
if c.PrimaryLLM.Concurrency <= 0 {
issues = append(issues, "primary llm concurrency must be greater than zero")
if c.TotalLLMConcurrency <= 0 {
issues = append(issues, "total llm concurrency must be greater than zero")
}
if c.ProposalLLMConcurrency <= 0 {
issues = append(issues, "proposal llm concurrency must be greater than zero")
}
if c.ProposalLLMConcurrency > c.TotalLLMConcurrency {
issues = append(issues, "proposal llm concurrency must be less than or equal to total llm concurrency")
}
if c.ValidationLLM.TimeoutSeconds != nil && *c.ValidationLLM.TimeoutSeconds <= 0 {
@@ -34,11 +40,11 @@ func (c Config) Validate() error {
if c.ValidationLLM.MaxRetries != nil && *c.ValidationLLM.MaxRetries < 0 {
issues = append(issues, "validation max retries must be zero or greater")
}
if c.ValidationLLM.Concurrency != nil && *c.ValidationLLM.Concurrency <= 0 {
if c.ValidationLLMConcurrency != nil && *c.ValidationLLMConcurrency <= 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.ValidationLLMConcurrency != nil && *c.ValidationLLMConcurrency > c.TotalLLMConcurrency {
issues = append(issues, "validation llm concurrency must be less than or equal to total llm concurrency")
}
if c.ValidationMaxPromptTokens <= 0 {