From 0d5bdd1c84a646dc0e88bf984dca28dd55375ca6 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Wed, 29 Apr 2026 22:33:45 -0500 Subject: [PATCH] Added a configuration flag to set the maximum number of tokens in validation prompts --- README.md | 1 + src/audita/cli.py | 6 +++ src/audita/core/config.py | 12 ++++++ src/audita/validators/llm.py | 2 +- tests/test_llm_validators.py | 73 +++++++++++++++++++++++++++++++++++- tests/test_new_cli.py | 5 +++ tests/test_new_config.py | 24 ++++++++++++ 7 files changed, 121 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 0181613..e13d571 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,7 @@ Useful configuration can be supplied by CLI flag or environment variable. CLI fl | `AUDITA_VALIDATION_BASE_URL` | `--validation-base-url` | unset | Validation-phase OpenAI-compatible API base URL; defaults to `AUDITA_BASE_URL` | | `AUDITA_LLM_TIMEOUT_SECONDS` | `--llm-timeout-seconds` | `600` | Per-request timeout in seconds for LLM calls to the configured OpenAI-compatible endpoint | | `AUDITA_VALIDATION_LLM_TIMEOUT_SECONDS` | `--validation-llm-timeout-seconds` | unset | Validation-phase per-request timeout in seconds; defaults to `AUDITA_LLM_TIMEOUT_SECONDS` | +| `AUDITA_VALIDATION_MAX_PROMPT_TOKENS` | `--validation-max-prompt-tokens` | `2048` | Maximum estimated tokens per validation-phase LLM prompt batch | | `AUDITA_TARGET_SECTIONS` | `--target-sections` | unset | Exact number of contiguous proposal-stage transcript sections; errors if min/max token bounds cannot be satisfied | | `AUDITA_MAX_RETRIES` | `--max-retries` | `3` | Maximum Instructor retries for structured responses | | `AUDITA_VALIDATION_MAX_RETRIES` | `--validation-max-retries` | unset | Validation-phase structured-output retries; defaults to `AUDITA_MAX_RETRIES` | diff --git a/src/audita/cli.py b/src/audita/cli.py index a7795cc..fe40519 100644 --- a/src/audita/cli.py +++ b/src/audita/cli.py @@ -58,6 +58,11 @@ def _build_parser() -> argparse.ArgumentParser: type=int, help="maximum structured-output retries for validation phases; defaults to --max-retries", ) + process.add_argument( + "--validation-max-prompt-tokens", + type=int, + help="maximum estimated tokens per validation-phase LLM prompt batch", + ) process.add_argument( "--target-sections", type=int, @@ -138,6 +143,7 @@ def _process(args: argparse.Namespace) -> int: validation_model=args.validation_model, validation_base_url=args.validation_base_url, validation_max_retries=args.validation_max_retries, + validation_max_prompt_tokens=args.validation_max_prompt_tokens, target_sections=args.target_sections, module_keys=args.modules, model=args.model, diff --git a/src/audita/core/config.py b/src/audita/core/config.py index 7dc5219..94f6c58 100644 --- a/src/audita/core/config.py +++ b/src/audita/core/config.py @@ -14,6 +14,7 @@ DEFAULT_BASE_URL = "https://openrouter.ai/api/v1" DEFAULT_LLM_CONCURRENCY = 1 DEFAULT_LLM_TIMEOUT_SECONDS = 600 DEFAULT_MAX_RETRIES = 3 +DEFAULT_VALIDATION_MAX_PROMPT_TOKENS = 2048 DEFAULT_MAX_SECTION_TOKENS = 8192 DEFAULT_MIN_SECTION_TOKENS = 2048 DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD = 0.80 @@ -39,6 +40,7 @@ class ConfigOverrides: validation_model: Optional[str] = None validation_base_url: Optional[str] = None validation_max_retries: Optional[int] = None + validation_max_prompt_tokens: Optional[int] = None target_sections: Optional[int] = None module_keys: Optional[Union[str, Sequence[str]]] = None model: Optional[str] = None @@ -69,6 +71,7 @@ class AuditaConfig: validation_model: Optional[str] = None validation_base_url: Optional[str] = None validation_max_retries: Optional[int] = None + validation_max_prompt_tokens: int = DEFAULT_VALIDATION_MAX_PROMPT_TOKENS target_sections: Optional[int] = None module_keys: Tuple[str, ...] = DEFAULT_MODULE_KEYS model: str = DEFAULT_MODEL @@ -140,6 +143,12 @@ class AuditaConfig: source.get("AUDITA_VALIDATION_MAX_RETRIES"), "AUDITA_VALIDATION_MAX_RETRIES", ), + validation_max_prompt_tokens=_select_int( + selected.validation_max_prompt_tokens, + source.get("AUDITA_VALIDATION_MAX_PROMPT_TOKENS"), + DEFAULT_VALIDATION_MAX_PROMPT_TOKENS, + "AUDITA_VALIDATION_MAX_PROMPT_TOKENS", + ), target_sections=_select_optional_int( selected.target_sections, source.get("AUDITA_TARGET_SECTIONS"), @@ -247,6 +256,8 @@ class AuditaConfig: raise AuditaConfigError("AUDITA_VALIDATION_LLM_TIMEOUT_SECONDS must be greater than zero.") if self.validation_max_retries is not None and self.validation_max_retries < 0: raise AuditaConfigError("AUDITA_VALIDATION_MAX_RETRIES must be greater than or equal to zero.") + if self.validation_max_prompt_tokens <= 0: + raise AuditaConfigError("AUDITA_VALIDATION_MAX_PROMPT_TOKENS must be greater than zero.") if self.target_sections is not None and self.target_sections <= 0: raise AuditaConfigError("AUDITA_TARGET_SECTIONS must be greater than zero.") if not self.model.strip(): @@ -355,6 +366,7 @@ class AuditaConfig: "validation_model": self.validation_model, "validation_base_url": self.validation_base_url, "validation_max_retries": self.validation_max_retries, + "validation_max_prompt_tokens": self.validation_max_prompt_tokens, "target_sections": self.target_sections, "module_keys": list(self.module_keys), "model": self.model, diff --git a/src/audita/validators/llm.py b/src/audita/validators/llm.py index e261b79..bc62f08 100644 --- a/src/audita/validators/llm.py +++ b/src/audita/validators/llm.py @@ -75,7 +75,7 @@ class _BaseLLMValidator: raise AuditaLLMError(f"Validator '{self.name}' requires a structured LLM client.") batches = chunk_payload_items( previewable, - context.config.max_section_tokens, + context.config.validation_max_prompt_tokens, payload_fn=lambda item: item.to_prompt_payload(), empty_error_message="Validation input must contain at least one proposal.", ) diff --git a/tests/test_llm_validators.py b/tests/test_llm_validators.py index aeb9eb7..163d6f1 100644 --- a/tests/test_llm_validators.py +++ b/tests/test_llm_validators.py @@ -1092,7 +1092,10 @@ def test_llm_validators_use_shared_token_batching_helper(monkeypatch, tmp_path): assert [decision.approved for decision in result.decisions] == [True, True, True] assert len(client.calls) == 2 assert len(chunk_calls) == 1 - assert chunk_calls[0]["max_tokens"] == AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}).max_section_tokens + assert ( + chunk_calls[0]["max_tokens"] + == AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}).validation_max_prompt_tokens + ) assert all("corrected_segment_text" in payload for payload in chunk_calls[0]["payloads"]) @@ -1272,6 +1275,74 @@ def test_validator_uses_validation_llm_config(tmp_path): assert validation_config.llm_concurrency == 3 +def test_validator_batches_use_validation_max_prompt_tokens(tmp_path, monkeypatch): + transcript = parse_transcript_json( + """ + [ + {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "There were gestures at the temple."} + ] + """ + ) + proposals = [ + CorrectionProposal( + proposal_index=0, + module_instance="homophones", + module_key="homophones", + id=1, + original_text="gestures", + corrected_text="Jesters", + confidence=0.95, + ) + ] + client = FakeStructuredLLMClient( + [ + { + "validations": [ + { + "correction_index": 0, + "approved": True, + "confidence": 0.97, + "reason": "ok", + } + ] + } + ] + ) + chunk_calls = [] + + def fake_chunk_payload_items(items, max_tokens, payload_fn, empty_error_message): + chunk_calls.append( + { + "count": len(items), + "max_tokens": max_tokens, + } + ) + return [TokenBatch(batch_index=0, items=list(items), token_count=1)] + + monkeypatch.setattr(llm_module, "chunk_payload_items", fake_chunk_payload_items) + config = AuditaConfig.from_sources( + env={"OPENROUTER_API_KEY": "test-key"}, + overrides=ConfigOverrides(validation_max_prompt_tokens=1024), + ) + context = ValidationContext( + proposals=proposals, + transcript=transcript, + glossary=_glossary(), + config=config.validation_llm_config(), + run_spec=ModuleRunSpec( + instance_name="homophones", + module_key="homophones", + module=_Module("require_unique"), + ), + run_dir=tmp_path, + llm_client=client, + ) + + SpokenFormPlausibilityValidator("spoken_form_plausibility_review").validate(context) + + assert chunk_calls[0]["max_tokens"] == 1024 + + def test_validator_uses_validation_llm_concurrency_override(tmp_path, monkeypatch): transcript = parse_transcript_json( """ diff --git a/tests/test_new_cli.py b/tests/test_new_cli.py index 0b06900..b1b778a 100644 --- a/tests/test_new_cli.py +++ b/tests/test_new_cli.py @@ -29,6 +29,7 @@ def test_process_help_exposes_framework_flags(capsys): assert "--validation-model" in output assert "--validation-base-url" in output assert "--validation-max-retries" in output + assert "--validation-max-prompt-tokens" in output assert "--target-sections" in output assert "--modules" in output assert "--model" in output @@ -351,6 +352,7 @@ def test_cli_process_passes_validation_llm_overrides_to_config(monkeypatch, tmp_ captured["validation_model"] = overrides.validation_model captured["validation_base_url"] = overrides.validation_base_url captured["validation_max_retries"] = overrides.validation_max_retries + captured["validation_max_prompt_tokens"] = overrides.validation_max_prompt_tokens return object() monkeypatch.setattr("audita.cli.AuditaConfig.from_sources", _fake_from_sources) @@ -376,6 +378,8 @@ def test_cli_process_passes_validation_llm_overrides_to_config(monkeypatch, tmp_ "http://localhost:9000/v1", "--validation-max-retries", "2", + "--validation-max-prompt-tokens", + "1024", ] ) @@ -387,6 +391,7 @@ def test_cli_process_passes_validation_llm_overrides_to_config(monkeypatch, tmp_ "validation_model": "validator-model", "validation_base_url": "http://localhost:9000/v1", "validation_max_retries": 2, + "validation_max_prompt_tokens": 1024, } diff --git a/tests/test_new_config.py b/tests/test_new_config.py index c758e2b..5c4829e 100644 --- a/tests/test_new_config.py +++ b/tests/test_new_config.py @@ -12,6 +12,7 @@ from audita.core.config import ( DEFAULT_MIN_SECTION_TOKENS, DEFAULT_NORMALIZE_MAX_SEGMENT_GAP, DEFAULT_SPOKEN_WORD_CONFIDENCE_THRESHOLD, + DEFAULT_VALIDATION_MAX_PROMPT_TOKENS, DEFAULT_WORK_DIR_RETENTION, ) from audita.core.errors import AuditaConfigError @@ -30,6 +31,7 @@ def test_default_config_allows_missing_api_key(): assert config.validation_model is None assert config.validation_base_url is None assert config.validation_max_retries is None + assert config.validation_max_prompt_tokens == DEFAULT_VALIDATION_MAX_PROMPT_TOKENS assert config.target_sections is None assert config.max_section_tokens == DEFAULT_MAX_SECTION_TOKENS assert config.min_section_tokens == DEFAULT_MIN_SECTION_TOKENS @@ -138,6 +140,21 @@ def test_validation_max_retries_env_is_parsed(): assert config.validation_max_retries == 7 +def test_validation_max_prompt_tokens_cli_override_takes_precedence(): + config = AuditaConfig.from_sources( + env={"AUDITA_VALIDATION_MAX_PROMPT_TOKENS": "1024"}, + overrides=ConfigOverrides(validation_max_prompt_tokens=4096), + ) + + assert config.validation_max_prompt_tokens == 4096 + + +def test_validation_max_prompt_tokens_env_is_parsed(): + config = AuditaConfig.from_sources(env={"AUDITA_VALIDATION_MAX_PROMPT_TOKENS": "1024"}) + + assert config.validation_max_prompt_tokens == 1024 + + def test_target_sections_cli_override_takes_precedence(): config = AuditaConfig.from_sources( env={"AUDITA_TARGET_SECTIONS": "2"}, @@ -296,6 +313,7 @@ def test_report_dict_includes_effective_validation_llm_config(): report = config.to_report_dict() assert report["validation_model"] == "validation-model" + assert report["validation_max_prompt_tokens"] == DEFAULT_VALIDATION_MAX_PROMPT_TOKENS assert report["effective_validation_llm"]["api_key_configured"] is True assert report["effective_validation_llm"]["model"] == "validation-model" @@ -379,6 +397,12 @@ def test_invalid_validation_max_retries_is_rejected(value): AuditaConfig.from_sources(env={"AUDITA_VALIDATION_MAX_RETRIES": value}) +@pytest.mark.parametrize("value", ["0", "-1", "many"]) +def test_invalid_validation_max_prompt_tokens_is_rejected(value): + with pytest.raises(AuditaConfigError, match="AUDITA_VALIDATION_MAX_PROMPT_TOKENS"): + AuditaConfig.from_sources(env={"AUDITA_VALIDATION_MAX_PROMPT_TOKENS": value}) + + @pytest.mark.parametrize("value", ["0", "-1", "many"]) def test_invalid_target_sections_is_rejected(value): with pytest.raises(AuditaConfigError, match="AUDITA_TARGET_SECTIONS"):