diff --git a/src/wikimind/engine/compiler.py b/src/wikimind/engine/compiler.py
index 86e158a6..fda2fbdc 100644
--- a/src/wikimind/engine/compiler.py
+++ b/src/wikimind/engine/compiler.py
@@ -14,7 +14,7 @@
import structlog
from slugify import slugify
from sqlalchemy.exc import SQLAlchemyError
-from sqlmodel import select
+from sqlmodel import col, select
from wikimind._datetime import utcnow_naive
from wikimind.config import get_settings
@@ -49,6 +49,7 @@
ReinforcementEvent,
RelationType,
Source,
+ SourceSpan,
TaskType,
TypedBacklinkSuggestion,
)
@@ -70,6 +71,8 @@
log = structlog.get_logger()
+_SPAN_PREVIEW_MAX_CHARS = 120
+
def _normalize_backlink_suggestions(raw: list[str | dict]) -> list[str]:
"""Normalize typed backlink suggestions to plain strings for CompilationResult.
@@ -203,6 +206,33 @@ def __init__(self, user_id: str):
# Compilation monitoring — set during compile(), read during save_article().
self._last_compilation_duration_ms: int | None = None
self._last_compilation_tokens: int | None = None
+ # Source spans loaded for the current compilation (issue #450 Phase 2).
+ # Populated by compile/compile_with_guidance, consumed by _persist_claims.
+ self._source_spans: list[SourceSpan] = []
+
+ async def _load_source_spans(
+ self,
+ source_id: str,
+ session: AsyncSession,
+ ) -> list[SourceSpan]:
+ """Load SourceSpan rows for a source, for inclusion in the compiler prompt.
+
+ Args:
+ source_id: The source UUID whose spans to load.
+ session: Async database session.
+
+ Returns:
+ List of SourceSpan instances ordered by creation time.
+ """
+ result = await session.execute(
+ select(SourceSpan)
+ .where(
+ SourceSpan.source_id == source_id,
+ SourceSpan.user_id == self.user_id,
+ )
+ .order_by(col(SourceSpan.created_at))
+ )
+ return list(result.scalars().all())
async def extract_takeaways(
self,
@@ -278,7 +308,13 @@ async def compile_with_guidance(
await session.commit()
return await self._compile_chunked(doc, session, progress_callback)
- user_prompt = self._build_user_prompt(doc)
+ # Load source spans for claim-level citation (issue #450 Phase 2).
+ spans: list[SourceSpan] = []
+ if doc.raw_source_id:
+ spans = await self._load_source_spans(doc.raw_source_id, session)
+ self._source_spans = spans
+
+ user_prompt = self._build_user_prompt(doc, spans=spans)
safe_guidance = _sanitize_guidance(guidance)
user_prompt += (
f"\n\nUSER GUIDANCE — weight the article toward these priorities:\n{safe_guidance}"
@@ -347,7 +383,13 @@ async def compile(
if doc.estimated_tokens > 80_000:
return await self._compile_chunked(doc, session, progress_callback)
- user_prompt = self._build_user_prompt(doc)
+ # Load source spans for claim-level citation (issue #450 Phase 2).
+ spans: list[SourceSpan] = []
+ if doc.raw_source_id:
+ spans = await self._load_source_spans(doc.raw_source_id, session)
+ self._source_spans = spans
+
+ user_prompt = self._build_user_prompt(doc, spans=spans)
# Concept ID registry injection: prevents concept fragmentation by
# telling the LLM which concepts already exist (issue #143, Phase 2).
@@ -425,8 +467,18 @@ async def compile(
)
return None
- def _build_user_prompt(self, doc: NormalizedDocument) -> str:
- """Build the user prompt for the LLM compiler."""
+ def _build_user_prompt(
+ self,
+ doc: NormalizedDocument,
+ spans: list[SourceSpan] | None = None,
+ ) -> str:
+ """Build the user prompt for the LLM compiler.
+
+ Args:
+ doc: Normalized document to compile.
+ spans: Optional source spans to include for claim-level citation.
+ When present, the LLM is instructed to cite span IDs per claim.
+ """
max_chars = get_settings().compiler.source_text_max_chars
meta = f"Title: {doc.title}"
if doc.author:
@@ -436,15 +488,29 @@ def _build_user_prompt(self, doc: NormalizedDocument) -> str:
if doc.raw_source_id:
meta += f"\nSource ID: {doc.raw_source_id}"
- return f"""{meta}
+ prompt = f"""{meta}
---
{doc.clean_text[:max_chars]}
----
+---"""
-Compile this into a wiki article following the JSON schema exactly."""
+ if spans:
+ max_span_chars = get_settings().compiler.source_text_max_chars // 4
+ span_section = "\n\n## Source Spans\n\nCite these span IDs in key_claims.source_span_ids:\n"
+ used_chars = 0
+ for span in spans:
+ preview = span.text[:_SPAN_PREVIEW_MAX_CHARS]
+ line = f'- {span.id}: "{preview}"\n'
+ if used_chars + len(line) > max_span_chars:
+ break
+ span_section += line
+ used_chars += len(line)
+ prompt += span_section
+
+ prompt += "\n\nCompile this into a wiki article following the JSON schema exactly."
+ return prompt
async def _compile_chunked(
self,
@@ -838,10 +904,29 @@ async def _persist_claims(
Each claim receives a numeric ``confidence_score`` computed from its
categorical confidence label and the number of backing sources
(issue #465).
+
+ Source span IDs returned by the LLM are validated against the actual
+ spans loaded during compilation. Invalid span IDs are silently
+ dropped to prevent hallucinated citations (issue #450 Phase 2).
"""
+ # Build set of valid span IDs for validation.
+ valid_span_ids = {s.id for s in self._source_spans}
+
for dto in result.key_claims:
claim_source_ids = dto.source_ids or [source.id]
confidence_level = dto.confidence.value if hasattr(dto.confidence, "value") else str(dto.confidence)
+
+ # Validate span IDs: keep only those that exist in the source.
+ raw_span_ids = dto.source_span_ids or []
+ validated_span_ids = [sid for sid in raw_span_ids if sid in valid_span_ids]
+ if len(validated_span_ids) < len(raw_span_ids):
+ rejected = set(raw_span_ids) - set(validated_span_ids)
+ log.warning(
+ "Rejected invalid span IDs from LLM",
+ article_id=article_id,
+ rejected_count=len(rejected),
+ )
+
claim = CompiledClaim(
article_id=article_id,
user_id=self.user_id,
@@ -855,6 +940,7 @@ async def _persist_claims(
),
quote=dto.quote,
source_ids=json.dumps(claim_source_ids),
+ source_span_ids=json.dumps(validated_span_ids),
)
session.add(claim)
await session.commit()
diff --git a/src/wikimind/engine/prompts.py b/src/wikimind/engine/prompts.py
index eefa886f..f790798e 100644
--- a/src/wikimind/engine/prompts.py
+++ b/src/wikimind/engine/prompts.py
@@ -40,7 +40,8 @@
"confidence": "sourced|inferred|opinion",
"subjects": ["canonical-subject-name"],
"source_ids": [""],
- "quote": "Optional direct quote under 15 words if the exact wording matters"
+ "quote": "Optional direct quote under 15 words if the exact wording matters",
+ "source_span_ids": ["span-uuid-1"]
}
],
"concepts": ["concept-name-1", "concept-name-2"],
@@ -63,6 +64,7 @@
- article_body must be substantive -- at least 300 words
- Never fabricate quotes or statistics not in the source
- For concepts: reuse existing concept names when they match your intent -- do not invent synonyms or near-duplicates
+- For source_span_ids: if the source material includes a "## Source Spans" section with span IDs, cite the span IDs that support each claim. Only use span IDs listed in that section. If no spans are provided, omit source_span_ids or use an empty list
Rich content preservation:
- Math: if the source contains mathematical expressions, reproduce them in LaTeX using $...$ for inline math and $$...$$ for display math blocks. Copy formulas verbatim from the source -- do not simplify or rewrite them.
diff --git a/src/wikimind/models/dto/compilation.py b/src/wikimind/models/dto/compilation.py
index a5aae073..37cecfd1 100644
--- a/src/wikimind/models/dto/compilation.py
+++ b/src/wikimind/models/dto/compilation.py
@@ -33,6 +33,7 @@ class CompiledClaimDTO(BaseModel):
predicate: str | None = None # LLM-extracted predicate
quote: str | None = None # Direct quote < 15 words if critical
source_ids: list[str] = [] # Source UUIDs supporting this claim
+ source_span_ids: list[str] = [] # SourceSpan UUIDs anchoring this claim (issue #450)
# ---------------------------------------------------------------------------
diff --git a/tests/unit/test_citations.py b/tests/unit/test_citations.py
index a0f9a36a..c337acbd 100644
--- a/tests/unit/test_citations.py
+++ b/tests/unit/test_citations.py
@@ -1,7 +1,8 @@
"""Tests for span-level citations (issue #450).
Covers the SourceSpan model, the CompiledClaim.source_span_ids field,
-the CitationService, and the GET /api/wiki/articles/{id}/citations endpoint.
+the CitationService, the GET /api/wiki/articles/{id}/citations endpoint,
+and Phase 2 claim-span linkage in the compiler.
"""
from __future__ import annotations
@@ -9,15 +10,24 @@
import hashlib
import json
import uuid
+from types import SimpleNamespace
+from unittest.mock import patch
import pytest
+from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession
from tests.conftest import TEST_USER_ID
+from wikimind.engine import base_compiler as base_compiler_mod
+from wikimind.engine.compiler import Compiler
from wikimind.models import (
Article,
+ CompilationResult,
CompiledClaim,
+ CompiledClaimDTO,
+ ConfidenceLevel,
LocatorKind,
+ NormalizedDocument,
Source,
SourceSpan,
)
@@ -398,3 +408,260 @@ async def test_citations_endpoint_with_spans(self, client, async_engine) -> None
assert len(data["claims"][0]["source_spans"]) == 1
assert data["claims"][0]["source_spans"][0]["id"] == span_id
assert data["claims"][0]["source_spans"][0]["locator_kind"] == "html-paragraph-offset"
+
+
+# ---------------------------------------------------------------------------
+# Phase 2: Compiler claim-span linkage tests (issue #450)
+# ---------------------------------------------------------------------------
+
+
+def _fake_settings() -> SimpleNamespace:
+ return SimpleNamespace(
+ data_dir="/tmp/wm-test",
+ compiler=SimpleNamespace(
+ max_tokens=8192,
+ source_text_max_chars=60000,
+ guidance_max_length=2000,
+ slug_max_attempts=1000,
+ ),
+ )
+
+
+def _make_compiler():
+ with (
+ patch.object(base_compiler_mod, "get_llm_router"),
+ patch.object(base_compiler_mod, "get_settings", return_value=_fake_settings()),
+ ):
+ return Compiler(user_id=TEST_USER_ID)
+
+
+class TestBuildUserPromptWithSpans:
+ """Verify _build_user_prompt includes span IDs when available."""
+
+ def test_prompt_without_spans(self) -> None:
+ c = _make_compiler()
+ doc = NormalizedDocument(
+ raw_source_id="src-1",
+ clean_text="Hello world",
+ title="Test",
+ estimated_tokens=10,
+ )
+ prompt = c._build_user_prompt(doc)
+ assert "Source Spans" not in prompt
+ assert "Compile this into a wiki article" in prompt
+
+ def test_prompt_with_spans(self) -> None:
+ c = _make_compiler()
+ doc = NormalizedDocument(
+ raw_source_id="src-1",
+ clean_text="Hello world",
+ title="Test",
+ estimated_tokens=10,
+ )
+ span = SourceSpan(
+ id="span-abc",
+ source_id="src-1",
+ user_id=TEST_USER_ID,
+ locator_kind=LocatorKind.TEXT_BYTE_RANGE,
+ locator={"start": 0, "end": 11},
+ text="Hello world",
+ fingerprint=_fingerprint("Hello world"),
+ )
+ prompt = c._build_user_prompt(doc, spans=[span])
+ assert "## Source Spans" in prompt
+ assert "span-abc" in prompt
+ assert "Hello world" in prompt
+
+
+class TestCompiledClaimDTOSpanIds:
+ """Verify CompiledClaimDTO accepts source_span_ids."""
+
+ def test_default_empty(self) -> None:
+ dto = CompiledClaimDTO(claim="X", confidence=ConfidenceLevel.SOURCED)
+ assert dto.source_span_ids == []
+
+ def test_with_span_ids(self) -> None:
+ dto = CompiledClaimDTO(
+ claim="X",
+ confidence=ConfidenceLevel.SOURCED,
+ source_span_ids=["span-1", "span-2"],
+ )
+ assert dto.source_span_ids == ["span-1", "span-2"]
+
+ def test_parsed_from_json(self) -> None:
+ data = {
+ "claim": "X",
+ "confidence": "sourced",
+ "source_span_ids": ["span-a"],
+ }
+ dto = CompiledClaimDTO(**data)
+ assert dto.source_span_ids == ["span-a"]
+
+
+class TestPersistClaimsWithSpanValidation:
+ """Verify _persist_claims validates and stores span IDs."""
+
+ @pytest.mark.asyncio
+ async def test_valid_span_ids_are_persisted(self, db_session: AsyncSession) -> None:
+ # Set up source, article, and spans
+ source_id = _uid()
+ source = Source(id=source_id, user_id=TEST_USER_ID, source_type="text", title="S")
+ db_session.add(source)
+ await db_session.flush()
+
+ article_id = _uid()
+ article = Article(
+ id=article_id,
+ user_id=TEST_USER_ID,
+ slug="persist-test",
+ title="Persist Test",
+ file_path="wiki/persist-test.md",
+ )
+ db_session.add(article)
+ await db_session.flush()
+
+ span_id = _uid()
+ span = SourceSpan(
+ id=span_id,
+ source_id=source_id,
+ user_id=TEST_USER_ID,
+ locator_kind=LocatorKind.TEXT_BYTE_RANGE,
+ locator={"start": 0, "end": 10},
+ text="Some text.",
+ fingerprint=_fingerprint("Some text."),
+ )
+ db_session.add(span)
+ await db_session.flush()
+
+ # Create compiler with valid spans loaded
+ c = _make_compiler()
+ c._source_spans = [span]
+
+ result = CompilationResult(
+ title="T",
+ summary="S. S.",
+ key_claims=[
+ CompiledClaimDTO(
+ claim="Claim one",
+ confidence=ConfidenceLevel.SOURCED,
+ source_span_ids=[span_id],
+ ),
+ ],
+ concepts=[],
+ backlink_suggestions=[],
+ open_questions=[],
+ article_body="body",
+ )
+
+ await c._persist_claims(article_id, result, source, db_session)
+
+ # Verify the claim was persisted with span IDs
+ stmt = select(CompiledClaim).where(CompiledClaim.article_id == article_id)
+ claims = (await db_session.exec(stmt)).all()
+ assert len(claims) == 1
+ assert json.loads(claims[0].source_span_ids) == [span_id]
+
+ @pytest.mark.asyncio
+ async def test_invalid_span_ids_are_rejected(self, db_session: AsyncSession) -> None:
+ source_id = _uid()
+ source = Source(id=source_id, user_id=TEST_USER_ID, source_type="text", title="S")
+ db_session.add(source)
+ await db_session.flush()
+
+ article_id = _uid()
+ article = Article(
+ id=article_id,
+ user_id=TEST_USER_ID,
+ slug="reject-test",
+ title="Reject Test",
+ file_path="wiki/reject-test.md",
+ )
+ db_session.add(article)
+ await db_session.flush()
+
+ valid_span_id = _uid()
+ span = SourceSpan(
+ id=valid_span_id,
+ source_id=source_id,
+ user_id=TEST_USER_ID,
+ locator_kind=LocatorKind.TEXT_BYTE_RANGE,
+ locator={"start": 0, "end": 5},
+ text="Valid",
+ fingerprint=_fingerprint("Valid"),
+ )
+ db_session.add(span)
+ await db_session.flush()
+
+ c = _make_compiler()
+ c._source_spans = [span]
+
+ fake_span_id = _uid()
+ result = CompilationResult(
+ title="T",
+ summary="S. S.",
+ key_claims=[
+ CompiledClaimDTO(
+ claim="Claim with mixed IDs",
+ confidence=ConfidenceLevel.SOURCED,
+ source_span_ids=[valid_span_id, fake_span_id],
+ ),
+ ],
+ concepts=[],
+ backlink_suggestions=[],
+ open_questions=[],
+ article_body="body",
+ )
+
+ await c._persist_claims(article_id, result, source, db_session)
+
+ stmt = select(CompiledClaim).where(CompiledClaim.article_id == article_id)
+ claims = (await db_session.exec(stmt)).all()
+ assert len(claims) == 1
+ persisted_span_ids = json.loads(claims[0].source_span_ids)
+ # Only the valid span ID should be persisted
+ assert persisted_span_ids == [valid_span_id]
+ assert fake_span_id not in persisted_span_ids
+
+ @pytest.mark.asyncio
+ async def test_no_spans_loaded_all_rejected(self, db_session: AsyncSession) -> None:
+ source_id = _uid()
+ source = Source(id=source_id, user_id=TEST_USER_ID, source_type="text", title="S")
+ db_session.add(source)
+ await db_session.flush()
+
+ article_id = _uid()
+ article = Article(
+ id=article_id,
+ user_id=TEST_USER_ID,
+ slug="empty-spans-test",
+ title="Empty Spans Test",
+ file_path="wiki/empty-spans-test.md",
+ )
+ db_session.add(article)
+ await db_session.flush()
+
+ c = _make_compiler()
+ c._source_spans = [] # No spans loaded
+
+ result = CompilationResult(
+ title="T",
+ summary="S. S.",
+ key_claims=[
+ CompiledClaimDTO(
+ claim="Claim with hallucinated spans",
+ confidence=ConfidenceLevel.SOURCED,
+ source_span_ids=[_uid()],
+ ),
+ ],
+ concepts=[],
+ backlink_suggestions=[],
+ open_questions=[],
+ article_body="body",
+ )
+
+ await c._persist_claims(article_id, result, source, db_session)
+
+ stmt = select(CompiledClaim).where(CompiledClaim.article_id == article_id)
+ claims = (await db_session.exec(stmt)).all()
+ assert len(claims) == 1
+ assert json.loads(claims[0].source_span_ids) == []
diff --git a/tests/unit/test_concept_recompilation.py b/tests/unit/test_concept_recompilation.py
index a72c03ff..a1c8f543 100644
--- a/tests/unit/test_concept_recompilation.py
+++ b/tests/unit/test_concept_recompilation.py
@@ -286,6 +286,7 @@ async def test_replace_calls_maybe_trigger(self, db_session, tmp_path):
compiler._last_typed_suggestions = {}
compiler._last_compilation_duration_ms = None
compiler._last_compilation_tokens = None
+ compiler._source_spans = []
with (
patch(