From c7ea8ab546bd412df1d17f77f3312910ff4313e5 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 17:23:42 +0200 Subject: [PATCH] fix(agent): stream LLM completions to avoid gateway idle-timeout Corporate LLM gateways (e.g. AI.proxy on AWS) enforce an idle timeout on buffered (non-streaming) requests, so a long-running compile step can hit a Gateway Timeout even though the provider would have eventually finished. Switch _llm_call() and _llm_call_async() in openkb/agent/compiler.py to litellm.completion()/acompletion() with stream=True: streaming keeps bytes flowing over the connection, so idle-timeout gateways never see a silent connection. Chunks are merged back into the existing response shape via a new _merge_stream_chunks() helper, using LiteLLM's own litellm.stream_chunk_builder() for genuine multi-chunk streams. An exception raised mid-stream propagates as a complete failure (list() never returns a partial buffer), matching prior all-or-nothing behavior. Adapts the compiler test mocks (_mock_completion/_mock_acompletion and a handful of inline mocks) to return a single-chunk fake stream, plus the litellm.completion/acompletion mocks in test_llm_timeout.py. Resolves #235. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- openkb/agent/compiler.py | 49 ++++++++++++++++-- tests/test_compiler.py | 106 ++++++++++++-------------------------- tests/test_llm_timeout.py | 19 ++++--- 3 files changed, 91 insertions(+), 83 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878d..2f9df71e1 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -397,6 +397,23 @@ 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 _llm_call( model: str, messages: list[dict], @@ -406,7 +423,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 +442,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 +451,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 = list(stream) + 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 +479,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 +493,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 = [chunk async for chunk in stream] + else: + chunks = list(stream) + 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..e17ad47d7 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1112,36 +1112,45 @@ 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 @@ -1342,15 +1351,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 +1508,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 +1577,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 +1711,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 +1801,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 +1831,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 +1894,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])) 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)