fix(compiler): guard concept/entity generation against malformed & truncated LLM output

openkb/agent/compiler.py

@@ -392,7 +392,14 @@ def _fmt_messages(messages: list[dict], max_content: int = 200) -> str:
     return "\n".join(parts)
 
 
-def _llm_call(model: str, messages: list[dict], step_name: str, **kwargs) -> str:
+class TruncatedResponseError(Exception):
+    """Raised when an LLM response hit the length cap and the caller asked to
+    treat truncation as a failure (so a partial page is skipped, not written)."""
+
+
+def _llm_call(
+    model: str, messages: list[dict], step_name: str, raise_on_truncation: bool = False, **kwargs
+) -> str:
     """Single LLM call with animated progress and debug logging."""
     messages = _prepare_messages(model, messages)
     extra_headers = get_extra_headers()
@@ -411,16 +418,22 @@ def _llm_call(model: str, messages: list[dict], step_name: str, **kwargs) -> str
 
     response = litellm.completion(model=model, messages=messages, **kwargs)
     content = response.choices[0].message.content or ""
-    _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
+    truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
 
     spinner.stop(_format_usage(time.time() - t0, response.usage))
     logger.debug(
         "LLM response [%s]:\n%s", step_name, content[:500] + ("..." if len(content) > 500 else "")
     )
+    if raise_on_truncation and truncated:
+        raise TruncatedResponseError(
+            f"LLM [{step_name}] hit the length limit; skipping to avoid a truncated page"
+        )
     return content.strip()
 
 
-async def _llm_call_async(model: str, messages: list[dict], step_name: str, **kwargs) -> str:
+async def _llm_call_async(
+    model: str, messages: list[dict], step_name: str, raise_on_truncation: bool = False, **kwargs
+) -> str:
     """Async LLM call with timing output and debug logging."""
     messages = _prepare_messages(model, messages)
     extra_headers = get_extra_headers()
@@ -437,17 +450,32 @@ async def _llm_call_async(model: str, messages: list[dict], step_name: str, **kw
 
     response = await litellm.acompletion(model=model, messages=messages, **kwargs)
     content = response.choices[0].message.content or ""
-    _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
+    truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
 
     elapsed = time.time() - t0
     sys.stdout.write(f"    {step_name}... {_format_usage(elapsed, response.usage)}\n")
     sys.stdout.flush()
     logger.debug(
         "LLM response [%s]:\n%s", step_name, content[:500] + ("..." if len(content) > 500 else "")
     )
+    if raise_on_truncation and truncated:
+        raise TruncatedResponseError(
+            f"LLM [{step_name}] hit the length limit; skipping to avoid a truncated page"
+        )
     return content.strip()
 
 
+async def _llm_call_page_async(model: str, messages: list[dict], step_name: str, **kwargs) -> str:
+    """``_llm_call_async`` for a step that writes a wiki page from the response.
+
+    Hard-codes ``raise_on_truncation=True`` so a truncated response skips the
+    write instead of silently persisting a partial page (#148). Use this for
+    every page-generating call so the guarantee can't be forgotten at a new
+    call site.
+    """
+    return await _llm_call_async(model, messages, step_name, raise_on_truncation=True, **kwargs)
+
+
 async def _close_async_llm_clients() -> None:
     """Close LiteLLM's cached async (aiohttp) clients for the current loop.
 
@@ -465,22 +493,26 @@ async def _close_async_llm_clients() -> None:
         logger.debug("litellm async client cleanup failed", exc_info=True)
 
 
-def _warn_if_truncated(response, step_name: str, max_tokens: int | None) -> None:
-    """Emit a warning when the LLM hit the max_tokens cap.
+def _warn_if_truncated(response, step_name: str, max_tokens: int | None) -> bool:
+    """Warn when the LLM hit the max_tokens cap; return True if it did.
 
     ``json_repair`` will silently salvage the truncated prefix, so without
-    this the caller can't tell a short response from a cut-off one.
+    this the caller can't tell a short response from a cut-off one. Callers
+    that write a page from the response can pass ``raise_on_truncation=True``
+    to ``_llm_call``/`_llm_call_async`` to turn a truncated response into a
+    skip instead of persisting partial content.
     """
     try:
         finish_reason = response.choices[0].finish_reason
     except (AttributeError, IndexError):
-        return
+        return False
     if finish_reason != "length":
-        return
+        return False
     cap = f" (max_tokens={max_tokens})" if max_tokens else ""
     logger.warning("LLM [%s] hit length limit%s — output may be truncated.", step_name, cap)
     sys.stdout.write(f"    [WARN] {step_name} hit length limit{cap} — output may be truncated.\n")
     sys.stdout.flush()
+    return True
 
 
 def _parse_json(text: str) -> list | dict:
@@ -499,6 +531,46 @@ def _parse_json(text: str) -> list | dict:
     return result
 
 
+def _parse_page_json(text: str) -> dict | None:
+    """Parse an LLM page response into a single JSON object.
+
+    Unwraps a single-element ``[{...}]`` array (some models wrap the object in
+    a list). Returns ``None`` when the response is valid JSON of the wrong
+    shape (empty/multi-element array, list of scalars) so callers skip the page
+    rather than persisting the raw JSON text as its body. Propagates the
+    json/ValueError from ``_parse_json`` when the text isn't JSON at all, which
+    callers catch to fall back to treating ``raw`` as a prose-markdown body.
+    """
+    parsed = _parse_json(text)
+    if isinstance(parsed, list) and len(parsed) == 1 and isinstance(parsed[0], dict):
+        parsed = parsed[0]
+    return parsed if isinstance(parsed, dict) else None
+
+
+def _page_fields(raw: str) -> tuple[str, str, dict | None]:
+    """Map a page LLM response to ``(brief, content, obj)``.
+
+    - JSON object (or a single-element ``[{...}]`` array): brief/content come
+      from it and ``obj`` is the dict (entity callers read ``type`` from it).
+    - Valid JSON of the wrong shape (multi/empty array, scalar): ``("", "",
+      None)`` — the empty content makes ``_require_nonempty_content`` skip the
+      page rather than persisting the raw JSON text as its body.
+    - Not JSON at all: ``("", raw, None)`` — ``raw`` is written as a
+      prose-markdown body (the legitimate fallback for models that emit
+      markdown instead of JSON).
+
+    Shared by all four page-generation closures so a new edge case is handled
+    in one place instead of four near-identical blocks.
+    """
+    try:
+        obj = _parse_page_json(raw)
+    except (json.JSONDecodeError, ValueError):
+        return "", raw, None
+    if obj is None:
+        return "", "", None
+    return obj.get("description", ""), (obj.get("content") or ""), obj
+
+
 def _filter_concept_items(items: list, label: str) -> list[dict]:
     """Keep only dicts that carry a non-empty ``name``; warn about anything else."""
     if not isinstance(items, list):
@@ -1740,7 +1812,7 @@ async def _gen_create(concept: dict) -> tuple[str, str, bool, str]:
         name = concept["name"]
         title = concept.get("title", name)
         async with semaphore:
-            raw = await _llm_call_async(
+            raw = await _llm_call_page_async(
                 model,
                 [
                     system_msg,
@@ -1759,17 +1831,7 @@ async def _gen_create(concept: dict) -> tuple[str, str, bool, str]:
                 f"concept: {name}",
                 response_format=_JSON_RESPONSE_FORMAT,
             )
-        try:
-            parsed = _parse_json(raw)
-            brief = parsed.get("description", "")
-            # Parse succeeded: do NOT fall back to ``raw`` (the JSON string).
-            # An empty/None ``content`` field yields "" so
-            # ``_require_nonempty_content`` raises and the page is skipped,
-            # rather than writing the raw JSON as the markdown body.
-            content = parsed.get("content") or ""
-        except (json.JSONDecodeError, ValueError):
-            # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback.
-            brief, content = "", raw
+        brief, content, _ = _page_fields(raw)
         _require_nonempty_content(content, name)
         return name, content, False, brief
 
@@ -1784,7 +1846,7 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool, str]:
         else:
             existing_content = "(page not found — create from scratch)"
         async with semaphore:
-            raw = await _llm_call_async(
+            raw = await _llm_call_page_async(
                 model,
                 [
                     system_msg,
@@ -1803,14 +1865,7 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool, str]:
                 f"update: {name}",
                 response_format=_JSON_RESPONSE_FORMAT,
             )
-        try:
-            parsed = _parse_json(raw)
-            brief = parsed.get("description", "")
-            # Parse succeeded: do NOT fall back to ``raw`` (the JSON string).
-            content = parsed.get("content") or ""
-        except (json.JSONDecodeError, ValueError):
-            # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback.
-            brief, content = "", raw
+        brief, content, _ = _page_fields(raw)
         _require_nonempty_content(content, name)
         return name, content, True, brief
 
@@ -1819,7 +1874,7 @@ async def _gen_entity_create(ent: dict) -> tuple[str, str, str, str]:
         title = ent.get("title", name)
         etype = ent.get("type", "other")
         async with semaphore:
-            raw = await _llm_call_async(
+            raw = await _llm_call_page_async(
                 model,
                 [
                     system_msg,
@@ -1838,15 +1893,8 @@ async def _gen_entity_create(ent: dict) -> tuple[str, str, str, str]:
                 f"entity: {name}",
                 response_format=_JSON_RESPONSE_FORMAT,
             )
-        try:
-            parsed = _parse_json(raw)
-            brief = parsed.get("description", "")
-            etype_out = parsed.get("type") if parsed.get("type") in valid_types else etype
-            # Parse succeeded: do NOT fall back to ``raw`` (the JSON string).
-            content = parsed.get("content") or ""
-        except (json.JSONDecodeError, ValueError):
-            # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback.
-            brief, etype_out, content = "", etype, raw
+        brief, content, obj = _page_fields(raw)
+        etype_out = obj.get("type") if obj and obj.get("type") in valid_types else etype
         _require_nonempty_content(content, name)
         return name, content, brief, etype_out
 
@@ -1862,7 +1910,7 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]:
         else:
             existing_content = "(page not found — create from scratch)"
         async with semaphore:
-            raw = await _llm_call_async(
+            raw = await _llm_call_page_async(
                 model,
                 [
                     system_msg,
@@ -1882,15 +1930,8 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]:
                 f"entity-update: {name}",
                 response_format=_JSON_RESPONSE_FORMAT,
             )
-        try:
-            parsed = _parse_json(raw)
-            brief = parsed.get("description", "")
-            etype_out = parsed.get("type") if parsed.get("type") in valid_types else etype
-            # Parse succeeded: do NOT fall back to ``raw`` (the JSON string).
-            content = parsed.get("content") or ""
-        except (json.JSONDecodeError, ValueError):
-            # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback.
-            brief, etype_out, content = "", etype, raw
+        brief, content, obj = _page_fields(raw)
+        etype_out = obj.get("type") if obj and obj.get("type") in valid_types else etype
         _require_nonempty_content(content, name)
         return name, content, brief, etype_out