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

[sebastianbraun25]

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) 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.

Update: skip non-retryable errors and oversized documents instead of wasting retries

The fixed 3-attempt retry floor this PR introduces in _llm_call/_llm_call_async
is a good resilience net for transient stream/connection errors, but it blindly
retried every exception — including litellm.ContextWindowExceededError, which
means the prompt itself exceeds the model's context window. Retrying an identical
request in that case is guaranteed to fail again; it only wastes attempts (and,
for a stream, the connection time to discover the failure again).

Observed in practice: converting a large spreadsheet/text attachment to Markdown
produced a ~570k-token prompt against a 200k-token model context window, and the
short-doc compile path retried it 3 times per attempt across both retry layers
(6 doomed attempts total) before giving up and rolling back the whole add for
that file.

Changes:

• _llm_call() / _llm_call_async(): a new _NON_RETRYABLE_LLM_ERRORS tuple
  (currently just litellm.ContextWindowExceededError) is checked in the retry
  loop — a matching exception raises immediately instead of retrying.
• compile_short_doc(): preflights the summary prompt's token count via
  litellm.token_counter() against the model's context window (via a new
  _max_input_tokens() helper, itself a best-effort litellm.get_model_info()
  lookup that returns None — and disables the check — for a model litellm
  doesn't recognize). A prompt that's already known to exceed the window skips
  the LLM call entirely.
• As a safety net for when the preflight can't determine the window (unmapped
  model) or its token estimate came in under the real one, the summary call
  itself is now wrapped to catch _NON_RETRYABLE_LLM_ERRORS.
• In either case, a new _write_unprocessable_stub() helper writes a plain
  placeholder summary page (reusing _write_summary()) noting that the document
  was too large for the model's context window. The raw source stays in the
  knowledge base as a plain reference — no concept/entity extraction is
  attempted for it — mirroring how an unreadable/undecodable image is handled,
  rather than aborting the add for that file or discarding it.
• tests/test_compiler.py: new TestNonRetryableLLMErrors (no-retry
  classification for both call sites), TestMaxInputTokens (known vs.
  unrecognized model), and TestOversizedDocumentSkip (preflight skip and the
  ContextWindowExceededError safety-net fallback, both writing the stub
  summary and producing no concept pages).

Update: fix the stub summary being an undiscoverable orphan

_write_unprocessable_stub() only wrote the stub summary page itself —
_update_index() (which normally appends the ## Documents entry in
index.md) is only called from inside _compile_concepts(), which this path
deliberately skips. The stub page therefore had no index.md entry: it was
technically on disk, but undiscoverable through the wiki's own navigation, and
openkb lint's check_index_sync() would flag it as a sync error
(summaries/<doc>.md not mentioned in index.md).

_write_unprocessable_stub() now also calls _update_index(wiki_dir, doc_name, [], doc_brief=description) — same as every other short-doc path, just with an
empty concept list. tests/test_compiler.py's TestOversizedDocumentSkip
cases now also assert the index.md entry is present.

Update: also handle a context-window overflow in the concepts-plan step

The doc-too-large handling above only guards the summary call. Separately, a
document can still overflow the context window one step later, at
concepts-plan — even when the document itself fit comfortably:
_read_concept_briefs()/_read_entity_briefs() read every concept/entity
page in the wiki with no cap, and that full index is appended on top of the
already-cached document for the concepts-plan call. As a KB accumulates more
concepts/entities over time, this index grows without bound, so the same doc
that summarized fine can push concepts-plan over the limit purely because
the KB has grown — independent of the document's own size.

Observed in practice: a document summarized successfully at 128,692 tokens,
then failed concepts-plan at 220,670 tokens once the existing concept/entity
index was added on top — and (before this fix) the resulting
ContextWindowExceededError propagated uncaught out of _compile_concepts(),
through compile_short_doc(), and was retried wastefully by cli.py's outer
retry (including re-running the summary call) before ultimately failing the
whole add for that file.

Unlike the doc-too-large case, the summary already exists by the time
concepts-plan runs, so discarding it for a generic "too large" stub would be
wasteful. _compile_concepts() now catches _NON_RETRYABLE_LLM_ERRORS around
the concepts-plan call and reuses the same fallback already used for an
unparseable/empty plan: write the real v1 summary (ghost-wikilink-stripped)
and the normal index.md entry, just skipping concept/entity generation for
this document. Applies to compile_long_doc() too, since it shares
_compile_concepts().

Note: this only fixes the crash/retry-waste for this specific step. The
underlying scaling problem — an unbounded concept/entity index injected in
full into every concepts-plan call — remains and will resurface for larger
KBs; a follow-up (e.g. bounding/retrieving only the relevant subset of the
index instead of dumping all of it) is tracked separately.

tests/test_compiler.py: new test_concepts_plan_context_window_exceeded_keeps_real_summary
verifies the real summary (not a stub) is kept, ghost-stripped, indexed, and
that no concept pages are produced, with no wasted retries on the doomed
concepts-plan request.

Issues

Resolves #235.

[sebastianbraun25]

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.