From 4c899122d96755e53bbcdeef7cfa5091ced9f29f Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Wed, 29 Apr 2026 20:10:00 -0500 Subject: [PATCH] Added a configuration flag to set the target number of sections, and updated default min and max section token limits --- README.md | 5 +- src/audita/cli.py | 6 +++ src/audita/core/chunking.py | 56 +++++++++++++++++++- src/audita/core/config.py | 25 ++++++++- src/audita/framework/runner.py | 3 +- tests/test_framework_chunking.py | 90 ++++++++++++++++++++++++++++++++ tests/test_framework_runner.py | 51 +++++++++++++++++- tests/test_new_cli.py | 55 +++++++++++++++++++ tests/test_new_config.py | 28 ++++++++++ 9 files changed, 312 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 6e53b6a..a380e7c 100644 --- a/README.md +++ b/README.md @@ -88,9 +88,10 @@ Useful configuration can be supplied by CLI flag or environment variable. CLI fl | `AUDITA_MODEL` | `--model` | `openrouter/google/gemma-4-31b-it` | LLM model name sent to the configured OpenAI-compatible endpoint | | `AUDITA_BASE_URL` | `--base-url` | `https://openrouter.ai/api/v1` | OpenAI-compatible API 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_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_MAX_SECTION_TOKENS` | `--max-section-tokens` | `16384` | Maximum estimated tokens per proposal-stage transcript section | -| `AUDITA_MIN_SECTION_TOKENS` | `--min-section-tokens` | `4096` | Minimum estimated tokens per proposal-stage transcript section when balancing for concurrency | +| `AUDITA_MAX_SECTION_TOKENS` | `--max-section-tokens` | `8192` | Maximum estimated tokens per proposal-stage transcript section | +| `AUDITA_MIN_SECTION_TOKENS` | `--min-section-tokens` | `2048` | Minimum estimated tokens per proposal-stage transcript section when balancing for concurrency | | `AUDITA_GLOSSARY_CONFIDENCE_THRESHOLD` | `--glossary-confidence-threshold` | `0.8` | Minimum confidence required for glossary proposals to survive validation | | `AUDITA_GRAMMAR_CONFIDENCE_THRESHOLD` | `--grammar-confidence-threshold` | `0.8` | Minimum confidence required for grammar proposals to survive validation | | `AUDITA_HOMOPHONES_CONFIDENCE_THRESHOLD` | `--homophones-confidence-threshold` | `0.8` | Minimum confidence required for homophone proposals to survive validation | diff --git a/src/audita/cli.py b/src/audita/cli.py index a7c27a1..c4b1f3d 100644 --- a/src/audita/cli.py +++ b/src/audita/cli.py @@ -34,6 +34,11 @@ def _build_parser() -> argparse.ArgumentParser: type=float, help="per-request timeout in seconds for LLM calls to the configured OpenAI-compatible endpoint", ) + process.add_argument( + "--target-sections", + type=int, + help="exact number of contiguous proposal-stage transcript sections to create", + ) process.add_argument( "--modules", help="comma-separated module keys to run, for example: grammar or glossary,homophones,glossary", @@ -103,6 +108,7 @@ def _process(args: argparse.Namespace) -> int: llm_api_key=args.llm_api_key, llm_concurrency=args.llm_concurrency, llm_timeout_seconds=args.llm_timeout_seconds, + target_sections=args.target_sections, module_keys=args.modules, model=args.model, base_url=args.base_url, diff --git a/src/audita/core/chunking.py b/src/audita/core/chunking.py index 8b5e995..cc578db 100644 --- a/src/audita/core/chunking.py +++ b/src/audita/core/chunking.py @@ -119,6 +119,7 @@ def chunk_transcript( max_section_tokens: int, min_section_tokens: int = 1, target_section_count: Optional[int] = None, + exact_target_section_count: Optional[int] = None, estimator: Optional[TokenEstimatorProtocol] = None, ) -> List[TranscriptSection]: indexed = [IndexedSegment(index=index, segment=segment) for index, segment in enumerate(segments)] @@ -127,6 +128,7 @@ def chunk_transcript( max_section_tokens, min_section_tokens=min_section_tokens, target_section_count=target_section_count, + exact_target_section_count=exact_target_section_count, estimator=estimator, ) @@ -136,9 +138,18 @@ def chunk_indexed_segments( max_section_tokens: int, min_section_tokens: int = 1, target_section_count: Optional[int] = None, + exact_target_section_count: Optional[int] = None, estimator: Optional[TokenEstimatorProtocol] = None, ) -> List[TranscriptSection]: - if target_section_count is None: + if exact_target_section_count is not None: + batches = _chunk_indexed_segments_for_exact_target_count( + indexed_segments, + max_section_tokens=max_section_tokens, + min_section_tokens=min_section_tokens, + target_section_count=exact_target_section_count, + estimator=estimator, + ) + elif target_section_count is None: batches = chunk_payload_items( indexed_segments, max_section_tokens, @@ -214,6 +225,49 @@ def _chunk_indexed_segments_for_target_count( return _build_balanced_batches(indexed_segments, single_tokens, section_count, token_estimator) +def _chunk_indexed_segments_for_exact_target_count( + indexed_segments: List[IndexedSegment], + *, + max_section_tokens: int, + min_section_tokens: int, + target_section_count: int, + estimator: Optional[TokenEstimatorProtocol], +) -> List[TokenBatch[IndexedSegment]]: + if max_section_tokens <= 0: + raise AuditaValidationError("Maximum section token count must be greater than zero.") + if min_section_tokens <= 0: + raise AuditaValidationError("Minimum section token count must be greater than zero.") + if min_section_tokens > max_section_tokens: + raise AuditaValidationError( + "Minimum section token count must be less than or equal to maximum section token count." + ) + if not indexed_segments: + raise AuditaValidationError("Transcript must contain at least one segment.") + if target_section_count <= 0: + raise AuditaValidationError("Target section count must be greater than zero.") + if target_section_count > len(indexed_segments): + raise AuditaValidationError( + "Target section count exceeds the number of transcript segments. " + "Lower AUDITA_TARGET_SECTIONS or pre-merge the transcript." + ) + + token_estimator = TokenEstimator() if estimator is None else estimator + single_tokens = [token_estimator.estimate_json([item.prompt_payload()]) for item in indexed_segments] + if any(tokens > max_section_tokens for tokens in single_tokens): + raise AuditaValidationError( + "A single transcript segment exceeds the maximum section token limit. " + "Raise the limit or pre-split the transcript." + ) + + batches = _build_balanced_batches(indexed_segments, single_tokens, target_section_count, token_estimator) + if not _batches_within_bounds(batches, min_section_tokens=min_section_tokens, max_section_tokens=max_section_tokens): + raise AuditaValidationError( + "AUDITA_TARGET_SECTIONS cannot produce contiguous transcript sections within " + "AUDITA_MIN_SECTION_TOKENS and AUDITA_MAX_SECTION_TOKENS." + ) + return batches + + def _resolve_section_count( *, indexed_segments: List[IndexedSegment], diff --git a/src/audita/core/config.py b/src/audita/core/config.py index 30bcfe2..9021e48 100644 --- a/src/audita/core/config.py +++ b/src/audita/core/config.py @@ -14,8 +14,8 @@ DEFAULT_BASE_URL = "https://openrouter.ai/api/v1" DEFAULT_LLM_CONCURRENCY = 1 DEFAULT_LLM_TIMEOUT_SECONDS = 600 DEFAULT_MAX_RETRIES = 3 -DEFAULT_MAX_SECTION_TOKENS = 16384 -DEFAULT_MIN_SECTION_TOKENS = 4096 +DEFAULT_MAX_SECTION_TOKENS = 8192 +DEFAULT_MIN_SECTION_TOKENS = 2048 DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD = 0.80 DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD = 0.80 DEFAULT_HOMOPHONES_CONFIDENCE_THRESHOLD = 0.80 @@ -33,6 +33,7 @@ class ConfigOverrides: llm_api_key: Optional[str] = None llm_concurrency: Optional[int] = None llm_timeout_seconds: Optional[float] = None + target_sections: Optional[int] = None module_keys: Optional[Union[str, Sequence[str]]] = None model: Optional[str] = None base_url: Optional[str] = None @@ -56,6 +57,7 @@ class AuditaConfig: api_key: Optional[str] = None llm_concurrency: int = DEFAULT_LLM_CONCURRENCY llm_timeout_seconds: float = DEFAULT_LLM_TIMEOUT_SECONDS + target_sections: Optional[int] = None module_keys: Tuple[str, ...] = DEFAULT_MODULE_KEYS model: str = DEFAULT_MODEL base_url: str = DEFAULT_BASE_URL @@ -99,6 +101,11 @@ class AuditaConfig: DEFAULT_LLM_TIMEOUT_SECONDS, "AUDITA_LLM_TIMEOUT_SECONDS", ), + target_sections=_select_optional_int( + selected.target_sections, + source.get("AUDITA_TARGET_SECTIONS"), + "AUDITA_TARGET_SECTIONS", + ), module_keys=_select_module_keys( selected.module_keys, source.get("AUDITA_MODULES"), @@ -193,6 +200,8 @@ class AuditaConfig: raise AuditaConfigError("AUDITA_LLM_TIMEOUT_SECONDS must be finite.") if self.llm_timeout_seconds <= 0: raise AuditaConfigError("AUDITA_LLM_TIMEOUT_SECONDS 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(): raise AuditaConfigError("AUDITA_MODEL must not be empty.") if not self.base_url.strip(): @@ -241,6 +250,7 @@ class AuditaConfig: "api_key_configured": bool(self.api_key), "llm_concurrency": self.llm_concurrency, "llm_timeout_seconds": self.llm_timeout_seconds, + "target_sections": self.target_sections, "module_keys": list(self.module_keys), "model": self.model, "base_url": self.base_url, @@ -290,6 +300,17 @@ def _select_int(cli_value: Optional[int], env_value: Optional[str], default: int raise AuditaConfigError(f"{name} must be an integer.") from exc +def _select_optional_int(cli_value: Optional[int], env_value: Optional[str], name: str) -> Optional[int]: + if cli_value is not None: + return cli_value + if env_value is None: + return None + try: + return int(env_value) + except ValueError as exc: + raise AuditaConfigError(f"{name} must be an integer.") from exc + + def _select_float(cli_value: Optional[float], env_value: Optional[str], default: float, name: str) -> float: if cli_value is not None: return cli_value diff --git a/src/audita/framework/runner.py b/src/audita/framework/runner.py index e1ded5a..ce4bb7a 100644 --- a/src/audita/framework/runner.py +++ b/src/audita/framework/runner.py @@ -92,7 +92,8 @@ class PipelineRunner: working, config.max_section_tokens, min_section_tokens=config.min_section_tokens, - target_section_count=config.llm_concurrency, + target_section_count=config.llm_concurrency if config.target_sections is None else None, + exact_target_section_count=config.target_sections, ) if progress is not None: progress( diff --git a/tests/test_framework_chunking.py b/tests/test_framework_chunking.py index 3ceaf6f..6935856 100644 --- a/tests/test_framework_chunking.py +++ b/tests/test_framework_chunking.py @@ -1,4 +1,7 @@ +import pytest + from audita.core.chunking import TokenEstimatorProtocol, chunk_transcript +from audita.core.errors import AuditaValidationError from audita.core.schemas import parse_transcript_json @@ -100,6 +103,93 @@ def test_chunk_transcript_reduces_section_count_when_target_sections_fall_below_ assert sections[0].token_count >= 8 +def test_chunk_transcript_exact_target_sections_returns_exact_count_when_feasible(): + transcript = parse_transcript_json( + """ + [ + {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"}, + {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"}, + {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}, + {"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"} + ] + """ + ) + + sections = chunk_transcript( + transcript, + max_section_tokens=8, + min_section_tokens=4, + exact_target_section_count=2, + estimator=FakeEstimator(), + ) + + assert len(sections) == 2 + assert [segment.segment.id for segment in sections[0].segments] == [1, 2] + assert [segment.segment.id for segment in sections[1].segments] == [3, 4] + + +def test_chunk_transcript_exact_target_sections_errors_when_target_exceeds_segment_count(): + transcript = parse_transcript_json( + """ + [ + {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"}, + {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"} + ] + """ + ) + + with pytest.raises(AuditaValidationError, match="Target section count exceeds the number of transcript segments"): + chunk_transcript( + transcript, + max_section_tokens=8, + min_section_tokens=4, + exact_target_section_count=3, + estimator=FakeEstimator(), + ) + + +def test_chunk_transcript_exact_target_sections_errors_when_section_would_exceed_max(): + transcript = parse_transcript_json( + """ + [ + {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"}, + {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"}, + {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"} + ] + """ + ) + + with pytest.raises(AuditaValidationError, match="AUDITA_TARGET_SECTIONS cannot produce contiguous transcript sections"): + chunk_transcript( + transcript, + max_section_tokens=8, + min_section_tokens=4, + exact_target_section_count=1, + estimator=FakeEstimator(), + ) + + +def test_chunk_transcript_exact_target_sections_errors_when_section_would_fall_below_min(): + transcript = parse_transcript_json( + """ + [ + {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"}, + {"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"}, + {"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"} + ] + """ + ) + + with pytest.raises(AuditaValidationError, match="AUDITA_TARGET_SECTIONS cannot produce contiguous transcript sections"): + chunk_transcript( + transcript, + max_section_tokens=12, + min_section_tokens=8, + exact_target_section_count=3, + estimator=FakeEstimator(), + ) + + def test_chunk_transcript_prompt_payload_includes_categories_when_present(): transcript = parse_transcript_json( """ diff --git a/tests/test_framework_runner.py b/tests/test_framework_runner.py index 09a59f2..0531b28 100644 --- a/tests/test_framework_runner.py +++ b/tests/test_framework_runner.py @@ -498,7 +498,7 @@ def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_s module = ConcurrentRecordingModule(seen, threading.Barrier(2, timeout=1.0)) monkeypatch.setattr( "audita.framework.runner.chunk_transcript", - lambda working, max_tokens, min_section_tokens=1, target_section_count=None: sections, + lambda working, max_tokens, min_section_tokens=1, target_section_count=None, exact_target_section_count=None: sections, ) runner = PipelineRunner() @@ -517,6 +517,55 @@ def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_s assert result.transcript[1].text == "Beta revised." +def test_pipeline_runner_passes_exact_target_sections_to_chunker(tmp_path, monkeypatch): + transcript = parse_transcript_json( + """ + [ + {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."} + ] + """ + ) + glossary = parse_glossary_yaml( + """ + glossary: + - name: "Alpha" + category: noun + summary: "Alpha." + """ + ) + seen_args = {} + sections = [ + TranscriptSection( + section_index=0, + start_index=0, + segments=[IndexedSegment(index=0, segment=transcript[0])], + token_count=1, + ) + ] + module = RecordingModule("noop", [], [], []) + + def _fake_chunk_transcript(working, max_tokens, min_section_tokens=1, target_section_count=None, exact_target_section_count=None): + seen_args["target_section_count"] = target_section_count + seen_args["exact_target_section_count"] = exact_target_section_count + return sections + + monkeypatch.setattr("audita.framework.runner.chunk_transcript", _fake_chunk_transcript) + + runner = PipelineRunner() + runner.run( + transcript=transcript, + glossary=glossary, + module_specs=[ModuleRunSpec(instance_name="noop", module_key="noop", module=module)], + config=AuditaConfig.from_sources(env={}, overrides=ConfigOverrides(target_sections=3)), + run_dir=tmp_path / "run", + ) + + assert seen_args == { + "target_section_count": None, + "exact_target_section_count": 3, + } + + def test_pipeline_runner_reports_first_llm_validator_rejection_in_chain_order(tmp_path): transcript = parse_transcript_json( """ diff --git a/tests/test_new_cli.py b/tests/test_new_cli.py index 436ef53..8cb84e2 100644 --- a/tests/test_new_cli.py +++ b/tests/test_new_cli.py @@ -23,6 +23,7 @@ def test_process_help_exposes_framework_flags(capsys): assert "--llm-api-key" in output assert "--llm-concurrency" in output assert "--llm-timeout-seconds" in output + assert "--target-sections" in output assert "--modules" in output assert "--model" in output assert "--base-url" in output @@ -307,6 +308,60 @@ def test_cli_process_passes_llm_timeout_seconds_override_to_config(monkeypatch, assert captured["llm_timeout_seconds"] == 900.0 +def test_cli_process_passes_target_sections_override_to_config(monkeypatch, tmp_path): + captured = {} + transcript = parse_transcript_json( + """ + [ + {"id": 1, "speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Fixed."} + ] + """ + ) + report = RunReport( + status="success", + config={"model": "m", "base_url": "b"}, + normalization={"source_segment_count": 1, "normalized_segment_count": 1, "merge_count": 0}, + pipeline=["grammar"], + modules=[], + applied_changes=[], + skipped_corrections=[], + totals={"output_segment_count": 1, "applied_change_count": 0, "skipped_correction_count": 0}, + work_dir_retention="auto", + work_dir_retained=False, + work_dir=None, + error=None, + ) + result = ProcessResult( + transcript=transcript, + report=report, + run_dir=tmp_path / "run", + work_dir_retained=False, + ) + + def _fake_from_sources(*, overrides=None): + captured["target_sections"] = overrides.target_sections + return object() + + monkeypatch.setattr("audita.cli.AuditaConfig.from_sources", _fake_from_sources) + monkeypatch.setattr("audita.cli.load_transcript", lambda path: []) + monkeypatch.setattr("audita.cli.load_glossary", lambda path: object()) + monkeypatch.setattr("audita.cli.process_transcript_result", lambda *args, **kwargs: result) + + exit_code = main( + [ + "process", + "transcript.json", + "--glossary", + "glossary.yaml", + "--target-sections", + "4", + ] + ) + + assert exit_code == 0 + assert captured["target_sections"] == 4 + + def test_cli_process_passes_min_section_tokens_override_to_config(monkeypatch, tmp_path): captured = {} transcript = parse_transcript_json( diff --git a/tests/test_new_config.py b/tests/test_new_config.py index 252ecca..74c3f38 100644 --- a/tests/test_new_config.py +++ b/tests/test_new_config.py @@ -24,6 +24,7 @@ def test_default_config_allows_missing_api_key(): assert config.api_key is None assert config.llm_concurrency == DEFAULT_LLM_CONCURRENCY assert config.llm_timeout_seconds == DEFAULT_LLM_TIMEOUT_SECONDS + assert config.target_sections is None assert config.max_section_tokens == DEFAULT_MAX_SECTION_TOKENS assert config.min_section_tokens == DEFAULT_MIN_SECTION_TOKENS assert config.module_keys == DEFAULT_MODULE_KEYS @@ -83,6 +84,21 @@ def test_llm_timeout_seconds_env_is_parsed(): assert config.llm_timeout_seconds == 120.5 +def test_target_sections_cli_override_takes_precedence(): + config = AuditaConfig.from_sources( + env={"AUDITA_TARGET_SECTIONS": "2"}, + overrides=ConfigOverrides(target_sections=5), + ) + + assert config.target_sections == 5 + + +def test_target_sections_env_is_parsed(): + config = AuditaConfig.from_sources(env={"AUDITA_TARGET_SECTIONS": "3"}) + + assert config.target_sections == 3 + + def test_min_section_tokens_env_is_parsed(): config = AuditaConfig.from_sources(env={"AUDITA_MIN_SECTION_TOKENS": "3000"}) @@ -153,6 +169,12 @@ def test_report_dict_includes_llm_timeout_seconds(): assert config.to_report_dict()["llm_timeout_seconds"] == 321.0 +def test_report_dict_includes_target_sections(): + config = AuditaConfig.from_sources(env={"AUDITA_TARGET_SECTIONS": "7"}) + + assert config.to_report_dict()["target_sections"] == 7 + + def test_threshold_overrides_take_precedence(): config = AuditaConfig.from_sources( env={ @@ -214,6 +236,12 @@ def test_invalid_llm_timeout_seconds_is_rejected(value): AuditaConfig.from_sources(env={"AUDITA_LLM_TIMEOUT_SECONDS": value}) +@pytest.mark.parametrize("value", ["0", "-1", "many"]) +def test_invalid_target_sections_is_rejected(value): + with pytest.raises(AuditaConfigError, match="AUDITA_TARGET_SECTIONS"): + AuditaConfig.from_sources(env={"AUDITA_TARGET_SECTIONS": value}) + + @pytest.mark.parametrize("value", ["0", "-1", "many"]) def test_invalid_min_section_tokens_is_rejected(value): with pytest.raises(AuditaConfigError, match="AUDITA_MIN_SECTION_TOKENS"):