diff --git a/config.yaml.example b/config.yaml.example index 45b4b8c3..7efb5dd1 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -17,3 +17,7 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex # - organization # - dataset # - model + +# Optional: per-request LLM timeout in seconds, forwarded to LiteLLM. +# Defaults to LiteLLM's 600s; raise it for slow local backends (e.g. Ollama). +# timeout: 1200 diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index c2b92619..480258bc 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -29,7 +29,12 @@ import litellm from openkb import frontmatter -from openkb.config import DEFAULT_ENTITY_TYPES, get_extra_headers, resolve_entity_types +from openkb.config import ( + DEFAULT_ENTITY_TYPES, + get_extra_headers, + get_timeout, + resolve_entity_types, +) from openkb.lint import list_existing_wiki_targets, strip_ghost_wikilinks from openkb.schema import INDEX_SEED, get_agents_md @@ -326,6 +331,9 @@ def _llm_call(model: str, messages: list[dict], step_name: str, **kwargs) -> str extra_headers = get_extra_headers() if extra_headers: kwargs.setdefault("extra_headers", extra_headers) + timeout = get_timeout() + if timeout is not None: + kwargs.setdefault("timeout", timeout) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -348,6 +356,9 @@ async def _llm_call_async(model: str, messages: list[dict], step_name: str, **kw extra_headers = get_extra_headers() if extra_headers: kwargs.setdefault("extra_headers", extra_headers) + timeout = get_timeout() + if timeout is not None: + kwargs.setdefault("timeout", timeout) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) diff --git a/openkb/agent/linter.py b/openkb/agent/linter.py index 7e9ae8e7..0519003a 100644 --- a/openkb/agent/linter.py +++ b/openkb/agent/linter.py @@ -7,7 +7,7 @@ from agents.model_settings import ModelSettings from openkb.agent.tools import list_wiki_files, read_wiki_file -from openkb.config import get_extra_headers +from openkb.config import get_extra_headers, get_timeout_extra_args MAX_TURNS = 50 from openkb.schema import get_agents_md @@ -81,7 +81,10 @@ def read_file(path: str) -> str: instructions=instructions, tools=[list_files, read_file], model=f"litellm/{model}", - model_settings=ModelSettings(extra_headers=get_extra_headers() or None), + model_settings=ModelSettings( + extra_headers=get_extra_headers() or None, + extra_args=get_timeout_extra_args(), + ), ) diff --git a/openkb/agent/query.py b/openkb/agent/query.py index a8a97ce5..5a755d76 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -6,7 +6,7 @@ from agents import Agent, Runner, function_tool from agents import ToolOutputImage, ToolOutputText -from openkb.config import get_extra_headers +from openkb.config import get_extra_headers, get_timeout_extra_args from openkb.agent.tools import ( get_wiki_page_content, read_wiki_file, @@ -98,6 +98,7 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: model_settings=ModelSettings( parallel_tool_calls=False, extra_headers=get_extra_headers() or None, + extra_args=get_timeout_extra_args(), ), ) diff --git a/openkb/cli.py b/openkb/cli.py index a6236b7d..a34a0b95 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -43,7 +43,7 @@ def filter(self, record: logging.LogRecord) -> bool: from openkb.config import ( DEFAULT_CONFIG, load_config, save_config, load_global_config, register_kb, - resolve_extra_headers, set_extra_headers, + resolve_extra_headers, set_extra_headers, resolve_timeout, set_timeout, ) from openkb.converter import _registry_path, convert_document from openkb.locks import atomic_write_json, atomic_write_text, kb_ingest_lock, kb_read_lock @@ -108,9 +108,11 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: api_key = os.environ.get("LLM_API_KEY", "") - # Try to resolve the active provider and extra headers from the KB config + # Try to resolve the active provider, extra headers, and request timeout + # from the KB config provider: str | None = None extra_headers: dict[str, str] = {} + timeout: float | None = None if kb_dir is not None: config_path = kb_dir / ".openkb" / "config.yaml" if config_path.exists(): @@ -118,7 +120,9 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: model = config.get("model", "") provider = _extract_provider(str(model)) extra_headers = resolve_extra_headers(config) + timeout = resolve_timeout(config) set_extra_headers(extra_headers) + set_timeout(timeout) if not api_key: # Check if any provider key is already set. OAuth-based providers diff --git a/openkb/config.py b/openkb/config.py index 9d0d6cd4..d5489688 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -2,6 +2,7 @@ import contextlib import logging +import math import re from pathlib import Path from typing import Any, Iterator @@ -136,6 +137,41 @@ def resolve_extra_headers(config: dict) -> dict[str, str]: return headers +def resolve_timeout(config: dict) -> float | None: + """Resolve the optional ``timeout:`` key to a finite positive number of seconds. + + Returns ``None`` (use LiteLLM's default) when absent or invalid; rejects + bools and ``nan``/``inf``, warning when present but unusable. + """ + raw = config.get("timeout") + if raw is None: + return None + if isinstance(raw, bool) or not isinstance(raw, (int, float, str)): + logger.warning( + "config: 'timeout' must be a positive number of seconds, got %s — " + "ignoring it.", + type(raw).__name__, + ) + return None + try: + value = float(raw) + except (TypeError, ValueError): + logger.warning( + "config: 'timeout' must be a positive number of seconds, got %r — " + "ignoring it.", + raw, + ) + return None + if not math.isfinite(value) or value <= 0: + logger.warning( + "config: 'timeout' must be a finite positive number of seconds, got " + "%s — ignoring it.", + value, + ) + return None + return value + + # Process-wide extra headers for LLM requests, resolved from the active KB's # config by the CLI entry points (cli._setup_llm_key). LLM call sites read it # via get_extra_headers() so the value doesn't have to be threaded through @@ -155,6 +191,29 @@ def get_extra_headers() -> dict[str, str]: return dict(_runtime_extra_headers) +# Process-wide LLM request timeout (seconds), set from config by the CLI and +# read at the call sites via get_timeout(). None = use LiteLLM's default. +_runtime_timeout: float | None = None + + +def set_timeout(timeout: float | None) -> None: + """Set the process-wide LLM request timeout in seconds; ``None`` clears it.""" + global _runtime_timeout + _runtime_timeout = timeout + + +def get_timeout() -> float | None: + """Return the process-wide LLM request timeout in seconds, or ``None``.""" + return _runtime_timeout + + +def get_timeout_extra_args() -> dict[str, float] | None: + """Timeout as Agents-SDK ``ModelSettings.extra_args`` (it has no ``timeout`` + field), or ``None``. The LiteLLM provider forwards it to the completion call. + """ + return {"timeout": _runtime_timeout} if _runtime_timeout is not None else None + + def load_config(config_path: Path) -> dict[str, Any]: """Load YAML config from config_path, merged with DEFAULT_CONFIG. diff --git a/openkb/skill/creator.py b/openkb/skill/creator.py index 7a9d68b1..339356f0 100644 --- a/openkb/skill/creator.py +++ b/openkb/skill/creator.py @@ -19,7 +19,7 @@ from agents import Agent, Runner, ToolOutputImage, ToolOutputText, function_tool from agents.model_settings import ModelSettings -from openkb.config import get_extra_headers +from openkb.config import get_extra_headers, get_timeout_extra_args from openkb.skill import skill_dir from openkb.skill.tools import ( get_skill_page_content as _get_page_content_impl, @@ -154,6 +154,7 @@ def done(summary: str) -> str: model_settings=ModelSettings( parallel_tool_calls=True, extra_headers=get_extra_headers() or None, + extra_args=get_timeout_extra_args(), ), ) diff --git a/openkb/skill/evaluator.py b/openkb/skill/evaluator.py index 9150ea54..223966cb 100644 --- a/openkb/skill/evaluator.py +++ b/openkb/skill/evaluator.py @@ -46,7 +46,7 @@ from agents.exceptions import MaxTurnsExceeded from agents.model_settings import ModelSettings -from openkb.config import get_extra_headers +from openkb.config import get_extra_headers, get_timeout_extra_args from openkb.skill import extract_body, extract_frontmatter @@ -243,7 +243,10 @@ async def generate_eval_set( name="eval-set-generator", instructions=instructions, model=f"litellm/{model}", - model_settings=ModelSettings(extra_headers=get_extra_headers() or None), + model_settings=ModelSettings( + extra_headers=get_extra_headers() or None, + extra_args=get_timeout_extra_args(), + ), ) try: result = await Runner.run(agent, "Generate the eval set now.", max_turns=3) @@ -302,7 +305,10 @@ async def grade_one( name="trigger-grader", instructions=instructions, model=f"litellm/{model}", - model_settings=ModelSettings(extra_headers=get_extra_headers() or None), + model_settings=ModelSettings( + extra_headers=get_extra_headers() or None, + extra_args=get_timeout_extra_args(), + ), ) try: result = await Runner.run(agent, f"Question: {question}", max_turns=2) @@ -351,7 +357,10 @@ async def grade_coverage( name="coverage-grader", instructions=instructions, model=f"litellm/{model}", - model_settings=ModelSettings(extra_headers=get_extra_headers() or None), + model_settings=ModelSettings( + extra_headers=get_extra_headers() or None, + extra_args=get_timeout_extra_args(), + ), ) try: result = await Runner.run(agent, f"Question: {question}", max_turns=2) diff --git a/tests/conftest.py b/tests/conftest.py index 3c556551..f4f55065 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,10 +4,11 @@ @pytest.fixture(autouse=True) def _reset_extra_headers(): - """Keep the process-wide LLM extra-headers stash from leaking across tests.""" - from openkb.config import set_extra_headers + """Keep the process-wide LLM extra-headers / timeout stashes from leaking across tests.""" + from openkb.config import set_extra_headers, set_timeout yield set_extra_headers({}) + set_timeout(None) @pytest.fixture diff --git a/tests/test_cli.py b/tests/test_cli.py index a94963ff..3f727138 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -373,12 +373,14 @@ class TestSetupLlmKey: """_setup_llm_key: OAuth-provider warning skip + extra-headers stash.""" @staticmethod - def _make_kb(tmp_path, model, extra_headers=None): + def _make_kb(tmp_path, model, extra_headers=None, timeout=None): openkb_dir = tmp_path / ".openkb" openkb_dir.mkdir() config = {"model": model} if extra_headers is not None: config["extra_headers"] = extra_headers + if timeout is not None: + config["timeout"] = timeout (openkb_dir / "config.yaml").write_text( yaml.safe_dump(config), encoding="utf-8" ) @@ -445,3 +447,20 @@ def test_extra_headers_reset_when_config_has_none(self, tmp_path): kb = self._make_kb(tmp_path, "gpt-5.4-mini") _setup_llm_key(kb) assert get_extra_headers() == {} + + def test_timeout_stashed_from_config(self, tmp_path): + from openkb.cli import _setup_llm_key + from openkb.config import get_timeout + + kb = self._make_kb(tmp_path, "gpt-5.4-mini", timeout=1200) + _setup_llm_key(kb) + assert get_timeout() == 1200.0 + + def test_timeout_reset_when_config_has_none(self, tmp_path): + from openkb.cli import _setup_llm_key + from openkb.config import get_timeout, set_timeout + + set_timeout(999.0) + kb = self._make_kb(tmp_path, "gpt-5.4-mini") + _setup_llm_key(kb) + assert get_timeout() is None diff --git a/tests/test_config.py b/tests/test_config.py index 22b720f7..3fd8edc6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,10 +1,13 @@ from openkb.config import ( DEFAULT_CONFIG, get_extra_headers, + get_timeout, load_config, resolve_extra_headers, + resolve_timeout, save_config, set_extra_headers, + set_timeout, ) @@ -102,3 +105,48 @@ def test_extra_headers_stash_roundtrip_and_isolation(): assert get_extra_headers() == {"A": "1"} set_extra_headers({}) assert get_extra_headers() == {} + + +# --- timeout ----------------------------------------------------------------- + +def test_resolve_timeout_absent_returns_none(): + assert resolve_timeout({}) is None + + +def test_resolve_timeout_int_and_float(): + assert resolve_timeout({"timeout": 1200}) == 1200.0 + assert resolve_timeout({"timeout": 0.5}) == 0.5 + + +def test_resolve_timeout_numeric_string_coerced(): + assert resolve_timeout({"timeout": "1200"}) == 1200.0 + + +def test_resolve_timeout_rejects_non_positive(): + assert resolve_timeout({"timeout": 0}) is None + assert resolve_timeout({"timeout": -10}) is None + + +def test_resolve_timeout_rejects_bool(): + # bool is a subclass of int; True/False are not durations. + assert resolve_timeout({"timeout": True}) is None + + +def test_resolve_timeout_rejects_non_numeric(): + assert resolve_timeout({"timeout": "soon"}) is None + assert resolve_timeout({"timeout": [1200]}) is None + + +def test_resolve_timeout_rejects_nan_and_inf(): + # nan/inf pass a naive `<= 0` check; YAML's .nan/.inf yield real floats. + assert resolve_timeout({"timeout": float("inf")}) is None + assert resolve_timeout({"timeout": float("nan")}) is None + assert resolve_timeout({"timeout": "inf"}) is None + assert resolve_timeout({"timeout": "nan"}) is None + + +def test_timeout_stash_roundtrip_and_reset(): + set_timeout(1200.0) + assert get_timeout() == 1200.0 + set_timeout(None) + assert get_timeout() is None diff --git a/tests/test_llm_timeout.py b/tests/test_llm_timeout.py new file mode 100644 index 00000000..1ad93c8e --- /dev/null +++ b/tests/test_llm_timeout.py @@ -0,0 +1,80 @@ +"""The configurable LLM request timeout is forwarded to LiteLLM. + +`timeout:` in config.yaml is resolved into a process-wide stash (see +test_config.py) and read at the LiteLLM call sites in openkb.agent.compiler. +These tests pin the call-site behavior: a configured timeout is forwarded to +`litellm.(a)completion`, and nothing is forwarded when it is unset (so LiteLLM +keeps applying its own default). +""" +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +from openkb.agent.compiler import _llm_call, _llm_call_async +from openkb.config import set_timeout + + +def _fake_response(): + choice = MagicMock() + choice.message.content = "ok" + choice.finish_reason = "stop" + resp = MagicMock() + resp.choices = [choice] + return resp + + +def test_llm_call_forwards_configured_timeout(): + set_timeout(1200.0) + with patch("openkb.agent.compiler.litellm.completion", + return_value=_fake_response()) as completion: + _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") + assert completion.call_args.kwargs["timeout"] == 1200.0 + + +def test_llm_call_omits_timeout_when_unset(): + set_timeout(None) + with patch("openkb.agent.compiler.litellm.completion", + return_value=_fake_response()) as completion: + _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") + assert "timeout" not in completion.call_args.kwargs + + +def test_llm_call_does_not_override_explicit_timeout(): + # An explicit per-call timeout kwarg wins over the configured default. + set_timeout(1200.0) + with patch("openkb.agent.compiler.litellm.completion", + return_value=_fake_response()) as completion: + _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30) + assert completion.call_args.kwargs["timeout"] == 30 + + +def test_llm_call_async_forwards_configured_timeout(): + set_timeout(900.0) + with patch("openkb.agent.compiler.litellm.acompletion", + new_callable=AsyncMock, return_value=_fake_response()) as acompletion: + asyncio.run( + _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step") + ) + assert acompletion.call_args.kwargs["timeout"] == 900.0 + + +def test_llm_call_async_omits_timeout_when_unset(): + set_timeout(None) + with patch("openkb.agent.compiler.litellm.acompletion", + new_callable=AsyncMock, return_value=_fake_response()) as acompletion: + asyncio.run( + _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step") + ) + assert "timeout" not in acompletion.call_args.kwargs + + +def test_llm_call_async_does_not_override_explicit_timeout(): + set_timeout(900.0) + with patch("openkb.agent.compiler.litellm.acompletion", + new_callable=AsyncMock, return_value=_fake_response()) as acompletion: + asyncio.run( + _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step", + timeout=30) + ) + assert acompletion.call_args.kwargs["timeout"] == 30 diff --git a/tests/test_query.py b/tests/test_query.py index 4b19b17d..a4d29387 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -163,3 +163,22 @@ def test_extra_headers_applied_from_stash(self, tmp_path): def test_no_extra_headers_by_default(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") assert agent.model_settings.extra_headers is None + + +class TestQueryAgentTimeout: + """Config-driven timeout reaches the agents-SDK model settings via extra_args. + + ModelSettings has no ``timeout`` field, so it is forwarded through + ``extra_args`` (which the LiteLLM provider passes on to the completion call). + """ + + def test_timeout_applied_from_stash(self, tmp_path): + from openkb.config import set_timeout + + set_timeout(1200.0) + agent = build_query_agent(str(tmp_path), "gpt-4o-mini") + assert agent.model_settings.extra_args == {"timeout": 1200.0} + + def test_no_timeout_by_default(self, tmp_path): + agent = build_query_agent(str(tmp_path), "gpt-4o-mini") + assert agent.model_settings.extra_args is None