-
Notifications
You must be signed in to change notification settings - Fork 470
Expand file tree
/
Copy pathcompiler.py
More file actions
2136 lines (1810 loc) · 85.5 KB
/
Copy pathcompiler.py
File metadata and controls
2136 lines (1810 loc) · 85.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""Wiki compilation pipeline for OpenKB.
Pipeline leveraging LLM prompt caching:
Step 1: Build base context A (schema + document content).
Step 2: A → generate summary.
Step 3: A + summary → concepts plan (create/update/related).
Step 4: Concurrent LLM calls (A cached) → generate new + rewrite updated concepts.
Step 5: Code adds cross-ref links to related concepts, updates index.
Anthropic prompt caching is enabled via ``cache_control`` markers at two
breakpoints: end of the document message (caches system + doc across all
N+M+2 calls) and end of the assistant summary message (caches the additional
summary prefix across N+M concept-generation calls). Providers that do not
support cache_control receive a normalized list-of-blocks content payload,
which LiteLLM passes through cleanly.
"""
from __future__ import annotations
import asyncio
import json
import logging
import re
import sys
import threading
import time
import unicodedata
from pathlib import Path
import litellm
from openkb import frontmatter
from openkb.config import (
DEFAULT_ENTITY_TYPES,
get_extra_headers,
get_timeout,
resolve_entity_types,
)
from openkb.lint import list_existing_wiki_targets, strip_ghost_wikilinks
from openkb.locks import atomic_write_text
from openkb.schema import INDEX_SEED, get_agents_md
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Prompt templates
# ---------------------------------------------------------------------------
# DeepSeek/Qwen require the prompt itself to mention "json" when this kwarg
# is set; the templates below already do.
_JSON_RESPONSE_FORMAT = {"type": "json_object"}
_SYSTEM_TEMPLATE = """\
You are OpenKB's wiki compilation agent for a personal knowledge base.
{schema_md}
Write all content in {language} language.
Use [[wikilinks]] to connect related pages (e.g. [[concepts/attention]]).
"""
_SUMMARY_USER = """\
New document: {doc_name}
Full text:
{content}
Write a summary page for this document in Markdown.
Return a JSON object with two keys:
- "description": A single sentence (under 100 chars) describing the document's main contribution
- "content": The full summary in Markdown. Include key concepts, findings, ideas, \
and [[wikilinks]] to concepts that could become cross-document concept pages
Return ONLY valid JSON, no fences.
"""
# 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)
_CONCEPTS_PLAN_USER = """\
Based on the summary above, decide how to update the wiki's CONCEPT pages and
ENTITY pages.
A CONCEPT is an abstract, recurring idea/pattern/mechanism (e.g. "agentic
systems"). An ENTITY is a specific named thing — a person, organization,
place, product, named work, or event (e.g. "Anthropic"). Each name goes in
exactly ONE group. A topic may have both (entity "NVIDIA" and concept
"ai-infrastructure-demand"); they cross-link, they do not merge.
Existing concept pages:
{concept_briefs}
Existing entity pages (with source counts = how many docs already cite them):
{entity_briefs}
Return a JSON object with two top-level keys, "concepts" and "entities".
"concepts" is an object with:
1. "create" — new concepts. Array of {{"name": "concept-slug", "title": "Title"}}
2. "update" — existing concepts with significant new info. Same shape.
3. "related" — existing concept slugs to cross-link only. Array of strings.
"entities" is an object with the same three keys, but create/update objects
add a "type" field, one of: __ENTITY_TYPES__. Example:
{{"name": "anthropic", "title": "Anthropic", "type": "organization"}}
Rules:
- For the first few documents, create 2-3 foundational concepts at most.
- Create an ENTITY page only when the entity is (a) central to this document
or (b) likely to recur across sources. Do NOT page proper nouns mentioned
only in passing. Roughly 5-15 entities per document is typical; fewer for
sparse documents.
- Prefer "update" over "create" for any concept or entity already listed above.
- Do NOT create a concept/entity that overlaps an existing one — use "update".
- Do NOT create concepts that are just the document topic itself.
- "related" is lightweight cross-linking only, no content rewrite.
Return ONLY valid JSON, no fences, no explanation.
"""
_KNOWN_TARGETS_USER = """\
The wiki currently contains these pages, and they are the COMPLETE list of \
valid [[wikilink]] targets you may use in the responses that follow:
{known_targets}
Rules for [[wikilinks]] in all subsequent responses:
- For [[concepts/X]]: X must appear in the whitelist above.
- For [[summaries/Y]]: Y must appear in the whitelist above.
- For [[entities/Z]]: Z must appear in the whitelist above.
- Do NOT invent new wikilink targets. If you want to mention a concept \
or entity that is not in the whitelist, write it as plain text without brackets.
"""
_CONCEPT_PAGE_USER = """\
Write the concept page for: {title}
This concept relates to the document "{doc_name}" summarized above.
{update_instruction}
Return a JSON object with two keys:
- "description": A single sentence (under 100 chars) defining this concept
- "content": The full concept page in Markdown. Include clear explanation, \
key details from the source document, and [[wikilinks]] to related concepts \
and [[summaries/{doc_name}]] — subject to the wikilink rules from the \
whitelist message above.
Return ONLY valid JSON, no fences.
"""
_CONCEPT_UPDATE_USER = """\
Update the concept page for: {title}
Current content of this page:
{existing_content}
New information from document "{doc_name}" (summarized above) should be \
integrated into this page. Rewrite the full page incorporating the new \
information naturally — do not just append. Preserve the existing structure \
and intent of the page.
For [[wikilinks]] in the rewrite, follow the whitelist rules from the \
message above: keep links whose target is in the whitelist, convert any \
existing links whose target is NOT in the whitelist to plain text, and do \
not invent new wikilink targets.
Return a JSON object with two keys:
- "description": A single sentence (under 100 chars) defining this concept (may differ from before)
- "content": The rewritten full concept page in Markdown
Return ONLY valid JSON, no fences.
"""
_ENTITY_PAGE_USER = """\
Write the entity page for: {title} (type: {type})
This entity relates to the document "{doc_name}" summarized above.
Return a JSON object with three keys:
- "description": A single sentence (under 100 chars) identifying this entity
- "type": one of __ENTITY_TYPES__
- "content": The full entity page in Markdown — what this entity is, the key
facts about it from this document, and [[wikilinks]] to related concepts,
other [[entities/...]], and [[summaries/{doc_name}]] — subject to the
whitelist rules from the message above.
Return ONLY valid JSON, no fences.
"""
_ENTITY_UPDATE_USER = """\
Update the entity page for: {title} (type: {type})
Current content of this page:
{existing_content}
Integrate the new facts about this entity from document "{doc_name}"
(summarized above). Rewrite the full page — do not just append. Preserve the
existing structure and intent. Follow the whitelist rules from the message
above for all [[wikilinks]].
Return a JSON object with three keys:
- "description": A single sentence (under 100 chars) identifying this entity
- "type": one of __ENTITY_TYPES__
- "content": The rewritten full entity page in Markdown
Return ONLY valid JSON, no fences.
"""
# 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``
# 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
# the ``{{ }}`` JSON braces these templates feed to ``str.format``.
_SUMMARY_REWRITE_USER = """\
Task: Rewrite the summary you wrote above into a final version that is \
consistent with the concept pages now in the wiki (per the whitelist message \
above).
STRICT rules:
- Preserve every factual claim, finding, and detail from your draft. Do \
NOT add or remove technical content, examples, or claims.
- For [[wikilinks]], follow the whitelist message above: keep valid links, \
replace targets not in the whitelist with plain text, do not invent new \
wikilink targets.
- You MAY upgrade plain-text mentions to [[wikilinks]] when the concept \
appears in the whitelist — this is encouraged.
- Keep the headings, paragraph structure, and approximately the same length \
as the draft.
Return ONLY the rewritten Markdown content (no JSON, no fences, no frontmatter).
"""
_LONG_DOC_SUMMARY_USER = """\
This is a PageIndex summary for long document "{doc_name}" (doc_id: {doc_id}):
{content}
Based on this structured summary, write a concise overview that captures \
the key themes and findings. This will be used to generate concept pages.
Return ONLY the Markdown content (no frontmatter, no code fences).
"""
# ---------------------------------------------------------------------------
# LLM helpers
# ---------------------------------------------------------------------------
def _cached_text(text: str) -> list[dict]:
"""Wrap a text payload into a content-block list with an Anthropic
ephemeral cache_control marker.
LiteLLM passes the marker through to Anthropic (and OpenRouter →
Anthropic). For other providers the marker is stripped at the request
egress (see :func:`_strip_cache_control`, applied in :func:`_llm_call`),
because not every provider merely *ignores* it — Gemini in particular
turns it into a 400. The list-of-blocks payload that remains is a valid
OpenAI-compatible content shape.
"""
return [{"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}]
def _accepts_cache_control(model: str) -> bool:
"""Whether ``model`` honours Anthropic-style ``cache_control`` markers.
The markers emitted by :func:`_cached_text` are an Anthropic feature.
LiteLLM forwards them to Anthropic directly, and to Anthropic (Claude)
models served via OpenRouter, Bedrock and Vertex. For other providers —
notably Gemini — LiteLLM instead translates the marker into a
provider-native cached-content object that conflicts with
``system_instruction``/``tools`` and makes *every* request fail with
``400 CachedContent can not be used with ...``. Detect the provider so the
marker can be dropped before it reaches such a backend.
"""
# Import the real symbol rather than going through the module-level
# ``litellm`` reference: provider detection must stay correct even when a
# caller patches ``openkb.agent.compiler.litellm`` to stub out completion.
from litellm import get_llm_provider
try:
provider = get_llm_provider(model)[1]
except Exception:
provider = ""
lowered = model.lower()
if provider == "anthropic":
return True
if provider in ("openrouter", "bedrock", "vertex_ai") and (
"claude" in lowered or "anthropic" in lowered
):
return True
return False
def _strip_cache_control(messages: list[dict]) -> list[dict]:
"""Return ``messages`` with every ``cache_control`` key removed.
Only list-of-blocks contents (see :func:`_cached_text`) can carry the
marker; plain-string contents pass through untouched. The input is not
mutated.
"""
cleaned: list[dict] = []
for msg in messages:
content = msg.get("content")
if isinstance(content, list):
blocks = [
{k: v for k, v in block.items() if k != "cache_control"}
if isinstance(block, dict)
else block
for block in content
]
msg = {**msg, "content": blocks}
cleaned.append(msg)
return cleaned
def _prepare_messages(model: str, messages: list[dict]) -> list[dict]:
"""Drop cache_control markers when ``model`` would reject them."""
if _accepts_cache_control(model):
return messages
return _strip_cache_control(messages)
class _Spinner:
"""Animated dots spinner that runs in a background thread."""
def __init__(self, label: str):
self._label = label
self._stop = threading.Event()
self._thread: threading.Thread | None = None
def start(self) -> None:
sys.stdout.write(f" {self._label}")
sys.stdout.flush()
self._thread = threading.Thread(target=self._run, daemon=True)
self._thread.start()
def _run(self) -> None:
while not self._stop.wait(timeout=1.0):
sys.stdout.write(".")
sys.stdout.flush()
def stop(self, suffix: str = "") -> None:
self._stop.set()
if self._thread:
self._thread.join()
sys.stdout.write(f" {suffix}\n")
sys.stdout.flush()
def _format_usage(elapsed: float, usage) -> str:
"""Format timing and token usage into a short summary string."""
cached = getattr(usage, "prompt_tokens_details", None)
cache_info = ""
if cached and hasattr(cached, "cached_tokens") and cached.cached_tokens:
cache_info = f", cached={cached.cached_tokens}"
return f"{elapsed:.1f}s (in={usage.prompt_tokens}, out={usage.completion_tokens}{cache_info})"
def _fmt_messages(messages: list[dict], max_content: int = 200) -> str:
"""Format messages for debug output, truncating long content.
Accepts both plain-string content and the list-of-blocks shape used by
cache_control-tagged messages (joins all text blocks for preview).
"""
parts = []
for msg in messages:
role = msg["role"]
raw = msg["content"]
if isinstance(raw, list):
text = "".join(b.get("text", "") for b in raw if isinstance(b, dict))
else:
text = raw
if len(text) > max_content:
preview = text[:max_content] + f"... ({len(text)} chars)"
else:
preview = text
parts.append(f" [{role}] {preview}")
return "\n".join(parts)
def _llm_call(model: str, messages: list[dict], step_name: str, **kwargs) -> str:
"""Single LLM call with animated progress and debug logging."""
messages = _prepare_messages(model, messages)
extra_headers = get_extra_headers()
if extra_headers:
kwargs.setdefault("extra_headers", extra_headers)
timeout = get_timeout()
if timeout is not None:
kwargs.setdefault("timeout", timeout)
logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages))
if kwargs:
logger.debug("LLM kwargs [%s]: %s", step_name, kwargs)
spinner = _Spinner(step_name)
spinner.start()
t0 = time.time()
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"))
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 ""))
return content.strip()
async def _llm_call_async(model: str, messages: list[dict], step_name: str, **kwargs) -> str:
"""Async LLM call with timing output and debug logging."""
messages = _prepare_messages(model, messages)
extra_headers = get_extra_headers()
if extra_headers:
kwargs.setdefault("extra_headers", extra_headers)
timeout = get_timeout()
if timeout is not None:
kwargs.setdefault("timeout", timeout)
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)
content = response.choices[0].message.content or ""
_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 ""))
return content.strip()
async def _close_async_llm_clients() -> None:
"""Close LiteLLM's cached async (aiohttp) clients for the current loop.
LiteLLM caches its async clients per event loop. ``add_single_file`` runs
each doc in its own ``asyncio.run`` loop, so without this the clients are
orphaned when the loop is torn down and their connections pile up in
CLOSE-WAIT, leaking sockets/FDs across a long ingest. Call this from a
``finally`` inside the compile coroutines so the clients are closed in the
same loop that created them. Best-effort: never raises, so cleanup can't
mask a real compilation error or break ingest.
"""
try:
await litellm.close_litellm_async_clients()
except Exception:
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.
``json_repair`` will silently salvage the truncated prefix, so without
this the caller can't tell a short response from a cut-off one.
"""
try:
finish_reason = response.choices[0].finish_reason
except (AttributeError, IndexError):
return
if finish_reason != "length":
return
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()
def _parse_json(text: str) -> list | dict:
"""Parse JSON from LLM response, handling fences, prose, and malformed JSON."""
from json_repair import repair_json
cleaned = text.strip()
if cleaned.startswith("```"):
first_nl = cleaned.find("\n")
cleaned = cleaned[first_nl + 1:] if first_nl != -1 else cleaned[3:]
if cleaned.endswith("```"):
cleaned = cleaned[:-3]
result = json.loads(repair_json(cleaned.strip()))
if not isinstance(result, (dict, list)):
raise ValueError(f"Expected JSON object or array, got {type(result).__name__}")
return result
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):
logger.warning("concepts plan: %s was %s, expected list — dropping",
label, type(items).__name__)
return []
valid = [c for c in items if isinstance(c, dict) and isinstance(c.get("name"), str) and c["name"].strip()]
if len(valid) < len(items):
reasons: list[str] = []
for c in items:
if not isinstance(c, dict):
reasons.append(type(c).__name__)
elif not isinstance(c.get("name"), str) or not c["name"].strip():
reasons.append("dict-missing-name")
logger.warning(
"concepts plan: dropped %d malformed %s item(s) (reasons: %s)",
len(items) - len(valid), label, ", ".join(sorted(set(reasons))),
)
return valid
def _require_nonempty_content(content, name: str) -> None:
"""Raise if a concept body is missing or whitespace-only."""
if not isinstance(content, str) or not content.strip():
raise ValueError(f"LLM returned empty content for concept {name!r}")
def _filter_related_slugs(items: list) -> list[str]:
"""Keep only non-empty string slugs; warn about anything else."""
if not isinstance(items, list):
logger.warning("concepts plan: related was %s, expected list — dropping",
type(items).__name__)
return []
valid = [s for s in items if isinstance(s, str) and s.strip()]
if len(valid) < len(items):
bad_types = sorted({type(s).__name__ for s in items if not (isinstance(s, str) and s.strip())})
logger.warning(
"concepts plan: dropped %d malformed related item(s) (types: %s)",
len(items) - len(valid), ", ".join(bad_types),
)
return valid
def _filter_entity_items(
items: object, valid_types: frozenset | None = None
) -> list[dict]:
"""Validate entity create/update objects: require name+title, coerce type.
Each kept item is normalized to ``{"name", "title", "type"}`` where
``type`` falls back to ``"other"`` when missing or outside ``valid_types``
and ``title`` falls back to ``name``. ``valid_types`` defaults to the
module-level ``_ENTITY_TYPES`` so callers that don't thread a config-driven
set keep today's behavior.
"""
if valid_types is None:
valid_types = _ENTITY_TYPES
out: list[dict] = []
if not isinstance(items, list):
return out
for it in items:
if not isinstance(it, dict):
continue
name = it.get("name")
if not isinstance(name, str) or not name.strip():
continue
title = it.get("title") if isinstance(it.get("title"), str) else name
etype = it.get("type")
if not isinstance(etype, str) or etype not in valid_types:
etype = "other"
out.append({"name": name, "title": title, "type": etype})
return out
def _parse_entities_plan(parsed: object, valid_types: frozenset | None = None) -> dict:
"""Extract the entities group from a plan dict, with graceful fallback.
Returns ``{"create": [...], "update": [...], "related": [...]}``. A
missing/malformed ``entities`` key yields empty lists, so older or
partial LLM responses never raise.
"""
empty = {"create": [], "update": [], "related": []}
if not isinstance(parsed, dict):
return empty
group = parsed.get("entities")
if not isinstance(group, dict):
return empty
return {
"create": _filter_entity_items(group.get("create", []), valid_types),
"update": _filter_entity_items(group.get("update", []), valid_types),
"related": _filter_related_slugs(group.get("related", [])),
}
# ---------------------------------------------------------------------------
# File I/O helpers
# ---------------------------------------------------------------------------
def _read_wiki_context(wiki_dir: Path) -> tuple[str, list[str]]:
"""Read current index.md content and list of existing concept slugs."""
index_path = wiki_dir / "index.md"
index_content = index_path.read_text(encoding="utf-8") if index_path.exists() else ""
concepts_dir = wiki_dir / "concepts"
existing = sorted(p.stem for p in concepts_dir.glob("*.md")) if concepts_dir.exists() else []
return index_content, existing
def _resolve_description(fm: dict) -> str:
"""Return a non-empty description string from a frontmatter dict.
Checks ``description`` first, then the legacy ``brief`` key. Returns
an empty string when neither key holds a non-blank string value.
"""
for key in ("description", "brief"):
v = fm.get(key)
if isinstance(v, str) and v.strip():
return v.strip()
return ""
def _read_concept_briefs(wiki_dir: Path) -> str:
"""Read existing concept pages and return compact one-line summaries.
For each concept, reads the ``description:`` field (falling back to legacy
``brief:``) from YAML frontmatter if present; otherwise falls back to
truncating the first 150 chars of the body (newlines collapsed to spaces).
Formats each as ``- {slug}: {description}``.
Returns "(none yet)" if the concepts directory is missing or empty.
"""
concepts_dir = wiki_dir / "concepts"
if not concepts_dir.exists():
return "(none yet)"
md_files = sorted(concepts_dir.glob("*.md"))
if not md_files:
return "(none yet)"
lines: list[str] = []
for path in md_files:
text = path.read_text(encoding="utf-8")
fm_dict = frontmatter.parse(text)
brief = _resolve_description(fm_dict)
if not brief:
parts = frontmatter.split(text)
body = parts[1] if parts is not None else text
brief = body.strip().replace("\n", " ")[:150]
if brief:
lines.append(f"- {path.stem}: {brief}")
return "\n".join(lines) or "(none yet)"
def _read_entity_briefs(wiki_dir: Path) -> str:
"""Read existing entity pages as compact lines for the plan call.
Formats each as ``- {slug} ({type}, {n} sources) — {brief}``. The source
count is the cross-document recurrence signal the LLM uses to decide
create-vs-update and salience. Returns "(none yet)" when empty.
"""
entities_dir = wiki_dir / "entities"
if not entities_dir.exists():
return "(none yet)"
md_files = sorted(entities_dir.glob("*.md"))
if not md_files:
return "(none yet)"
lines: list[str] = []
for path in md_files:
text = path.read_text(encoding="utf-8")
fm_dict = frontmatter.parse(text)
brief = _resolve_description(fm_dict)
etype = str(fm_dict.get("type") or "").strip().lower() or "other"
n_sources = len(fm_dict["sources"]) if isinstance(fm_dict.get("sources"), list) else 0
if not brief:
parts = frontmatter.split(text)
body = parts[1] if parts is not None else text
brief = body.strip().replace("\n", " ")[:150]
suffix = f" — {brief}" if brief else ""
lines.append(f"- {path.stem} ({etype}, {n_sources} sources){suffix}")
return "\n".join(lines) or "(none yet)"
def _iter_h2_headings(lines: list[str]) -> list[tuple[int, str]]:
"""Return ``[(line_index, normalized_heading), ...]`` for every ATX H2.
A line counts as H2 when it starts with ``"## "`` (two hashes + space).
``normalized_heading`` is the line with trailing whitespace stripped, so
``"## Documents "`` normalizes to ``"## Documents"`` — letting callers
use exact-string comparison without tripping on stray whitespace.
Used by ``_get_section_bounds`` so heading lookup and the next-section
boundary share one scan and one normalization rule.
"""
return [
(i, line.rstrip())
for i, line in enumerate(lines)
if line.startswith("## ")
]
def _get_section_bounds(lines: list[str], heading: str) -> tuple[int, int] | None:
"""Return the [start, end) bounds for a Markdown H2 section.
Uses ``_iter_h2_headings`` so the same H2 detection that finds the
target heading also determines the section's end (the next H2). A
drifted ``"## Documents "`` matches ``"## Documents"`` because both
sides are normalized.
"""
headings = _iter_h2_headings(lines)
for k, (idx, normalized) in enumerate(headings):
if normalized == heading:
start = idx + 1
end = headings[k + 1][0] if k + 1 < len(headings) else len(lines)
return start, end
return None
def _ensure_h2_section(lines: list[str], heading: str, *, quiet: bool = False) -> None:
"""Ensure an H2 section ``heading`` exists in ``lines``; append if missing.
Recovers from hand-edited or drifted index.md files where the expected
section was removed or renamed — without this, downstream inserts would
silently no-op and entries would be dropped.
``quiet=True`` suppresses the drift warning. Use it when adding a section
is the normal, expected operation (e.g. a backlink helper creating a
``## Related Documents`` / ``## Entities`` section on a page for the first
time), as opposed to repairing a drifted index.
"""
if _get_section_bounds(lines, heading) is not None:
return
if not quiet:
logger.warning(
"Wiki page is missing %r section; appending it. "
"Check whether the file was hand-edited away from the canonical layout.",
heading,
)
while lines and lines[-1] == "":
lines.pop()
if lines:
lines.append("")
lines.append(heading)
lines.append("")
def _ensure_h2_section_before(
lines: list[str], heading: str, before: str,
) -> None:
"""Ensure H2 ``heading`` exists, inserting it just before ``before``.
If ``heading`` is already present, no-op. If ``before`` is absent, fall
back to :func:`_ensure_h2_section` (append at end). This keeps the
canonical index order (e.g. ``## Entities`` ahead of ``## Explorations``)
when recovering an older index.md that predates the section.
"""
if _get_section_bounds(lines, heading) is not None:
return
before_bounds = _get_section_bounds(lines, before)
if before_bounds is None:
_ensure_h2_section(lines, heading)
return
# ``start`` is the line after the ``before`` heading; insert the new
# section (heading + blank line) right before that heading line.
insert_at = before_bounds[0] - 1
logger.warning(
"Wiki index is missing %r section; inserting it before %r. "
"Check whether the file was hand-edited away from the canonical layout.",
heading, before,
)
lines[insert_at:insert_at] = [heading, ""]
def _section_contains_link(lines: list[str], heading: str, link: str) -> bool:
"""Check whether an index entry already exists inside the named section."""
bounds = _get_section_bounds(lines, heading)
if bounds is None:
return False
start, end = bounds
entry_prefix = f"- {link}"
return any(line.startswith(entry_prefix) for line in lines[start:end])
def _replace_section_entry(lines: list[str], heading: str, link: str, entry: str) -> bool:
"""Replace the first matching entry within a specific section."""
bounds = _get_section_bounds(lines, heading)
if bounds is None:
return False
start, end = bounds
entry_prefix = f"- {link}"
for i in range(start, end):
if lines[i].startswith(entry_prefix):
lines[i] = entry
return True
return False
def _insert_section_entry(lines: list[str], heading: str, entry: str) -> bool:
"""Insert a new entry at the top of a specific section."""
bounds = _get_section_bounds(lines, heading)
if bounds is None:
return False
start, _ = bounds
lines.insert(start, entry)
return True
def _remove_section_entry(lines: list[str], heading: str, link: str) -> bool:
"""Remove the first entry whose line starts with ``- {link}`` in the named
section. Returns True if an entry was removed.
Matching is intentionally strict (prefix-only, matching the canonical
bullet form written by ``_insert_section_entry`` and friends). An earlier
substring fallback could wrongly delete sibling bullets whose brief text
referenced the removed link.
"""
bounds = _get_section_bounds(lines, heading)
if bounds is None:
return False
start, end = bounds
entry_prefix = f"- {link}"
for i in range(start, end):
if lines[i].startswith(entry_prefix):
del lines[i]
return True
return False
def _write_summary(wiki_dir: Path, doc_name: str, summary: str,
doc_type: str = "short", description: str = "") -> None:
"""Write summary page with frontmatter."""
parts = frontmatter.split(summary)
if parts is not None:
_, summary = parts
summary = summary.lstrip("\n")
summaries_dir = wiki_dir / "summaries"
summaries_dir.mkdir(parents=True, exist_ok=True)
ext = "md" if doc_type == "short" else "json"
fm_lines = [_yaml_kv_line("type", "Summary")]
if description:
fm_lines.append(_yaml_kv_line("description", description))
fm_lines.append(f"doc_type: {doc_type}")
fm_lines.append(_yaml_kv_line("full_text", f"sources/{doc_name}.{ext}"))
fm_block = "---\n" + "\n".join(fm_lines) + "\n---\n\n"
atomic_write_text(summaries_dir / f"{doc_name}.md", fm_block + summary)
_SAFE_NAME_RE = re.compile(r'[^\w\-]')
def _sanitize_concept_name(name: str) -> str:
"""Sanitize a concept name for safe use as a filename."""
name = unicodedata.normalize("NFKC", name)
sanitized = _SAFE_NAME_RE.sub("-", name).strip("-")
return sanitized or "unnamed-concept"
_yaml_kv_line = frontmatter.kv_line
_yaml_list_line = frontmatter.list_line
_parse_yaml_list_value = frontmatter.parse_list_value
def _write_concept(wiki_dir: Path, name: str, content: str, source_file: str, is_update: bool, brief: str = "") -> None:
"""Write or update a concept page, managing the sources frontmatter."""
concepts_dir = wiki_dir / "concepts"
concepts_dir.mkdir(parents=True, exist_ok=True)
safe_name = _sanitize_concept_name(name)
path = (concepts_dir / f"{safe_name}.md").resolve()
if not path.is_relative_to(concepts_dir.resolve()):
logger.warning("Concept name escapes concepts dir: %s", name)
return
if is_update and path.exists():
existing = path.read_text(encoding="utf-8")
if source_file not in existing:
existing = _prepend_source_to_frontmatter(existing, source_file)
# Strip frontmatter from LLM content to avoid duplicate blocks
clean_parts = frontmatter.split(content)
clean = clean_parts[1].lstrip("\n") if clean_parts is not None else content
# Replace body with LLM rewrite (prompt asks for full rewrite, not delta)
ex_parts = frontmatter.split(existing)
if ex_parts is not None:
fm_block, _ = ex_parts
existing = fm_block + "\n" + clean
else:
# Malformed/absent frontmatter (opening ``---`` with no closing
# delimiter, or no frontmatter at all): rebuild valid frontmatter
# rather than writing a bare body. Recover any sources already
# listed in the broken block first.
recovered: list[str] = []
for ln in existing.split("\n"):
if ln.lstrip().startswith("sources:"):
parsed = _parse_yaml_list_value(ln)
if parsed:
recovered = parsed
break
merged = [source_file] + [s for s in recovered if s != source_file]
fm_lines = [
_yaml_kv_line("type", "Concept"),
_yaml_list_line("sources", merged),
]
if brief:
fm_lines.append(_yaml_kv_line("description", brief))
existing = frontmatter.block(fm_lines) + clean
atomic_write_text(path, existing)
return
# Guarantee type + refresh description on update; remove legacy brief:.
ex_parts2 = frontmatter.split(existing)
if ex_parts2 is not None:
fm_block, body = ex_parts2
fm_block = _set_fm_line(fm_block, "type", "Concept")
if brief:
fm_block = _set_fm_line(fm_block, "description", brief)
# Drop legacy brief: lines (migrated to description:).
fm_block = frontmatter.drop_line(fm_block, "brief")
existing = fm_block + body
atomic_write_text(path, existing)
else:
clean_parts = frontmatter.split(content)
if clean_parts is not None:
content = clean_parts[1].lstrip("\n")
fm_lines = [
_yaml_kv_line("type", "Concept"),
_yaml_list_line("sources", [source_file]),
]
if brief:
fm_lines.append(_yaml_kv_line("description", brief))
fm_block = "---\n" + "\n".join(fm_lines) + "\n---\n\n"
atomic_write_text(path, fm_block + content)
def _write_entity(
wiki_dir: Path, name: str, content: str, source_file: str,
is_update: bool, brief: str = "", type_: str = "other",
aliases: list[str] | None = None,
) -> None:
"""Write or update an entity page in entities/, managing frontmatter.
Frontmatter fields: ``sources`` (list), ``type`` (one of the entity
enum, capitalized on write), ``description`` (one-liner), and optional
``aliases`` (list, omitted when empty). On update the new source is prepended and the body replaced
with the LLM rewrite; ``type`` is preserved from the new write.
"""
entities_dir = wiki_dir / "entities"
entities_dir.mkdir(parents=True, exist_ok=True)
safe_name = _sanitize_concept_name(name)
path = (entities_dir / f"{safe_name}.md").resolve()
if not path.is_relative_to(entities_dir.resolve()):
logger.warning("Entity name escapes entities dir: %s", name)
return
# Strip any frontmatter the LLM body may carry.
clean_parts = frontmatter.split(content)
clean = clean_parts[1].lstrip("\n") if clean_parts is not None else content
def _build_entity_frontmatter(sources: list[str]) -> str:
fm_lines = [_yaml_list_line("sources", sources)]
fm_lines.append(_yaml_kv_line("type", (type_ or "other").title()))
if brief:
fm_lines.append(_yaml_kv_line("description", brief))
if aliases:
fm_lines.append(_yaml_list_line("aliases", aliases))
return "---\n" + "\n".join(fm_lines) + "\n---\n\n"
if is_update and path.exists():
existing = path.read_text(encoding="utf-8")
if source_file not in existing:
existing = _prepend_source_to_frontmatter(existing, source_file)
ex_parts = frontmatter.split(existing)
if ex_parts is not None:
fm_block, _ = ex_parts
fm_block = _set_fm_line(fm_block, "description", brief) if brief else fm_block
fm_block = _set_fm_line(fm_block, "type", type_.title()) if type_ else fm_block
# Drop any legacy ``brief:`` key (migrated to ``description:``),
# mirroring _write_concept's update path.
fm_block = frontmatter.drop_line(fm_block, "brief")
existing = fm_block + "\n" + clean
else:
# Malformed/absent frontmatter (opening ``---`` with no closing
# delimiter, or no frontmatter at all): rebuild valid frontmatter
# rather than writing a body-only page. Recover any sources already
# listed in the broken block first — otherwise a multi-source
# entity would be truncated to just this document.
recovered: list[str] = []
for ln in existing.split("\n"):
if ln.lstrip().startswith("sources:"):
parsed = _parse_yaml_list_value(ln)
if parsed:
recovered = parsed
break
merged = [source_file] + [s for s in recovered if s != source_file]
existing = _build_entity_frontmatter(merged) + clean
atomic_write_text(path, existing)
return
atomic_write_text(path, _build_entity_frontmatter([source_file]) + clean)