fix(agent): stream LLM completions to avoid gateway idle-timeout - #236
Open
sebastianbraun25 wants to merge 1 commit into
Open
fix(agent): stream LLM completions to avoid gateway idle-timeout#236sebastianbraun25 wants to merge 1 commit into
sebastianbraun25 wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
openkb's wiki compiler (openkb/agent/compiler.py) callslitellm.completion()/litellm.acompletion()in buffered (non-streaming) mode. On a long-running compilestep (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.proxygateway (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()inopenkb/agent/compiler.pycalllitellm.completion()/litellm.acompletion()withoutstream=True. A bufferedrequest 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 calllitellm.completion()/litellm.acompletion()withstream=True, so bytes keep flowing over theconnection as tokens arrive.
_merge_stream_chunks()helper merges the streamed chunks back into thesame 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 ownlitellm.stream_chunk_builder()for genuine multi-chunk streams.complete failure —
list(stream)never returns a partial buffer, matchingthe prior all-or-nothing behavior of a failed
litellm.completion()call.stream=True/Falseconfig toggle.
tests/test_compiler.py:_mock_completion()/_mock_acompletion()(and thehandful 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=Truecall signature.tests/test_llm_timeout.py: same adjustment for itslitellm.completion/acompletionmocks.Backward compatible: no config/CLI surface changes, callers of
_llm_call/_llm_call_asyncsee the same return type and response shape as before.Issues
Resolves #235.