fix(agent): stream LLM completions to avoid gateway idle-timeout - #236
Open
sebastianbraun25 wants to merge 3 commits into
Open
fix(agent): stream LLM completions to avoid gateway idle-timeout#236sebastianbraun25 wants to merge 3 commits into
sebastianbraun25 wants to merge 3 commits 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>
Add per-chunk debug logging around streamed LiteLLM completion consumption in `openkb.agent.compiler`. This diagnostic instrumentation is enabled via the existing `openkb -v` flag and helps narrow the still-observed ~60s production cutoff to either a no-first-byte case (for example slow time-to-first-token) or a proxy/gateway path that buffers or drops streamed bytes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Per-chunk DEBUG logging (_log_chunk_timing) drowned out the rest of a log on a long response — hundreds of LLM stream chunk [...] #N lines for a single LLM call, e.g. one concepts-plan request logging 205 individual chunk lines. Replaces it with exactly two log lines per LLM call: - _log_stream_start: logged once, when the first chunk arrives (time-to-first-token). - _log_stream_end: logged once, when the stream finishes cleanly (total chunk count + elapsed time for the last chunk). - _log_stream_interrupted: logged once instead of _log_stream_end if the stream raises mid-iteration — reports how many chunks were successfully received and when, right before the exception is re-raised (still a complete failure, no partial buffer). Special-cases zero chunks (failure before any byte arrived) with dedicated wording instead of an inapplicable chunk number. Updated tests/test_compiler.py::TestLLMStreamTimingDebugLogging to match: asserts exactly one start + one end/interrupted line, and the explicit absence of the old per-chunk lines.
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
1 participant
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.
Note
This PR was created in collaboration between a human and AI: implementation, tests, and
PR text were created by an AI assistant under the guidance and review of the human author.
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.