feat: wire brief+content JSON through compile pipeline to index and frontmatter

Author: KylinMountain

openkb/agent/compiler.py

@@ -497,6 +497,7 @@ async def _compile_concepts(
     summary: str,
     doc_name: str,
     max_concurrency: int,
+    doc_brief: str = "",
 ) -> None:
     """Shared Steps 2-4: concepts plan → generate/update → index.
 
@@ -546,11 +547,11 @@ async def _compile_concepts(
     # --- Step 3: Generate/update concept pages concurrently (A cached) ---
     semaphore = asyncio.Semaphore(max_concurrency)
 
-    async def _gen_create(concept: dict) -> tuple[str, str, bool]:
+    async def _gen_create(concept: dict) -> tuple[str, str, bool, str]:
         name = concept["name"]
         title = concept.get("title", name)
         async with semaphore:
-            page_content = await _llm_call_async(model, [
+            raw = await _llm_call_async(model, [
                 system_msg,
                 doc_msg,
                 {"role": "assistant", "content": summary},
@@ -559,9 +560,15 @@ async def _gen_create(concept: dict) -> tuple[str, str, bool]:
                     update_instruction="",
                 )},
             ], f"concept:{name}")
-        return name, page_content, False
-
-    async def _gen_update(concept: dict) -> tuple[str, str, bool]:
+        try:
+            parsed = _parse_json(raw)
+            brief = parsed.get("brief", "")
+            content = parsed.get("content", raw)
+        except (json.JSONDecodeError, ValueError):
+            brief, content = "", raw
+        return name, content, False, brief
+
+    async def _gen_update(concept: dict) -> tuple[str, str, bool, str]:
         name = concept["name"]
         title = concept.get("title", name)
         concept_path = wiki_dir / "concepts" / f"{name}.md"
@@ -575,7 +582,7 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool]:
         else:
             existing_content = "(page not found — create from scratch)"
         async with semaphore:
-            page_content = await _llm_call_async(model, [
+            raw = await _llm_call_async(model, [
                 system_msg,
                 doc_msg,
                 {"role": "assistant", "content": summary},
@@ -584,13 +591,20 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool]:
                     existing_content=existing_content,
                 )},
             ], f"update:{name}")
-        return name, page_content, True
+        try:
+            parsed = _parse_json(raw)
+            brief = parsed.get("brief", "")
+            content = parsed.get("content", raw)
+        except (json.JSONDecodeError, ValueError):
+            brief, content = "", raw
+        return name, content, True, brief
 
     tasks = []
     tasks.extend(_gen_create(c) for c in create_items)
     tasks.extend(_gen_update(c) for c in update_items)
 
     concept_names: list[str] = []
+    concept_briefs_map: dict[str, str] = {}
 
     if tasks:
         total = len(tasks)
@@ -603,9 +617,11 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool]:
             if isinstance(r, Exception):
                 logger.warning("Concept generation failed: %s", r)
                 continue
-            name, page_content, is_update = r
-            _write_concept(wiki_dir, name, page_content, source_file, is_update)
+            name, page_content, is_update, brief = r
+            _write_concept(wiki_dir, name, page_content, source_file, is_update, brief=brief)
             concept_names.append(name)
+            if brief:
+                concept_briefs_map[name] = brief
 
     # --- Step 3b: Process related items (code only, no LLM) ---
     for slug in related_items:
@@ -618,7 +634,8 @@ async def _gen_update(concept: dict) -> tuple[str, str, bool]:
         _backlink_concepts(wiki_dir, doc_name, all_concept_slugs)
 
     # --- Step 4: Update index (code only) ---
-    _update_index(wiki_dir, doc_name, concept_names)
+    _update_index(wiki_dir, doc_name, concept_names,
+                  doc_brief=doc_brief, concept_briefs=concept_briefs_map)
 
 
 async def compile_short_doc(
@@ -653,13 +670,20 @@ async def compile_short_doc(
     )}
 
     # --- Step 1: Generate summary ---
-    summary = _llm_call(model, [system_msg, doc_msg], "summary")
-    _write_summary(wiki_dir, doc_name, source_file, summary)
+    summary_raw = _llm_call(model, [system_msg, doc_msg], "summary")
+    try:
+        summary_parsed = _parse_json(summary_raw)
+        doc_brief = summary_parsed.get("brief", "")
+        summary = summary_parsed.get("content", summary_raw)
+    except (json.JSONDecodeError, ValueError):
+        doc_brief = ""
+        summary = summary_raw
+    _write_summary(wiki_dir, doc_name, source_file, summary, brief=doc_brief)
 
     # --- Steps 2-4: Concept plan → generate/update → index ---
     await _compile_concepts(
         wiki_dir, kb_dir, model, system_msg, doc_msg,
-        summary, doc_name, max_concurrency,
+        summary, doc_name, max_concurrency, doc_brief=doc_brief,
     )
 
 
@@ -669,6 +693,7 @@ async def compile_long_doc(
     doc_id: str,
     kb_dir: Path,
     model: str,
+    doc_description: str = "",
     max_concurrency: int = DEFAULT_COMPILE_CONCURRENCY,
 ) -> None:
     """Compile a long (PageIndex) document's concepts and index.
@@ -700,5 +725,5 @@ async def compile_long_doc(
     # --- Steps 2-4: Concept plan → generate/update → index ---
     await _compile_concepts(
         wiki_dir, kb_dir, model, system_msg, doc_msg,
-        overview, doc_name, max_concurrency,
+        overview, doc_name, max_concurrency, doc_brief=doc_description,
     )

openkb/cli.py

@@ -132,7 +132,8 @@ def _add_single_file(file_path: Path, kb_dir: Path) -> None:
         for attempt in range(2):
             try:
                 asyncio.run(
-                    compile_long_doc(doc_name, summary_path, index_result.doc_id, kb_dir, model)
+                    compile_long_doc(doc_name, summary_path, index_result.doc_id, kb_dir, model,
+                                     doc_description=index_result.description)
                 )
                 break
             except Exception as exc:

tests/test_compiler.py

@@ -487,13 +487,19 @@ async def test_full_pipeline(self, tmp_path):
         (tmp_path / "raw").mkdir()
         (tmp_path / "raw" / "test-doc.pdf").write_bytes(b"fake")
 
-        summary_response = "# Summary\n\nThis document discusses transformers."
+        summary_response = json.dumps({
+            "brief": "Discusses transformers",
+            "content": "# Summary\n\nThis document discusses transformers.",
+        })
         concepts_list_response = json.dumps({
             "create": [{"name": "transformer", "title": "Transformer"}],
             "update": [],
             "related": [],
         })
-        concept_page_response = "# Transformer\n\nA neural network architecture."
+        concept_page_response = json.dumps({
+            "brief": "NN architecture using self-attention",
+            "content": "# Transformer\n\nA neural network architecture.",
+        })
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(
@@ -534,7 +540,7 @@ async def test_handles_bad_json(self, tmp_path):
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(
-                side_effect=_mock_completion(["Summary text", "not valid json"])
+                side_effect=_mock_completion(["Plain summary text", "not valid json"])
             )
             # Should not raise
             await compile_short_doc("doc", source_path, tmp_path, "gpt-4o-mini")
@@ -567,7 +573,10 @@ async def test_full_pipeline(self, tmp_path):
             "update": [],
             "related": [],
         })
-        concept_page_response = "# Deep Learning\n\nA subfield of ML."
+        concept_page_response = json.dumps({
+            "brief": "Subfield of ML using neural networks",
+            "content": "# Deep Learning\n\nA subfield of ML.",
+        })
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(
@@ -624,8 +633,14 @@ async def test_create_and_update_flow(self, tmp_path):
             "update": [{"name": "attention", "title": "Attention"}],
             "related": [],
         })
-        create_page_response = "# Flash Attention\n\nAn efficient attention algorithm."
-        update_page_response = "# Attention\n\nUpdated content with new info."
+        create_page_response = json.dumps({
+            "brief": "Efficient attention algorithm",
+            "content": "# Flash Attention\n\nAn efficient attention algorithm.",
+        })
+        update_page_response = json.dumps({
+            "brief": "Updated attention mechanism",
+            "content": "# Attention\n\nUpdated content with new info.",
+        })
 
         system_msg = {"role": "system", "content": "You are a wiki agent."}
         doc_msg = {"role": "user", "content": "Document about attention mechanisms."}
@@ -720,7 +735,10 @@ async def test_fallback_list_format(self, tmp_path):
         plan_response = json.dumps([
             {"name": "attention", "title": "Attention"},
         ])
-        concept_page_response = "# Attention\n\nA mechanism for focusing."
+        concept_page_response = json.dumps({
+            "brief": "A mechanism for focusing",
+            "content": "# Attention\n\nA mechanism for focusing.",
+        })
 
         system_msg = {"role": "system", "content": "You are a wiki agent."}
         doc_msg = {"role": "user", "content": "Document content."}
@@ -744,3 +762,57 @@ async def test_fallback_list_format(self, tmp_path):
         att_text = att_path.read_text()
         assert "sources: [test-doc.pdf]" in att_text
         assert "Attention" in att_text
+
+
+class TestBriefIntegration:
+    @pytest.mark.asyncio
+    async def test_short_doc_briefs_in_index_and_frontmatter(self, tmp_path):
+        wiki = tmp_path / "wiki"
+        (wiki / "sources").mkdir(parents=True)
+        (wiki / "summaries").mkdir(parents=True)
+        (wiki / "concepts").mkdir(parents=True)
+        (wiki / "index.md").write_text(
+            "# Index\n\n## Documents\n\n## Concepts\n\n## Explorations\n",
+            encoding="utf-8",
+        )
+        source_path = wiki / "sources" / "test-doc.md"
+        source_path.write_text("# Test Doc\n\nContent.", encoding="utf-8")
+        (tmp_path / ".openkb").mkdir()
+        (tmp_path / "raw").mkdir()
+        (tmp_path / "raw" / "test-doc.pdf").write_bytes(b"fake")
+
+        summary_resp = json.dumps({
+            "brief": "A paper about transformers",
+            "content": "# Summary\n\nThis paper discusses transformers.",
+        })
+        plan_resp = json.dumps({
+            "create": [{"name": "transformer", "title": "Transformer"}],
+            "update": [],
+            "related": [],
+        })
+        concept_resp = json.dumps({
+            "brief": "NN architecture using self-attention",
+            "content": "# Transformer\n\nA neural network architecture.",
+        })
+
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.completion = MagicMock(
+                side_effect=_mock_completion([summary_resp, plan_resp])
+            )
+            mock_litellm.acompletion = AsyncMock(
+                side_effect=_mock_acompletion([concept_resp])
+            )
+            await compile_short_doc("test-doc", source_path, tmp_path, "gpt-4o-mini")
+
+        # Summary frontmatter has brief
+        summary_text = (wiki / "summaries" / "test-doc.md").read_text()
+        assert "brief: A paper about transformers" in summary_text
+
+        # Concept frontmatter has brief
+        concept_text = (wiki / "concepts" / "transformer.md").read_text()
+        assert "brief: NN architecture using self-attention" in concept_text
+
+        # Index has briefs
+        index_text = (wiki / "index.md").read_text()
+        assert "— A paper about transformers" in index_text
+        assert "— NN architecture using self-attention" in index_text