Skip to content

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

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.

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