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>

Author: Sebastian Braun

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

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]))

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)