eric/audita
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

Added a configurable LLM timeout

Browse Source
This commit is contained in:
Eric Rakestraw
2026-04-29 18:08:39 -05:00
parent d8bc84934e
commit f841f7eb71
7 changed files with 143 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
README.md
View File

@@ -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 |

6
src/audita/cli.py
View File

@@ -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,

14
src/audita/core/config.py
View File

@@ -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,

10
src/audita/framework/llm.py
View File

@@ -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

36
tests/test_framework_llm.py
View File

@@ -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}]

55
tests/test_new_cli.py
View File

@@ -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(

29
tests/test_new_config.py
View File

@@ -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"):