feat: add official record entity matching fields (#24)

Co-authored-by: Codex <noreply@openai.com>

Author: Pigbibi

src/political_event_tracking_research/official_event_import.py

@@ -34,6 +34,10 @@
     }
 )
 
+ALLOWED_ENTITY_RELATIONSHIP_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 = "unverified"
 
 
 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", "unverified"),
+                match_evidence=row.get("match_evidence", ""),
+                relationship_type=row.get("relationship_type", "unverified"),
             )
         )
     return records
@@ -87,6 +97,12 @@ 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")
+    for field_name, value in (
+        ("entity_match_type", record.entity_match_type),
+        ("relationship_type", record.relationship_type),
+    ):
+        if value not in ALLOWED_ENTITY_RELATIONSHIP_TYPES:
+            raise ValueError(f"{record.record_id}: unsupported {field_name} {value!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:

tests/test_official_event_import.py

@@ -7,6 +7,7 @@
 from political_event_tracking_research.official_event_import import (
     OfficialRecord,
     import_official_events,
+    load_official_records,
     normalize_records,
 )
 
@@ -28,6 +29,66 @@ 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 set(rows[0]) == {
+        "event_id",
+        "event_date",
+        "symbol",
+        "event_type",
+        "direction",
+        "confidence",
+        "source_url",
+        "notes",
+    }
+
+
+def test_load_official_records_defaults_entity_fields_for_legacy_csv(tmp_path: Path) -> None:
+    input_path = tmp_path / "legacy.csv"
+    input_path.write_text(
+        "record_id,record_date,symbol,source_type,event_type,direction,source_url,summary\n"
+        "legacy-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish,"
+        "https://www.sec.gov/example/legacy-1,Legacy record.\n",
+        encoding="utf-8",
+    )
+
+    record = load_official_records(input_path)[0]
+
+    assert record.entity_match_type == "unverified"
+    assert record.match_evidence == ""
+    assert record.relationship_type == "unverified"
+
+
+@pytest.mark.parametrize("value", ["issuer", "direct_beneficiary", "industry_context", "unverified"])
+def test_new_entity_fields_accept_allowed_values(tmp_path: Path, value: str) -> None:
+    input_path = tmp_path / "entity-fields.csv"
+    input_path.write_text(
+        "record_id,record_date,symbol,source_type,event_type,direction,source_url,summary,"
+        "entity_match_type,match_evidence,relationship_type\n"
+        "entity-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish,"
+        f"https://www.sec.gov/example/entity-1,Entity record.,{value},SEC filing,{value}\n",
+        encoding="utf-8",
+    )
+
+    rows = normalize_records(load_official_records(input_path))
+
+    assert rows[0]["event_id"] == "official-government-filing-entity-1"
+
+
+@pytest.mark.parametrize("field", ["entity_match_type", "relationship_type"])
+@pytest.mark.parametrize("value", ["", "   ", "not-allowed"])
+def test_new_entity_fields_fail_closed_for_blank_or_invalid_values(tmp_path: Path, field: str, value: str) -> None:
+    input_path = tmp_path / f"invalid-{field}.csv"
+    entity_match_type = value if field == "entity_match_type" else "issuer"
+    relationship_type = value if field == "relationship_type" else "issuer"
+    input_path.write_text(
+        "record_id,record_date,symbol,source_type,event_type,direction,source_url,summary,"
+        "entity_match_type,match_evidence,relationship_type\n"
+        f"entity-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish,"
+        f"https://www.sec.gov/example/entity-1,Entity record.,{entity_match_type},SEC filing,{relationship_type}\n",
+        encoding="utf-8",
+    )
+
+    with pytest.raises(ValueError):
+        normalize_records(load_official_records(input_path))
 
 
 def test_government_records_reject_non_gov_urls() -> None: