diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index dd580696..fc712b15 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,9 +458,24 @@ 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() +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 diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 4e0b69cd..95a57cc4 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1770,6 +1770,189 @@ 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" + + 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