diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878d..756f66ec7 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -397,6 +397,134 @@ class TruncatedResponseError(Exception): treat truncation as a failure (so a partial page is skipped, not written).""" +def _merge_stream_chunks(chunks: list, messages: list[dict]): + """Merge streamed LLM chunks back into a single, non-streaming response. + + Genuine LiteLLM stream chunks only ever carry a ``.delta`` (never a + ``.message``), so a real multi-chunk stream is merged via LiteLLM's own + :func:`litellm.stream_chunk_builder`. A single chunk that already looks + like a complete, non-streaming ``ModelResponse`` (exposing ``.message``) + is used as-is — there's nothing left to merge, and it lets test doubles + fake a one-shot response without simulating LiteLLM's internal delta + format. + """ + choices = getattr(chunks[0], "choices", None) or [] + if len(chunks) == 1 and choices and hasattr(choices[0], "message"): + return chunks[0] + return litellm.stream_chunk_builder(chunks, messages=messages) + + +def _log_stream_start(step_name: str, t0: float, first_chunk_t: float) -> None: + """Debug-log the time-to-first-chunk (TTFT) once a stream's first chunk arrives. + + Marks the start of a "chunk phase" in the log. The counterpart is + :func:`_log_stream_end` (clean finish) or :func:`_log_stream_interrupted` + (mid-stream failure) — together these replace a debug line per chunk + (which used to drown out the rest of the log on a long response, e.g. + hundreds of lines for one LLM call) with exactly one line at the start + and exactly one more at the end/interruption. + """ + logger.debug( + "LLM stream started [%s]: first chunk after %.2fs", + step_name, + first_chunk_t - t0, + ) + + +def _log_stream_end(step_name: str, chunk_count: int, t0: float, last_chunk_t: float) -> None: + """Debug-log a stream's clean completion: total chunk count and elapsed time.""" + logger.debug( + "LLM stream finished [%s]: %d chunk(s), last chunk after %.2fs total", + step_name, + chunk_count, + last_chunk_t - t0, + ) + + +def _log_stream_interrupted( + step_name: str, chunk_count: int, t0: float, last_chunk_t: float +) -> None: + """Debug-log a stream that raised mid-iteration, right before it is re-raised. + + ``chunk_count`` is how many chunks were successfully received before the + failure (0 if the very first chunk never arrived). The exception itself + (with traceback) is attached via ``exc_info=True`` so the failure and the + chunk-phase summary land in a single log record. + """ + now = time.time() + if chunk_count == 0: + logger.debug( + "LLM stream [%s] interrupted unexpectedly before any chunk arrived (%.2fs total)", + step_name, + now - t0, + exc_info=True, + ) + return + logger.debug( + "LLM stream [%s] interrupted unexpectedly after chunk %d " + "(last chunk after %.2fs, failure after %.2fs total)", + step_name, + chunk_count, + last_chunk_t - t0, + now - t0, + exc_info=True, + ) + + +def _consume_stream(stream, step_name: str, t0: float) -> list: + """Collect a sync LiteLLM stream into a list, debug-logging the chunk phase. + + Logs exactly one line when the first chunk arrives (time-to-first-token) + and exactly one more line when the stream ends — either + :func:`_log_stream_end` on a clean finish or :func:`_log_stream_interrupted` + if it raises mid-iteration. A mid-stream exception (e.g. the gateway + idle-timeout firing) propagates after being logged, so callers still see + a complete failure — no partial buffer is ever returned. + """ + if not logger.isEnabledFor(logging.DEBUG): + return list(stream) + + chunks: list = [] + last_t = t0 + try: + for chunk in stream: + now = time.time() + if not chunks: + _log_stream_start(step_name, t0, now) + chunks.append(chunk) + last_t = now + except Exception: + _log_stream_interrupted(step_name, len(chunks), t0, last_t) + raise + _log_stream_end(step_name, len(chunks), t0, last_t) + return chunks + + +async def _consume_stream_async(stream, step_name: str, t0: float) -> list: + """Collect an async LiteLLM stream into a list, debug-logging the chunk phase. + + Mirrors :func:`_consume_stream`, including the start/end-or-interrupted + logging and the no-partial-buffer invariant on failure. + """ + if not logger.isEnabledFor(logging.DEBUG): + return [chunk async for chunk in stream] + + chunks: list = [] + last_t = t0 + try: + async for chunk in stream: + now = time.time() + if not chunks: + _log_stream_start(step_name, t0, now) + chunks.append(chunk) + last_t = now + except Exception: + _log_stream_interrupted(step_name, len(chunks), t0, last_t) + raise + _log_stream_end(step_name, len(chunks), t0, last_t) + return chunks + + def _llm_call( model: str, messages: list[dict], @@ -406,7 +534,15 @@ def _llm_call( bundle=None, **kwargs, ) -> str: - """Single LLM call with animated progress and debug logging.""" + """Single LLM call with animated progress and debug logging. + + Uses ``stream=True``: some corporate LLM gateways enforce an idle + timeout on buffered (non-streaming) requests, which a long-running + completion can hit before the response is ever sent. Streaming keeps + bytes flowing over the connection so that timeout never fires; the + chunks are merged back into a single response via + :func:`_merge_stream_chunks` so callers see the same shape as before. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -417,6 +553,7 @@ def _llm_call( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + kwargs.setdefault("stream_options", {"include_usage": True}) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -425,7 +562,11 @@ def _llm_call( spinner.start() t0 = time.time() - response = litellm.completion(model=model, messages=messages, **kwargs) + stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs) + chunks = _consume_stream(stream, step_name, t0) + if not chunks: + raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") + response = _merge_stream_chunks(chunks, messages) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) @@ -449,7 +590,10 @@ async def _llm_call_async( bundle=None, **kwargs, ) -> str: - """Async LLM call with timing output and debug logging.""" + """Async LLM call with timing output and debug logging. + + See ``_llm_call`` for why ``stream=True`` is used. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -460,13 +604,21 @@ async def _llm_call_async( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + kwargs.setdefault("stream_options", {"include_usage": True}) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) t0 = time.time() - response = await litellm.acompletion(model=model, messages=messages, **kwargs) + stream = await litellm.acompletion(model=model, messages=messages, stream=True, **kwargs) + if hasattr(stream, "__aiter__"): + chunks = await _consume_stream_async(stream, step_name, t0) + else: + chunks = _consume_stream(stream, step_name, t0) + if not chunks: + raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") + response = _merge_stream_chunks(chunks, messages) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 95a57cc4c..34a94ac6c 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1112,40 +1113,96 @@ def test_frontmatter_without_sources_line_gets_one_inserted(self, tmp_path): assert "[[summaries/new-doc]]" in text +def _mock_response(content, finish_reason: str = "stop") -> MagicMock: + """Build a fake, already-complete LLM response (single-chunk stream). + + ``_llm_call``/``_llm_call_async`` now call ``litellm.completion``/ + ``acompletion`` with ``stream=True`` and merge the resulting chunks back + into one response (see ``_merge_stream_chunks``). Exposing ``.message`` + (rather than the ``.delta`` a genuine stream chunk carries) tells + ``_merge_stream_chunks`` this single chunk *is* the final response, so it + is used as-is without needing to fake LiteLLM's internal delta format. + """ + mock_resp = MagicMock() + mock_resp.choices = [MagicMock()] + mock_resp.choices[0].message.content = content + mock_resp.choices[0].finish_reason = finish_reason + mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) + mock_resp.usage.prompt_tokens_details = None + return mock_resp + + def _mock_completion(responses: list[str]): - """Create a mock for litellm.completion that returns responses in order.""" + """Create a mock for litellm.completion returning a single-chunk stream.""" call_count = {"n": 0} def side_effect(*args, **kwargs): idx = min(call_count["n"], len(responses) - 1) call_count["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(responses[idx])] return side_effect def _mock_acompletion(responses: list[str]): - """Create an async mock for litellm.acompletion.""" + """Create an async mock for litellm.acompletion returning a single-chunk stream.""" call_count = {"n": 0} async def side_effect(*args, **kwargs): idx = min(call_count["n"], len(responses) - 1) call_count["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(responses[idx])] return side_effect +class _NoOpSpinner: + """Test double that disables spinner side effects.""" + + def __init__(self, *_args, **_kwargs): + pass + + def start(self) -> None: + pass + + def stop(self, _suffix: str = "") -> None: + pass + + +class _AsyncStream: + """Simple async iterator for exercising streamed LiteLLM responses in tests.""" + + def __init__( + self, + chunks: list[object], + *, + error: Exception | None = None, + raise_after: int | None = None, + ) -> None: + self._chunks = chunks + self._error = error + self._raise_after = raise_after + self._index = 0 + + def __aiter__(self) -> _AsyncStream: + return self + + async def __anext__(self) -> object: + if self._raise_after is not None and self._index == self._raise_after: + raise self._error or RuntimeError("stream exploded") + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +def _stream_then_raise(chunks: list[object], error: Exception): + """Yield all chunks, then raise ``error`` on the next iteration.""" + yield from chunks + raise error + + class TestCompileShortDoc: @pytest.mark.asyncio async def test_full_pipeline(self, tmp_path): @@ -1342,15 +1399,7 @@ def sync_side_effect(*args, **kwargs): sync_call_count["n"] += 1 if idx == 2: # the summary-rewrite call raise RuntimeError("simulated API failure") - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = [ - summary_response, - plan_response, - ][idx] - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response([summary_response, plan_response][idx])] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1507,21 +1556,11 @@ async def test_short_doc_marks_doc_and_summary(self, tmp_path): def sync_side_effect(*args, **kwargs): captured_sync_calls.append(kwargs["messages"]) idx = min(len(captured_sync_calls) - 1, len(sync_responses) - 1) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = sync_responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(sync_responses[idx])] async def async_side_effect(*args, **kwargs): captured_async_calls.append(kwargs["messages"]) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = concept_response - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(concept_response)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1586,15 +1625,9 @@ async def test_long_doc_marks_doc_message(self, tmp_path): def sync_side_effect(*args, **kwargs): captured.append(kwargs["messages"]) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] # First call: overview (plain text); second: plan (JSON). - mock_resp.choices[0].message.content = ( - "Overview text" if len(captured) == 1 else plan_response - ) - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + content = "Overview text" if len(captured) == 1 else plan_response + return [_mock_response(content)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1726,16 +1759,9 @@ async def test_create_and_update_flow(self, tmp_path): async def ordered_acompletion(*args, **kwargs): idx = call_order["n"] call_order["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] # create tasks come first, then update tasks - if idx == 0: - mock_resp.choices[0].message.content = create_page_response - else: - mock_resp.choices[0].message.content = update_page_response - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + content = create_page_response if idx == 0 else update_page_response + return [_mock_response(content)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1823,13 +1849,7 @@ async def test_truncated_update_preserves_existing_page(self, tmp_path): ) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1859,13 +1879,7 @@ async def test_truncated_create_skips_partial_page(self, tmp_path): truncated_page = json.dumps({"brief": "x", "content": "# Ghost\n\nPartial"}) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1928,13 +1942,7 @@ async def test_truncated_entity_update_preserves_existing_page(self, tmp_path): ) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -2701,6 +2709,142 @@ async def test_llm_call_async_injects_extra_headers(self): assert kwargs["extra_headers"] == {"Copilot-Integration-Id": "vscode-chat"} +class TestLLMStreamTimingDebugLogging: + """Chunk-phase debug logging should be visible when verbose logging is + enabled: one line when the first chunk arrives, one more when the stream + ends cleanly or is interrupted — never one line per chunk (see #).""" + + def test_llm_call_logs_stream_start_and_end_at_debug(self, caplog): + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock( + return_value=iter(["chunk-1", "chunk-2", "chunk-3"]) + ) + mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok")) + + out = _llm_call("m", [{"role": "user", "content": "hi"}], "sync-step") + + assert out == "ok" + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [sync-step]" in m] + end_logs = [m for m in messages if "LLM stream finished [sync-step]" in m] + # Exactly one start line and one end line — never a line per chunk. + assert len(start_logs) == 1 + assert len(end_logs) == 1 + assert "3 chunk(s)" in end_logs[0] + assert not any("LLM stream chunk [sync-step]" in m for m in messages) + + @pytest.mark.asyncio + async def test_llm_call_async_logs_stream_start_and_end_at_debug(self, caplog): + from openkb.agent.compiler import _llm_call_async + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.acompletion = AsyncMock( + return_value=_AsyncStream(["chunk-1", "chunk-2", "chunk-3"]) + ) + mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok")) + + out = await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-step") + + assert out == "ok" + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [async-step]" in m] + end_logs = [m for m in messages if "LLM stream finished [async-step]" in m] + assert len(start_logs) == 1 + assert len(end_logs) == 1 + assert "3 chunk(s)" in end_logs[0] + assert not any("LLM stream chunk [async-step]" in m for m in messages) + + def test_llm_call_logs_interruption_before_reraising(self, caplog): + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("stream exploded") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock( + return_value=_stream_then_raise(["chunk-1", "chunk-2"], error) + ) + + with pytest.raises(RuntimeError, match="stream exploded"): + _llm_call("m", [{"role": "user", "content": "hi"}], "sync-fail-step") + + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [sync-fail-step]" in m] + interrupted_logs = [ + m for m in messages if "LLM stream [sync-fail-step] interrupted unexpectedly" in m + ] + # Exactly one start line and one interruption line, no end-of-stream line, + # and no per-chunk lines in between. + assert len(start_logs) == 1 + assert len(interrupted_logs) == 1 + assert "after chunk 2" in interrupted_logs[0] + assert not any("LLM stream finished [sync-fail-step]" in m for m in messages) + assert not any("LLM stream chunk [sync-fail-step]" in m for m in messages) + + @pytest.mark.asyncio + async def test_llm_call_async_logs_interruption_before_reraising(self, caplog): + from openkb.agent.compiler import _llm_call_async + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("stream exploded") + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.acompletion = AsyncMock( + return_value=_AsyncStream( + ["chunk-1", "chunk-2"], + error=error, + raise_after=2, + ) + ) + + with pytest.raises(RuntimeError, match="stream exploded"): + await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-fail-step") + + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [async-fail-step]" in m] + interrupted_logs = [ + m for m in messages if "LLM stream [async-fail-step] interrupted unexpectedly" in m + ] + assert len(start_logs) == 1 + assert len(interrupted_logs) == 1 + assert "after chunk 2" in interrupted_logs[0] + assert not any("LLM stream finished [async-fail-step]" in m for m in messages) + assert not any("LLM stream chunk [async-fail-step]" in m for m in messages) + + def test_llm_call_logs_interruption_before_any_chunk(self, caplog): + """No chunk ever arrives (e.g. a proxy silently buffering despite + stream=True): no start line, and the interruption line says so + instead of an inapplicable chunk number.""" + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("connect timeout") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock(return_value=_stream_then_raise([], error)) + + with pytest.raises(RuntimeError, match="connect timeout"): + _llm_call("m", [{"role": "user", "content": "hi"}], "sync-no-chunk-step") + + messages = [record.getMessage() for record in caplog.records] + interrupted_logs = [ + m for m in messages if "LLM stream [sync-no-chunk-step] interrupted unexpectedly" in m + ] + assert len(interrupted_logs) == 1 + assert "before any chunk arrived" in interrupted_logs[0] + assert not any("LLM stream started [sync-no-chunk-step]" in m for m in messages) + + class TestCacheControlStripping: """cache_control markers must only reach providers that honour them. diff --git a/tests/test_llm_timeout.py b/tests/test_llm_timeout.py index ca7d80e68..db119df3c 100644 --- a/tests/test_llm_timeout.py +++ b/tests/test_llm_timeout.py @@ -17,6 +17,13 @@ def _fake_response(): + """A fake, already-complete LLM response (single-chunk stream). + + See ``openkb.agent.compiler._merge_stream_chunks``: a chunk exposing + ``.message`` (as this one does) is treated as already-complete and used + as-is, so callers of ``litellm.completion``/``acompletion`` with + ``stream=True`` can be mocked to just return a one-item list. + """ choice = MagicMock() choice.message.content = "ok" choice.finish_reason = "stop" @@ -28,7 +35,7 @@ def _fake_response(): def test_llm_call_forwards_configured_timeout(): set_timeout(1200.0) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "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 @@ -37,7 +44,7 @@ def test_llm_call_forwards_configured_timeout(): def test_llm_call_omits_timeout_when_unset(): set_timeout(None) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "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 @@ -47,7 +54,7 @@ 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() + "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 @@ -58,7 +65,7 @@ def test_llm_call_async_forwards_configured_timeout(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + 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 @@ -69,7 +76,7 @@ def test_llm_call_async_omits_timeout_when_unset(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + 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 @@ -80,7 +87,7 @@ def test_llm_call_async_does_not_override_explicit_timeout(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + return_value=[_fake_response()], ) as acompletion: asyncio.run( _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30)