feat: add official event entity schema fields

Co-Authored-By: Codex <noreply@openai.com>

Author: Pigbibi

src/political_event_tracking_research/official_event_import.py

@@ -34,6 +34,10 @@
     }
 )
 
+ALLOWED_ENTITY_MATCH_TYPES = frozenset(
+    {"issuer", "direct_beneficiary", "industry_context", "unverified"}
+)
+
 
 @dataclass(frozen=True)
 class OfficialRecord:
@@ -45,6 +49,9 @@ class OfficialRecord:
     direction: str
     source_url: str
     summary: str
+    entity_match_type: str = "unverified"
+    match_evidence: str = ""
+    relationship_type: str = ""
 
 
 def slug(value: str) -> str:
@@ -76,6 +83,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") or "unverified",
+                match_evidence=row.get("match_evidence", ""),
+                relationship_type=row.get("relationship_type", ""),
             )
         )
     return records
@@ -89,6 +99,10 @@ def validate_record(record: OfficialRecord) -> None:
         raise ValueError(f"{record.record_id}: symbol is required")
     if record.event_type not in ALLOWED_EVENT_TYPES:
         raise ValueError(f"{record.record_id}: unsupported event_type {record.event_type!r}")
+    if record.entity_match_type not in ALLOWED_ENTITY_MATCH_TYPES:
+        raise ValueError(
+            f"{record.record_id}: unsupported entity_match_type {record.entity_match_type!r}"
+        )
     if record.source_type in GOVERNMENT_SOURCE_TYPES:
         if not is_government_url(record.source_url):
             raise ValueError(f"{record.record_id}: government source URLs must be https .gov URLs")
@@ -125,6 +139,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 +153,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

tests/test_official_event_import.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import csv
 from pathlib import Path
 
 import pytest
@@ -30,6 +31,94 @@ def test_import_official_events_normalizes_to_event_schema(tmp_path: Path) -> No
     assert rows[-1]["confidence"] == "low"
 
 
+def test_legacy_input_defaults_entity_schema_fields_and_serializes_stably(tmp_path: Path) -> None:
+    output = tmp_path / "events.csv"
+
+    rows = import_official_events(ROOT / "examples/official_records.example.csv", output)
+
+    assert all(row["entity_match_type"] == "unverified" for row in rows)
+    assert all(row["match_evidence"] == "" for row in rows)
+    assert all(row["relationship_type"] == "" for row in rows)
+    with output.open(newline="", encoding="utf-8") as handle:
+        assert next(csv.reader(handle)) == [
+            "event_id",
+            "event_date",
+            "symbol",
+            "event_type",
+            "direction",
+            "confidence",
+            "source_url",
+            "notes",
+            "entity_match_type",
+            "match_evidence",
+            "relationship_type",
+        ]
+
+
+def test_entity_schema_fields_are_imported_and_normalized() -> None:
+    record = OfficialRecord(
+        record_id="entity-record",
+        record_date="2026-01-10",
+        symbol="EVT",
+        source_type="government_filing",
+        event_type="disclosure_buy",
+        direction="bullish",
+        source_url="https://www.sec.gov/example/entity-record",
+        summary="Entity evidence.",
+        entity_match_type="direct_beneficiary",
+        match_evidence="Named beneficiary in filing.",
+        relationship_type="supplier",
+    )
+
+    assert normalize_records([record])[0] == {
+        "event_id": "official-government-filing-entity-record",
+        "event_date": "2026-01-10",
+        "symbol": "EVT",
+        "event_type": "disclosure_buy",
+        "direction": "bullish",
+        "confidence": "high",
+        "source_url": "https://www.sec.gov/example/entity-record",
+        "notes": "Entity evidence.",
+        "entity_match_type": "direct_beneficiary",
+        "match_evidence": "Named beneficiary in filing.",
+        "relationship_type": "supplier",
+    }
+
+
+@pytest.mark.parametrize("entity_match_type", ["issuer", "direct_beneficiary", "industry_context", "unverified"])
+def test_entity_match_type_accepts_only_contract_values(entity_match_type: str) -> None:
+    record = OfficialRecord(
+        record_id="entity-record",
+        record_date="2026-01-10",
+        symbol="EVT",
+        source_type="government_filing",
+        event_type="disclosure_buy",
+        direction="bullish",
+        source_url="https://www.sec.gov/example/entity-record",
+        summary="Entity evidence.",
+        entity_match_type=entity_match_type,
+    )
+
+    normalize_records([record])
+
+
+def test_entity_match_type_rejects_unknown_values() -> None:
+    record = OfficialRecord(
+        record_id="entity-record",
+        record_date="2026-01-10",
+        symbol="EVT",
+        source_type="government_filing",
+        event_type="disclosure_buy",
+        direction="bullish",
+        source_url="https://www.sec.gov/example/entity-record",
+        summary="Entity evidence.",
+        entity_match_type="inferred",
+    )
+
+    with pytest.raises(ValueError, match="unsupported entity_match_type"):
+        normalize_records([record])
+
+
 def test_government_records_reject_non_gov_urls() -> None:
     record = OfficialRecord(
         record_id="bad-media-record",