feat: retain entity relationship evidence in extracted events

Pigbibi · codex · Pigbibi · commit 1d21f8b7b437 · 2026-07-13T16:42:05.000+08:00
Co-Authored-By: Codex &lt;noreply@openai.com&gt;
diff --git a/src/political_event_tracking_research/official_event_import.py b/src/political_event_tracking_research/official_event_import.py
@@ -45,6 +45,9 @@ class OfficialRecord:
     direction: str
     source_url: str
     summary: str
+    entity_match_type: str = "unverified"
+    match_evidence: str = ""
+    relationship_type: str = "unverified"
 
 
 def slug(value: str) -> str:
@@ -76,6 +79,9 @@ def load_official_records(path: str | Path) -> list[OfficialRecord]:
                 direction=row.get("direction", ""),
                 source_url=row["source_url"],
                 summary=row.get("summary", ""),
+                entity_match_type=row.get("entity_match_type", "unverified") or "unverified",
+                match_evidence=row.get("match_evidence", ""),
+                relationship_type=row.get("relationship_type", "unverified") or "unverified",
             )
         )
     return records
@@ -87,6 +93,11 @@ def validate_record(record: OfficialRecord) -> None:
         raise ValueError("record_id is required")
     if not record.symbol.strip():
         raise ValueError(f"{record.record_id}: symbol is required")
+    allowed_match_types = {"issuer", "direct_beneficiary", "industry_context", "unverified"}
+    if record.entity_match_type not in allowed_match_types:
+        raise ValueError(f"{record.record_id}: unsupported entity_match_type {record.entity_match_type!r}")
+    if record.relationship_type not in allowed_match_types:
+        raise ValueError(f"{record.record_id}: unsupported relationship_type {record.relationship_type!r}")
     if record.event_type not in ALLOWED_EVENT_TYPES:
         raise ValueError(f"{record.record_id}: unsupported event_type {record.event_type!r}")
     if record.source_type in GOVERNMENT_SOURCE_TYPES:
@@ -125,6 +136,9 @@ def normalize_records(records: list[OfficialRecord]) -> list[dict[str, str]]:
                 "confidence": confidence_for_source(record),
                 "source_url": record.source_url,
                 "notes": record.summary,
+                "entity_match_type": record.entity_match_type,
+                "match_evidence": record.match_evidence,
+                "relationship_type": record.relationship_type,
             }
         )
     rows.sort(key=lambda row: (row["event_date"], row["symbol"], row["event_id"]))
@@ -136,7 +150,19 @@ def import_official_events(input_path: str | Path, output_path: str | Path) -> l
     rows = normalize_records(records)
     write_csv_rows(
         output_path,
-        ["event_id", "event_date", "symbol", "event_type", "direction", "confidence", "source_url", "notes"],
+        [
+            "event_id",
+            "event_date",
+            "symbol",
+            "event_type",
+            "direction",
+            "confidence",
+            "source_url",
+            "notes",
+            "entity_match_type",
+            "match_evidence",
+            "relationship_type",
+        ],
         rows,
     )
     return rows
diff --git a/src/political_event_tracking_research/source_mention_extract.py b/src/political_event_tracking_research/source_mention_extract.py
@@ -10,6 +10,32 @@
 from .official_event_import import OfficialRecord, normalize_records
 
 
+GENERIC_ALIAS_DENYLIST = frozenset(
+    {
+        "advanced nuclear",
+        "bitcoin treasury",
+        "chips act",
+        "crypto assets",
+        "cybersecurity",
+        "digital assets",
+        "energy-related infrastructure",
+        "foundry",
+        "nuclear energy",
+        "nuclear reactor",
+        "semiconductor manufacturing",
+        "strategy",
+        "tokenization",
+    }
+)
+
+RELATIONSHIP_PRIORITY = {
+    "unverified": 0,
+    "industry_context": 1,
+    "issuer": 2,
+    "direct_beneficiary": 3,
+}
+
+
 @dataclass(frozen=True)
 class RawSourceItem:
     item_id: str
@@ -24,6 +50,7 @@ class RawSourceItem:
 class MentionAlias:
     symbol: str
     aliases: tuple[str, ...]
+    name: str = ""
 
 
 def load_raw_items(path: str | Path) -> list[RawSourceItem]:
@@ -51,10 +78,13 @@ def load_aliases(path: str | Path) -> list[MentionAlias]:
     records: list[MentionAlias] = []
     for row in read_csv_rows(path):
         symbol = row["symbol"].upper()
+        name = row.get("name", "").strip()
         aliases = split_aliases(row.get("aliases", ""))
+        if name and name.casefold() not in {alias.casefold() for alias in aliases}:
+            aliases = (name, *aliases)
         if symbol not in aliases:
             aliases = (symbol, *aliases)
-        records.append(MentionAlias(symbol=symbol, aliases=aliases))
+        records.append(MentionAlias(symbol=symbol, aliases=aliases, name=name))
     return records
 
 
@@ -88,6 +118,40 @@ def match_symbols(text: str, aliases: list[MentionAlias]) -> list[str]:
     return sorted(dict.fromkeys(matches))
 
 
+def _is_generic_alias(alias_record: MentionAlias, alias: str) -> bool:
+    normalized = normalize_match_text(alias).strip()
+    if normalized.upper() == alias_record.symbol.upper():
+        return False
+    return normalized.casefold() in GENERIC_ALIAS_DENYLIST
+
+
+def match_evidence(text: str, alias_record: MentionAlias, source_type: str = "") -> tuple[str, str] | None:
+    normalized_text = normalize_match_text(text)
+    matches: list[tuple[str, str]] = []
+    for alias in alias_record.aliases:
+        normalized_alias = normalize_match_text(alias)
+        if not alias_pattern(alias).search(normalized_text):
+            continue
+        direct_pattern = re.compile(
+            rf"(?:awarded to|contract with|funding for|benefit(?:s|ed)? from)\s+"
+            rf"(?:the\s+)?{re.escape(normalized_alias)}",
+            re.IGNORECASE,
+        )
+        is_canonical_name = bool(alias_record.name) and alias.casefold() == alias_record.name.casefold()
+        if direct_pattern.search(normalized_text):
+            relationship = "direct_beneficiary"
+        elif source_type == "issuer_release" and is_canonical_name:
+            relationship = "issuer"
+        elif _is_generic_alias(alias_record, alias):
+            relationship = "industry_context"
+        else:
+            relationship = "issuer"
+        matches.append((relationship, normalized_alias))
+    if not matches:
+        return None
+    return max(matches, key=lambda match: (RELATIONSHIP_PRIORITY[match[0]], match[1].casefold()))
+
+
 def infer_event_type(item: RawSourceItem) -> str:
     source_type = item.source_type
     text = item.text.lower()
@@ -133,24 +197,42 @@ def extract_source_records(raw_items_path: str | Path, aliases_path: str | Path,
     aliases = load_aliases(aliases_path)
     records: list[OfficialRecord] = []
     for item in raw_items:
-        symbols = match_symbols(item.text, aliases)
-        for symbol in symbols:
+        for alias_record in aliases:
+            evidence = match_evidence(item.text, alias_record, item.source_type)
+            if evidence is None:
+                continue
+            entity_match_type, matched_text = evidence
             records.append(
                 OfficialRecord(
-                    record_id=f"{item.item_id}-{symbol.lower()}",
+                    record_id=f"{item.item_id}-{alias_record.symbol.lower()}",
                     record_date=item_date(item),
-                    symbol=symbol,
+                    symbol=alias_record.symbol,
                     source_type=item.source_type,
                     event_type=infer_event_type(item),
                     direction=infer_direction(item.text),
                     source_url=item.source_url,
                     summary=f"{item.author}: {item.text}".strip(": "),
+                    entity_match_type=entity_match_type,
+                    match_evidence=matched_text,
+                    relationship_type=entity_match_type,
                 )
             )
     rows = normalize_records(records)
     write_csv_rows(
         output_path,
-        ["event_id", "event_date", "symbol", "event_type", "direction", "confidence", "source_url", "notes"],
+        [
+            "event_id",
+            "event_date",
+            "symbol",
+            "event_type",
+            "direction",
+            "confidence",
+            "source_url",
+            "notes",
+            "entity_match_type",
+            "match_evidence",
+            "relationship_type",
+        ],
         rows,
     )
     return rows
diff --git a/tests/test_official_event_import.py b/tests/test_official_event_import.py
@@ -28,6 +28,8 @@ def test_import_official_events_normalizes_to_event_schema(tmp_path: Path) -> No
     assert rows[-2]["event_id"] == "official-issuer-release-demo-issuer-evt3"
     assert rows[-1]["event_id"] == "official-financial-media-demo-media-evt5"
     assert rows[-1]["confidence"] == "low"
+    assert rows[0]["entity_match_type"] == "unverified"
+    assert rows[0]["relationship_type"] == "unverified"
 
 
 def test_government_records_reject_non_gov_urls() -> None:
diff --git a/tests/test_source_mention_extract.py b/tests/test_source_mention_extract.py
@@ -51,3 +51,106 @@ def test_extract_source_records_outputs_confidence_by_source_type(tmp_path: Path
     assert by_symbol["EVT3"]["confidence"] == "low"
     assert by_symbol["EVT4"]["event_type"] == "procurement"
     assert output.exists()
+
+
+def test_entity_evidence_prefers_strongest_alias_match(tmp_path: Path) -> None:
+    raw_items = tmp_path / "source_items.csv"
+    raw_items.write_text(
+        "item_id,published_at,source_type,source_url,author,text\n"
+        "mixed,2026-04-01T00:00:00Z,government_procurement,https://www.govinfo.gov/example,Agency,"
+        'Funding for Current Company supports cybersecurity.\n',
+        encoding="utf-8",
+    )
+    aliases = tmp_path / "aliases.csv"
+    aliases.write_text(
+        "symbol,name,aliases\nTEST,Current Company,cybersecurity|Current Company\n",
+        encoding="utf-8",
+    )
+
+    rows = extract_source_records(raw_items, aliases, tmp_path / "events.csv")
+
+    assert rows[0]["entity_match_type"] == "direct_beneficiary"
+    assert rows[0]["match_evidence"] == "Current Company"
+    assert rows[0]["relationship_type"] == "direct_beneficiary"
+
+
+def test_generic_context_does_not_become_company_evidence(tmp_path: Path) -> None:
+    raw_items = tmp_path / "source_items.csv"
+    raw_items.write_text(
+        "item_id,published_at,source_type,source_url,author,text\n"
+        "nist,2026-04-01T00:00:00Z,government_policy,https://www.nist.gov/example,NIST,"
+        'NIST guidance discusses nuclear reactor safety and tokenization.\n',
+        encoding="utf-8",
+    )
+    aliases = tmp_path / "aliases.csv"
+    aliases.write_text(
+        "symbol,name,aliases\nOKLO,Oklo,Oklo|OKLO|nuclear reactor\n"
+        "COIN,Coinbase,Coinbase|COIN|tokenization\n",
+        encoding="utf-8",
+    )
+
+    rows = extract_source_records(raw_items, aliases, tmp_path / "events.csv")
+
+    assert {row["symbol"] for row in rows} == {"OKLO", "COIN"}
+    assert {row["entity_match_type"] for row in rows} == {"industry_context"}
+
+
+def test_canonical_name_is_trusted_only_in_issuer_release(tmp_path: Path) -> None:
+    raw_items = tmp_path / "source_items.csv"
+    raw_items.write_text(
+        "item_id,published_at,source_type,source_url,author,text\n"
+        "issuer,2026-04-01T00:00:00Z,issuer_release,https://example.com/release,Issuer,"
+        'Strategy announced a new product.\n'
+        "policy,2026-04-01T00:00:00Z,government_policy,https://www.nist.gov/example,NIST,"
+        'Strategy guidance was published.\n',
+        encoding="utf-8",
+    )
+    aliases = tmp_path / "aliases.csv"
+    aliases.write_text("symbol,name,aliases\nMSTR,Strategy,Strategy|MSTR\n", encoding="utf-8")
+
+    rows = extract_source_records(raw_items, aliases, tmp_path / "events.csv")
+    by_id = {row["event_id"]: row for row in rows}
+
+    assert by_id["official-issuer-release-issuer-mstr"]["entity_match_type"] == "issuer"
+    assert by_id["official-government-policy-policy-mstr"]["entity_match_type"] == "industry_context"
+
+
+def test_direct_beneficiary_overrides_generic_alias(tmp_path: Path) -> None:
+    raw_items = tmp_path / "source_items.csv"
+    raw_items.write_text(
+        "item_id,published_at,source_type,source_url,author,text\n"
+        "award,2026-04-01T00:00:00Z,government_procurement,https://www.govinfo.gov/example,Agency,"
+        'Funding for Strategy was announced.\n',
+        encoding="utf-8",
+    )
+    aliases = tmp_path / "aliases.csv"
+    aliases.write_text("symbol,name,aliases\nMSTR,Strategy,Strategy|MSTR\n", encoding="utf-8")
+
+    rows = extract_source_records(raw_items, aliases, tmp_path / "events.csv")
+
+    assert rows[0]["entity_match_type"] == "direct_beneficiary"
+
+
+def test_known_ticker_collisions_remain_context_only(tmp_path: Path) -> None:
+    raw_items = tmp_path / "source_items.csv"
+    raw_items.write_text(
+        "item_id,published_at,source_type,source_url,author,text\n"
+        "mstr,2026-04-01T00:00:00Z,government_policy,https://www.nist.gov/mstr,NIST,Strategy guidance\n"
+        "oklo,2026-04-01T00:00:00Z,government_policy,https://www.nist.gov/oklo,NIST,nuclear reactor safety\n"
+        "panw,2026-04-01T00:00:00Z,government_policy,https://www.nist.gov/panw,NIST,generic cybersecurity guidance\n"
+        "coin,2026-04-01T00:00:00Z,government_policy,https://www.federalreserve.gov/coin,Fed,tokenization policy\n"
+        "intc,2026-04-01T00:00:00Z,government_policy,https://www.commerce.gov/intc,Commerce,CHIPS Act funding\n",
+        encoding="utf-8",
+    )
+    aliases = tmp_path / "aliases.csv"
+    aliases.write_text(
+        "symbol,name,aliases\nMSTR,Strategy,Strategy|MSTR\nOKLO,Oklo,Oklo|OKLO|nuclear reactor\n"
+        "PANW,Palo Alto Networks,Palo Alto Networks|PANW|cybersecurity\n"
+        "COIN,Coinbase,Coinbase|COIN|tokenization\nINTC,Intel,Intel|INTC|CHIPS Act\n",
+        encoding="utf-8",
+    )
+
+    rows = extract_source_records(raw_items, aliases, tmp_path / "events.csv")
+
+    assert {row["symbol"] for row in rows} == {"MSTR", "OKLO", "PANW", "COIN", "INTC"}
+    assert {row["entity_match_type"] for row in rows} == {"industry_context"}
