Skip to content

fix(agent): stream LLM completions to avoid gateway idle-timeout on long compiles #235

Open

Description

@sebastianbraun25

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.proxy gateway (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 invoke litellm.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

  1. sebastianbraun25 commented on Sep 10, 2026

    @sebastianbraun25
    Author

    Follow-up note (scope clarification): the "oversized document" handling added on this branch (_write_unprocessable_stub for 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:

    • kept as a plain stub (raw file + a "not processed by the LLM" placeholder summary, no concepts/entities) when even the summary call can't fit, or
    • kept with a real summary but with concepts/entities derived from that summary instead of the original file, when only the concepts-plan step (summary + existing concept/entity briefs) overflows.

    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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions