Repository navigation
fix(agent): stream LLM completions to avoid gateway idle-timeout on long compiles #235
Description
Problem
openkb's wiki compiler (openkb/agent/compiler.py) calls litellm.completion() /
litellm.acompletion() in buffered (non-streaming) mode. On a long-running compile
(large documents, many concepts, verbose models), some corporate LLM gateways sit
in front of the actual provider and enforce an idle timeout on buffered requests
— if no bytes are sent back to the client for N seconds while the gateway waits for
the upstream provider to finish generating the full response, the gateway kills the
connection with a Gateway Timeout (502/504), even though the underlying provider
call would have eventually succeeded.
This has been observed against a corporate AI.proxy gateway deployed on AWS: any
step whose completion takes longer than the gateway's idle window (large concept
pages, long summaries, etc.) fails with a timeout, forcing a retry (or, without a
retry mechanism, aborting the whole compile step) purely because of buffering, not
because the model itself was slow to start responding.
Reproduction
- Command:
openkb add <large-document>(or any compile step that produces a long
LLM response) - Provider: corporate
AI.proxygateway (OpenAI-compatible, deployed at AWS) in
front of the real LLM provider - Observed: the LLM call hangs until the gateway's idle timeout fires, then raises a
Gateway Timeout style error (502/504) even though the model would have completed
the response given enough time — the request is only ever a single buffered
round-trip, so no bytes flow until the entire response is ready. - Expected: as long as the provider is actively generating tokens, the call should
not be killed by an idle-connection timeout.
Context
openkb/agent/compiler.py:_llm_call()(sync) and_llm_call_async()(async)
are the two call sites that invokelitellm.completion()/litellm.acompletion()
for every wiki-compilation LLM call (summaries, concept plans, concept/entity
pages, overviews).- LiteLLM pinned at
1.87.2(pyproject.toml).
Proposed solution
Switch both call sites to litellm.completion(..., stream=True) /
litellm.acompletion(..., stream=True). Streaming keeps the HTTP connection's bytes
flowing as tokens arrive, so an idle-timeout gateway never observes a silent
connection and doesn't kill the request. The streamed chunks are then merged back
into the same response shape the rest of the compiler already expects, via
LiteLLM's own litellm.stream_chunk_builder() — so downstream code (response.choices[0].message.content,
response.usage, response.choices[0].finish_reason) doesn't need to change.
This is a straight swap to streaming-only (no stream=True/False config toggle);
an exception raised mid-stream must be treated as a complete failure — no partial
buffer is used — matching today's all-or-nothing behavior for a failed
litellm.completion() call.
This issue was drafted with the assistance of an AI assistant.
Activity
sebastianbraun25 commentedon Sep 10, 2026
Follow-up note (scope clarification): the "oversized document" handling added on this branch (
_write_unprocessable_stubfor a doc whose content alone already exceeds the model's context window, and the concepts-plan fallback for a doc whose plan-step prompt only exceeds the window once the existing concept/entity index is added on top) is an intentionally simplified interim solution, not the end state.For short/plain-text documents there is currently no PageIndex-style chunked ingestion path (unlike long PDFs via
index_long_document/PageIndex). Until such a path exists for arbitrary large short documents too, it is accepted that a document too large for a single LLM call gets ingested in a degraded form:This is a deliberate, accepted trade-off for now (simplicity + no dropped documents), not a design goal — filing this as a known limitation here rather than a new issue, since it's directly adjacent to the context-window handling already introduced on this branch.