feat: honor canonical lexemes in consolidated export (#657)

Co-authored-by: TarahAssistant <265369510+ArdeleanLucas@users.noreply.github.com>

Author: TarahAssistant

python/ai/tools/export_tools.py

@@ -653,7 +653,8 @@ def export_lingpy_tsv(tools: "ParseChatTools", args: Dict[str, Any]) -> Dict[str
             try:
                 from compare.consolidated_matrix import build_wordlist_rows, wordlist_rows_to_tsv
                 cognate_sets, meta, id_to_key, _speakers, allowed_ids = _consolidated_sets(tools, concept_tag)
-                rows = build_wordlist_rows(tools.annotations_dir, cognate_sets, id_to_key, allowed_ids)
+                canonical_lexemes = (_read_json_file(tools.enrichments_path, {}).get("manual_overrides") or {}).get("canonical_lexemes")
+                rows = build_wordlist_rows(tools.annotations_dir, cognate_sets, id_to_key, allowed_ids, canonical_lexemes)
                 content = wordlist_rows_to_tsv(rows)
             except ChatToolError:
                 raise

python/ai/tools/test_consolidated_export_wiring.py

@@ -6,12 +6,14 @@
 import re
 import sys
 import xml.etree.ElementTree as ET
+from typing import Any
 
 import pytest
 
 sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2]))
 
 from ai.chat_tools import ParseChatTools  # noqa: E402  (import first; avoids export_tools load cycle)
+from ai.workflow_tools import WorkflowTools  # noqa: E402
 
 
 def _seed_workspace(root: pathlib.Path) -> None:
@@ -53,6 +55,73 @@ def write(speaker: str, concept_id: str, ipa: str) -> None:
     write("Spk3", "150", "kalan")  # tagged under the duplicate JBIL id
 
 
+def _write_annotation(root: pathlib.Path, speaker: str, intervals: list[dict[str, Any]]) -> None:
+    ann = root / "annotations"
+    ann.mkdir(exist_ok=True)
+    concept_intervals = []
+    ipa_intervals = []
+    for index, item in enumerate(intervals):
+        start = float(index)
+        end = start + 0.5
+        concept_intervals.append(
+            {"text": item.get("label", ""), "concept_id": item["concept_id"], "start": start, "end": end}
+        )
+        ipa_intervals.append({"text": item["ipa"], "start": start, "end": end})
+    (ann / f"{speaker}.parse.json").write_text(
+        json.dumps({"speaker": speaker, "tiers": {"concept": {"intervals": concept_intervals}, "ipa": {"intervals": ipa_intervals}}}),
+        encoding="utf-8",
+    )
+
+
+def _seed_roundtrip_workspace(root: pathlib.Path) -> None:
+    (root / "concepts.csv").write_text(
+        "id,concept_en,source_survey,source_item\n"
+        "101,hand,KLQ,1.1\n"
+        "201,hand,JBIL,101\n"
+        "102,foot,KLQ,1.2\n"
+        "202,foot,JBIL,102\n"
+        "103,salt,KLQ,1.3\n"
+        "203,salt,JBIL,103\n",
+        encoding="utf-8",
+    )
+    groups = {"A": ["Spk1", "Spk2"], "B": ["Spk3"]}
+    (root / "parse-enrichments.json").write_text(
+        json.dumps({
+            "manual_overrides": {
+                "cognate_sets": {
+                    "101": groups,
+                    "201": dict(groups),
+                    "102": {"A": ["Spk1"], "B": ["Spk2", "Spk3"]},
+                    "202": {"A": ["Spk1"], "B": ["Spk2", "Spk3"]},
+                    # Disagreement safety: this overlapping gloss must stay split.
+                    "103": {"A": ["Spk1", "Spk2"], "B": ["Spk3"]},
+                    "203": {"A": ["Spk1"], "B": ["Spk2", "Spk3"]},
+                },
+                # Spk1 has both survey forms for hand; the explicit canonical
+                # choice should select row 201 rather than the earlier 101 form.
+                "canonical_lexemes": {"bundle:hand": {"Spk1": {"csv_row_id": "201", "source": "manual"}}},
+            }
+        }),
+        encoding="utf-8",
+    )
+    (root / "project.json").write_text(
+        json.dumps({"speakers": {"Spk1": {}, "Spk2": {}, "Spk3": {}}}), encoding="utf-8"
+    )
+    _write_annotation(
+        root,
+        "Spk1",
+        [
+            {"concept_id": "101", "label": "hand", "ipa": "old-hand"},
+            {"concept_id": "201", "label": "hand", "ipa": "chosen-hand"},
+            {"concept_id": "102", "label": "foot", "ipa": "foot-one"},
+            {"concept_id": "103", "label": "salt", "ipa": "salt-one"},
+            {"concept_id": "203", "label": "salt", "ipa": "salt-two"},
+        ],
+    )
+    _write_annotation(root, "Spk2", [{"concept_id": "201", "label": "hand", "ipa": "hand-two"}, {"concept_id": "202", "label": "foot", "ipa": "foot-two"}])
+    _write_annotation(root, "Spk3", [{"concept_id": "101", "label": "hand", "ipa": "hand-three"}, {"concept_id": "202", "label": "foot", "ipa": "foot-three"}])
+
+
 def test_export_nexus_consolidates_and_tag_filters(tmp_path: pathlib.Path) -> None:
     _seed_workspace(tmp_path)
     tools = ParseChatTools(project_root=tmp_path)
@@ -115,3 +184,35 @@ def test_export_beast2_xml_rejects_bad_chain_length(tmp_path: pathlib.Path) -> N
     tools = ParseChatTools(project_root=tmp_path)
     with pytest.raises(Exception):
         tools._tool_export_beast2_xml({"consolidate": True, "chainLength": 0, "dryRun": True})
+
+
+def test_complete_export_consolidate_roundtrip_collapses_n_not_2n(tmp_path: pathlib.Path) -> None:
+    _seed_roundtrip_workspace(tmp_path)
+    workflow = WorkflowTools(project_root=tmp_path)
+
+    payload = workflow.execute(
+        "export_complete_lingpy_dataset",
+        {"with_contact_lexemes": False, "consolidate": True, "outputDir": "exports/roundtrip", "dryRun": False},
+    )["result"]
+
+    nexus_text = (tmp_path / "exports" / "roundtrip" / "dataset.nex").read_text(encoding="utf-8")
+    wordlist_text = (tmp_path / "exports" / "roundtrip" / "wordlist.tsv").read_text(encoding="utf-8")
+    nchar = int(re.search(r"NCHAR=(\d+)", nexus_text).group(1))
+    concept_values = [line.split("\t")[1] for line in wordlist_text.splitlines()[1:]]
+    hand_rows = [line.split("\t") for line in wordlist_text.splitlines()[1:] if line.split("\t")[1] == "hand"]
+
+    assert payload["consolidated"] is True
+    assert payload["stages"][0]["payload"]["consolidation"]["collapsed_groups"] == 2
+    assert payload["stages"][0]["payload"]["consolidation"]["concept_count"] == 4
+    assert payload["stages"][0]["payload"]["consolidation"]["needs_recluster_groups"] == 1
+    assert any("kept separate" in warning for warning in payload["warnings"])
+    # N=2 safe overlapping meanings (hand + foot) produce N consolidated concept
+    # blocks, not 2N per-survey duplicates; disagreeing salt remains split.
+    assert "hand#" not in nexus_text
+    assert "foot#" not in nexus_text
+    assert "salt#103" in nexus_text and "salt#203" in nexus_text
+    assert nchar == 8  # hand(2) + foot(2) + safe split salt#103(2) + salt#203(2)
+    assert concept_values.count("hand") == 3
+    assert len({row[2] for row in hand_rows}) == len(hand_rows)
+    assert any(row[3] == "chosen-hand" for row in hand_rows)
+    assert all(row[3] != "old-hand" for row in hand_rows)

python/compare/consolidated_matrix.py

@@ -238,6 +238,7 @@ def build_wordlist_rows(
     cognate_sets: CognateSets,
     id_to_key: Mapping[str, str],
     allowed_ids: Optional[Set[str]] = None,
+    canonical_lexemes: Optional[Mapping[str, Any]] = None,
 ) -> List[Tuple[Any, ...]]:
     """Build LingPy wordlist rows remapped onto consolidated canonical concepts.
 
@@ -248,6 +249,22 @@ def build_wordlist_rows(
     forms_by_concept, _speakers = load_annotations(annotations_dir)
     cogid_lookup, cogid_group_by_speaker = _build_cogid_lookup(cognate_sets)
 
+    selected_row_by_key_speaker: Dict[Tuple[str, str], str] = {}
+    if isinstance(canonical_lexemes, Mapping):
+        for speaker_map in canonical_lexemes.values():
+            if not isinstance(speaker_map, Mapping):
+                continue
+            for speaker, selection in speaker_map.items():
+                if not isinstance(selection, Mapping):
+                    continue
+                row_id = str(selection.get("csv_row_id") or "").strip()
+                if not row_id:
+                    continue
+                key = id_to_key.get(row_id)
+                if not key:
+                    continue
+                selected_row_by_key_speaker[(key, _speaker_lookup_key(str(speaker)))] = row_id
+
     # Remap and de-dup forms onto canonical gloss keys.
     by_key_speaker: Dict[Tuple[str, str], Any] = {}
     for concept_id, records in forms_by_concept.items():
@@ -258,7 +275,17 @@ def build_wordlist_rows(
             sk = _speaker_lookup_key(record.speaker)
             slot = (key, sk)
             existing = by_key_speaker.get(slot)
-            if existing is None or record.start_sec < existing.start_sec:
+            selected_row_id = selected_row_by_key_speaker.get(slot)
+            if existing is None:
+                by_key_speaker[slot] = record
+            elif selected_row_id:
+                record_is_selected = str(record.concept_id) == selected_row_id
+                existing_is_selected = str(existing.concept_id) == selected_row_id
+                if record_is_selected and not existing_is_selected:
+                    by_key_speaker[slot] = record
+                elif record_is_selected == existing_is_selected and record.start_sec < existing.start_sec:
+                    by_key_speaker[slot] = record
+            elif record.start_sec < existing.start_sec:
                 by_key_speaker[slot] = record
 
     forms_by_key: Dict[str, List[Any]] = {}