From f834ffad970616519de13cd8b7a4b9a02c2feb2f Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Tue, 28 Apr 2026 21:45:38 -0500 Subject: [PATCH] Implemented provider-neutral OpenAI-compatible LLM endpoint support --- README.md | 30 ++++++- src/audita/cli.py | 6 +- src/audita/core/config.py | 20 ++++- src/audita/framework/llm.py | 44 ++++++--- src/audita/pipeline.py | 4 +- tests/test_framework_llm.py | 174 ++++++++++++++++++++++++++++++++++++ tests/test_new_cli.py | 55 ++++++++++++ tests/test_new_config.py | 38 ++++++++ tests/test_new_pipeline.py | 3 +- 9 files changed, 353 insertions(+), 21 deletions(-) create mode 100644 tests/test_framework_llm.py diff --git a/README.md b/README.md index e2bbc8e..3c8a2da 100644 --- a/README.md +++ b/README.md @@ -79,12 +79,13 @@ audita process transcript.json --glossary glossary.yaml --output corrected.json Without `--output`, Audita writes the corrected transcript JSON to stdout and progress logs to stderr. `--report-json` writes a separate machine-readable run report and never mixes report data into stdout. -Useful configuration can be supplied by CLI flag or environment variable. CLI flags take precedence over environment variables. Normal runs now require `OPENROUTER_API_KEY`, because the `glossary`, `homophones`, `spoken_word`, and `grammar` modules make real LLM calls. +Useful configuration can be supplied by CLI flag or environment variable. CLI flags take precedence over environment variables. Normal runs now require LLM API credentials, because the `glossary`, `homophones`, `spoken_word`, and `grammar` modules make real LLM calls. `AUDITA_LLM_API_KEY` and `--llm-api-key` are the preferred provider-neutral credential surfaces, while `OPENROUTER_API_KEY` remains supported as a backward-compatible fallback. | 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_LLM_API_KEY` | `--llm-api-key` | unset | Preferred provider-neutral LLM API credential; CLI overrides both environment-key variants | +| `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_MAX_RETRIES` | `--max-retries` | `3` | Maximum Instructor retries for structured responses | | `AUDITA_MAX_SECTION_TOKENS` | `--max-section-tokens` | `6144` | Maximum estimated tokens per transcript batch | @@ -101,6 +102,31 @@ Useful configuration can be supplied by CLI flag or environment variable. CLI fl Set `AUDITA_MODULES=grammar` to run only the grammar module by default, or override it per command with `--modules`. +OpenRouter remains the default out of the box: + +```sh +export AUDITA_LLM_API_KEY=your-openrouter-key +audita process transcript.json --glossary glossary.yaml --output corrected.json +``` + +You can point Audita at any OpenAI-compatible endpoint by changing `AUDITA_BASE_URL` and, if needed, `AUDITA_MODEL`. For example, a local vLLM server: + +```sh +export AUDITA_LLM_API_KEY=local-dev-key +export AUDITA_BASE_URL=http://localhost:8000/v1 +export AUDITA_MODEL=meta-llama/Llama-3.1-8B-Instruct +audita process transcript.json --glossary glossary.yaml --output corrected.json +``` + +Or the actual OpenAI API: + +```sh +export AUDITA_LLM_API_KEY=your-openai-key +export AUDITA_BASE_URL=https://api.openai.com/v1 +export AUDITA_MODEL=gpt-4.1-mini +audita process transcript.json --glossary glossary.yaml --output corrected.json +``` + `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. diff --git a/src/audita/cli.py b/src/audita/cli.py index 9a065cf..d3398e3 100644 --- a/src/audita/cli.py +++ b/src/audita/cli.py @@ -27,12 +27,13 @@ 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("--llm-api-key", help="LLM API key for the configured OpenAI-compatible endpoint") 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("--model", help="LLM model name for the configured OpenAI-compatible endpoint") + process.add_argument("--base-url", help="OpenAI-compatible API base URL for Audita LLM stages") process.add_argument("--max-retries", type=int, help="maximum structured-output retries for LLM stages") process.add_argument("--max-section-tokens", type=int, help="maximum estimated tokens per transcript batch") process.add_argument( @@ -88,6 +89,7 @@ def _process(args: argparse.Namespace) -> int: try: config = AuditaConfig.from_sources( overrides=ConfigOverrides( + llm_api_key=args.llm_api_key, module_keys=args.modules, model=args.model, base_url=args.base_url, diff --git a/src/audita/core/config.py b/src/audita/core/config.py index 5b16b86..32d6850 100644 --- a/src/audita/core/config.py +++ b/src/audita/core/config.py @@ -27,6 +27,7 @@ DEFAULT_NORMALIZE_MAX_SEGMENT_TOKENS = 2048 @dataclass(frozen=True) class ConfigOverrides: + llm_api_key: Optional[str] = None module_keys: Optional[Union[str, Sequence[str]]] = None model: Optional[str] = None base_url: Optional[str] = None @@ -72,7 +73,11 @@ class AuditaConfig: source = os.environ if env is None else env selected = ConfigOverrides() if overrides is None else overrides config = cls( - api_key=_select_optional_string(source.get("OPENROUTER_API_KEY")), + api_key=_select_api_key( + selected.llm_api_key, + source.get("AUDITA_LLM_API_KEY"), + source.get("OPENROUTER_API_KEY"), + ), module_keys=_select_module_keys( selected.module_keys, source.get("AUDITA_MODULES"), @@ -219,6 +224,19 @@ def _select_optional_string(value: Optional[str]) -> Optional[str]: return stripped or None +def _select_api_key( + cli_value: Optional[str], + generic_env_value: Optional[str], + legacy_env_value: Optional[str], +) -> Optional[str]: + if cli_value is not None: + return _select_optional_string(cli_value) + generic = _select_optional_string(generic_env_value) + if generic is not None: + return generic + return _select_optional_string(legacy_env_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 diff --git a/src/audita/framework/llm.py b/src/audita/framework/llm.py index 4144830..ae8faf4 100644 --- a/src/audita/framework/llm.py +++ b/src/audita/framework/llm.py @@ -2,11 +2,11 @@ from __future__ import annotations from typing import Any, Optional, Sequence, Tuple -from audita.core.config import AuditaConfig +from audita.core.config import AuditaConfig, DEFAULT_BASE_URL from audita.core.errors import AuditaLLMError -class OpenRouterStructuredLLMClient: +class OpenAICompatibleStructuredLLMClient: def __init__(self) -> None: self._client = None self._client_identity: Optional[Tuple[str, str]] = None @@ -21,23 +21,25 @@ class OpenRouterStructuredLLMClient: ) -> Any: if not config.api_key: raise AuditaLLMError( - f"Structured LLM stage '{stage_name}' requires OPENROUTER_API_KEY to be configured." + f"Structured LLM stage '{stage_name}' requires LLM API credentials to be configured " + "via --llm-api-key, AUDITA_LLM_API_KEY, or OPENROUTER_API_KEY." ) client = self._get_client(config) - model = _normalize_openrouter_model(config.model) + request = { + "model": _request_model_name(config), + "messages": list(messages), + "response_model": response_model, + "max_retries": config.max_retries, + } + if _uses_openrouter_shape(config): + request["extra_body"] = {"provider": {"require_parameters": True}} try: - return client.chat.completions.create( - model=model, - messages=list(messages), - response_model=response_model, - max_retries=config.max_retries, - extra_body={"provider": {"require_parameters": True}}, - ) + return client.chat.completions.create(**request) except Exception as exc: raise AuditaLLMError( - f"Structured LLM stage '{stage_name}' failed. Confirm the configured OpenRouter model " - "supports tool calling or structured outputs." + f"Structured LLM stage '{stage_name}' failed. Confirm the configured model and " + "OpenAI-compatible endpoint support tool calling or structured outputs." ) from exc def _get_client(self, config: AuditaConfig) -> Any: @@ -64,3 +66,19 @@ def _normalize_openrouter_model(model: str) -> str: if model.startswith(prefix): return model[len(prefix) :] return model + + +def _request_model_name(config: AuditaConfig) -> str: + if _uses_openrouter_shape(config): + return _normalize_openrouter_model(config.model) + return config.model + + +def _uses_openrouter_shape(config: AuditaConfig) -> bool: + return _normalized_base_url(config.base_url) == _normalized_base_url(DEFAULT_BASE_URL) or config.model.startswith( + "openrouter/" + ) + + +def _normalized_base_url(base_url: str) -> str: + return base_url.rstrip("/") diff --git a/src/audita/pipeline.py b/src/audita/pipeline.py index 34d6662..8f2a38b 100644 --- a/src/audita/pipeline.py +++ b/src/audita/pipeline.py @@ -10,7 +10,7 @@ from .core.errors import AuditaError from .core.normalization import normalize_transcript from .core.reporting import AppliedChange, ModuleRunReport, ProcessResult, ReportedSkip, RunReport from .core.schemas import Glossary, SourceTranscriptSegment, TranscriptSegment, source_transcript_to_json, transcript_to_json -from .framework.llm import OpenRouterStructuredLLMClient +from .framework.llm import OpenAICompatibleStructuredLLMClient from .framework.models import StructuredLLMClient from .framework.runner import PipelineRunError, PipelineRunner from .modules import resolve_module_specs @@ -70,7 +70,7 @@ def process_transcript_result( ) pipeline_runner = PipelineRunner() - effective_llm_client = OpenRouterStructuredLLMClient() if llm_client is None else llm_client + effective_llm_client = OpenAICompatibleStructuredLLMClient() if llm_client is None else llm_client pipeline_result = pipeline_runner.run( transcript=normalized_transcript, glossary=glossary, diff --git a/tests/test_framework_llm.py b/tests/test_framework_llm.py new file mode 100644 index 0000000..8f094ab --- /dev/null +++ b/tests/test_framework_llm.py @@ -0,0 +1,174 @@ +import sys +import types + +import pytest + +from audita.core.config import AuditaConfig +from audita.core.errors import AuditaLLMError +from audita.framework.llm import OpenAICompatibleStructuredLLMClient + + +class DummyResponseModel: + pass + + +def _config(**overrides): + base = AuditaConfig.from_sources(env={}) + data = { + "api_key": "test-key", + "module_keys": base.module_keys, + "model": base.model, + "base_url": base.base_url, + "max_retries": base.max_retries, + "max_section_tokens": base.max_section_tokens, + "glossary_confidence_threshold": base.glossary_confidence_threshold, + "grammar_confidence_threshold": base.grammar_confidence_threshold, + "homophones_confidence_threshold": base.homophones_confidence_threshold, + "spoken_word_confidence_threshold": base.spoken_word_confidence_threshold, + "normalize_max_segment_gap": base.normalize_max_segment_gap, + "normalize_ellipsis_gap": base.normalize_ellipsis_gap, + "normalize_max_segment_duration": base.normalize_max_segment_duration, + "normalize_max_segment_tokens": base.normalize_max_segment_tokens, + "work_dir": base.work_dir, + "work_dir_retention": base.work_dir_retention, + } + data.update(overrides) + return AuditaConfig(**data) + + +def _install_fake_llm_modules(monkeypatch): + create_calls = [] + openai_inits = [] + + class FakePatchedClient: + def __init__(self): + self.chat = types.SimpleNamespace(completions=types.SimpleNamespace(create=self._create)) + + def _create(self, **kwargs): + create_calls.append(kwargs) + return {"ok": True} + + class FakeOpenAI: + def __init__(self, *, api_key, base_url): + openai_inits.append({"api_key": api_key, "base_url": base_url}) + + fake_instructor = types.SimpleNamespace( + Mode=types.SimpleNamespace(TOOLS="TOOLS"), + patch=lambda client, mode: FakePatchedClient(), + ) + fake_openai = types.SimpleNamespace(OpenAI=FakeOpenAI) + monkeypatch.setitem(sys.modules, "instructor", fake_instructor) + monkeypatch.setitem(sys.modules, "openai", fake_openai) + return create_calls, openai_inits + + +def test_openrouter_requests_strip_prefix_and_include_extra_body(monkeypatch): + create_calls, _ = _install_fake_llm_modules(monkeypatch) + client = OpenAICompatibleStructuredLLMClient() + + client.run_structured( + stage_name="test-stage", + messages=[{"role": "user", "content": "Hello"}], + response_model=DummyResponseModel, + config=_config(model="openrouter/google/gemma-4-31b-it"), + ) + + assert create_calls == [ + { + "model": "google/gemma-4-31b-it", + "messages": [{"role": "user", "content": "Hello"}], + "response_model": DummyResponseModel, + "max_retries": 3, + "extra_body": {"provider": {"require_parameters": True}}, + } + ] + + +def test_openrouter_default_base_url_uses_openrouter_request_shape(monkeypatch): + create_calls, _ = _install_fake_llm_modules(monkeypatch) + client = OpenAICompatibleStructuredLLMClient() + + client.run_structured( + stage_name="test-stage", + messages=[{"role": "user", "content": "Hello"}], + response_model=DummyResponseModel, + config=_config(model="google/gemma-4-31b-it"), + ) + + assert create_calls == [ + { + "model": "google/gemma-4-31b-it", + "messages": [{"role": "user", "content": "Hello"}], + "response_model": DummyResponseModel, + "max_retries": 3, + "extra_body": {"provider": {"require_parameters": True}}, + } + ] + + +def test_generic_endpoint_requests_keep_model_and_omit_extra_body(monkeypatch): + create_calls, _ = _install_fake_llm_modules(monkeypatch) + client = OpenAICompatibleStructuredLLMClient() + + client.run_structured( + stage_name="test-stage", + messages=[{"role": "user", "content": "Hello"}], + response_model=DummyResponseModel, + config=_config( + model="meta-llama/Llama-3.1-8B-Instruct", + base_url="http://localhost:8000/v1", + ), + ) + + assert create_calls == [ + { + "model": "meta-llama/Llama-3.1-8B-Instruct", + "messages": [{"role": "user", "content": "Hello"}], + "response_model": DummyResponseModel, + "max_retries": 3, + } + ] + + +def test_client_cache_identity_uses_api_key_and_base_url(monkeypatch): + _, openai_inits = _install_fake_llm_modules(monkeypatch) + client = OpenAICompatibleStructuredLLMClient() + first = _config(api_key="key-1", base_url="http://localhost:8000/v1") + second = _config(api_key="key-1", base_url="http://localhost:8000/v1") + third = _config(api_key="key-1", base_url="https://api.openai.com/v1") + + client.run_structured( + stage_name="one", + messages=[{"role": "user", "content": "Hello"}], + response_model=DummyResponseModel, + config=first, + ) + client.run_structured( + stage_name="two", + messages=[{"role": "user", "content": "Hello"}], + response_model=DummyResponseModel, + config=second, + ) + client.run_structured( + stage_name="three", + messages=[{"role": "user", "content": "Hello"}], + response_model=DummyResponseModel, + config=third, + ) + + assert openai_inits == [ + {"api_key": "key-1", "base_url": "http://localhost:8000/v1"}, + {"api_key": "key-1", "base_url": "https://api.openai.com/v1"}, + ] + + +def test_missing_api_key_error_is_provider_neutral(): + client = OpenAICompatibleStructuredLLMClient() + + with pytest.raises(AuditaLLMError, match="AUDITA_LLM_API_KEY"): + client.run_structured( + stage_name="test-stage", + messages=[{"role": "user", "content": "Hello"}], + response_model=DummyResponseModel, + config=_config(api_key=None), + ) diff --git a/tests/test_new_cli.py b/tests/test_new_cli.py index 4d8824a..f7284e7 100644 --- a/tests/test_new_cli.py +++ b/tests/test_new_cli.py @@ -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 "--llm-api-key" in output assert "--modules" in output assert "--model" in output assert "--base-url" in output @@ -139,3 +140,57 @@ def test_cli_process_passes_modules_override_to_config(monkeypatch, tmp_path): assert exit_code == 0 assert captured["module_keys"] == "grammar" + + +def test_cli_process_passes_llm_api_key_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["llm_api_key"] = overrides.llm_api_key + 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", + "--llm-api-key", + "cli-key", + ] + ) + + assert exit_code == 0 + assert captured["llm_api_key"] == "cli-key" diff --git a/tests/test_new_config.py b/tests/test_new_config.py index 2d88e23..07dff15 100644 --- a/tests/test_new_config.py +++ b/tests/test_new_config.py @@ -36,6 +36,44 @@ def test_cli_overrides_take_precedence(): assert config.max_section_tokens == 2000 +def test_generic_llm_api_key_env_is_read(): + config = AuditaConfig.from_sources(env={"AUDITA_LLM_API_KEY": "generic-key"}) + + assert config.api_key == "generic-key" + + +def test_generic_llm_api_key_takes_precedence_over_openrouter_env(): + config = AuditaConfig.from_sources( + env={ + "AUDITA_LLM_API_KEY": "generic-key", + "OPENROUTER_API_KEY": "legacy-key", + } + ) + + assert config.api_key == "generic-key" + + +def test_llm_api_key_cli_override_takes_precedence_over_env(): + config = AuditaConfig.from_sources( + env={ + "AUDITA_LLM_API_KEY": "generic-key", + "OPENROUTER_API_KEY": "legacy-key", + }, + overrides=ConfigOverrides(llm_api_key="cli-key"), + ) + + assert config.api_key == "cli-key" + + +def test_blank_llm_api_key_override_resolves_to_none(): + config = AuditaConfig.from_sources( + env={"OPENROUTER_API_KEY": "legacy-key"}, + overrides=ConfigOverrides(llm_api_key=" "), + ) + + assert config.api_key is None + + def test_module_key_overrides_take_precedence(): config = AuditaConfig.from_sources( env={"AUDITA_MODULES": "grammar"}, diff --git a/tests/test_new_pipeline.py b/tests/test_new_pipeline.py index 52b3db7..791af82 100644 --- a/tests/test_new_pipeline.py +++ b/tests/test_new_pipeline.py @@ -267,7 +267,7 @@ def test_process_transcript_result_missing_api_key_writes_failed_report(tmp_path work_dir_retention="never", ) - with pytest.raises(AuditaLLMError, match="OPENROUTER_API_KEY"): + with pytest.raises(AuditaLLMError, match="AUDITA_LLM_API_KEY"): process_transcript_result(_transcript(), _glossary(), config) run_dir = next((tmp_path / "work").iterdir()) @@ -287,6 +287,7 @@ def test_process_transcript_result_missing_api_key_writes_failed_report(tmp_path assert report["skipped_corrections"] == [] assert report["work_dir_retained"] is True assert report["work_dir"] == str(run_dir) + assert "AUDITA_LLM_API_KEY" in report["error"] assert "OPENROUTER_API_KEY" in report["error"]