diff --git a/.gitignore b/.gitignore index 36b13f1..d1e3e3e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,176 +1,11 @@ -# ---> Python -# Byte-compiled / optimized / DLL files +.DS_Store +.venv/ __pycache__/ *.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py,cover -.hypothesis/ .pytest_cache/ -cover/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -.pybuilder/ -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: -# .python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# UV -# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -#uv.lock - -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control -#poetry.lock - -# pdm -# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -#pdm.lock -# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it -# in version control. -# https://pdm.fming.dev/latest/usage/project/#working-with-version-control -.pdm.toml -.pdm-python -.pdm-build/ - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Environments -.env -.venv -env/ -venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ - -# pytype static type analyzer -.pytype/ - -# Cython debug symbols -cython_debug/ - -# PyCharm -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ - -# Ruff stuff: .ruff_cache/ - -# PyPI configuration file -.pypirc +.mypy_cache/ +dist/ +build/ +*.egg-info/ diff --git a/README.md b/README.md index bf42373..e241da7 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,34 @@ # Audita -Audita takes raw audio transcripts and uses an LLM to identify and fix misheard words, jargon, and domain-specific terms. \ No newline at end of file +Audita takes raw audio transcripts and uses an LLM to identify and fix misheard words, jargon, and domain-specific terms. + +## Development + +This project is set up for `uv`. + +```sh +uv sync --extra dev +uv run pytest +``` + +## Usage + +Set an OpenRouter API key, then process a transcript with a glossary: + +```sh +export OPENROUTER_API_KEY=... +uv run audia process transcript.json --glossary glossary.yaml --output corrected.json +``` + +Without `--output`, Audita writes the corrected transcript JSON to stdout and progress logs to stderr. + +Useful configuration can be supplied by CLI flag or environment variable: + +- `AUDITA_MODEL`, default `openrouter/mistralai/mistral-small-3.2-24b-instruct` +- `AUDITA_BASE_URL`, default `https://openrouter.ai/api/v1` +- `AUDITA_MAX_SECTION_TOKENS`, default `16000` +- `AUDITA_CONFIDENCE_THRESHOLD`, default `0.80` +- `AUDITA_MAX_RETRIES`, default `3` +- `AUDITA_WORK_DIR`, default `/tmp/audita` + +`AUDITA_WORK_DIR` stores per-run diagnostics while processing. Successful runs clean up their run directory; failed runs preserve it for debugging. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..a140d68 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,32 @@ +[project] +name = "audita" +version = "0.1.0" +description = "Correct audio transcripts with glossary-guided LLM passes." +readme = "README.md" +requires-python = ">=3.9" +license = { text = "BSD-3-Clause" } +dependencies = [ + "instructor>=1.7", + "openai>=1.55", + "pydantic>=2.8", + "PyYAML>=6.0", + "tiktoken>=0.8", + "eval-type-backport>=0.2.0; python_version < '3.10'", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", +] + +[project.scripts] +audia = "audita.cli:main" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.pytest.ini_options] +pythonpath = ["src"] +testpaths = ["tests"] +addopts = "-ra" diff --git a/src/audita/__init__.py b/src/audita/__init__.py new file mode 100644 index 0000000..5f528c6 --- /dev/null +++ b/src/audita/__init__.py @@ -0,0 +1,6 @@ +"""Audita transcript correction package.""" + +__all__ = ["__version__"] + +__version__ = "0.1.0" + diff --git a/src/audita/__main__.py b/src/audita/__main__.py new file mode 100644 index 0000000..0b6ae7c --- /dev/null +++ b/src/audita/__main__.py @@ -0,0 +1,6 @@ +from .cli import main + + +if __name__ == "__main__": + raise SystemExit(main()) + diff --git a/src/audita/chunking.py b/src/audita/chunking.py new file mode 100644 index 0000000..821c133 --- /dev/null +++ b/src/audita/chunking.py @@ -0,0 +1,124 @@ +import json +from dataclasses import dataclass +from math import ceil +from typing import Any, List, Optional, Protocol + +from .errors import AuditaValidationError +from .schemas import TranscriptSegment, parse_transcript_json + + +class TokenEstimatorProtocol(Protocol): + def estimate_json(self, value: Any) -> int: + ... + + +class TokenEstimator: + def __init__(self, fallback_chars_per_token: int = 4) -> None: + self._fallback_chars_per_token = fallback_chars_per_token + self._encoding = None + try: + import tiktoken + + self._encoding = tiktoken.get_encoding("cl100k_base") + except Exception: + self._encoding = None + + def estimate_text(self, text: str) -> int: + if self._encoding is not None: + return len(self._encoding.encode(text)) + return max(1, ceil(len(text) / self._fallback_chars_per_token)) + + def estimate_json(self, value: Any) -> int: + return self.estimate_text(json.dumps(value, ensure_ascii=False, separators=(",", ":"))) + + +@dataclass(frozen=True) +class IndexedSegment: + index: int + segment: TranscriptSegment + + def transcript_payload(self) -> dict: + return self.segment.model_dump(mode="json") + + def prompt_payload(self) -> dict: + payload = self.transcript_payload() + payload["segment_index"] = self.index + return payload + + +@dataclass(frozen=True) +class TranscriptSection: + section_index: int + start_index: int + segments: List[IndexedSegment] + token_count: int + + def transcript_payload(self) -> List[dict]: + return [item.transcript_payload() for item in self.segments] + + def prompt_payload(self) -> List[dict]: + return [item.prompt_payload() for item in self.segments] + + def transcript_json(self) -> str: + return json.dumps(self.transcript_payload(), ensure_ascii=False, indent=2) + "\n" + + +def chunk_transcript( + segments: List[TranscriptSegment], + max_section_tokens: int, + estimator: Optional[TokenEstimatorProtocol] = None, +) -> List[TranscriptSection]: + if max_section_tokens <= 0: + raise AuditaValidationError("Maximum section token count must be greater than zero.") + if not segments: + raise AuditaValidationError("Transcript must contain at least one segment.") + + token_estimator = TokenEstimator() if estimator is None else estimator + indexed = [IndexedSegment(index=index, segment=segment) for index, segment in enumerate(segments)] + + sections: List[TranscriptSection] = [] + current: List[IndexedSegment] = [] + current_tokens = 0 + + for item in indexed: + single_payload = [item.transcript_payload()] + single_tokens = token_estimator.estimate_json(single_payload) + if single_tokens > max_section_tokens: + raise AuditaValidationError( + "A single transcript segment exceeds the maximum section token limit. " + "Raise the limit or pre-split the transcript." + ) + + candidate = current + [item] + candidate_tokens = token_estimator.estimate_json( + [candidate_item.transcript_payload() for candidate_item in candidate] + ) + if current and candidate_tokens > max_section_tokens: + sections.append(_make_section(len(sections), current, current_tokens)) + current = [item] + current_tokens = single_tokens + else: + current = candidate + current_tokens = candidate_tokens + + if current: + sections.append(_make_section(len(sections), current, current_tokens)) + + for section in sections: + parse_transcript_json(section.transcript_json()) + + return sections + + +def _make_section( + section_index: int, + segments: List[IndexedSegment], + token_count: int, +) -> TranscriptSection: + return TranscriptSection( + section_index=section_index, + start_index=segments[0].index, + segments=list(segments), + token_count=token_count, + ) + diff --git a/src/audita/cli.py b/src/audita/cli.py new file mode 100644 index 0000000..bb98d0e --- /dev/null +++ b/src/audita/cli.py @@ -0,0 +1,70 @@ +import argparse +import sys +from pathlib import Path +from typing import Optional, Sequence + +from .config import AuditaConfig, ConfigOverrides +from .errors import AuditaError +from .io import load_glossary, load_transcript, write_transcript +from .pipeline import process_transcript +from .schemas import transcript_to_json + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + + if args.command == "process": + return _process(args) + + parser.print_help(sys.stderr) + return 2 + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="audia") + subparsers = parser.add_subparsers(dest="command", required=True) + + process = subparsers.add_parser("process", help="correct a transcript using a glossary") + process.add_argument("transcript", type=Path, help="path to the input transcript JSON") + 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("--model", help="OpenRouter model to use") + process.add_argument("--base-url", help="OpenAI-compatible API base URL") + process.add_argument("--max-section-tokens", type=int, help="maximum estimated tokens per transcript section") + process.add_argument("--confidence-threshold", type=float, help="minimum confidence required to apply a correction") + process.add_argument("--max-retries", type=int, help="maximum Instructor retries for structured response validation") + process.add_argument("--work-dir", type=Path, help="directory for per-run scratch diagnostics") + return parser + + +def _process(args: argparse.Namespace) -> int: + try: + config = AuditaConfig.from_sources( + overrides=ConfigOverrides( + model=args.model, + base_url=args.base_url, + max_section_tokens=args.max_section_tokens, + confidence_threshold=args.confidence_threshold, + max_retries=args.max_retries, + work_dir=args.work_dir, + ) + ) + transcript = load_transcript(args.transcript) + glossary = load_glossary(args.glossary) + revised = process_transcript( + transcript, + glossary, + config, + progress=lambda message: print(message, file=sys.stderr), + ) + + if args.output is not None: + write_transcript(args.output, revised) + else: + sys.stdout.write(transcript_to_json(revised)) + return 0 + except AuditaError as exc: + print(f"audia: error: {exc}", file=sys.stderr) + return 1 + diff --git a/src/audita/config.py b/src/audita/config.py new file mode 100644 index 0000000..01068db --- /dev/null +++ b/src/audita/config.py @@ -0,0 +1,128 @@ +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Mapping, Optional + +from .errors import AuditaConfigError + + +DEFAULT_MODEL = "openrouter/mistralai/mistral-small-3.2-24b-instruct" +DEFAULT_BASE_URL = "https://openrouter.ai/api/v1" +DEFAULT_MAX_SECTION_TOKENS = 16000 +DEFAULT_CONFIDENCE_THRESHOLD = 0.80 +DEFAULT_MAX_RETRIES = 3 +DEFAULT_WORK_DIR = "/tmp/audita" + + +@dataclass(frozen=True) +class ConfigOverrides: + model: Optional[str] = None + base_url: Optional[str] = None + max_section_tokens: Optional[int] = None + confidence_threshold: Optional[float] = None + max_retries: Optional[int] = None + work_dir: Optional[Path] = None + + +@dataclass(frozen=True) +class AuditaConfig: + api_key: str + model: str = DEFAULT_MODEL + base_url: str = DEFAULT_BASE_URL + max_section_tokens: int = DEFAULT_MAX_SECTION_TOKENS + confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD + max_retries: int = DEFAULT_MAX_RETRIES + work_dir: Path = Path(DEFAULT_WORK_DIR) + + @classmethod + def from_sources( + cls, + env: Optional[Mapping[str, str]] = None, + overrides: Optional[ConfigOverrides] = None, + ) -> "AuditaConfig": + source = os.environ if env is None else env + selected = ConfigOverrides() if overrides is None else overrides + + api_key = _get_required_env(source, "OPENROUTER_API_KEY") + 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_section_tokens = _select_int( + selected.max_section_tokens, + source.get("AUDITA_MAX_SECTION_TOKENS"), + DEFAULT_MAX_SECTION_TOKENS, + "AUDITA_MAX_SECTION_TOKENS", + ) + confidence_threshold = _select_float( + selected.confidence_threshold, + source.get("AUDITA_CONFIDENCE_THRESHOLD"), + DEFAULT_CONFIDENCE_THRESHOLD, + "AUDITA_CONFIDENCE_THRESHOLD", + ) + max_retries = _select_int( + selected.max_retries, + source.get("AUDITA_MAX_RETRIES"), + DEFAULT_MAX_RETRIES, + "AUDITA_MAX_RETRIES", + ) + work_dir_value = selected.work_dir or Path(source.get("AUDITA_WORK_DIR") or DEFAULT_WORK_DIR) + + config = cls( + api_key=api_key, + model=model, + base_url=base_url, + max_section_tokens=max_section_tokens, + confidence_threshold=confidence_threshold, + max_retries=max_retries, + work_dir=Path(work_dir_value), + ) + config.validate() + return config + + def validate(self) -> None: + if not self.api_key.strip(): + raise AuditaConfigError("OPENROUTER_API_KEY must not be empty.") + if not self.model.strip(): + raise AuditaConfigError("AUDITA_MODEL must not be empty.") + if not self.base_url.strip(): + raise AuditaConfigError("AUDITA_BASE_URL must not be empty.") + if self.max_section_tokens <= 0: + raise AuditaConfigError("AUDITA_MAX_SECTION_TOKENS must be greater than zero.") + if not 0.0 <= self.confidence_threshold <= 1.0: + raise AuditaConfigError("AUDITA_CONFIDENCE_THRESHOLD must be between 0.0 and 1.0.") + if self.max_retries < 0: + raise AuditaConfigError("AUDITA_MAX_RETRIES must be greater than or equal to zero.") + + +def _get_required_env(env: Mapping[str, str], name: str) -> str: + value = env.get(name) + if value is None or not value.strip(): + raise AuditaConfigError(f"{name} is required.") + return value + + +def _select_int(cli_value: Optional[int], env_value: Optional[str], default: int, name: str) -> int: + if cli_value is not None: + return cli_value + if env_value is None: + return default + 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 + if env_value is None: + return default + try: + return float(env_value) + except ValueError as exc: + raise AuditaConfigError(f"{name} must be a number.") from exc + diff --git a/src/audita/corrections.py b/src/audita/corrections.py new file mode 100644 index 0000000..adcb225 --- /dev/null +++ b/src/audita/corrections.py @@ -0,0 +1,53 @@ +from typing import Dict, Iterable, List + +from .errors import AuditaValidationError +from .schemas import CorrectionCandidate, TranscriptSegment + + +def apply_corrections( + transcript: List[TranscriptSegment], + corrections: Iterable[CorrectionCandidate], + confidence_threshold: float, +) -> List[TranscriptSegment]: + if not 0.0 <= confidence_threshold <= 1.0: + raise AuditaValidationError("Confidence threshold must be between 0.0 and 1.0.") + + correction_by_index: Dict[int, CorrectionCandidate] = {} + for correction in corrections: + if correction.segment_index in correction_by_index: + raise AuditaValidationError( + f"Duplicate corrections returned for segment {correction.segment_index}." + ) + _validate_target(transcript, correction) + correction_by_index[correction.segment_index] = correction + + revised: List[TranscriptSegment] = [] + for original_index, segment in enumerate(transcript): + correction = correction_by_index.get(original_index) + if correction is not None and correction.confidence >= confidence_threshold: + revised.append(segment.model_copy(update={"text": correction.corrected_text})) + else: + revised.append(segment) + + indexed_revised = list(enumerate(revised)) + indexed_revised.sort(key=lambda item: (item[1].start, item[1].end, item[0])) + return [segment for _, segment in indexed_revised] + + +def _validate_target(transcript: List[TranscriptSegment], correction: CorrectionCandidate) -> None: + if correction.segment_index >= len(transcript): + raise AuditaValidationError( + f"Correction targets missing segment {correction.segment_index}." + ) + + segment = transcript[correction.segment_index] + if ( + correction.speaker != segment.speaker + or correction.start != segment.start + or correction.end != segment.end + or correction.original_text != segment.text + ): + raise AuditaValidationError( + f"Correction target does not exactly match segment {correction.segment_index}." + ) + diff --git a/src/audita/errors.py b/src/audita/errors.py new file mode 100644 index 0000000..f017433 --- /dev/null +++ b/src/audita/errors.py @@ -0,0 +1,15 @@ +class AuditaError(Exception): + """Base exception for user-facing Audita failures.""" + + +class AuditaValidationError(AuditaError): + """Raised when input data does not match Audita's expected schema.""" + + +class AuditaConfigError(AuditaError): + """Raised when runtime configuration is invalid or incomplete.""" + + +class AuditaLLMError(AuditaError): + """Raised when an LLM request or structured response fails.""" + diff --git a/src/audita/io.py b/src/audita/io.py new file mode 100644 index 0000000..fbae020 --- /dev/null +++ b/src/audita/io.py @@ -0,0 +1,18 @@ +from pathlib import Path +from typing import List + +from .schemas import Glossary, TranscriptSegment +from .schemas import parse_glossary_yaml, parse_transcript_json, transcript_to_json + + +def load_transcript(path: Path) -> List[TranscriptSegment]: + return parse_transcript_json(path.read_text(encoding="utf-8")) + + +def load_glossary(path: Path) -> Glossary: + return parse_glossary_yaml(path.read_text(encoding="utf-8")) + + +def write_transcript(path: Path, segments: List[TranscriptSegment]) -> None: + path.write_text(transcript_to_json(segments), encoding="utf-8") + diff --git a/src/audita/llm.py b/src/audita/llm.py new file mode 100644 index 0000000..9bad628 --- /dev/null +++ b/src/audita/llm.py @@ -0,0 +1,45 @@ +from typing import List + +from .config import AuditaConfig +from .errors import AuditaLLMError +from .prompts import Message +from .schemas import CorrectionSet + + +class InstructorLLMClient: + def __init__(self, config: AuditaConfig) -> None: + try: + import instructor + from openai import OpenAI + except ImportError as exc: + raise AuditaLLMError( + "The LLM dependencies are not installed. Run `uv sync` before using audia." + ) from exc + + self._instructor = instructor + openai_client = OpenAI(api_key=config.api_key, base_url=config.base_url) + self._client = instructor.patch(openai_client, mode=instructor.Mode.TOOLS) + + def create_corrections(self, messages: List[Message], config: AuditaConfig) -> CorrectionSet: + model = _normalize_openrouter_model(config.model) + try: + return self._client.chat.completions.create( + model=model, + messages=messages, + response_model=CorrectionSet, + max_retries=config.max_retries, + extra_body={"provider": {"require_parameters": True}}, + ) + except Exception as exc: + raise AuditaLLMError( + "LLM correction request failed. Confirm the configured OpenRouter model " + "supports tool calling or structured outputs." + ) from exc + + +def _normalize_openrouter_model(model: str) -> str: + prefix = "openrouter/" + if model.startswith(prefix): + return model[len(prefix) :] + return model + diff --git a/src/audita/passes.py b/src/audita/passes.py new file mode 100644 index 0000000..d2c1c88 --- /dev/null +++ b/src/audita/passes.py @@ -0,0 +1,46 @@ +import json +from pathlib import Path +from typing import List, Protocol + +from .chunking import TranscriptSection +from .config import AuditaConfig +from .prompts import build_glossary_correction_messages +from .schemas import CorrectionCandidate, CorrectionSet, Glossary + + +class LLMClient(Protocol): + def create_corrections(self, messages: List[dict], config: AuditaConfig) -> CorrectionSet: + ... + + +class CorrectionPass(Protocol): + def run( + self, + section: TranscriptSection, + glossary: Glossary, + config: AuditaConfig, + run_dir: Path, + ) -> List[CorrectionCandidate]: + ... + + +class GlossaryCorrectionPass: + def __init__(self, llm_client: LLMClient) -> None: + self._llm_client = llm_client + + def run( + self, + section: TranscriptSection, + glossary: Glossary, + config: AuditaConfig, + run_dir: Path, + ) -> List[CorrectionCandidate]: + messages = build_glossary_correction_messages(section, glossary) + prompt_path = run_dir / f"prompt-{section.section_index:04d}.json" + prompt_path.write_text(json.dumps(messages, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + response = self._llm_client.create_corrections(messages, config) + response_path = run_dir / f"corrections-{section.section_index:04d}.json" + response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8") + return list(response.corrections) + diff --git a/src/audita/pipeline.py b/src/audita/pipeline.py new file mode 100644 index 0000000..f3ef352 --- /dev/null +++ b/src/audita/pipeline.py @@ -0,0 +1,105 @@ +import json +import shutil +from datetime import datetime +from pathlib import Path +from typing import Callable, List, Optional +from uuid import uuid4 + +from .chunking import TranscriptSection, chunk_transcript +from .config import AuditaConfig +from .corrections import apply_corrections +from .errors import AuditaError +from .passes import GlossaryCorrectionPass, LLMClient +from .schemas import Glossary, TranscriptSegment, parse_transcript_json + + +ProgressCallback = Callable[[str], None] + + +def process_transcript( + transcript: List[TranscriptSegment], + glossary: Glossary, + config: AuditaConfig, + llm_client: Optional[LLMClient] = None, + progress: Optional[ProgressCallback] = None, +) -> List[TranscriptSegment]: + run_dir = _create_run_dir(config.work_dir) + try: + _log(progress, f"Created work directory {run_dir}") + sections = chunk_transcript(transcript, config.max_section_tokens) + _write_run_metadata(run_dir, sections, config) + + if llm_client is None: + from .llm import InstructorLLMClient + + llm_client = InstructorLLMClient(config) + + correction_pass = GlossaryCorrectionPass(llm_client) + corrections = [] + for section in sections: + _write_and_validate_section(run_dir, section) + _log( + progress, + f"Processing section {section.section_index + 1}/{len(sections)} " + f"({len(section.segments)} segments, estimated {section.token_count} tokens)", + ) + corrections.extend(correction_pass.run(section, glossary, config, run_dir)) + + revised = apply_corrections(transcript, corrections, config.confidence_threshold) + except Exception as exc: + message = f"{exc} Diagnostics preserved at {run_dir}" + if isinstance(exc, AuditaError): + raise type(exc)(message) from exc + raise AuditaError(message) from exc + + shutil.rmtree(run_dir) + _log(progress, "Removed work directory after successful run") + return revised + + +def _create_run_dir(work_dir: Path) -> Path: + work_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ") + run_dir = work_dir / f"run-{timestamp}-{uuid4().hex[:8]}" + run_dir.mkdir() + return run_dir + + +def _write_run_metadata( + run_dir: Path, + sections: List[TranscriptSection], + config: AuditaConfig, +) -> None: + metadata = { + "model": config.model, + "base_url": config.base_url, + "max_section_tokens": config.max_section_tokens, + "confidence_threshold": config.confidence_threshold, + "max_retries": config.max_retries, + "sections": [ + { + "section_index": section.section_index, + "start_index": section.start_index, + "segment_count": len(section.segments), + "estimated_tokens": section.token_count, + } + for section in sections + ], + } + (run_dir / "metadata.json").write_text( + json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def _write_and_validate_section(run_dir: Path, section: TranscriptSection) -> None: + section_path = run_dir / f"section-{section.section_index:04d}.json" + section_json = section.transcript_json() + section_path.write_text(section_json, encoding="utf-8") + parse_transcript_json(section_json) + + +def _log(progress: Optional[ProgressCallback], message: str) -> None: + if progress is not None: + progress(message) + diff --git a/src/audita/prompts.py b/src/audita/prompts.py new file mode 100644 index 0000000..9f2f8d0 --- /dev/null +++ b/src/audita/prompts.py @@ -0,0 +1,34 @@ +import json +from typing import Dict, List + +from .chunking import TranscriptSection +from .schemas import Glossary + + +Message = Dict[str, str] + + +def build_glossary_correction_messages(section: TranscriptSection, glossary: Glossary) -> List[Message]: + glossary_json = json.dumps(glossary.model_dump(mode="json"), ensure_ascii=False, indent=2) + section_json = json.dumps(section.prompt_payload(), ensure_ascii=False, indent=2) + + system = ( + "You are Audita, a careful transcript correction assistant. " + "Identify only transcription errors that are strongly supported by the glossary. " + "Do not make generic grammar, spelling, capitalization, or style edits. " + "Do not rewrite unchanged transcript segments. " + "Preserve speaker names, timestamps, and meaning." + ) + user = ( + "Review this transcript section and return only corrections that should be applied.\n\n" + "Rules:\n" + "- Correct domain-specific names, aliases, jargon, deities, locations, NPCs, players, and similar terms when the glossary supports the correction.\n" + "- Use the exact segment_index, speaker, start, end, and original_text from the input segment.\n" + "- corrected_text must contain the full corrected text for that segment.\n" + "- confidence must be between 0.0 and 1.0.\n" + "- If no corrections are needed, return an empty corrections list.\n\n" + f"Glossary:\n{glossary_json}\n\n" + f"Transcript section:\n{section_json}" + ) + return [{"role": "system", "content": system}, {"role": "user", "content": user}] + diff --git a/src/audita/py.typed b/src/audita/py.typed new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/audita/py.typed @@ -0,0 +1 @@ + diff --git a/src/audita/schemas.py b/src/audita/schemas.py new file mode 100644 index 0000000..d88ff5d --- /dev/null +++ b/src/audita/schemas.py @@ -0,0 +1,161 @@ +import json +import math +from typing import Any, List + +from pydantic import BaseModel, ConfigDict, Field, StrictStr, TypeAdapter +from pydantic import ValidationError, field_validator, model_validator + +from .errors import AuditaValidationError + + +class TranscriptSegment(BaseModel): + model_config = ConfigDict(extra="forbid") + + speaker: StrictStr + start: float + end: float + text: StrictStr + + @field_validator("speaker", "text") + @classmethod + def require_non_empty_text(cls, value: str) -> str: + if not value.strip(): + raise ValueError("must not be empty") + return value + + @field_validator("start", "end", mode="before") + @classmethod + def require_number(cls, value: Any) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError("must be a JSON number") + number = float(value) + if not math.isfinite(number): + raise ValueError("must be finite") + if number < 0: + raise ValueError("must be non-negative") + return number + + @model_validator(mode="after") + def require_ordered_times(self) -> "TranscriptSegment": + if self.end < self.start: + raise ValueError("end must be greater than or equal to start") + return self + + +class GlossaryEntry(BaseModel): + model_config = ConfigDict(extra="forbid") + + name: StrictStr + aliases: List[StrictStr] = Field(default_factory=list) + category: StrictStr + summary: StrictStr + + @field_validator("name", "category", "summary") + @classmethod + def require_non_empty_text(cls, value: str) -> str: + if not value.strip(): + raise ValueError("must not be empty") + return value + + @field_validator("aliases") + @classmethod + def require_non_empty_aliases(cls, aliases: List[str]) -> List[str]: + for alias in aliases: + if not alias.strip(): + raise ValueError("aliases must not contain empty strings") + return aliases + + +class Glossary(BaseModel): + model_config = ConfigDict(extra="forbid") + + glossary: List[GlossaryEntry] + + @field_validator("glossary") + @classmethod + def require_entries(cls, entries: List[GlossaryEntry]) -> List[GlossaryEntry]: + if not entries: + raise ValueError("glossary must contain at least one entry") + return entries + + +class CorrectionCandidate(BaseModel): + model_config = ConfigDict(extra="forbid") + + segment_index: int = Field(ge=0) + speaker: StrictStr + start: float + end: float + original_text: StrictStr + corrected_text: StrictStr + confidence: float = Field(ge=0.0, le=1.0) + + @field_validator("segment_index", mode="before") + @classmethod + def require_integer_index(cls, value: Any) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError("must be an integer") + return value + + @field_validator("start", "end", "confidence", mode="before") + @classmethod + def require_number(cls, value: Any) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError("must be a JSON number") + number = float(value) + if not math.isfinite(number): + raise ValueError("must be finite") + return number + + +class CorrectionSet(BaseModel): + model_config = ConfigDict(extra="forbid") + + corrections: List[CorrectionCandidate] = Field(default_factory=list) + + +_TRANSCRIPT_ADAPTER = TypeAdapter(List[TranscriptSegment]) + + +def validate_transcript_data(data: Any) -> List[TranscriptSegment]: + if not isinstance(data, list): + raise AuditaValidationError("Transcript must be a JSON array.") + if not data: + raise AuditaValidationError("Transcript must contain at least one segment.") + try: + return _TRANSCRIPT_ADAPTER.validate_python(data) + except ValidationError as exc: + raise AuditaValidationError(f"Transcript schema validation failed: {exc}") from exc + + +def parse_transcript_json(raw: str) -> List[TranscriptSegment]: + try: + data = json.loads(raw) + except json.JSONDecodeError as exc: + raise AuditaValidationError(f"Transcript is not valid JSON: {exc}") from exc + return validate_transcript_data(data) + + +def parse_glossary_yaml(raw: str) -> Glossary: + try: + import yaml + except ImportError as exc: + raise AuditaValidationError("PyYAML is required to read glossary files.") from exc + + try: + data = yaml.safe_load(raw) + except yaml.YAMLError as exc: + raise AuditaValidationError(f"Glossary is not valid YAML: {exc}") from exc + + if data is None: + data = {} + try: + return Glossary.model_validate(data) + except ValidationError as exc: + raise AuditaValidationError(f"Glossary schema validation failed: {exc}") from exc + + +def transcript_to_json(segments: List[TranscriptSegment]) -> str: + payload = [segment.model_dump(mode="json") for segment in segments] + return json.dumps(payload, ensure_ascii=False, indent=2) + "\n" + diff --git a/tests/test_chunking.py b/tests/test_chunking.py new file mode 100644 index 0000000..c71ca9a --- /dev/null +++ b/tests/test_chunking.py @@ -0,0 +1,40 @@ +import pytest + +from audita.chunking import chunk_transcript +from audita.errors import AuditaValidationError +from audita.schemas import parse_transcript_json + + +class CountEstimator: + def estimate_json(self, value): + return len(value) * 10 + + +def _segments(count): + payload = [ + {"speaker": "Eric", "start": float(i), "end": float(i + 1), "text": f"Segment {i}"} + for i in range(count) + ] + import json + + return parse_transcript_json(json.dumps(payload)) + + +def test_chunk_transcript_splits_on_segment_boundaries(): + sections = chunk_transcript(_segments(5), max_section_tokens=20, estimator=CountEstimator()) + + assert [len(section.segments) for section in sections] == [2, 2, 1] + assert [section.start_index for section in sections] == [0, 2, 4] + + +def test_chunk_transcript_allows_exact_limit(): + sections = chunk_transcript(_segments(2), max_section_tokens=20, estimator=CountEstimator()) + + assert len(sections) == 1 + assert len(sections[0].segments) == 2 + + +def test_chunk_transcript_rejects_oversized_single_segment(): + with pytest.raises(AuditaValidationError): + chunk_transcript(_segments(1), max_section_tokens=9, estimator=CountEstimator()) + diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..28f6137 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,65 @@ +from pathlib import Path + +import pytest + +from audita.config import AuditaConfig, ConfigOverrides +from audita.config import DEFAULT_MAX_RETRIES, DEFAULT_MAX_SECTION_TOKENS, DEFAULT_WORK_DIR +from audita.errors import AuditaConfigError + + +def test_config_uses_defaults_with_api_key(): + config = AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "key"}) + + assert config.max_section_tokens == DEFAULT_MAX_SECTION_TOKENS + assert config.max_retries == DEFAULT_MAX_RETRIES + assert config.work_dir == Path(DEFAULT_WORK_DIR) + + +def test_config_env_overrides_defaults(): + config = AuditaConfig.from_sources( + env={ + "OPENROUTER_API_KEY": "key", + "AUDITA_MAX_SECTION_TOKENS": "42", + "AUDITA_CONFIDENCE_THRESHOLD": "0.9", + "AUDITA_MAX_RETRIES": "5", + "AUDITA_WORK_DIR": "/tmp/custom-audita", + } + ) + + assert config.max_section_tokens == 42 + assert config.confidence_threshold == 0.9 + assert config.max_retries == 5 + assert config.work_dir == Path("/tmp/custom-audita") + + +def test_config_cli_overrides_env(): + config = AuditaConfig.from_sources( + env={ + "OPENROUTER_API_KEY": "key", + "AUDITA_MAX_SECTION_TOKENS": "42", + "AUDITA_MAX_RETRIES": "5", + "AUDITA_WORK_DIR": "/tmp/env-audita", + }, + overrides=ConfigOverrides( + max_section_tokens=100, + max_retries=3, + work_dir=Path("/tmp/cli-audita"), + ), + ) + + assert config.max_section_tokens == 100 + assert config.max_retries == 3 + assert config.work_dir == Path("/tmp/cli-audita") + + +def test_config_requires_api_key(): + with pytest.raises(AuditaConfigError): + AuditaConfig.from_sources(env={}) + + +def test_config_rejects_bad_env_int(): + with pytest.raises(AuditaConfigError): + AuditaConfig.from_sources( + env={"OPENROUTER_API_KEY": "key", "AUDITA_MAX_SECTION_TOKENS": "many"} + ) + diff --git a/tests/test_corrections.py b/tests/test_corrections.py new file mode 100644 index 0000000..97e0e41 --- /dev/null +++ b/tests/test_corrections.py @@ -0,0 +1,88 @@ +import pytest + +from audita.corrections import apply_corrections +from audita.errors import AuditaValidationError +from audita.schemas import CorrectionCandidate, parse_transcript_json + + +def _transcript(): + return parse_transcript_json( + """ + [ + {"speaker": "Eric", "start": 10.0, "end": 11.0, "text": "I ask Chontia."}, + {"speaker": "Mike", "start": 0.0, "end": 1.0, "text": "Then Lyra."} + ] + """ + ) + + +def test_apply_corrections_uses_threshold_and_sorts_chronologically(): + transcript = _transcript() + corrections = [ + CorrectionCandidate( + segment_index=0, + speaker="Eric", + start=10.0, + end=11.0, + original_text="I ask Chontia.", + corrected_text="I ask Chauntea.", + confidence=0.8, + ) + ] + + revised = apply_corrections(transcript, corrections, confidence_threshold=0.8) + + assert [segment.speaker for segment in revised] == ["Mike", "Eric"] + assert revised[1].text == "I ask Chauntea." + + +def test_apply_corrections_ignores_below_threshold(): + transcript = _transcript() + corrections = [ + CorrectionCandidate( + segment_index=0, + speaker="Eric", + start=10.0, + end=11.0, + original_text="I ask Chontia.", + corrected_text="I ask Chauntea.", + confidence=0.79, + ) + ] + + revised = apply_corrections(transcript, corrections, confidence_threshold=0.8) + + assert revised[1].text == "I ask Chontia." + + +def test_apply_corrections_rejects_duplicate_targets(): + transcript = _transcript() + correction = CorrectionCandidate( + segment_index=0, + speaker="Eric", + start=10.0, + end=11.0, + original_text="I ask Chontia.", + corrected_text="I ask Chauntea.", + confidence=0.8, + ) + + with pytest.raises(AuditaValidationError): + apply_corrections(transcript, [correction, correction], confidence_threshold=0.8) + + +def test_apply_corrections_rejects_mismatched_original_text(): + transcript = _transcript() + correction = CorrectionCandidate( + segment_index=0, + speaker="Eric", + start=10.0, + end=11.0, + original_text="Different text.", + corrected_text="I ask Chauntea.", + confidence=0.8, + ) + + with pytest.raises(AuditaValidationError): + apply_corrections(transcript, [correction], confidence_threshold=0.8) + diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..1ab3ecb --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,99 @@ +from pathlib import Path + +import pytest + +from audita.config import AuditaConfig +from audita.errors import AuditaValidationError +from audita.pipeline import process_transcript +from audita.schemas import CorrectionCandidate, CorrectionSet, parse_glossary_yaml, parse_transcript_json + + +class FakeLLMClient: + def __init__(self, responses): + self.responses = list(responses) + self.calls = 0 + + def create_corrections(self, messages, config): + self.calls += 1 + return self.responses.pop(0) + + +def _config(tmp_path): + return AuditaConfig( + api_key="key", + max_section_tokens=16000, + confidence_threshold=0.8, + max_retries=3, + work_dir=tmp_path / "work", + ) + + +def _glossary(): + return parse_glossary_yaml( + """ + glossary: + - name: "Chauntea" + category: deity + summary: "Chauntea is a deity." + """ + ) + + +def _transcript(): + return parse_transcript_json( + """ + [ + {"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "I ask Chontia."} + ] + """ + ) + + +def test_pipeline_processes_with_fake_llm_and_cleans_work_dir(tmp_path): + correction = CorrectionCandidate( + segment_index=0, + speaker="Eric", + start=0.0, + end=1.0, + original_text="I ask Chontia.", + corrected_text="I ask Chauntea.", + confidence=0.95, + ) + fake_client = FakeLLMClient([CorrectionSet(corrections=[correction])]) + + revised = process_transcript( + _transcript(), + _glossary(), + _config(tmp_path), + llm_client=fake_client, + ) + + assert revised[0].text == "I ask Chauntea." + assert fake_client.calls == 1 + assert list((tmp_path / "work").iterdir()) == [] + + +def test_pipeline_preserves_work_dir_on_failure(tmp_path): + correction = CorrectionCandidate( + segment_index=0, + speaker="Eric", + start=0.0, + end=1.0, + original_text="Different text.", + corrected_text="I ask Chauntea.", + confidence=0.95, + ) + fake_client = FakeLLMClient([CorrectionSet(corrections=[correction])]) + + with pytest.raises(AuditaValidationError): + process_transcript( + _transcript(), + _glossary(), + _config(tmp_path), + llm_client=fake_client, + ) + + preserved = list((tmp_path / "work").iterdir()) + assert len(preserved) == 1 + assert (Path(preserved[0]) / "section-0000.json").exists() + diff --git a/tests/test_validation.py b/tests/test_validation.py new file mode 100644 index 0000000..aa399dc --- /dev/null +++ b/tests/test_validation.py @@ -0,0 +1,76 @@ +import pytest + +from audita.errors import AuditaValidationError +from audita.schemas import parse_glossary_yaml, parse_transcript_json + + +def test_valid_transcript_parses(): + segments = parse_transcript_json( + """ + [ + {"speaker": "Eric", "start": 0.0, "end": 1.25, "text": "Then Lyra."} + ] + """ + ) + + assert len(segments) == 1 + assert segments[0].speaker == "Eric" + + +def test_transcript_rejects_extra_fields(): + with pytest.raises(AuditaValidationError): + parse_transcript_json( + """ + [ + {"speaker": "Eric", "start": 0.0, "end": 1.25, "text": "Hi", "extra": true} + ] + """ + ) + + +def test_transcript_rejects_bad_timestamps(): + with pytest.raises(AuditaValidationError): + parse_transcript_json( + """ + [ + {"speaker": "Eric", "start": 2.0, "end": 1.0, "text": "Hi"} + ] + """ + ) + + +def test_transcript_rejects_empty_input(): + with pytest.raises(AuditaValidationError): + parse_transcript_json("[]") + + +def test_valid_glossary_parses(): + glossary = parse_glossary_yaml( + """ + glossary: + - name: "Lyra" + category: npc + summary: "Lyra is a hostile NPC." + """ + ) + + assert glossary.glossary[0].name == "Lyra" + + +def test_glossary_rejects_empty_entries(): + with pytest.raises(AuditaValidationError): + parse_glossary_yaml("glossary: []") + + +def test_glossary_rejects_extra_fields(): + with pytest.raises(AuditaValidationError): + parse_glossary_yaml( + """ + glossary: + - name: "Lyra" + category: npc + summary: "Lyra is a hostile NPC." + extra: "nope" + """ + ) +