Implemented a --modules CLI flag to allow runtime selection of the modules to be run

This commit is contained in:
2026-04-25 07:48:21 -05:00
parent 2d1d21d314
commit 92c8c371a6
9 changed files with 277 additions and 36 deletions

View File

@@ -28,17 +28,31 @@ uv run audita process transcript.json --glossary glossary.yaml --output correcte
The framework currently runs this default module sequence:
1. `glossary_primary`
1. `glossary`
2. `homophones`
3. `glossary_secondary`
3. `glossary`
4. `spoken_word`
5. `grammar`
Resolved run instance names are auto-numbered for repeats, so the default report pipeline is:
1. `glossary_1`
2. `homophones`
3. `glossary_2`
4. `spoken_word`
5. `grammar`
The default module sequence is partially implemented today:
- `glossary_primary`, `homophones`, and `glossary_secondary` run real LLM-backed proposal and validation stages
- `glossary`, `homophones`, and the second `glossary` pass run real LLM-backed proposal and validation stages
- `spoken_word` and `grammar` remain stubs and currently propose no corrections
To run a custom module sequence, pass `--modules`:
```sh
uv run audita process transcript.json --glossary glossary.yaml --modules grammar --output corrected.json
```
To also write a structured JSON report:
```sh
@@ -67,6 +81,7 @@ Useful configuration can be supplied by CLI flag or environment variable. CLI fl
| Environment variable | CLI flag | Default | Purpose |
| --- | --- | --- | --- |
| `AUDITA_MODULES` | `--modules` | `glossary,homophones,glossary,spoken_word,grammar` | Comma-separated logical module keys to run; CLI overrides the environment value |
| `AUDITA_MODEL` | `--model` | `openrouter/google/gemma-4-31b-it` | LLM model used by glossary and homophones proposal/validation stages |
| `AUDITA_BASE_URL` | `--base-url` | `https://openrouter.ai/api/v1` | OpenAI-compatible API base URL |
| `AUDITA_MAX_RETRIES` | `--max-retries` | `3` | Maximum Instructor retries for structured responses |
@@ -80,6 +95,8 @@ Useful configuration can be supplied by CLI flag or environment variable. CLI fl
| `AUDITA_WORK_DIR` | `--work-dir` | `/tmp/audita` | Per-run scratch diagnostics directory |
| `AUDITA_WORK_DIR_RETENTION` | `--work-dir-retention` | `auto` | Whether to retain the per-run work directory: `auto`, `always`, or `never` |
Set `AUDITA_MODULES=grammar` to run only the grammar module by default, or override it per command with `--modules`.
`AUDITA_WORK_DIR` stores per-run diagnostics while processing. Under the default `AUDITA_WORK_DIR_RETENTION=auto`, clean successful runs are removed, while failed runs and successful runs with final skipped corrections are preserved. Use `always` to keep every run directory and `never` to remove successful run directories even when skips remain.
Failed runs always preserve the run directory and include an authoritative `report.json` alongside normalization and prompt/response diagnostics.

View File

@@ -27,6 +27,10 @@ def _build_parser() -> argparse.ArgumentParser:
process.add_argument("--glossary", type=Path, required=True, help="path to the glossary YAML")
process.add_argument("--output", type=Path, help="write corrected transcript JSON to this path")
process.add_argument("--report-json", type=Path, help="write structured run report JSON to this path")
process.add_argument(
"--modules",
help="comma-separated module keys to run, for example: grammar or glossary,homophones,glossary",
)
process.add_argument("--model", help="OpenRouter model for glossary and homophones LLM stages")
process.add_argument("--base-url", help="OpenAI-compatible API base URL for glossary and homophones LLM stages")
process.add_argument("--max-retries", type=int, help="maximum structured-output retries for LLM stages")
@@ -74,6 +78,7 @@ def _process(args: argparse.Namespace) -> int:
try:
config = AuditaConfig.from_sources(
overrides=ConfigOverrides(
module_keys=args.modules,
model=args.model,
base_url=args.base_url,
max_retries=args.max_retries,

View File

@@ -2,7 +2,9 @@ import math
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Mapping, Optional
from typing import Mapping, Optional, Sequence, Tuple, Union
from audita.modules import DEFAULT_MODULE_KEYS, normalize_module_keys
from .errors import AuditaConfigError
@@ -23,6 +25,7 @@ DEFAULT_NORMALIZE_MAX_SEGMENT_TOKENS = 2048
@dataclass(frozen=True)
class ConfigOverrides:
module_keys: Optional[Union[str, Sequence[str]]] = None
model: Optional[str] = None
base_url: Optional[str] = None
max_retries: Optional[int] = None
@@ -40,6 +43,7 @@ class ConfigOverrides:
@dataclass(frozen=True)
class AuditaConfig:
api_key: Optional[str] = None
module_keys: Tuple[str, ...] = DEFAULT_MODULE_KEYS
model: str = DEFAULT_MODEL
base_url: str = DEFAULT_BASE_URL
max_retries: int = DEFAULT_MAX_RETRIES
@@ -63,6 +67,12 @@ class AuditaConfig:
selected = ConfigOverrides() if overrides is None else overrides
config = cls(
api_key=_select_optional_string(source.get("OPENROUTER_API_KEY")),
module_keys=_select_module_keys(
selected.module_keys,
source.get("AUDITA_MODULES"),
DEFAULT_MODULE_KEYS,
"AUDITA_MODULES",
),
model=selected.model or source.get("AUDITA_MODEL") or DEFAULT_MODEL,
base_url=selected.base_url or source.get("AUDITA_BASE_URL") or DEFAULT_BASE_URL,
max_retries=_select_int(
@@ -126,6 +136,7 @@ class AuditaConfig:
return config
def validate(self) -> None:
normalize_module_keys(self.module_keys)
if not self.model.strip():
raise AuditaConfigError("AUDITA_MODEL must not be empty.")
if not self.base_url.strip():
@@ -162,6 +173,7 @@ class AuditaConfig:
def to_report_dict(self) -> dict:
return {
"api_key_configured": bool(self.api_key),
"module_keys": list(self.module_keys),
"model": self.model,
"base_url": self.base_url,
"max_retries": self.max_retries,
@@ -216,3 +228,24 @@ def _select_choice(
if value not in valid_choices:
raise AuditaConfigError(f"{name} must be one of {', '.join(valid_choices)}.")
return value
def _select_module_keys(
cli_value: Optional[Union[str, Sequence[str]]],
env_value: Optional[str],
default: Tuple[str, ...],
name: str,
) -> Tuple[str, ...]:
if cli_value is not None:
return _coerce_module_keys(cli_value, name)
if env_value is not None:
return _coerce_module_keys(env_value, name)
return default
def _coerce_module_keys(value: Union[str, Sequence[str]], name: str) -> Tuple[str, ...]:
try:
raw_items = value.split(",") if isinstance(value, str) else list(value)
return normalize_module_keys(raw_items)
except AuditaConfigError as exc:
raise AuditaConfigError(f"{name} is invalid: {exc}") from exc

View File

@@ -1,16 +1,80 @@
from audita.framework.models import ModuleRunSpec
from __future__ import annotations
from .glossary import GlossaryModule
from .grammar import GrammarModule
from .homophones import HomophonesModule
from .spoken_word import SpokenWordModule
from collections import Counter, defaultdict
from typing import Sequence
from audita.core.errors import AuditaConfigError
def default_module_specs() -> list[ModuleRunSpec]:
SUPPORTED_MODULE_KEYS = ("glossary", "homophones", "spoken_word", "grammar")
DEFAULT_MODULE_KEYS = ("glossary", "homophones", "glossary", "spoken_word", "grammar")
def normalize_module_keys(module_keys: Sequence[str]) -> tuple[str, ...]:
normalized = tuple(_normalize_module_key(item) for item in module_keys)
if not normalized:
raise AuditaConfigError("Module sequence must contain at least one module key.")
invalid = tuple(item for item in normalized if item not in SUPPORTED_MODULE_KEYS)
if invalid:
valid = ", ".join(SUPPORTED_MODULE_KEYS)
unknown = ", ".join(invalid)
raise AuditaConfigError(f"Unknown module key(s): {unknown}. Valid module keys: {valid}.")
return normalized
def resolve_module_specs(module_keys: Sequence[str]) -> list["ModuleRunSpec"]:
from audita.framework.models import ModuleRunSpec
normalized = normalize_module_keys(module_keys)
counts = Counter(normalized)
seen: defaultdict[str, int] = defaultdict(int)
return [
ModuleRunSpec(instance_name="glossary_primary", module_key="glossary", module=GlossaryModule()),
ModuleRunSpec(instance_name="homophones", module_key="homophones", module=HomophonesModule()),
ModuleRunSpec(instance_name="glossary_secondary", module_key="glossary", module=GlossaryModule()),
ModuleRunSpec(instance_name="spoken_word", module_key="spoken_word", module=SpokenWordModule()),
ModuleRunSpec(instance_name="grammar", module_key="grammar", module=GrammarModule()),
ModuleRunSpec(
instance_name=_instance_name(module_key, counts, seen),
module_key=module_key,
module=_instantiate_module(module_key),
)
for module_key in normalized
]
def default_module_specs() -> list["ModuleRunSpec"]:
return resolve_module_specs(DEFAULT_MODULE_KEYS)
def _instance_name(module_key: str, counts: Counter[str], seen: defaultdict[str, int]) -> str:
seen[module_key] += 1
if counts[module_key] == 1:
return module_key
return f"{module_key}_{seen[module_key]}"
def _instantiate_module(module_key: str):
if module_key == "glossary":
from .glossary import GlossaryModule
return GlossaryModule()
if module_key == "homophones":
from .homophones import HomophonesModule
return HomophonesModule()
if module_key == "spoken_word":
from .spoken_word import SpokenWordModule
return SpokenWordModule()
if module_key == "grammar":
from .grammar import GrammarModule
return GrammarModule()
raise AuditaConfigError(f"Unknown module key '{module_key}'.")
def _normalize_module_key(value: str) -> str:
if not isinstance(value, str):
raise AuditaConfigError("Module keys must be strings.")
normalized = value.strip()
if not normalized:
raise AuditaConfigError("Module keys must not contain empty values.")
return normalized

View File

@@ -2,7 +2,7 @@ import json
import shutil
from datetime import datetime
from pathlib import Path
from typing import Callable, List, Optional
from typing import Callable, List, Optional, Sequence
from uuid import uuid4
from .core.config import AuditaConfig
@@ -13,7 +13,7 @@ from .core.schemas import Glossary, SourceTranscriptSegment, TranscriptSegment,
from .framework.llm import OpenRouterStructuredLLMClient
from .framework.models import StructuredLLMClient
from .framework.runner import PipelineRunError, PipelineRunner
from .modules import default_module_specs
from .modules import resolve_module_specs
ProgressCallback = Callable[[str], None]
@@ -23,6 +23,7 @@ def process_transcript(
transcript: List[SourceTranscriptSegment],
glossary: Glossary,
config: AuditaConfig,
module_keys: Optional[Sequence[str]] = None,
llm_client: Optional[StructuredLLMClient] = None,
progress: Optional[ProgressCallback] = None,
) -> List[TranscriptSegment]:
@@ -30,6 +31,7 @@ def process_transcript(
transcript,
glossary,
config,
module_keys=module_keys,
llm_client=llm_client,
progress=progress,
).transcript
@@ -39,11 +41,12 @@ def process_transcript_result(
transcript: List[SourceTranscriptSegment],
glossary: Glossary,
config: AuditaConfig,
module_keys: Optional[Sequence[str]] = None,
llm_client: Optional[StructuredLLMClient] = None,
progress: Optional[ProgressCallback] = None,
) -> ProcessResult:
run_dir = _create_run_dir(config.work_dir)
module_specs = default_module_specs()
module_specs = resolve_module_specs(config.module_keys if module_keys is None else module_keys)
normalized_transcript: List[TranscriptSegment] = []
normalization_summary: Optional[dict] = None
report: Optional[RunReport] = None

View File

@@ -207,13 +207,13 @@ def test_process_transcript_result_uses_injected_fake_client_and_applies_sequent
assert result.transcript[0].text == "There were Jesters at the damn."
assert [call["stage_name"] for call in client.calls] == [
"glossary_primary:proposal",
"glossary_primary:spoken_form_plausibility_review",
"glossary_primary:meaning_reversal_review",
"glossary_1:proposal",
"glossary_1:spoken_form_plausibility_review",
"glossary_1:meaning_reversal_review",
"homophones:proposal",
"homophones:spoken_form_plausibility_review",
"homophones:meaning_reversal_review",
"glossary_secondary:proposal",
"glossary_2:proposal",
]
assert "There were Jesters at the dam." in client.calls[3]["messages"][1]["content"]
@@ -263,9 +263,9 @@ def test_process_transcript_result_rejects_below_threshold_proposals_before_llm_
assert result.transcript[0].text == "There were gestures at the temple."
assert [call["stage_name"] for call in client.calls] == [
"glossary_primary:proposal",
"glossary_1:proposal",
"homophones:proposal",
"glossary_secondary:proposal",
"glossary_2:proposal",
]
assert result.report.skipped_corrections[0].reason == "proposal confidence below threshold"
assert result.report.skipped_corrections[0].source == "validator:proposal_confidence_guard"

View File

@@ -20,6 +20,7 @@ def test_process_help_exposes_framework_flags(capsys):
assert exc.value.code == 0
output = capsys.readouterr().out
assert "--report-json" in output
assert "--modules" in output
assert "--model" in output
assert "--base-url" in output
assert "--max-retries" in output
@@ -43,7 +44,7 @@ def test_cli_process_writes_report_json(monkeypatch, tmp_path):
status="success",
config={"model": "m", "base_url": "b"},
normalization={"source_segment_count": 1, "normalized_segment_count": 1, "merge_count": 0},
pipeline=["glossary_primary", "homophones", "glossary_secondary", "spoken_word", "grammar"],
pipeline=["glossary_1", "homophones", "glossary_2", "spoken_word", "grammar"],
modules=[],
applied_changes=[],
skipped_corrections=[],
@@ -82,3 +83,57 @@ def test_cli_process_writes_report_json(monkeypatch, tmp_path):
assert exit_code == 0
assert report_path.exists()
def test_cli_process_passes_modules_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["module_keys"] = overrides.module_keys
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",
"--modules",
"grammar",
]
)
assert exit_code == 0
assert captured["module_keys"] == "grammar"

View File

@@ -9,12 +9,14 @@ from audita.core.config import (
DEFAULT_WORK_DIR_RETENTION,
)
from audita.core.errors import AuditaConfigError
from audita.modules import DEFAULT_MODULE_KEYS
def test_default_config_allows_missing_api_key():
config = AuditaConfig.from_sources(env={})
assert config.api_key is None
assert config.module_keys == DEFAULT_MODULE_KEYS
assert config.glossary_confidence_threshold == DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD
assert config.homophones_confidence_threshold == DEFAULT_HOMOPHONES_CONFIDENCE_THRESHOLD
assert config.normalize_max_segment_gap == DEFAULT_NORMALIZE_MAX_SEGMENT_GAP
@@ -30,6 +32,21 @@ def test_cli_overrides_take_precedence():
assert config.max_section_tokens == 2000
def test_module_key_overrides_take_precedence():
config = AuditaConfig.from_sources(
env={"AUDITA_MODULES": "grammar"},
overrides=ConfigOverrides(module_keys="homophones,grammar"),
)
assert config.module_keys == ("homophones", "grammar")
def test_module_key_env_is_parsed_and_trimmed():
config = AuditaConfig.from_sources(env={"AUDITA_MODULES": " glossary , grammar "})
assert config.module_keys == ("glossary", "grammar")
def test_invalid_work_dir_retention_is_rejected():
with pytest.raises(AuditaConfigError):
AuditaConfig.from_sources(env={"AUDITA_WORK_DIR_RETENTION": "sometimes"})
@@ -51,6 +68,19 @@ def test_threshold_overrides_take_precedence():
assert config.homophones_confidence_threshold == 0.9
@pytest.mark.parametrize(
"value",
[
"",
"grammar,,homophones",
"bogus",
],
)
def test_invalid_module_sequences_are_rejected(value):
with pytest.raises(AuditaConfigError, match="AUDITA_MODULES"):
AuditaConfig.from_sources(env={"AUDITA_MODULES": value})
@pytest.mark.parametrize(
"env_name",
[

View File

@@ -2,11 +2,11 @@ import json
import pytest
from audita.core.config import AuditaConfig
from audita.core.config import AuditaConfig, ConfigOverrides
from audita.core.errors import AuditaLLMError
from audita.core.io import write_report
from audita.core.schemas import parse_glossary_yaml, parse_source_transcript_json
from audita.modules import default_module_specs
from audita.modules import DEFAULT_MODULE_KEYS, default_module_specs, resolve_module_specs
from audita.pipeline import process_transcript, process_transcript_result
@@ -73,9 +73,9 @@ def test_process_transcript_runs_noop_framework(tmp_path):
assert revised[0].text == "Hello. Again."
assert revised[1].text == "Done."
assert [call["stage_name"] for call in llm_client.calls] == [
"glossary_primary:proposal",
"glossary_1:proposal",
"homophones:proposal",
"glossary_secondary:proposal",
"glossary_2:proposal",
]
@@ -111,9 +111,9 @@ def test_process_transcript_result_writes_report_and_preserves_skips_per_policy(
assert result.work_dir_retained is True
assert result.report.pipeline == [
"glossary_primary",
"glossary_1",
"homophones",
"glossary_secondary",
"glossary_2",
"spoken_word",
"grammar",
]
@@ -142,13 +142,21 @@ def test_external_report_can_be_written(tmp_path):
write_report(report_path, result.report)
payload = json.loads(report_path.read_text(encoding="utf-8"))
assert payload["pipeline"][0] == "glossary_primary"
assert payload["pipeline"][0] == "glossary_1"
assert payload["totals"]["applied_change_count"] == 0
def test_default_module_specs_expose_final_validator_order():
specs = default_module_specs()
assert DEFAULT_MODULE_KEYS == ("glossary", "homophones", "glossary", "spoken_word", "grammar")
assert [spec.instance_name for spec in specs] == [
"glossary_1",
"homophones",
"glossary_2",
"spoken_word",
"grammar",
]
assert [validator.name for validator in specs[0].module.validators()] == [
"proposal_confidence_guard",
"protected_glossary_guard",
@@ -201,9 +209,9 @@ def test_process_transcript_result_missing_api_key_writes_failed_report(tmp_path
assert report["status"] == "failed"
assert report["normalization"]["normalized_segment_count"] == 2
assert report["pipeline"] == [
"glossary_primary",
"glossary_1",
"homophones",
"glossary_secondary",
"glossary_2",
"spoken_word",
"grammar",
]
@@ -286,7 +294,7 @@ def test_process_transcript_result_preserves_partial_progress_when_later_module_
assert report["status"] == "failed"
assert report["normalization"]["normalized_segment_count"] == 1
assert [module["instance_name"] for module in report["modules"]] == ["glossary_primary"]
assert [module["instance_name"] for module in report["modules"]] == ["glossary_1"]
assert report["applied_changes"][0]["corrected_text"] == "Jesters"
assert report["applied_changes"][0]["segment_text_after"] == "There were Jesters at the dam."
assert report["totals"]["applied_change_count"] == 1
@@ -351,7 +359,7 @@ def test_process_transcript_result_preserves_partial_skips_and_validator_diagnos
run_dir = next((tmp_path / "work").iterdir())
report = json.loads((run_dir / "report.json").read_text(encoding="utf-8"))
validator_dir = run_dir / "glossary_primary"
validator_dir = run_dir / "glossary_1"
assert report["status"] == "failed"
assert report["modules"] == []
@@ -361,3 +369,29 @@ def test_process_transcript_result_preserves_partial_skips_and_validator_diagnos
assert "unknown correction_index" in report["error"]
assert (validator_dir / "spoken_form_plausibility_review-prompt-0000.json").exists()
assert (validator_dir / "spoken_form_plausibility_review-response-0000.json").exists()
def test_resolve_module_specs_numbers_repeated_keys():
specs = resolve_module_specs(["glossary", "homophones", "glossary"])
assert [spec.instance_name for spec in specs] == ["glossary_1", "homophones", "glossary_2"]
assert [spec.module_key for spec in specs] == ["glossary", "homophones", "glossary"]
def test_process_transcript_result_supports_grammar_only_module_override(tmp_path):
config = AuditaConfig.from_sources(
env={},
overrides=ConfigOverrides(work_dir=tmp_path / "work", work_dir_retention="always"),
)
result = process_transcript_result(
_transcript(),
_glossary(),
config,
module_keys=["grammar"],
llm_client=FakeStructuredLLMClient([]),
)
assert [segment.id for segment in result.transcript] == [1, 2]
assert result.report.pipeline == ["grammar"]
assert result.report.totals["applied_change_count"] == 0