From 0d58eb55e2c752c5e5dbfc8cac5eb646dd06718c Mon Sep 17 00:00:00 2001 From: Kylin Date: Thu, 2 Jul 2026 17:33:08 +0800 Subject: [PATCH 1/2] fix(compiler): guard concept/entity generation against malformed & truncated LLM output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two silent-data-loss bugs in _compile_concepts, both from trusting the per-page LLM response shape: - #158: a response returned as a JSON array (e.g. [{...}] or a multi-item list) reached .get() on a list and raised AttributeError, which the local except (JSONDecodeError, ValueError) did not catch — the page was dropped and, because the doc index records it as written, never retried. New _parse_page_json unwraps a single-element [{...}] array (recovering the common case) and returns None for other wrong shapes so the page is skipped cleanly instead of writing the raw JSON as its body. - #148: a response that hit finish_reason=length was repaired by json_repair and written anyway, overwriting an existing concept page with truncated content while still reporting [OK]. _warn_if_truncated now reports whether it truncated, and the four page-generation calls pass raise_on_truncation=True so a truncated response skips the write (existing page preserved). Other callers (plan, summary, overview) keep the warn-only behavior. Closes #158. Closes #148. Claude-Session: https://claude.ai/code/session_01UtbmJxjtw6FtP8fUXUKVtg --- openkb/agent/compiler.py | 120 ++++++++++++++++++++++++++++----------- tests/test_compiler.py | 112 ++++++++++++++++++++++++++++++++++++ 2 files changed, 198 insertions(+), 34 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index dd580696..55deab23 100644 --- a/openkb/agent/compiler.py +++ b/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,7 +450,7 @@ 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") @@ -445,6 +458,10 @@ async def _llm_call_async(model: str, messages: list[dict], step_name: str, **kw 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() @@ -465,22 +482,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 +520,22 @@ 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 _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): @@ -1758,18 +1795,21 @@ async def _gen_create(concept: dict) -> tuple[str, str, bool, str]: ], f"concept: {name}", response_format=_JSON_RESPONSE_FORMAT, + raise_on_truncation=True, ) 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 "" + obj = _parse_page_json(raw) except (json.JSONDecodeError, ValueError): - # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback. + # Not JSON at all: ``raw`` is the legitimate prose-markdown body. brief, content = "", raw + else: + if obj is None: + # Valid JSON, wrong shape (array/scalar): skip rather than write + # the JSON text as the body. Empty content trips the check below. + brief, content = "", "" + else: + brief = obj.get("description", "") + content = obj.get("content") or "" _require_nonempty_content(content, name) return name, content, False, brief @@ -1802,15 +1842,19 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool, str]: ], f"update: {name}", response_format=_JSON_RESPONSE_FORMAT, + raise_on_truncation=True, ) 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 "" + obj = _parse_page_json(raw) except (json.JSONDecodeError, ValueError): - # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback. + # Not JSON at all: ``raw`` is the legitimate prose-markdown body. brief, content = "", raw + else: + if obj is None: + brief, content = "", "" + else: + brief = obj.get("description", "") + content = obj.get("content") or "" _require_nonempty_content(content, name) return name, content, True, brief @@ -1837,16 +1881,20 @@ async def _gen_entity_create(ent: dict) -> tuple[str, str, str, str]: ], f"entity: {name}", response_format=_JSON_RESPONSE_FORMAT, + raise_on_truncation=True, ) 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 "" + obj = _parse_page_json(raw) except (json.JSONDecodeError, ValueError): - # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback. + # Not JSON at all: ``raw`` is the legitimate prose-markdown body. brief, etype_out, content = "", etype, raw + else: + if obj is None: + brief, etype_out, content = "", etype, "" + else: + brief = obj.get("description", "") + etype_out = obj.get("type") if obj.get("type") in valid_types else etype + content = obj.get("content") or "" _require_nonempty_content(content, name) return name, content, brief, etype_out @@ -1881,16 +1929,20 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: ], f"entity-update: {name}", response_format=_JSON_RESPONSE_FORMAT, + raise_on_truncation=True, ) 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 "" + obj = _parse_page_json(raw) except (json.JSONDecodeError, ValueError): - # Parse FAILED: ``raw`` is the legitimate non-JSON body fallback. + # Not JSON at all: ``raw`` is the legitimate prose-markdown body. brief, etype_out, content = "", etype, raw + else: + if obj is None: + brief, etype_out, content = "", etype, "" + else: + brief = obj.get("description", "") + etype_out = obj.get("type") if obj.get("type") in valid_types else etype + content = obj.get("content") or "" _require_nonempty_content(content, name) return name, content, brief, etype_out diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 4e0b69cd..577769fa 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1770,6 +1770,118 @@ async def ordered_acompletion(*args, **kwargs): assert "[[concepts/flash-attention]]" in index_text assert "[[concepts/attention]]" in index_text + def test_parse_page_json_unwraps_and_guards_shape(self): + """#158: _parse_page_json returns an object, unwraps a single-element + ``[{...}]`` array, and returns None for wrong-shaped-but-valid JSON.""" + from openkb.agent.compiler import _parse_page_json + + assert _parse_page_json('{"content": "x"}') == {"content": "x"} + assert _parse_page_json('[{"content": "x"}]') == {"content": "x"} # unwrapped + assert _parse_page_json("[]") is None + assert _parse_page_json('[{"a": 1}, {"b": 2}]') is None + assert _parse_page_json('["a", "b"]') is None + + @pytest.mark.asyncio + async def test_page_json_wrapped_in_single_array_is_recovered(self, tmp_path): + """#158: a page response the model wrapped as ``[{...}]`` (instead of a + bare object) is unwrapped and written, not dropped with an AttributeError.""" + wiki = self._setup_wiki(tmp_path) + plan_response = json.dumps( + {"create": [{"name": "attention", "title": "Attention"}], "update": [], "related": []} + ) + array_page = json.dumps([{"brief": "b", "content": "# Attention\n\nRecovered body."}]) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) + mock_litellm.acompletion = AsyncMock(side_effect=_mock_acompletion([array_page])) + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + ) + path = wiki / "concepts" / "attention.md" + assert path.exists(), "single-object array should be unwrapped and written" + text = path.read_text() + assert "Recovered body." in text + assert "[{" not in text # not the raw JSON array text + + @pytest.mark.asyncio + async def test_truncated_update_preserves_existing_page(self, tmp_path): + """#148: an update whose response hit finish_reason='length' must not + overwrite the existing (complete) page with truncated content.""" + original = "---\nsources: [old.pdf]\n---\n\n# Attention\n\nComplete original body." + wiki = self._setup_wiki(tmp_path, existing_concepts={"attention": original}) + plan_response = json.dumps( + {"create": [], "update": [{"name": "attention", "title": "Attention"}], "related": []} + ) + truncated_page = json.dumps( + {"brief": "x", "content": "# Attention\n\nTruncated tail cut off"} + ) + + 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 + + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) + mock_litellm.acompletion = AsyncMock(side_effect=truncated_acompletion) + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + ) + text = (wiki / "concepts" / "attention.md").read_text() + assert "Complete original body." in text, "existing page must survive a truncated update" + assert "Truncated tail" not in text + + @pytest.mark.asyncio + async def test_truncated_create_skips_partial_page(self, tmp_path): + """#148: a create whose response hit finish_reason='length' must not + write a partial page (which would be recorded as done and never retried).""" + wiki = self._setup_wiki(tmp_path) + plan_response = json.dumps( + {"create": [{"name": "ghost", "title": "Ghost"}], "update": [], "related": []} + ) + 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 + + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) + mock_litellm.acompletion = AsyncMock(side_effect=truncated_acompletion) + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + ) + assert not (wiki / "concepts" / "ghost.md").exists(), "truncated create must be skipped" + @pytest.mark.asyncio async def test_empty_content_skips_page_no_json_body(self, tmp_path): """#9: when the page LLM returns parseable JSON with empty content From 299f95942ba9abd3a2a5c30c9e56c291d5cd676c Mon Sep 17 00:00:00 2001 From: Kylin Date: Thu, 2 Jul 2026 17:53:15 +0800 Subject: [PATCH 2/2] refactor(compiler): extract _page_fields + _llm_call_page_async; cover entity/prose paths Addresses code-review findings on the #158/#148 fix, no behavior change: - Collapse the four near-identical parse/fallback blocks in the concept and entity generation closures into a shared _page_fields() helper, so a future edge case is fixed in one place instead of four copies that can drift. - Add _llm_call_page_async() which hard-codes raise_on_truncation=True, so a page-generating call site can't silently forget the truncation guard and regress #148. - Tests: a _page_fields shape unit test (object / single-element-array unwrap / wrong-shape skip / non-JSON prose fallback) and an entity-path truncation-skip integration test. Claude-Session: https://claude.ai/code/session_01UtbmJxjtw6FtP8fUXUKVtg --- openkb/agent/compiler.py | 101 +++++++++++++++++---------------------- tests/test_compiler.py | 71 +++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 56 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index 55deab23..fc712b15 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -465,6 +465,17 @@ async def _llm_call_async( 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. @@ -536,6 +547,30 @@ def _parse_page_json(text: str) -> dict | None: 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): @@ -1777,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, @@ -1795,21 +1830,8 @@ async def _gen_create(concept: dict) -> tuple[str, str, bool, str]: ], f"concept: {name}", response_format=_JSON_RESPONSE_FORMAT, - raise_on_truncation=True, ) - try: - obj = _parse_page_json(raw) - except (json.JSONDecodeError, ValueError): - # Not JSON at all: ``raw`` is the legitimate prose-markdown body. - brief, content = "", raw - else: - if obj is None: - # Valid JSON, wrong shape (array/scalar): skip rather than write - # the JSON text as the body. Empty content trips the check below. - brief, content = "", "" - else: - brief = obj.get("description", "") - content = obj.get("content") or "" + brief, content, _ = _page_fields(raw) _require_nonempty_content(content, name) return name, content, False, brief @@ -1824,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, @@ -1842,19 +1864,8 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool, str]: ], f"update: {name}", response_format=_JSON_RESPONSE_FORMAT, - raise_on_truncation=True, ) - try: - obj = _parse_page_json(raw) - except (json.JSONDecodeError, ValueError): - # Not JSON at all: ``raw`` is the legitimate prose-markdown body. - brief, content = "", raw - else: - if obj is None: - brief, content = "", "" - else: - brief = obj.get("description", "") - content = obj.get("content") or "" + brief, content, _ = _page_fields(raw) _require_nonempty_content(content, name) return name, content, True, brief @@ -1863,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, @@ -1881,20 +1892,9 @@ async def _gen_entity_create(ent: dict) -> tuple[str, str, str, str]: ], f"entity: {name}", response_format=_JSON_RESPONSE_FORMAT, - raise_on_truncation=True, ) - try: - obj = _parse_page_json(raw) - except (json.JSONDecodeError, ValueError): - # Not JSON at all: ``raw`` is the legitimate prose-markdown body. - brief, etype_out, content = "", etype, raw - else: - if obj is None: - brief, etype_out, content = "", etype, "" - else: - brief = obj.get("description", "") - etype_out = obj.get("type") if obj.get("type") in valid_types else etype - content = obj.get("content") or "" + 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 @@ -1910,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, @@ -1929,20 +1929,9 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: ], f"entity-update: {name}", response_format=_JSON_RESPONSE_FORMAT, - raise_on_truncation=True, ) - try: - obj = _parse_page_json(raw) - except (json.JSONDecodeError, ValueError): - # Not JSON at all: ``raw`` is the legitimate prose-markdown body. - brief, etype_out, content = "", etype, raw - else: - if obj is None: - brief, etype_out, content = "", etype, "" - else: - brief = obj.get("description", "") - etype_out = obj.get("type") if obj.get("type") in valid_types else etype - content = obj.get("content") or "" + 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 diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 577769fa..95a57cc4 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1882,6 +1882,77 @@ async def truncated_acompletion(*args, **kwargs): ) assert not (wiki / "concepts" / "ghost.md").exists(), "truncated create must be skipped" + def test_page_fields_maps_response_shapes(self): + """Shared mapping used by all four page closures: object, single-element + array unwrap, wrong-shape skip (empty content), and non-JSON prose + fallback (raw written as the body).""" + from openkb.agent.compiler import _page_fields + + brief, content, obj = _page_fields('{"description": "d", "content": "c", "type": "org"}') + assert (brief, content) == ("d", "c") + assert obj == {"description": "d", "content": "c", "type": "org"} + # Single-element [{...}] is unwrapped and used. + assert _page_fields('[{"description": "d", "content": "c"}]')[:2] == ("d", "c") + # Wrong shape (multi-element / empty array) → empty content so the + # caller's _require_nonempty_content skips the page. + assert _page_fields('[{"a": 1}, {"b": 2}]') == ("", "", None) + assert _page_fields("[]") == ("", "", None) + # Non-JSON prose → written verbatim as the markdown body. + prose = "# Heading\n\nJust markdown, not JSON." + assert _page_fields(prose) == ("", prose, None) + + @pytest.mark.asyncio + async def test_truncated_entity_update_preserves_existing_page(self, tmp_path): + """#148 (entity path): a truncated entity update must not overwrite the + existing entity page with cut-off content.""" + wiki = self._setup_wiki(tmp_path) + (wiki / "entities").mkdir() + (wiki / "entities" / "google.md").write_text( + "---\ntype: org\nsources: [old.pdf]\n---\n\n# Google\n\nComplete original entity body.", + encoding="utf-8", + ) + plan_response = json.dumps( + { + "create": [], + "update": [], + "related": [], + "entities": { + "create": [], + "update": [{"name": "google", "title": "Google", "type": "org"}], + "related": [], + }, + } + ) + truncated_page = json.dumps( + {"brief": "x", "content": "# Google\n\nTruncated entity tail cut"} + ) + + 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 + + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) + mock_litellm.acompletion = AsyncMock(side_effect=truncated_acompletion) + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + ) + text = (wiki / "entities" / "google.md").read_text() + assert "Complete original entity body." in text, "existing entity must survive truncation" + assert "Truncated entity tail" not in text + @pytest.mark.asyncio async def test_empty_content_skips_page_no_json_body(self, tmp_path): """#9: when the page LLM returns parseable JSON with empty content