diff --git a/src/openchronicle/config.py b/src/openchronicle/config.py index 21794c89..2fcf7ef8 100644 --- a/src/openchronicle/config.py +++ b/src/openchronicle/config.py @@ -18,6 +18,14 @@ class ModelConfig: api_key: str = "" api_key_env: str = "OPENAI_API_KEY" max_tokens: int | None = None + # Per-call timeout in seconds. ``None`` falls back to the writer + # default (writer.llm.DEFAULT_TIMEOUT_SECONDS) so a hung connection + # can never block a reducer / classifier thread indefinitely. + timeout_seconds: int | None = None + # Number of automatic retries on transient failures (429s, network + # blips). ``None`` falls back to writer.llm.DEFAULT_NUM_RETRIES. + # litellm honors Retry-After headers between attempts. + num_retries: int | None = None @dataclass @@ -207,6 +215,8 @@ def load(path: Path | None = None) -> Config: api_key_env = "OPENAI_API_KEY" # base_url = "" # api_key = "" # overrides api_key_env if set +# timeout_seconds = 120 # per-call timeout; raise this if you run a slow local model +# num_retries = 2 # automatic retries on transient errors (429, network blips) [models.compact] # Accuracy-sensitive — match or exceed the default. diff --git a/src/openchronicle/writer/llm.py b/src/openchronicle/writer/llm.py index 6e4317bd..ed87885e 100644 --- a/src/openchronicle/writer/llm.py +++ b/src/openchronicle/writer/llm.py @@ -11,6 +11,20 @@ logger = get("openchronicle.writer") +# litellm.completion() defaults to no client-side timeout and zero +# retries. Without these, a stuck connection (slow provider, partial +# response, dead TCP socket) blocks the reducer / classifier daemon +# thread *forever* — the daemon stays "alive" but no durable facts +# get written. Two minutes is generous for a long session reduce +# without being absurdly long; a slow local model can override per +# stage via ``[models.] timeout_seconds = N``. +DEFAULT_TIMEOUT_SECONDS = 120 +# Two retries gives 3 total attempts. litellm uses tenacity-style +# backoff between attempts and respects Retry-After headers from +# providers, so this is enough to absorb a brief 429 / 502 blip +# without piling on extra latency. +DEFAULT_NUM_RETRIES = 2 + def call_llm( cfg: Config, @@ -30,9 +44,21 @@ def call_llm( import litellm # imported lazily to keep CLI startup fast model_cfg = cfg.model_for(stage) + timeout = ( + model_cfg.timeout_seconds + if model_cfg.timeout_seconds is not None + else DEFAULT_TIMEOUT_SECONDS + ) + num_retries = ( + model_cfg.num_retries + if model_cfg.num_retries is not None + else DEFAULT_NUM_RETRIES + ) kwargs: dict[str, Any] = { "model": model_cfg.model, "messages": messages, + "timeout": timeout, + "num_retries": num_retries, } if model_cfg.base_url: kwargs["api_base"] = model_cfg.base_url @@ -47,7 +73,10 @@ def call_llm( if model_cfg.max_tokens: kwargs["max_tokens"] = model_cfg.max_tokens - logger.debug("llm call stage=%s model=%s", stage, model_cfg.model) + logger.debug( + "llm call stage=%s model=%s timeout=%ds retries=%d", + stage, model_cfg.model, timeout, num_retries, + ) return litellm.completion(**kwargs) diff --git a/tests/test_config.py b/tests/test_config.py index f5e742db..418e017e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -49,3 +49,54 @@ def test_api_key_precedence(tmp_path: Path, monkeypatch) -> None: assert config.resolve_api_key(cfg) == "direct" cfg2 = config.ModelConfig(api_key="", api_key_env="ENV_KEY") assert config.resolve_api_key(cfg2) == "from-env" + + +def test_timeout_and_retries_inherit_from_default(tmp_path: Path) -> None: + """[models.] without explicit timeout_seconds inherits the default.""" + path = tmp_path / "config.toml" + path.write_text( + """ +[models.default] +model = "gpt-5.4-nano" +timeout_seconds = 200 +num_retries = 4 + +[models.reducer] +model = "claude-sonnet-4" +""" + ) + cfg = config.load(path) + reducer = cfg.model_for("reducer") + assert reducer.model == "claude-sonnet-4" + # timeout_seconds / num_retries come down from [models.default] + assert reducer.timeout_seconds == 200 + assert reducer.num_retries == 4 + + +def test_timeout_per_stage_overrides_default(tmp_path: Path) -> None: + """A stage that explicitly sets timeout_seconds wins over the default.""" + path = tmp_path / "config.toml" + path.write_text( + """ +[models.default] +timeout_seconds = 120 + +[models.reducer] +timeout_seconds = 600 +""" + ) + cfg = config.load(path) + assert cfg.model_for("default").timeout_seconds == 120 + assert cfg.model_for("reducer").timeout_seconds == 600 + + +def test_missing_timeout_defaults_to_none(tmp_path: Path) -> None: + """An old config that doesn't mention timeout_seconds parses to None. + + The llm.py wrapper interprets None as "use module default" so a + user upgrading from a config that pre-dates this field gets the + DEFAULT_TIMEOUT_SECONDS automatically — no breakage. + """ + cfg = config.load(tmp_path / "missing.toml") + assert cfg.model_for("reducer").timeout_seconds is None + assert cfg.model_for("reducer").num_retries is None diff --git a/tests/test_writer_llm.py b/tests/test_writer_llm.py new file mode 100644 index 00000000..a4a46431 --- /dev/null +++ b/tests/test_writer_llm.py @@ -0,0 +1,107 @@ +"""Tests for the litellm wrapper: timeout / retries / per-stage override.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +from openchronicle.config import Config, ModelConfig +from openchronicle.writer import llm as llm_mod + + +def _cfg_with_models(models: dict[str, ModelConfig]) -> Config: + """Build a Config whose ``model_for(stage)`` returns the given mapping. + + The real ``Config.model_for`` falls back from stage → 'default' → + bare ``ModelConfig()``, so we make sure 'default' exists. + """ + if "default" not in models: + models["default"] = ModelConfig() + cfg = Config() + cfg.models = models + return cfg + + +def _stub_completion_response() -> Any: + """Return-shape doesn't matter for these tests — they only inspect kwargs.""" + msg = MagicMock() + msg.content = "" + msg.tool_calls = None + choice = MagicMock() + choice.message = msg + choice.finish_reason = "stop" + resp = MagicMock() + resp.choices = [choice] + return resp + + +def test_call_llm_passes_default_timeout_and_retries(monkeypatch) -> None: + """Stage with no explicit override → module defaults are forwarded to litellm.""" + monkeypatch.delenv("OPENCHRONICLE_LLM_MOCK", raising=False) + cfg = _cfg_with_models({}) + + with patch("litellm.completion", return_value=_stub_completion_response()) as mock: + llm_mod.call_llm( + cfg, "reducer", + messages=[{"role": "user", "content": "hi"}], + ) + kwargs = mock.call_args.kwargs + + assert kwargs["timeout"] == llm_mod.DEFAULT_TIMEOUT_SECONDS + assert kwargs["num_retries"] == llm_mod.DEFAULT_NUM_RETRIES + + +def test_call_llm_per_stage_timeout_override(monkeypatch) -> None: + """A stage-level timeout_seconds in the config overrides the default.""" + monkeypatch.delenv("OPENCHRONICLE_LLM_MOCK", raising=False) + cfg = _cfg_with_models( + { + "reducer": ModelConfig(timeout_seconds=300, num_retries=5), + } + ) + + with patch("litellm.completion", return_value=_stub_completion_response()) as mock: + llm_mod.call_llm( + cfg, "reducer", + messages=[{"role": "user", "content": "hi"}], + ) + kwargs = mock.call_args.kwargs + + assert kwargs["timeout"] == 300 + assert kwargs["num_retries"] == 5 + + +def test_call_llm_other_stages_keep_defaults_when_one_is_overridden(monkeypatch) -> None: + """Setting timeout for ``reducer`` must not leak into ``classifier``.""" + monkeypatch.delenv("OPENCHRONICLE_LLM_MOCK", raising=False) + cfg = _cfg_with_models( + { + "reducer": ModelConfig(timeout_seconds=300), + "classifier": ModelConfig(), # no override + } + ) + + with patch("litellm.completion", return_value=_stub_completion_response()) as mock: + llm_mod.call_llm( + cfg, "classifier", + messages=[{"role": "user", "content": "hi"}], + ) + kwargs = mock.call_args.kwargs + + assert kwargs["timeout"] == llm_mod.DEFAULT_TIMEOUT_SECONDS + + +def test_call_llm_mock_path_bypasses_litellm(monkeypatch) -> None: + """OPENCHRONICLE_LLM_MOCK=1 still short-circuits before importing litellm.""" + monkeypatch.setenv("OPENCHRONICLE_LLM_MOCK", "1") + cfg = _cfg_with_models({}) + + with patch("litellm.completion") as mock: + resp = llm_mod.call_llm( + cfg, "reducer", + messages=[{"role": "user", "content": "hi"}], + ) + + assert mock.call_count == 0 + # The mock response shape: choices[0].message.content is a string. + assert resp.choices[0].message.content