diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878d..58aacad15 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -397,6 +397,161 @@ class TruncatedResponseError(Exception): treat truncation as a failure (so a partial page is skipped, not written).""" +# Exceptions that are guaranteed to fail again on an identical retry (e.g. a +# prompt that already exceeds the model's context window) — retrying wastes +# whole attempts (and, for a stream, the connection time to discover the +# failure again) for zero chance of success. +_NON_RETRYABLE_LLM_ERRORS: tuple[type[Exception], ...] = (litellm.ContextWindowExceededError,) + + +def _max_input_tokens(model: str) -> int | None: + """Best-effort context-window lookup for ``model``; ``None`` if unknown. + + Used only to skip a call that's already known to be doomed before it's + even sent — never to second-guess a model litellm/the provider doesn't + also recognize, so an unmapped model just disables the preflight check. + """ + try: + max_input_tokens = litellm.get_model_info(model).get("max_input_tokens") + except Exception: + return None + return int(max_input_tokens) if isinstance(max_input_tokens, int | float) else None + + +# Reserved headroom (completion + rough token-counting slack) subtracted from +# a model's context window before comparing it to the prompt's token count — +# a prompt that just barely fits leaves no room for the model to respond. +_CONTEXT_WINDOW_HEADROOM_TOKENS = 4096 + + +def _merge_stream_chunks(chunks: list, messages: list[dict]): + """Merge streamed LLM chunks back into a single, non-streaming response. + + Genuine LiteLLM stream chunks only ever carry a ``.delta`` (never a + ``.message``), so a real multi-chunk stream is merged via LiteLLM's own + :func:`litellm.stream_chunk_builder`. A single chunk that already looks + like a complete, non-streaming ``ModelResponse`` (exposing ``.message``) + is used as-is — there's nothing left to merge, and it lets test doubles + fake a one-shot response without simulating LiteLLM's internal delta + format. + """ + choices = getattr(chunks[0], "choices", None) or [] + if len(chunks) == 1 and choices and hasattr(choices[0], "message"): + return chunks[0] + return litellm.stream_chunk_builder(chunks, messages=messages) + + +def _log_stream_start(step_name: str, t0: float, first_chunk_t: float) -> None: + """Debug-log the time-to-first-chunk (TTFT) once a stream's first chunk arrives. + + Marks the start of a "chunk phase" in the log. The counterpart is + :func:`_log_stream_end` (clean finish) or :func:`_log_stream_interrupted` + (mid-stream failure) — together these replace a debug line per chunk + (which used to drown out the rest of the log on a long response, e.g. + hundreds of lines for one LLM call) with exactly one line at the start + and exactly one more at the end/interruption. + """ + logger.debug( + "LLM stream started [%s]: first chunk after %.2fs", + step_name, + first_chunk_t - t0, + ) + + +def _log_stream_end(step_name: str, chunk_count: int, t0: float, last_chunk_t: float) -> None: + """Debug-log a stream's clean completion: total chunk count and elapsed time.""" + logger.debug( + "LLM stream finished [%s]: %d chunk(s), last chunk after %.2fs total", + step_name, + chunk_count, + last_chunk_t - t0, + ) + + +def _log_stream_interrupted( + step_name: str, chunk_count: int, t0: float, last_chunk_t: float +) -> None: + """Debug-log a stream that raised mid-iteration, right before it is re-raised. + + ``chunk_count`` is how many chunks were successfully received before the + failure (0 if the very first chunk never arrived). The exception itself + (with traceback) is attached via ``exc_info=True`` so the failure and the + chunk-phase summary land in a single log record. + """ + now = time.time() + if chunk_count == 0: + logger.debug( + "LLM stream [%s] interrupted unexpectedly before any chunk arrived (%.2fs total)", + step_name, + now - t0, + exc_info=True, + ) + return + logger.debug( + "LLM stream [%s] interrupted unexpectedly after chunk %d " + "(last chunk after %.2fs, failure after %.2fs total)", + step_name, + chunk_count, + last_chunk_t - t0, + now - t0, + exc_info=True, + ) + + +def _consume_stream(stream, step_name: str, t0: float) -> list: + """Collect a sync LiteLLM stream into a list, debug-logging the chunk phase. + + Logs exactly one line when the first chunk arrives (time-to-first-token) + and exactly one more line when the stream ends — either + :func:`_log_stream_end` on a clean finish or :func:`_log_stream_interrupted` + if it raises mid-iteration. A mid-stream exception (e.g. the gateway + idle-timeout firing) propagates after being logged, so callers still see + a complete failure — no partial buffer is ever returned. + """ + if not logger.isEnabledFor(logging.DEBUG): + return list(stream) + + chunks: list = [] + last_t = t0 + try: + for chunk in stream: + now = time.time() + if not chunks: + _log_stream_start(step_name, t0, now) + chunks.append(chunk) + last_t = now + except Exception: + _log_stream_interrupted(step_name, len(chunks), t0, last_t) + raise + _log_stream_end(step_name, len(chunks), t0, last_t) + return chunks + + +async def _consume_stream_async(stream, step_name: str, t0: float) -> list: + """Collect an async LiteLLM stream into a list, debug-logging the chunk phase. + + Mirrors :func:`_consume_stream`, including the start/end-or-interrupted + logging and the no-partial-buffer invariant on failure. + """ + if not logger.isEnabledFor(logging.DEBUG): + return [chunk async for chunk in stream] + + chunks: list = [] + last_t = t0 + try: + async for chunk in stream: + now = time.time() + if not chunks: + _log_stream_start(step_name, t0, now) + chunks.append(chunk) + last_t = now + except Exception: + _log_stream_interrupted(step_name, len(chunks), t0, last_t) + raise + _log_stream_end(step_name, len(chunks), t0, last_t) + return chunks + + def _llm_call( model: str, messages: list[dict], @@ -406,7 +561,15 @@ def _llm_call( bundle=None, **kwargs, ) -> str: - """Single LLM call with animated progress and debug logging.""" + """Single LLM call with animated progress and debug logging. + + Uses ``stream=True``: some corporate LLM gateways enforce an idle + timeout on buffered (non-streaming) requests, which a long-running + completion can hit before the response is ever sent. Streaming keeps + bytes flowing over the connection so that timeout never fires; the + chunks are merged back into a single response via + :func:`_merge_stream_chunks` so callers see the same shape as before. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -417,6 +580,7 @@ def _llm_call( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + kwargs.setdefault("stream_options", {"include_usage": True}) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -425,7 +589,29 @@ def _llm_call( spinner.start() t0 = time.time() - response = litellm.completion(model=model, messages=messages, **kwargs) + # Fixed 2 extra attempts for transient stream/LLM errors — not a tunable + # knob, just a resilience floor. The concept/entity sweep in + # _compile_concepts is the next retry tier above this one. + attempts = 3 + for attempt in range(attempts): + try: + stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs) + chunks = _consume_stream(stream, step_name, t0) + if not chunks: + raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") + response = _merge_stream_chunks(chunks, messages) + break + except Exception as exc: + if attempt == attempts - 1 or isinstance(exc, _NON_RETRYABLE_LLM_ERRORS): + spinner.stop("failed") + raise + logger.warning( + "LLM [%s] attempt %d/%d failed: %s; retrying...", + step_name, + attempt + 1, + attempts, + exc, + ) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) @@ -449,7 +635,10 @@ async def _llm_call_async( bundle=None, **kwargs, ) -> str: - """Async LLM call with timing output and debug logging.""" + """Async LLM call with timing output and debug logging. + + See ``_llm_call`` for why ``stream=True`` is used. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -460,13 +649,40 @@ async def _llm_call_async( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + kwargs.setdefault("stream_options", {"include_usage": True}) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) t0 = time.time() - response = await litellm.acompletion(model=model, messages=messages, **kwargs) + # Fixed 2 extra attempts for transient stream/LLM errors — not a tunable + # knob, just a resilience floor. The concept/entity sweep in + # _compile_concepts is the next retry tier above this one. + attempts = 3 + for attempt in range(attempts): + try: + stream = await litellm.acompletion( + model=model, messages=messages, stream=True, **kwargs + ) + if hasattr(stream, "__aiter__"): + chunks = await _consume_stream_async(stream, step_name, t0) + else: + chunks = _consume_stream(stream, step_name, t0) + if not chunks: + raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") + response = _merge_stream_chunks(chunks, messages) + break + except Exception as exc: + if attempt == attempts - 1 or isinstance(exc, _NON_RETRYABLE_LLM_ERRORS): + raise + logger.warning( + "LLM [%s] attempt %d/%d failed: %s; retrying...", + step_name, + attempt + 1, + attempts, + exc, + ) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) @@ -957,6 +1173,30 @@ def _write_summary( atomic_write_text(summaries_dir / f"{doc_name}.md", fm_block + summary) +def _write_unprocessable_stub(wiki_dir: Path, doc_name: str, reason: str) -> None: + """Write a placeholder summary for a doc whose content can't be fed to the LLM. + + Mirrors how an unreadable/undecodable image is handled: the source stays + in the knowledge base as a plain reference instead of aborting the whole + ``add`` or discarding it, but it's skipped for LLM ingestion entirely (no + summary/concept/entity generation), since a request this size is either + already known to exceed the model's context window or has just failed + with ``litellm.ContextWindowExceededError``. Also adds the usual + ``## Documents`` index.md entry (via ``_update_index``, with no concepts) + — without it the summary page would be an undiscoverable orphan, which + ``openkb lint`` flags as an index-sync error. + """ + description = "Not processed by the LLM \u2014 content too large for the context window." + body = ( + "This document was not processed by the LLM: its content is too " + f"large for the model's context window ({reason}). The raw source " + "is still kept in the knowledge base for reference, but no summary " + "or concept/entity extraction was generated for it." + ) + _write_summary(wiki_dir, doc_name, body, description=description) + _update_index(wiki_dir, doc_name, [], doc_brief=description) + + _SAFE_NAME_RE = re.compile(r"[^\w\-]") @@ -1631,25 +1871,6 @@ async def _compile_concepts( # (system + doc + summary) for the plan call and every concept call. summary_msg = {"role": "assistant", "content": _cached_text(summary)} - plan_raw = _llm_call( - model, - [ - system_msg, - doc_msg, - summary_msg, - { - "role": "user", - "content": _CONCEPTS_PLAN_USER.format( - concept_briefs=concept_briefs, - entity_briefs=entity_briefs, - ).replace("__ENTITY_TYPES__", types_str), - }, - ], - "concepts-plan", - response_format=_JSON_RESPONSE_FORMAT, - bundle=bundle, - ) - def _write_v1_summary_stripped() -> None: """Fallback writer for the v1 summary on early-return paths. @@ -1671,6 +1892,60 @@ def _write_v1_summary_stripped() -> None: ) _write_summary(wiki_dir, doc_name, cleaned, description=doc_brief) + concepts_plan_user_msg = { + "role": "user", + "content": _CONCEPTS_PLAN_USER.format( + concept_briefs=concept_briefs, + entity_briefs=entity_briefs, + ).replace("__ENTITY_TYPES__", types_str), + } + try: + plan_raw = _llm_call( + model, + [system_msg, doc_msg, summary_msg, concepts_plan_user_msg], + "concepts-plan", + response_format=_JSON_RESPONSE_FORMAT, + bundle=bundle, + ) + except _NON_RETRYABLE_LLM_ERRORS as exc: + # The existing concept/entity index grows with the KB (see + # _read_concept_briefs/_read_entity_briefs) and is added on top of the + # already-cached document — a doc that was fine for the summary call + # can still blow the window here once the index gets large enough. + # Retry once without the full document: the plan prompt already asks + # the model to work "based on the summary above" (see + # _CONCEPTS_PLAN_USER), so dropping doc_msg usually shrinks the + # prompt back under the window without losing the plan's intent. + logger.warning( + "concepts plan exceeded context window for %s: %s; retrying with summary as source", + doc_name, + exc, + ) + sys.stdout.write( + f" [WARN] concepts plan exceeded context window for {doc_name} — " + "retrying with the summary as source instead of the full document.\n" + ) + sys.stdout.flush() + try: + plan_raw = _llm_call( + model, + [system_msg, summary_msg, concepts_plan_user_msg], + "concepts-plan", + response_format=_JSON_RESPONSE_FORMAT, + bundle=bundle, + ) + except _NON_RETRYABLE_LLM_ERRORS as retry_exc: + # Still too large even with just the summary (e.g. the existing + # concept/entity index alone is huge) — the summary itself + # already exists (unlike the doc-too-large case above), so it's + # kept — same fallback as an unparseable/empty plan, just + # skipping concept/entity generation for this doc. + logger.warning("Skipping concept/entity extraction for %s: %s", doc_name, retry_exc) + if rewrite_summary: + _write_v1_summary_stripped() + _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) + return + try: parsed = _parse_json(plan_raw) except (json.JSONDecodeError, ValueError) as exc: @@ -2245,18 +2520,49 @@ async def compile_short_doc( ), } + # Preflight: skip the LLM entirely for a doc that's already known to + # exceed the model's context window (only when the model is recognized — + # see _max_input_tokens) instead of sending a request that's certain to + # fail. The doc stays in the KB as a plain reference, like an unreadable + # image would. + max_input_tokens = _max_input_tokens(model) + if max_input_tokens is not None: + prompt_tokens = litellm.token_counter(model=model, messages=[system_msg, doc_msg]) + if prompt_tokens > max_input_tokens - _CONTEXT_WINDOW_HEADROOM_TOKENS: + logger.warning( + "Skipping LLM ingestion for %s: %d prompt tokens > %s's %d-token context window", + doc_name, + prompt_tokens, + model, + max_input_tokens, + ) + _write_unprocessable_stub( + wiki_dir, + doc_name, + f"{prompt_tokens} tokens > {model}'s {max_input_tokens}-token context window", + ) + return + # --- Step 1: Generate summary (v1, held in memory) --- # The summary is NOT written to disk yet — it's used as cache context # for the plan + concept-generation calls, then rewritten into a final # v2 (with a whitelist of known wikilink targets) inside # _compile_concepts before being written to disk. - summary_raw = _llm_call( - model, - [system_msg, doc_msg], - "summary", - response_format=_JSON_RESPONSE_FORMAT, - bundle=bundle, - ) + try: + summary_raw = _llm_call( + model, + [system_msg, doc_msg], + "summary", + response_format=_JSON_RESPONSE_FORMAT, + bundle=bundle, + ) + except _NON_RETRYABLE_LLM_ERRORS as exc: + # The preflight check above is best-effort (unmapped model, or + # litellm's token_counter estimate came in under the real one) — this + # is the safety net for when it still slips through. + logger.warning("Skipping LLM ingestion for %s: %s", doc_name, exc) + _write_unprocessable_stub(wiki_dir, doc_name, str(exc)) + return try: summary_parsed = _parse_json(summary_raw) doc_brief = summary_parsed.get("description", "") diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 95a57cc4c..02851bf7c 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -3,8 +3,10 @@ from __future__ import annotations import json +import logging from unittest.mock import AsyncMock, MagicMock, patch +import litellm import pytest from openkb.agent.compiler import ( @@ -1112,40 +1114,96 @@ def test_frontmatter_without_sources_line_gets_one_inserted(self, tmp_path): assert "[[summaries/new-doc]]" in text +def _mock_response(content, finish_reason: str = "stop") -> MagicMock: + """Build a fake, already-complete LLM response (single-chunk stream). + + ``_llm_call``/``_llm_call_async`` now call ``litellm.completion``/ + ``acompletion`` with ``stream=True`` and merge the resulting chunks back + into one response (see ``_merge_stream_chunks``). Exposing ``.message`` + (rather than the ``.delta`` a genuine stream chunk carries) tells + ``_merge_stream_chunks`` this single chunk *is* the final response, so it + is used as-is without needing to fake LiteLLM's internal delta format. + """ + mock_resp = MagicMock() + mock_resp.choices = [MagicMock()] + mock_resp.choices[0].message.content = content + mock_resp.choices[0].finish_reason = finish_reason + mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) + mock_resp.usage.prompt_tokens_details = None + return mock_resp + + def _mock_completion(responses: list[str]): - """Create a mock for litellm.completion that returns responses in order.""" + """Create a mock for litellm.completion returning a single-chunk stream.""" call_count = {"n": 0} def side_effect(*args, **kwargs): idx = min(call_count["n"], len(responses) - 1) call_count["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(responses[idx])] return side_effect def _mock_acompletion(responses: list[str]): - """Create an async mock for litellm.acompletion.""" + """Create an async mock for litellm.acompletion returning a single-chunk stream.""" call_count = {"n": 0} async def side_effect(*args, **kwargs): idx = min(call_count["n"], len(responses) - 1) call_count["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(responses[idx])] return side_effect +class _NoOpSpinner: + """Test double that disables spinner side effects.""" + + def __init__(self, *_args, **_kwargs): + pass + + def start(self) -> None: + pass + + def stop(self, _suffix: str = "") -> None: + pass + + +class _AsyncStream: + """Simple async iterator for exercising streamed LiteLLM responses in tests.""" + + def __init__( + self, + chunks: list[object], + *, + error: Exception | None = None, + raise_after: int | None = None, + ) -> None: + self._chunks = chunks + self._error = error + self._raise_after = raise_after + self._index = 0 + + def __aiter__(self) -> _AsyncStream: + return self + + async def __anext__(self) -> object: + if self._raise_after is not None and self._index == self._raise_after: + raise self._error or RuntimeError("stream exploded") + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +def _stream_then_raise(chunks: list[object], error: Exception): + """Yield all chunks, then raise ``error`` on the next iteration.""" + yield from chunks + raise error + + class TestCompileShortDoc: @pytest.mark.asyncio async def test_full_pipeline(self, tmp_path): @@ -1342,15 +1400,7 @@ def sync_side_effect(*args, **kwargs): sync_call_count["n"] += 1 if idx == 2: # the summary-rewrite call raise RuntimeError("simulated API failure") - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = [ - summary_response, - plan_response, - ][idx] - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response([summary_response, plan_response][idx])] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1453,6 +1503,231 @@ async def test_scalar_plan_handled_gracefully(self, tmp_path): # No concept pages produced from the unusable plan. assert not list((wiki / "concepts").glob("*.md")) + @pytest.mark.asyncio + async def test_concepts_plan_context_window_exceeded_keeps_real_summary(self, tmp_path): + """The summary call already succeeded with real content by the time + concepts-plan runs (unlike TestOversizedDocumentSkip, where the doc + itself is too large) — this must NOT be replaced with a generic + "too large" stub. The concepts-plan call is retried once with the + summary as source (doc_msg dropped) before giving up; if that retry + ALSO hits the context window, same fallback as an unparseable/empty + plan: keep the real v1 summary (ghost-stripped), index it, skip + concept/entity generation for this doc.""" + wiki, source_path = self._setup_kb(tmp_path) + + v1_summary_content = "# Summary\n\nDiscusses [[concepts/nonexistent]] here." + summary_response = json.dumps( + {"description": "A real summary", "content": v1_summary_content} + ) + error = litellm.ContextWindowExceededError( + message="prompt is too long: 220670 tokens > 200000 maximum", + model="claude-sonnet-4-5", + llm_provider="anthropic", + ) + call_count = {"n": 0} + + def side_effect(*args, **kwargs): + idx = call_count["n"] + call_count["n"] += 1 + if idx == 0: + return [_mock_response(summary_response)] + raise error + + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=side_effect) + # Must not raise out of compile_short_doc. + await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5") + # summary + concepts-plan (full doc) + concepts-plan retry (summary + # only) — both plan attempts are doomed here, no further retries. + assert mock_litellm.completion.call_count == 3 + + # The retry attempt must not still include the full document message + # (_SUMMARY_USER's "Full text:" marker only ever appears in doc_msg). + retry_messages = mock_litellm.completion.call_args_list[2].kwargs["messages"] + retry_text = json.dumps(retry_messages) + assert "Full text:" not in retry_text + + summary_path = wiki / "summaries" / "doc.md" + assert summary_path.exists() + text = summary_path.read_text() + assert "Discusses" in text # real summary content kept, not a stub + assert "too large" not in text.lower() + assert "[[concepts/nonexistent]]" not in text # ghost link stripped + assert "nonexistent" in text # display text preserved + + index_text = (wiki / "index.md").read_text() + assert "[[summaries/doc]]" in index_text + assert not list((wiki / "concepts").glob("*.md")) + + @pytest.mark.asyncio + async def test_concepts_plan_context_window_retry_succeeds_with_summary_only(self, tmp_path): + """When the retry (summary-only) succeeds, the plan is processed + normally — concepts are generated from the summary-derived plan + instead of the doc being treated as incomplete.""" + wiki, source_path = self._setup_kb(tmp_path) + + v1_summary_content = "# Summary\n\nDiscusses transformers." + summary_response = json.dumps( + {"description": "A real summary", "content": v1_summary_content} + ) + error = litellm.ContextWindowExceededError( + message="prompt is too long: 220670 tokens > 200000 maximum", + model="claude-sonnet-4-5", + llm_provider="anthropic", + ) + plan_response = json.dumps( + { + "create": [{"name": "transformer", "title": "Transformer"}], + "update": [], + "related": [], + } + ) + concept_response = json.dumps({"description": "C", "content": "# T\n\nBody."}) + call_count = {"n": 0} + + def side_effect(*args, **kwargs): + idx = call_count["n"] + call_count["n"] += 1 + if idx == 0: + return [_mock_response(summary_response)] + if idx == 1: + raise error + return [_mock_response(plan_response)] + + with ( + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock(side_effect=side_effect) + mock_litellm.acompletion = AsyncMock(side_effect=_mock_acompletion([concept_response])) + await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5") + # summary + concepts-plan (full doc, fails) + concepts-plan retry + # (summary only, succeeds) + summary-rewrite (rewrite_summary=True). + assert mock_litellm.completion.call_count == 4 + + # The successful retry's messages must not include the full document + # message (_SUMMARY_USER's "Full text:" marker only ever appears in doc_msg). + retry_messages = mock_litellm.completion.call_args_list[2].kwargs["messages"] + retry_text = json.dumps(retry_messages) + assert "Full text:" not in retry_text + + concept_path = wiki / "concepts" / "transformer.md" + assert concept_path.exists() + + index_text = (wiki / "index.md").read_text() + assert "[[concepts/transformer]]" in index_text + + +class TestOversizedDocumentSkip: + """A document whose content can't fit an LLM call is treated like an + unreadable image: kept in the KB as a plain reference, but skipped for + LLM ingestion entirely rather than aborting the whole ``add`` or wasting + retries on a request that's guaranteed to fail again (#).""" + + @pytest.mark.asyncio + async def test_skips_llm_call_when_preflight_detects_oversized_prompt(self, tmp_path): + wiki, source_path = TestCompileShortDocFallbacks._setup_kb(tmp_path) + + with ( + patch("openkb.agent.compiler._max_input_tokens", return_value=1000), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.token_counter = MagicMock(return_value=999_999) + await compile_short_doc("doc", source_path, tmp_path, "gpt-4o-mini") + mock_litellm.completion.assert_not_called() + + summary_path = wiki / "summaries" / "doc.md" + assert summary_path.exists() + text = summary_path.read_text() + assert "too large" in text.lower() + assert not list((wiki / "concepts").glob("*.md")) + + # The stub must still get the usual index.md "## Documents" entry — + # otherwise it's an undiscoverable orphan (openkb lint's index-sync check). + index_text = (wiki / "index.md").read_text() + assert "[[summaries/doc]]" in index_text + + @pytest.mark.asyncio + async def test_writes_stub_when_summary_call_raises_context_window_exceeded(self, tmp_path): + wiki, source_path = TestCompileShortDocFallbacks._setup_kb(tmp_path) + + error = litellm.ContextWindowExceededError( + message="prompt is too long: 569887 tokens > 200000 maximum", + model="claude-sonnet-4-5", + llm_provider="anthropic", + ) + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock(side_effect=error) + # Must not raise out of compile_short_doc. + await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5") + # Only the summary call was attempted — no retries wasted on a + # deterministically doomed request, no concept-plan call either. + assert mock_litellm.completion.call_count == 1 + + summary_path = wiki / "summaries" / "doc.md" + assert summary_path.exists() + text = summary_path.read_text() + assert "too large" in text.lower() + assert not list((wiki / "concepts").glob("*.md")) + + index_text = (wiki / "index.md").read_text() + assert "[[summaries/doc]]" in index_text + + +class TestMaxInputTokens: + """``_max_input_tokens`` must only ever be a best-effort hint: an + unrecognized model disables the preflight check instead of guessing.""" + + def test_known_model_returns_context_window(self): + from openkb.agent.compiler import _max_input_tokens + + assert _max_input_tokens("claude-sonnet-4-5") == 200_000 + + def test_unknown_model_returns_none(self): + from openkb.agent.compiler import _max_input_tokens + + assert _max_input_tokens("totally-unknown-model-xyz") is None + + +class TestNonRetryableLLMErrors: + """litellm.ContextWindowExceededError means the exact same request will + fail again — retrying it just burns the fixed 3-attempt resilience floor + on a request that can never succeed.""" + + def test_llm_call_does_not_retry_context_window_exceeded(self): + from openkb.agent.compiler import _llm_call + + error = litellm.ContextWindowExceededError( + message="prompt is too long: 569887 tokens > 200000 maximum", + model="m", + llm_provider="anthropic", + ) + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock(side_effect=error) + with pytest.raises(litellm.ContextWindowExceededError): + _llm_call("m", [{"role": "user", "content": "hi"}], "step") + assert mock_litellm.completion.call_count == 1 + + @pytest.mark.asyncio + async def test_llm_call_async_does_not_retry_context_window_exceeded(self): + from openkb.agent.compiler import _llm_call_async + + error = litellm.ContextWindowExceededError( + message="prompt is too long: 569887 tokens > 200000 maximum", + model="m", + llm_provider="anthropic", + ) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.acompletion = AsyncMock(side_effect=error) + with pytest.raises(litellm.ContextWindowExceededError): + await _llm_call_async("m", [{"role": "user", "content": "hi"}], "step") + assert mock_litellm.acompletion.call_count == 1 + class TestCacheControl: """Verify cache_control breakpoints are emitted on the right messages @@ -1507,21 +1782,11 @@ async def test_short_doc_marks_doc_and_summary(self, tmp_path): def sync_side_effect(*args, **kwargs): captured_sync_calls.append(kwargs["messages"]) idx = min(len(captured_sync_calls) - 1, len(sync_responses) - 1) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = sync_responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(sync_responses[idx])] async def async_side_effect(*args, **kwargs): captured_async_calls.append(kwargs["messages"]) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = concept_response - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(concept_response)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1586,15 +1851,9 @@ async def test_long_doc_marks_doc_message(self, tmp_path): def sync_side_effect(*args, **kwargs): captured.append(kwargs["messages"]) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] # First call: overview (plain text); second: plan (JSON). - mock_resp.choices[0].message.content = ( - "Overview text" if len(captured) == 1 else plan_response - ) - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + content = "Overview text" if len(captured) == 1 else plan_response + return [_mock_response(content)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1726,16 +1985,9 @@ async def test_create_and_update_flow(self, tmp_path): async def ordered_acompletion(*args, **kwargs): idx = call_order["n"] call_order["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] # create tasks come first, then update tasks - if idx == 0: - mock_resp.choices[0].message.content = create_page_response - else: - mock_resp.choices[0].message.content = update_page_response - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + content = create_page_response if idx == 0 else update_page_response + return [_mock_response(content)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1823,13 +2075,7 @@ async def test_truncated_update_preserves_existing_page(self, tmp_path): ) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1859,13 +2105,7 @@ async def test_truncated_create_skips_partial_page(self, tmp_path): truncated_page = json.dumps({"brief": "x", "content": "# Ghost\n\nPartial"}) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1928,13 +2168,7 @@ async def test_truncated_entity_update_preserves_existing_page(self, tmp_path): ) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -2701,6 +2935,152 @@ async def test_llm_call_async_injects_extra_headers(self): assert kwargs["extra_headers"] == {"Copilot-Integration-Id": "vscode-chat"} +class TestLLMStreamTimingDebugLogging: + """Chunk-phase debug logging should be visible when verbose logging is + enabled: one line when the first chunk arrives, one more when the stream + ends cleanly or is interrupted — never one line per chunk (see #).""" + + def test_llm_call_logs_stream_start_and_end_at_debug(self, caplog): + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock( + return_value=iter(["chunk-1", "chunk-2", "chunk-3"]) + ) + mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok")) + + out = _llm_call("m", [{"role": "user", "content": "hi"}], "sync-step") + + assert out == "ok" + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [sync-step]" in m] + end_logs = [m for m in messages if "LLM stream finished [sync-step]" in m] + # Exactly one start line and one end line — never a line per chunk. + assert len(start_logs) == 1 + assert len(end_logs) == 1 + assert "3 chunk(s)" in end_logs[0] + assert not any("LLM stream chunk [sync-step]" in m for m in messages) + + @pytest.mark.asyncio + async def test_llm_call_async_logs_stream_start_and_end_at_debug(self, caplog): + from openkb.agent.compiler import _llm_call_async + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.acompletion = AsyncMock( + return_value=_AsyncStream(["chunk-1", "chunk-2", "chunk-3"]) + ) + mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok")) + + out = await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-step") + + assert out == "ok" + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [async-step]" in m] + end_logs = [m for m in messages if "LLM stream finished [async-step]" in m] + assert len(start_logs) == 1 + assert len(end_logs) == 1 + assert "3 chunk(s)" in end_logs[0] + assert not any("LLM stream chunk [async-step]" in m for m in messages) + + def test_llm_call_logs_interruption_before_reraising(self, caplog): + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("stream exploded") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + # Fresh generator per call: _llm_call now retries 3 times total + # (fixed resilience floor), and a real litellm.completion() call + # returns an independent stream every time, unlike reusing one + # already-exhausted mock generator across attempts. + mock_litellm.completion = MagicMock( + side_effect=lambda *a, **k: _stream_then_raise(["chunk-1", "chunk-2"], error) + ) + + with pytest.raises(RuntimeError, match="stream exploded"): + _llm_call("m", [{"role": "user", "content": "hi"}], "sync-fail-step") + + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [sync-fail-step]" in m] + interrupted_logs = [ + m for m in messages if "LLM stream [sync-fail-step] interrupted unexpectedly" in m + ] + # One start + one interruption line per attempt (3 attempts total), + # no end-of-stream line, and no per-chunk lines in between. + assert len(start_logs) == 3 + assert len(interrupted_logs) == 3 + assert all("after chunk 2" in m for m in interrupted_logs) + assert not any("LLM stream finished [sync-fail-step]" in m for m in messages) + assert not any("LLM stream chunk [sync-fail-step]" in m for m in messages) + + @pytest.mark.asyncio + async def test_llm_call_async_logs_interruption_before_reraising(self, caplog): + from openkb.agent.compiler import _llm_call_async + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("stream exploded") + with patch("openkb.agent.compiler.litellm") as mock_litellm: + # Fresh _AsyncStream per call (see the sync test above for why): + # _llm_call_async now retries 3 times total. + mock_litellm.acompletion = AsyncMock( + side_effect=lambda *a, **k: _AsyncStream( + ["chunk-1", "chunk-2"], + error=error, + raise_after=2, + ) + ) + + with pytest.raises(RuntimeError, match="stream exploded"): + await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-fail-step") + + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [async-fail-step]" in m] + interrupted_logs = [ + m for m in messages if "LLM stream [async-fail-step] interrupted unexpectedly" in m + ] + assert len(start_logs) == 3 + assert len(interrupted_logs) == 3 + assert all("after chunk 2" in m for m in interrupted_logs) + assert not any("LLM stream finished [async-fail-step]" in m for m in messages) + assert not any("LLM stream chunk [async-fail-step]" in m for m in messages) + + def test_llm_call_logs_interruption_before_any_chunk(self, caplog): + """No chunk ever arrives (e.g. a proxy silently buffering despite + stream=True): no start line, and the interruption line says so + instead of an inapplicable chunk number.""" + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("connect timeout") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + # Fresh generator per call (see test_llm_call_logs_interruption_before_reraising): + # _llm_call now retries 3 times total. + mock_litellm.completion = MagicMock( + side_effect=lambda *a, **k: _stream_then_raise([], error) + ) + + with pytest.raises(RuntimeError, match="connect timeout"): + _llm_call("m", [{"role": "user", "content": "hi"}], "sync-no-chunk-step") + + messages = [record.getMessage() for record in caplog.records] + interrupted_logs = [ + m for m in messages if "LLM stream [sync-no-chunk-step] interrupted unexpectedly" in m + ] + assert len(interrupted_logs) == 3 + assert all("before any chunk arrived" in m for m in interrupted_logs) + assert not any("LLM stream started [sync-no-chunk-step]" in m for m in messages) + + class TestCacheControlStripping: """cache_control markers must only reach providers that honour them. diff --git a/tests/test_llm_timeout.py b/tests/test_llm_timeout.py index ca7d80e68..db119df3c 100644 --- a/tests/test_llm_timeout.py +++ b/tests/test_llm_timeout.py @@ -17,6 +17,13 @@ def _fake_response(): + """A fake, already-complete LLM response (single-chunk stream). + + See ``openkb.agent.compiler._merge_stream_chunks``: a chunk exposing + ``.message`` (as this one does) is treated as already-complete and used + as-is, so callers of ``litellm.completion``/``acompletion`` with + ``stream=True`` can be mocked to just return a one-item list. + """ choice = MagicMock() choice.message.content = "ok" choice.finish_reason = "stop" @@ -28,7 +35,7 @@ def _fake_response(): def test_llm_call_forwards_configured_timeout(): set_timeout(1200.0) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()] ) as completion: _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") assert completion.call_args.kwargs["timeout"] == 1200.0 @@ -37,7 +44,7 @@ def test_llm_call_forwards_configured_timeout(): def test_llm_call_omits_timeout_when_unset(): set_timeout(None) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()] ) as completion: _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") assert "timeout" not in completion.call_args.kwargs @@ -47,7 +54,7 @@ def test_llm_call_does_not_override_explicit_timeout(): # An explicit per-call timeout kwarg wins over the configured default. set_timeout(1200.0) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()] ) as completion: _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30) assert completion.call_args.kwargs["timeout"] == 30 @@ -58,7 +65,7 @@ def test_llm_call_async_forwards_configured_timeout(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + return_value=[_fake_response()], ) as acompletion: asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step")) assert acompletion.call_args.kwargs["timeout"] == 900.0 @@ -69,7 +76,7 @@ def test_llm_call_async_omits_timeout_when_unset(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + return_value=[_fake_response()], ) as acompletion: asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step")) assert "timeout" not in acompletion.call_args.kwargs @@ -80,7 +87,7 @@ def test_llm_call_async_does_not_override_explicit_timeout(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + return_value=[_fake_response()], ) as acompletion: asyncio.run( _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30)