Skip to content

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

Open
sebastianbraun25 wants to merge 1 commit into
VectifyAI:mainfrom
sebastianbraun25:fix/issue-235-stream-llm-completions
Open

fix(agent): stream LLM completions to avoid gateway idle-timeout#236
sebastianbraun25 wants to merge 1 commit into
VectifyAI:mainfrom
sebastianbraun25:fix/issue-235-stream-llm-completions

Conversation

@sebastianbraun25

Copy link
Copy Markdown

Problem

openkb's wiki compiler (openkb/agent/compiler.py) calls litellm.completion() /
litellm.acompletion() in buffered (non-streaming) mode. On a long-running compile
step (large documents, many concepts, verbose models), some corporate LLM gateways
sitting in front of the actual provider 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, even though the underlying
provider call would have eventually succeeded.

This was observed against a corporate AI.proxy gateway (OpenAI-compatible,
deployed at AWS): any step whose completion took longer than the gateway's idle
window failed with a timeout purely because of buffering, not because the model
was slow to start responding.

Root Cause

_llm_call() and _llm_call_async() in openkb/agent/compiler.py call
litellm.completion() / litellm.acompletion() without stream=True. A buffered
request produces zero bytes on the wire until the entire response is ready, so an
idle-timeout gateway watching for connection activity can't distinguish "still
generating" from "connection died" and kills it.

Solution / Changes

  • openkb/agent/compiler.py:
    • _llm_call() and _llm_call_async() now call litellm.completion() /
      litellm.acompletion() with stream=True, so bytes keep flowing over the
      connection as tokens arrive.
    • New _merge_stream_chunks() helper merges the streamed chunks back into the
      same non-streaming response shape the rest of the compiler already expects
      (response.choices[0].message.content, response.usage,
      response.choices[0].finish_reason), using LiteLLM's own
      litellm.stream_chunk_builder() for genuine multi-chunk streams.
    • An exception raised mid-stream (e.g. a dropped connection) propagates as a
      complete failure — list(stream) never returns a partial buffer, matching
      the prior all-or-nothing behavior of a failed litellm.completion() call.
    • This is a straight swap to streaming-only; there is no stream=True/False
      config toggle.
  • tests/test_compiler.py: _mock_completion()/_mock_acompletion() (and the
    handful of inline mocks that built their own fake response) now return a
    single-chunk fake stream ([mock_resp]) instead of a bare response object,
    matching the new stream=True call signature.
  • tests/test_llm_timeout.py: same adjustment for its litellm.completion/
    acompletion mocks.

Backward compatible: no config/CLI surface changes, callers of _llm_call/
_llm_call_async see the same return type and response shape as before.

Issues

Resolves #235.

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 VectifyAI#235.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant