refactor: move entity-type resolution to config layer + co-locate remove-preview scan

Altitude cleanups from the review:
- Move resolve_entity_types + DEFAULT_ENTITY_TYPES into openkb/config.py (the
  config layer owns config validation/normalization; any command can reuse it
  without importing the heavy compiler module). compiler.py imports them;
  _ENTITY_TYPE_LIST/_ENTITY_TYPES remain as the default alias/validation set.
- Move the remove dry-run preview scan from cli.py into compiler.py as
  scan_affected_pages, beside remove_doc_from_*_pages and sharing
  _parse_yaml_list_value — so preview and executor can't drift on how the
  sources list is parsed (root cause of the earlier JSON-quote preview bug).

Author: KylinMountain

openkb/agent/compiler.py

@@ -29,6 +29,7 @@
 import litellm
 import yaml
 
+from openkb.config import DEFAULT_ENTITY_TYPES, resolve_entity_types
 from openkb.lint import list_existing_wiki_targets, strip_ghost_wikilinks
 from openkb.schema import INDEX_SEED, get_agents_md
 
@@ -68,56 +69,17 @@
 """
 
 
-# Default entity-type enum. The EFFECTIVE set is resolved per-KB from config
-# (see ``_resolve_entity_types``) and substituted into the plan + entity-page
-# prompts at call time inside ``_compile_concepts`` via the ``__ENTITY_TYPES__``
-# token. ``_ENTITY_TYPES`` is the default validation set used when no
-# config-driven set is threaded through.
-_ENTITY_TYPE_LIST = ("person", "organization", "place", "product", "work", "event", "other")
+# Default entity-type enum lives in the config layer (so config validation is
+# centralized there and reusable by any command). ``_ENTITY_TYPE_LIST`` /
+# ``_ENTITY_TYPES`` are the default name + validation set used when no
+# config-driven set is threaded through; the EFFECTIVE set is resolved per-KB
+# via ``resolve_entity_types(config)`` and substituted into the plan +
+# entity-page prompts at call time inside ``_compile_concepts`` via the
+# ``__ENTITY_TYPES__`` token.
+_ENTITY_TYPE_LIST = DEFAULT_ENTITY_TYPES
 _ENTITY_TYPES = frozenset(_ENTITY_TYPE_LIST)
 
 
-def _resolve_entity_types(config: dict) -> list[str]:
-    """Resolve the effective entity-type list from config.
-
-    If ``config["entity_types"]`` is a non-empty list, each item is cleaned
-    (``str(x).strip().lower()``, empties dropped); if anything survives, that
-    cleaned list is used (de-duped, order-preserving) with ``"other"`` always
-    appended when missing (it's the coercion fallback). Otherwise — the key is
-    absent, not a list, empty, or fully malformed — the default
-    ``_ENTITY_TYPE_LIST`` is returned, so behavior is byte-identical to today.
-    A warning is logged only when ``entity_types`` was present-but-malformed.
-    """
-    raw = config.get("entity_types")
-    if raw is None:
-        return list(_ENTITY_TYPE_LIST)
-    if not isinstance(raw, list):
-        logger.warning(
-            "config: 'entity_types' must be a list of strings, got %s — "
-            "falling back to the default entity types.",
-            type(raw).__name__,
-        )
-        return list(_ENTITY_TYPE_LIST)
-    cleaned: list[str] = []
-    for x in raw:
-        if not isinstance(x, str):
-            continue  # skip YAML nulls/numbers (str(None) would become "none")
-        # Restrict to a safe label charset so a stray '{'/'}' or punctuation
-        # can't leak into a prompt template or a frontmatter value.
-        s = re.sub(r"[^a-z0-9 _-]+", "", x.strip().lower()).strip()
-        if s and s not in cleaned:
-            cleaned.append(s)
-    if not cleaned:
-        logger.warning(
-            "config: 'entity_types' was present but yielded no usable values — "
-            "falling back to the default entity types.",
-        )
-        return list(_ENTITY_TYPE_LIST)
-    if "other" not in cleaned:
-        cleaned.append("other")
-    return cleaned
-
-
 _CONCEPTS_PLAN_USER = """\
 Based on the summary above, decide how to update the wiki's CONCEPT pages and
 ENTITY pages.
@@ -249,7 +211,7 @@ def _resolve_entity_types(config: dict) -> list[str]:
 
 # NOTE: the prompt templates intentionally KEEP the literal ``__ENTITY_TYPES__``
 # token at import time. The effective entity-type list is resolved per-compile
-# from config (see ``_resolve_entity_types``) and substituted via ``str.replace``
+# from config (see ``resolve_entity_types``) and substituted via ``str.replace``
 # at call time inside ``_compile_concepts``. This lets ``entity_types:`` in
 # ``.openkb/config.yaml`` override the default enum everywhere at once. The
 # token is a plain string (not a ``{}`` placeholder) so it does not collide with
@@ -1225,6 +1187,35 @@ def _remove_doc_from_pages(
     return {"modified": modified, "deleted": deleted}
 
 
+def scan_affected_pages(pages_dir: Path, source_file_marker: str) -> list[tuple[str, int]]:
+    """Return ``(slug, remaining_sources)`` for pages under ``pages_dir`` whose
+    frontmatter ``sources:`` list contains ``source_file_marker``.
+
+    Used by the ``openkb remove`` dry-run preview. Lives here, beside
+    ``remove_doc_from_concept_pages`` / ``remove_doc_from_entity_pages`` and
+    sharing ``_parse_yaml_list_value`` with them, so the preview and the
+    executor can't drift apart on how the sources list is parsed (a hand-rolled
+    comma-split here once kept the JSON quotes and matched nothing).
+    """
+    affected: list[tuple[str, int]] = []
+    if not pages_dir.is_dir():
+        return affected
+    for path in sorted(pages_dir.glob("*.md")):
+        text = path.read_text(encoding="utf-8")
+        if not text.startswith("---"):
+            continue
+        fm_end = text.find("---", 3)
+        if fm_end == -1:
+            continue
+        for line in text[:fm_end].split("\n"):
+            if line.lstrip().startswith("sources:"):
+                items = _parse_yaml_list_value(line)
+                if items is not None and source_file_marker in items:
+                    affected.append((path.stem, max(len(items) - 1, 0)))
+                break
+    return affected
+
+
 def remove_doc_from_concept_pages(
     wiki_dir: Path,
     doc_name: str,
@@ -1958,7 +1949,7 @@ async def compile_short_doc(
     openkb_dir = kb_dir / ".openkb"
     config = load_config(openkb_dir / "config.yaml")
     language: str = config.get("language", "en")
-    entity_types = _resolve_entity_types(config)
+    entity_types = resolve_entity_types(config)
 
     wiki_dir = kb_dir / "wiki"
     schema_md = get_agents_md(wiki_dir)
@@ -2016,7 +2007,7 @@ async def compile_long_doc(
     openkb_dir = kb_dir / ".openkb"
     config = load_config(openkb_dir / "config.yaml")
     language: str = config.get("language", "en")
-    entity_types = _resolve_entity_types(config)
+    entity_types = resolve_entity_types(config)
 
     wiki_dir = kb_dir / "wiki"
     schema_md = get_agents_md(wiki_dir)

openkb/cli.py

@@ -739,37 +739,6 @@ def _cleanup_pageindex(
     return True, f"deleted PageIndex doc ({doc_id[:12]}…)"
 
 
-def _scan_affected_pages(pages_dir: Path, source_file_marker: str) -> list[tuple[str, int]]:
-    """Return ``(slug, remaining_sources)`` for pages whose frontmatter
-    ``sources:`` list contains ``source_file_marker``.
-
-    Uses the same ``_parse_yaml_list_value`` parser the executor uses, so
-    JSON-quoted values (``sources: ["summaries/x.md"]`` — exactly how the
-    compiler writes them) are matched correctly. A hand-rolled comma-split
-    here previously kept the surrounding quotes, so the marker never matched
-    and the remove preview silently reported 0 affected pages.
-    """
-    from openkb.agent.compiler import _parse_yaml_list_value
-
-    affected: list[tuple[str, int]] = []
-    if not pages_dir.is_dir():
-        return affected
-    for path in sorted(pages_dir.glob("*.md")):
-        text = path.read_text(encoding="utf-8")
-        if not text.startswith("---"):
-            continue
-        fm_end = text.find("---", 3)
-        if fm_end == -1:
-            continue
-        for line in text[:fm_end].split("\n"):
-            if line.lstrip().startswith("sources:"):
-                items = _parse_yaml_list_value(line)
-                if items is not None and source_file_marker in items:
-                    affected.append((path.stem, max(len(items) - 1, 0)))
-                break
-    return affected
-
-
 def _resolve_doc_identifier(registry, identifier: str) -> list[tuple[str, dict]]:
     """Find registry entries matching ``identifier``.
 
@@ -833,6 +802,7 @@ def remove(ctx, identifier, keep_raw, keep_empty, dry_run, yes):
         remove_doc_from_concept_pages,
         remove_doc_from_entity_pages,
         remove_doc_from_index,
+        scan_affected_pages,
     )
     from openkb.lint import fix_broken_links
     from openkb.state import HashRegistry
@@ -894,7 +864,7 @@ def remove(ctx, identifier, keep_raw, keep_empty, dry_run, yes):
     # affect the delete/edit classification, so the plan reflects what
     # the executor will actually do.
     source_file_marker = f"summaries/{doc_name}.md"
-    affected_concepts = _scan_affected_pages(wiki_dir / "concepts", source_file_marker)
+    affected_concepts = scan_affected_pages(wiki_dir / "concepts", source_file_marker)
 
     concept_deletes = [s for s, r in affected_concepts if r == 0 and not keep_empty]
     concept_edits = [s for s, r in affected_concepts if r > 0 or keep_empty]
@@ -906,7 +876,7 @@ def remove(ctx, identifier, keep_raw, keep_empty, dry_run, yes):
     # Scan entity pages with the same frontmatter logic as concepts. The
     # executor calls ``remove_doc_from_entity_pages``; this only makes the
     # preview/summary truthful about what it will delete vs. edit.
-    affected_entities = _scan_affected_pages(wiki_dir / "entities", source_file_marker)
+    affected_entities = scan_affected_pages(wiki_dir / "entities", source_file_marker)
 
     entity_deletes = [s for s, r in affected_entities if r == 0 and not keep_empty]
     entity_edits = [s for s, r in affected_entities if r > 0 or keep_empty]

openkb/config.py

@@ -1,20 +1,71 @@
 from __future__ import annotations
 
+import logging
+import re
 from pathlib import Path
 from typing import Any
 
 import yaml
 
+logger = logging.getLogger(__name__)
+
 DEFAULT_CONFIG: dict[str, Any] = {
     "model": "gpt-5.4-mini",
     "language": "en",
     "pageindex_threshold": 20,
 }
 
+# Default entity-type vocabulary. Overridable per-KB via the optional
+# ``entity_types:`` config key (see ``resolve_entity_types``).
+DEFAULT_ENTITY_TYPES: tuple[str, ...] = (
+    "person", "organization", "place", "product", "work", "event", "other",
+)
+
 GLOBAL_CONFIG_DIR = Path.home() / ".config" / "openkb"
 GLOBAL_CONFIG_PATH = GLOBAL_CONFIG_DIR / "global.yaml"
 
 
+def resolve_entity_types(config: dict) -> list[str]:
+    """Resolve the effective entity-type list from a loaded config dict.
+
+    If ``config["entity_types"]`` is a non-empty list, each string item is
+    cleaned (lowercased, trimmed, restricted to ``[a-z0-9 _-]`` so a stray
+    brace/punctuation can't leak into a prompt template or frontmatter value);
+    non-string items (YAML nulls, numbers) are skipped. The cleaned list is
+    de-duped (order preserving) and ``"other"`` is always appended when missing
+    (it is the coercion fallback). Otherwise — key absent, not a list, empty,
+    or fully malformed — :data:`DEFAULT_ENTITY_TYPES` is returned, so behavior
+    is byte-identical to the default. A warning is logged only when
+    ``entity_types`` was present-but-malformed.
+    """
+    raw = config.get("entity_types")
+    if raw is None:
+        return list(DEFAULT_ENTITY_TYPES)
+    if not isinstance(raw, list):
+        logger.warning(
+            "config: 'entity_types' must be a list of strings, got %s — "
+            "falling back to the default entity types.",
+            type(raw).__name__,
+        )
+        return list(DEFAULT_ENTITY_TYPES)
+    cleaned: list[str] = []
+    for x in raw:
+        if not isinstance(x, str):
+            continue  # skip YAML nulls/numbers (str(None) would become "none")
+        s = re.sub(r"[^a-z0-9 _-]+", "", x.strip().lower()).strip()
+        if s and s not in cleaned:
+            cleaned.append(s)
+    if not cleaned:
+        logger.warning(
+            "config: 'entity_types' was present but yielded no usable values — "
+            "falling back to the default entity types.",
+        )
+        return list(DEFAULT_ENTITY_TYPES)
+    if "other" not in cleaned:
+        cleaned.append("other")
+    return cleaned
+
+
 def load_config(config_path: Path) -> dict[str, Any]:
     """Load YAML config from config_path, merged with DEFAULT_CONFIG.
 

tests/test_compiler.py

@@ -26,10 +26,10 @@
     _backlink_entities,
     _parse_entities_plan,
     _filter_entity_items,
-    _resolve_entity_types,
     _ENTITY_TYPE_LIST,
     remove_doc_from_entity_pages,
 )
+from openkb.config import resolve_entity_types
 
 
 class TestParseJson:
@@ -97,37 +97,37 @@ def test_bad_type_falls_back_to_other(self):
 
 class TestResolveEntityTypes:
     def test_default_when_key_absent(self):
-        assert _resolve_entity_types({}) == list(_ENTITY_TYPE_LIST)
+        assert resolve_entity_types({}) == list(_ENTITY_TYPE_LIST)
 
     def test_custom_list_is_used_and_normalized(self):
-        out = _resolve_entity_types({"entity_types": ["Person", " Dataset ", "MODEL"]})
+        out = resolve_entity_types({"entity_types": ["Person", " Dataset ", "MODEL"]})
         assert out == ["person", "dataset", "model", "other"]
 
     def test_always_includes_other(self):
-        out = _resolve_entity_types({"entity_types": ["person", "dataset"]})
+        out = resolve_entity_types({"entity_types": ["person", "dataset"]})
         assert "other" in out
         # already-present "other" is not duplicated
-        out2 = _resolve_entity_types({"entity_types": ["dataset", "other"]})
+        out2 = resolve_entity_types({"entity_types": ["dataset", "other"]})
         assert out2.count("other") == 1
 
     def test_dedupes_preserving_order(self):
-        out = _resolve_entity_types({"entity_types": ["a", "a", "b"]})
+        out = resolve_entity_types({"entity_types": ["a", "a", "b"]})
         assert out == ["a", "b", "other"]
 
     def test_malformed_string_falls_back_to_default(self):
-        assert _resolve_entity_types({"entity_types": "person"}) == list(_ENTITY_TYPE_LIST)
+        assert resolve_entity_types({"entity_types": "person"}) == list(_ENTITY_TYPE_LIST)
 
     def test_empty_list_falls_back_to_default(self):
-        assert _resolve_entity_types({"entity_types": []}) == list(_ENTITY_TYPE_LIST)
+        assert resolve_entity_types({"entity_types": []}) == list(_ENTITY_TYPE_LIST)
 
     def test_all_empty_strings_falls_back_to_default(self):
-        assert _resolve_entity_types({"entity_types": ["", "  "]}) == list(_ENTITY_TYPE_LIST)
+        assert resolve_entity_types({"entity_types": ["", "  "]}) == list(_ENTITY_TYPE_LIST)
 
     def test_sanitizes_punctuation_and_skips_non_strings(self):
         # '{'/'}' and other punctuation are stripped (so they can't leak into a
         # prompt template's .format()); non-string items (YAML null, ints) are
         # skipped (str(None) must NOT become the type "none").
-        out = _resolve_entity_types(
+        out = resolve_entity_types(
             {"entity_types": ["Per{son}", None, 123, "data set!"]}
         )
         assert out == ["person", "data set", "other"]
@@ -1931,7 +1931,7 @@ async def fake_llm_async(model, messages, label, **kw):
     @pytest.mark.asyncio
     async def test_brace_in_entity_type_does_not_crash_format(self, tmp_path, monkeypatch):
         """Defense-in-depth: even if a '{'/'}' reaches types_str (bypassing
-        _resolve_entity_types sanitization), the prompt build must not raise —
+        resolve_entity_types sanitization), the prompt build must not raise —
         the token is substituted AFTER .format(), so braces are inert."""
         wiki = tmp_path / "wiki"
         (wiki / "summaries").mkdir(parents=True)