fix(agent): retry summary-rewrite with summary as source before giving up - #257
Open
sebastianbraun25 wants to merge 10 commits into
Open
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.
_llm_call/_llm_call_async now wrap the stream fetch/consume/merge step in a fixed 3-attempt loop. A dropped connection or other transient error during the now-streamed completion (VectifyAI#236) previously failed the whole concept/entity generation immediately; it now gets 2 automatic retries before giving up. Not a tunable knob, just a resilience floor sitting below the end-of-batch sweep retry added for insert_mode (VectifyAI#241).
_llm_call/_llm_call_async no longer waste their 3-attempt retry floor on litellm.ContextWindowExceededError, since an identical retry is guaranteed to fail again. compile_short_doc additionally preflights the summary prompt's token count (when the model is recognized by litellm) and, on either that check or a ContextWindowExceededError from the summary call itself, skips LLM ingestion for the document entirely: the raw source stays in the KB as a plain reference (like an unreadable image would), with a stub summary noting why, instead of aborting the add or discarding the file.
_write_unprocessable_stub() previously only wrote the stub summary page, never calling _update_index() (that only happened inside _compile_concepts(), which the oversized-doc path deliberately skips). The stub page was therefore an undiscoverable orphan: missing from index.md's ## Documents section, which openkb lint's check_index_sync() flags as an index-sync error. Now calls _update_index(wiki_dir, doc_name, [], doc_brief=description) too, same as every other short-doc path.
…ding the summary The existing concept/entity index (_read_concept_briefs/_read_entity_briefs) grows unbounded with the KB and is appended on top of the already-cached document for the concepts-plan call -- a document whose summary call fit comfortably can still blow the context window once that index gets large enough. Previously this exception propagated uncaught out of _compile_concepts(), through compile_short_doc(), retried the whole doc (including a wasted repeat summary call) via cli.py's outer retry, and ultimately failed the add. _compile_concepts() now catches _NON_RETRYABLE_LLM_ERRORS around the concepts-plan call. Since the summary already succeeded by this point (unlike the doc-too-large case in compile_short_doc's own preflight/summary-call catch), it 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 doc -- instead of discarding a summary that was already successfully generated. Applies to compile_long_doc() too, since it shares _compile_concepts().
…g up known_targets_msg (the whitelist of every existing concept/entity page) grows with the KB, same as concepts-plan's brief lists (VectifyAI#226) - a document whose concepts-plan call still fit can push summary-rewrite over the context window once that whitelist gets large enough. Previously any exception from this call, including litellm.ContextWindowExceededError, fell straight through to the v1-summary fallback, silently forfeiting the improved cross-linking summary-rewrite exists to produce. Mirrors the existing concepts-plan fix: on _NON_RETRYABLE_LLM_ERRORS, retry once with doc_msg dropped (the rewrite prompt only needs summary_msg and known_targets_msg, not the original document) before falling through to the unchanged v1 fallback. Resolves VectifyAI#256 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
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
_compile_concepts()'ssummary-rewritecall (openkb/agent/compiler.py) sends[system_msg, doc_msg, summary_msg, known_targets_msg, ...].known_targets_msg— the whitelistof every existing
[[wikilink]]-able concept/entity page — grows with the size of the knowledgebase (see #226). Combined with a sizeable source document (
doc_msg), this can push the call overthe model's context window even for documents whose
concepts-plancall (a smaller prompt, nowhitelist) still fits fine.
Today any exception from
summary-rewrite— includinglitellm.ContextWindowExceededError— iscaught by a blanket
except Exceptionthat silently falls back to writing the unmodified v1summary. This avoids a crash, but permanently forfeits the improved cross-linking
summary-rewriteexists to produce, for every document from that point on as the KB keepsgrowing.
Solution / Changes
Mirrors the existing
concepts-planretry fix (from #235's branch, which this PR depends on — seebelow):
openkb/agent/compiler.py: on_NON_RETRYABLE_LLM_ERRORS(i.e.litellm.ContextWindowExceededError) from the firstsummary-rewritecall, retry once withdoc_msgdropped —[system_msg, summary_msg, known_targets_msg, user]— since the rewriteprompt only asks the model to reconcile the already-generated summary (
summary_msg) againstthe whitelist, not the original document.
except Exception→ v1-summary fallback still applies, unchanged.tests/test_compiler.py: two new tests mirroring the existingconcepts-plancontext-windowtests — retry succeeds without
doc_msg(rewritten content used), and retry also fails(existing v1 fallback still applies).
Dependencies
Depends on #235— this PR is built on top of that branch (_NON_RETRYABLE_LLM_ERRORSand theconcepts-planretry-with-summary-as-source fix it introduces). The diff below includes fix(agent): stream LLM completions to avoid gateway idle-timeout on long compiles #235'scommits; once fix(agent): stream LLM completions to avoid gateway idle-timeout on long compiles #235 merges first, this PR's diff will shrink to just the
summary-rewriteretrycommit.
Issues
Resolves #256