politics-tracking/backend/politrack/enrich/linker.py

"""Link documents (headlines, press releases, government news) to the politicians and parties they mention.

Conservative by design: full names (>= 2 tokens) only, plus "office title + surname" of *current* office
holders ("Bundeskanzler Stocker"); an ambiguous name is only linked if exactly one candidate is active.
False links are worse than missed ones on a site like this.
"""

from __future__ import annotations

import re
from collections import defaultdict
from collections.abc import Iterator
from datetime import datetime, timezone

from psycopg import Connection
from psycopg.types.json import Jsonb

from politrack.core import Enricher, RunContext, register
from politrack.core.records import Claim, Evidence, OrgRef, PersonRef
from politrack.core.spec import DERIVED
from politrack.text import normalize, strip_titles

TOKEN = re.compile(r"[\w'-]+")
OFFICE_TITLE = re.compile(
    r"\b(?:Bundeskanzler(?:in)?|Vizekanzler(?:in)?|[A-ZÄÖÜ][a-zäöüß]*minister(?:in)?|Minister(?:in)?|"
    r"Staatssekretär(?:in)?|Nationalratspräsident(?:in)?|Bundesratspräsident(?:in)?|Bundespräsident(?:in)?|"
    r"Landeshauptmann|Landeshauptfrau|Landesrät(?:in)?|Landesrat|Bürgermeister(?:in)?|Klubobmann|Klubobfrau|"
    r"Klubchef(?:in)?|Parteichef(?:in)?|Generalsekretär(?:in)?|Abgeordnete[rn]?|EU-Abgeordnete[rn]?|"
    r"EU-Kommissar(?:in)?|Präsident(?:in)?)\s+(?:(?:Dr|Mag|DI|Ing)\.\s+)?(?P<surname>[A-ZÄÖÜ][a-zäöüß]+(?:-[A-ZÄÖÜ][a-zäöüß]+)?)\b"
    # "Nationalratspräsident Walter Rosenkranz": a following capitalised word means we matched a first name.
    r"(?!\s+(?:von\s+)?[A-ZÄÖÜ][a-zäöüß]+)"
)
CONTEXT_PATTERNS = [
    ("interview", re.compile(r"\binterview\b", re.I)),
    ("tv", re.compile(r"\b(ZIB ?2|ZiB ?2|Pressestunde|Im Zentrum|Sommergespräch|Report|Puls ?24|ServusTV|Talk im Hangar|"
                      r"Milborn|Pro und Contra)\b", re.I)),
    ("radio", re.compile(r"\b(Ö1|Morgenjournal|Mittagsjournal|Journal zu Gast)\b", re.I)),
    ("op_ed", re.compile(r"\b(Gastkommentar|Kommentar der anderen)\b", re.I)),
]


def tokens(text: str) -> list[str]:
    return TOKEN.findall(normalize(text))


def detect_context(text: str) -> str | None:
    for name, pattern in CONTEXT_PATTERNS:
        if pattern.search(text):
            return name
    return None


class MentionLinker:
    def __init__(self, conn: Connection) -> None:
        self.person_index: dict[tuple[str, ...], set[int]] = defaultdict(set)
        for row in conn.execute("SELECT entity_id, alias FROM alias WHERE entity_type = 'person'"):
            key = tuple(tokens(strip_titles(row["alias"])))
            if len(key) >= 2:
                self.person_index[key].add(row["entity_id"])
        self.active = {r["person_id"] for r in conn.execute("SELECT person_id FROM person_status WHERE status = 'active'")}
        self.active_surnames: dict[str, set[int]] = defaultdict(set)
        for row in conn.execute(
            "SELECT p.id, p.family_name FROM person p JOIN person_status s ON s.person_id = p.id "
            "WHERE s.status = 'active' AND p.family_name IS NOT NULL"
        ):
            self.active_surnames[normalize(row["family_name"])].add(row["id"])
        self.org_index: dict[tuple[str, ...], int] = {}
        for row in conn.execute(
            """SELECT o.id, a.alias FROM organization o JOIN alias a ON a.entity_type = 'organization' AND a.entity_id = o.id
               WHERE o.kind = 'party'
               UNION SELECT id, short_name FROM organization WHERE kind = 'party' AND short_name IS NOT NULL"""
        ):
            key = tuple(tokens(row["alias"]))
            if key and len("".join(key)) >= 3:
                self.org_index[key] = row["id"]
        self.max_len = max((len(k) for k in self.person_index), default=2)

    def people(self, text: str) -> set[int]:
        toks = tokens(text)
        found: set[int] = set()
        for n in range(2, self.max_len + 1):
            for i in range(len(toks) - n + 1):
                candidates = self.person_index.get(tuple(toks[i : i + n]))
                if not candidates:
                    continue
                if len(candidates) == 1:
                    found |= candidates
                elif len(active := candidates & self.active) == 1:
                    found |= active
        return found

    def people_by_title(self, text: str) -> set[int]:
        found: set[int] = set()
        for m in OFFICE_TITLE.finditer(text):
            candidates = self.active_surnames.get(normalize(m.group("surname")), set())
            if len(candidates) == 1:
                found |= candidates
        return found

    def all_people(self, text: str) -> set[int]:
        return self.people(text) | self.people_by_title(text)

    def parties(self, text: str) -> set[int]:
        toks = tokens(text)
        found = set()
        for n in (1, 2, 3):
            for i in range(len(toks) - n + 1):
                if org := self.org_index.get(tuple(toks[i : i + n])):
                    found.add(org)
        return found


@register
class Linker(Enricher):
    spec = DERIVED
    name = "enrich.linker"
    description = "Link news headlines, press releases and government news to politicians and parties"
    schedule = "*/10 * * * *"

    def run(self, ctx: RunContext) -> Iterator[Claim]:
        conn = ctx.conn
        linker = MentionLinker(conn)
        while True:
            docs = conn.execute(
                """SELECT d.id, d.url, d.title, d.excerpt, left(d.full_text, 3000) AS full_text, d.published_at,
                          s.slug AS source, s.name AS source_name, s.kind AS source_kind,
                          COALESCE((d.meta->>'tdm_reserved')::boolean, false) AS tdm_reserved
                   FROM document d JOIN source s ON s.id = d.source_id
                   WHERE (s.kind IN ('media', 'user') OR s.slug IN ('gdelt', 'bundesregierung', 'wksta'))
                     AND NOT (d.meta ? 'linked_at') AND d.doc_type IN ('article', 'press_release')
                   ORDER BY d.id LIMIT 500"""
            ).fetchall()
            if not docs:
                return
            for doc in docs:
                title = doc["title"] or ""
                text = title if doc["tdm_reserved"] else f"{title} {doc['full_text'] or doc['excerpt'] or ''}"
                context = detect_context(title)
                people, parties = linker.all_people(text), linker.parties(title)
                evidence = Evidence(url=doc["url"], source=doc["source"], title=title, quote=title, stance="mentions",
                                    confidence=0.8, published_at=doc["published_at"])
                for person_id in people:
                    yield Claim(claim_type="media_mention", person=PersonRef(ids={"politrack": str(person_id)}),
                                occurred_at=doc["published_at"], summary_de=f"{doc['source_name']}: {title}",
                                data={"context": context, "source": doc["source"]}, evidence=[evidence])
                for org_id in parties:
                    yield Claim(claim_type="media_mention", org=OrgRef(ids={"politrack_org": str(org_id)}, create=False),
                                occurred_at=doc["published_at"], summary_de=f"{doc['source_name']}: {title}",
                                data={"context": context, "source": doc["source"]}, evidence=[evidence])
                conn.execute("UPDATE document SET meta = meta || %s WHERE id = %s",
                             (Jsonb({"linked_at": datetime.now(timezone.utc).isoformat(), "people": sorted(people),
                                     "parties": sorted(parties)}), doc["id"]))
                ctx.stats["documents_linked"] += 1