diff --git a/README.md b/README.md index 4a1bb01..6e53b6a 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,7 @@ Useful configuration can be supplied by CLI flag or environment variable. CLI fl | `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_LLM_TIMEOUT_SECONDS` | `--llm-timeout-seconds` | `600` | Per-request timeout in seconds for LLM calls to the configured OpenAI-compatible endpoint | | `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 | diff --git a/src/audita/cli.py b/src/audita/cli.py index 61aceea..a7c27a1 100644 --- a/src/audita/cli.py +++ b/src/audita/cli.py @@ -29,6 +29,11 @@ def _build_parser() -> argparse.ArgumentParser: 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("--llm-concurrency", type=int, help="maximum concurrent LLM calls within a module stage") + process.add_argument( + "--llm-timeout-seconds", + type=float, + help="per-request timeout in seconds for LLM calls to the configured OpenAI-compatible endpoint", + ) process.add_argument( "--modules", help="comma-separated module keys to run, for example: grammar or glossary,homophones,glossary", @@ -97,6 +102,7 @@ def _process(args: argparse.Namespace) -> int: overrides=ConfigOverrides( llm_api_key=args.llm_api_key, llm_concurrency=args.llm_concurrency, + llm_timeout_seconds=args.llm_timeout_seconds, 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 f126b28..30bcfe2 100644 --- a/src/audita/core/config.py +++ b/src/audita/core/config.py @@ -12,6 +12,7 @@ from .errors import AuditaConfigError DEFAULT_MODEL = "openrouter/google/gemma-4-31b-it" 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 @@ -31,6 +32,7 @@ DEFAULT_NORMALIZE_MAX_SEGMENT_TOKENS = 2048 class ConfigOverrides: llm_api_key: Optional[str] = None llm_concurrency: Optional[int] = None + llm_timeout_seconds: Optional[float] = None module_keys: Optional[Union[str, Sequence[str]]] = None model: Optional[str] = None base_url: Optional[str] = None @@ -53,6 +55,7 @@ class ConfigOverrides: class AuditaConfig: api_key: Optional[str] = None llm_concurrency: int = DEFAULT_LLM_CONCURRENCY + llm_timeout_seconds: float = DEFAULT_LLM_TIMEOUT_SECONDS module_keys: Tuple[str, ...] = DEFAULT_MODULE_KEYS model: str = DEFAULT_MODEL base_url: str = DEFAULT_BASE_URL @@ -90,6 +93,12 @@ class AuditaConfig: DEFAULT_LLM_CONCURRENCY, "AUDITA_LLM_CONCURRENCY", ), + llm_timeout_seconds=_select_float( + selected.llm_timeout_seconds, + source.get("AUDITA_LLM_TIMEOUT_SECONDS"), + DEFAULT_LLM_TIMEOUT_SECONDS, + "AUDITA_LLM_TIMEOUT_SECONDS", + ), module_keys=_select_module_keys( selected.module_keys, source.get("AUDITA_MODULES"), @@ -180,6 +189,10 @@ class AuditaConfig: normalize_module_keys(self.module_keys) if self.llm_concurrency <= 0: raise AuditaConfigError("AUDITA_LLM_CONCURRENCY must be greater than zero.") + if not math.isfinite(self.llm_timeout_seconds): + 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 not self.model.strip(): raise AuditaConfigError("AUDITA_MODEL must not be empty.") if not self.base_url.strip(): @@ -227,6 +240,7 @@ class AuditaConfig: return { "api_key_configured": bool(self.api_key), "llm_concurrency": self.llm_concurrency, + "llm_timeout_seconds": self.llm_timeout_seconds, "module_keys": list(self.module_keys), "model": self.model, "base_url": self.base_url, diff --git a/src/audita/framework/llm.py b/src/audita/framework/llm.py index 034d80c..02b28a8 100644 --- a/src/audita/framework/llm.py +++ b/src/audita/framework/llm.py @@ -10,7 +10,7 @@ from audita.core.errors import AuditaLLMError class OpenAICompatibleStructuredLLMClient: def __init__(self) -> None: self._client = None - self._client_identity: Optional[Tuple[str, str]] = None + self._client_identity: Optional[Tuple[str, str, float]] = None self._client_lock = threading.Lock() def run_structured( @@ -45,7 +45,7 @@ class OpenAICompatibleStructuredLLMClient: ) from exc def _get_client(self, config: AuditaConfig) -> Any: - identity = (config.api_key or "", config.base_url) + identity = (config.api_key or "", config.base_url, config.llm_timeout_seconds) with self._client_lock: if self._client is not None and self._client_identity == identity: return self._client @@ -58,7 +58,11 @@ class OpenAICompatibleStructuredLLMClient: "The LLM dependencies are not installed. Run `uv sync` before using Audita LLM stages." ) from exc - openai_client = OpenAI(api_key=config.api_key, base_url=config.base_url) + openai_client = OpenAI( + api_key=config.api_key, + base_url=config.base_url, + timeout=config.llm_timeout_seconds, + ) self._client = instructor.patch(openai_client, mode=instructor.Mode.TOOLS) self._client_identity = identity return self._client diff --git a/tests/test_framework_llm.py b/tests/test_framework_llm.py index c828cb8..4ad6ec7 100644 --- a/tests/test_framework_llm.py +++ b/tests/test_framework_llm.py @@ -18,6 +18,7 @@ def _config(**overrides): data = { "api_key": "test-key", "llm_concurrency": base.llm_concurrency, + "llm_timeout_seconds": base.llm_timeout_seconds, "module_keys": base.module_keys, "model": base.model, "base_url": base.base_url, @@ -51,8 +52,8 @@ def _install_fake_llm_modules(monkeypatch): return {"ok": True} class FakeOpenAI: - def __init__(self, *, api_key, base_url): - openai_inits.append({"api_key": api_key, "base_url": base_url}) + def __init__(self, *, api_key, base_url, timeout): + openai_inits.append({"api_key": api_key, "base_url": base_url, "timeout": timeout}) fake_instructor = types.SimpleNamespace( Mode=types.SimpleNamespace(TOOLS="TOOLS"), @@ -159,8 +160,33 @@ def test_client_cache_identity_uses_api_key_and_base_url(monkeypatch): ) assert openai_inits == [ - {"api_key": "key-1", "base_url": "http://localhost:8000/v1"}, - {"api_key": "key-1", "base_url": "https://api.openai.com/v1"}, + {"api_key": "key-1", "base_url": "http://localhost:8000/v1", "timeout": 600}, + {"api_key": "key-1", "base_url": "https://api.openai.com/v1", "timeout": 600}, + ] + + +def test_client_cache_identity_uses_timeout(monkeypatch): + _, openai_inits = _install_fake_llm_modules(monkeypatch) + client = OpenAICompatibleStructuredLLMClient() + first = _config(api_key="key-1", base_url="http://localhost:8000/v1", llm_timeout_seconds=600) + second = _config(api_key="key-1", base_url="http://localhost:8000/v1", llm_timeout_seconds=1200) + + 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, + ) + + assert openai_inits == [ + {"api_key": "key-1", "base_url": "http://localhost:8000/v1", "timeout": 600}, + {"api_key": "key-1", "base_url": "http://localhost:8000/v1", "timeout": 1200}, ] @@ -194,4 +220,4 @@ def test_client_initialization_is_safe_under_concurrent_calls(monkeypatch): ) ) - assert openai_inits == [{"api_key": "key-1", "base_url": "http://localhost:8000/v1"}] + assert openai_inits == [{"api_key": "key-1", "base_url": "http://localhost:8000/v1", "timeout": 600}] diff --git a/tests/test_new_cli.py b/tests/test_new_cli.py index 7f1927f..436ef53 100644 --- a/tests/test_new_cli.py +++ b/tests/test_new_cli.py @@ -22,6 +22,7 @@ def test_process_help_exposes_framework_flags(capsys): assert "--report-json" in output assert "--llm-api-key" in output assert "--llm-concurrency" in output + assert "--llm-timeout-seconds" in output assert "--modules" in output assert "--model" in output assert "--base-url" in output @@ -252,6 +253,60 @@ def test_cli_process_passes_llm_concurrency_override_to_config(monkeypatch, tmp_ assert captured["llm_concurrency"] == 3 +def test_cli_process_passes_llm_timeout_seconds_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_timeout_seconds"] = overrides.llm_timeout_seconds + 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-timeout-seconds", + "900", + ] + ) + + assert exit_code == 0 + assert captured["llm_timeout_seconds"] == 900.0 + + 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 caf2e04..252ecca 100644 --- a/tests/test_new_config.py +++ b/tests/test_new_config.py @@ -7,6 +7,7 @@ from audita.core.config import ( DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD, DEFAULT_HOMOPHONES_CONFIDENCE_THRESHOLD, DEFAULT_LLM_CONCURRENCY, + DEFAULT_LLM_TIMEOUT_SECONDS, DEFAULT_MAX_SECTION_TOKENS, DEFAULT_MIN_SECTION_TOKENS, DEFAULT_NORMALIZE_MAX_SEGMENT_GAP, @@ -22,6 +23,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.max_section_tokens == DEFAULT_MAX_SECTION_TOKENS assert config.min_section_tokens == DEFAULT_MIN_SECTION_TOKENS assert config.module_keys == DEFAULT_MODULE_KEYS @@ -66,6 +68,21 @@ def test_llm_concurrency_env_is_parsed(): assert config.llm_concurrency == 3 +def test_llm_timeout_seconds_cli_override_takes_precedence(): + config = AuditaConfig.from_sources( + env={"AUDITA_LLM_TIMEOUT_SECONDS": "120"}, + overrides=ConfigOverrides(llm_timeout_seconds=900.0), + ) + + assert config.llm_timeout_seconds == 900.0 + + +def test_llm_timeout_seconds_env_is_parsed(): + config = AuditaConfig.from_sources(env={"AUDITA_LLM_TIMEOUT_SECONDS": "120.5"}) + + assert config.llm_timeout_seconds == 120.5 + + def test_min_section_tokens_env_is_parsed(): config = AuditaConfig.from_sources(env={"AUDITA_MIN_SECTION_TOKENS": "3000"}) @@ -130,6 +147,12 @@ def test_invalid_work_dir_retention_is_rejected(): AuditaConfig.from_sources(env={"AUDITA_WORK_DIR_RETENTION": "sometimes"}) +def test_report_dict_includes_llm_timeout_seconds(): + config = AuditaConfig.from_sources(env={"AUDITA_LLM_TIMEOUT_SECONDS": "321"}) + + assert config.to_report_dict()["llm_timeout_seconds"] == 321.0 + + def test_threshold_overrides_take_precedence(): config = AuditaConfig.from_sources( env={ @@ -185,6 +208,12 @@ def test_invalid_llm_concurrency_is_rejected(value): AuditaConfig.from_sources(env={"AUDITA_LLM_CONCURRENCY": value}) +@pytest.mark.parametrize("value", ["0", "-1", "many"]) +def test_invalid_llm_timeout_seconds_is_rejected(value): + with pytest.raises(AuditaConfigError, match="AUDITA_LLM_TIMEOUT_SECONDS"): + AuditaConfig.from_sources(env={"AUDITA_LLM_TIMEOUT_SECONDS": 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"):