Skip to content

fix(agent): stream LLM completions to avoid gateway idle-timeout #236

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sebastianbraun25 wants to merge 1 commit into
base: main
Choose a base branch
from
+91 −83
Open

fix(agent): stream LLM completions to avoid gateway idle-timeout #236

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 45 additions & 4 deletions openkb/agent/compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand All @@ -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:
Expand All @@ -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)
Expand All @@ -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"))

Expand All @@ -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:
Expand All @@ -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"))

Expand Down
106 changes: 33 additions & 73 deletions tests/test_compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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]))
Expand Down Expand Up @@ -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]))
Expand Down Expand Up @@ -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]))
Expand Down Expand Up @@ -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]))
Expand Down
19 changes: 13 additions & 6 deletions tests/test_llm_timeout.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down