Add opt-in event_entity KG builder, event_join partitioner and event extraction profiles

MANIFEST.in

@@ -1,3 +1,4 @@
 include requirements.txt
 include graphgen/configs/*.yaml
 recursive-include graphgen/configs *.yaml
+recursive-include graphgen/templates/event_profiles *

graphgen/models/__init__.py

@@ -23,13 +23,14 @@
         TrueFalseGenerator,
         VQAGenerator,
     )
-    from .kg_builder import LightRAGKGBuilder, MMKGBuilder
+    from .kg_builder import EventEntityKGBuilder, LightRAGKGBuilder, MMKGBuilder
     from .llm import HTTPClient, OllamaClient, OpenAIClient
     from .partitioner import (
         AnchorBFSPartitioner,
         BFSPartitioner,
         DFSPartitioner,
         ECEPartitioner,
+        EventJoinPartitioner,
         LeidenPartitioner,
         QuintuplePartitioner,
         TriplePartitioner,
@@ -80,6 +81,7 @@
     # KG Builder
     "LightRAGKGBuilder": ".kg_builder",
     "MMKGBuilder": ".kg_builder",
+    "EventEntityKGBuilder": ".kg_builder",
     # LLM
     "HTTPClient": ".llm",
     "OllamaClient": ".llm",
@@ -89,6 +91,7 @@
     "BFSPartitioner": ".partitioner",
     "DFSPartitioner": ".partitioner",
     "ECEPartitioner": ".partitioner",
+    "EventJoinPartitioner": ".partitioner",
     "LeidenPartitioner": ".partitioner",
     "TriplePartitioner": ".partitioner",
     "QuintuplePartitioner": ".partitioner",

graphgen/models/generator/support.py

@@ -0,0 +1,128 @@
+"""Parse support blocks and build a display-safe graph view for generators."""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Optional
+
+_SUPPORT_RE = re.compile(r"<support>(.*?)</support>", re.DOTALL)
+_EDGE_SEPARATOR = " - "
+
+
+def parse_support(response: str) -> Optional[dict[str, Any]]:
+    """Parse support XML; malformed or absent blocks are ignored."""
+    match = _SUPPORT_RE.search(response or "")
+    if not match:
+        return None
+    try:
+        payload = json.loads(match.group(1).strip())
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(payload, dict) or set(payload) != {"cited"}:
+        return None
+    cited = payload.get("cited")
+    if not isinstance(cited, list) or not all(isinstance(value, str) and value.strip() for value in cited):
+        return None
+    return {"cited": [value.strip() for value in cited]}
+
+
+def _is_event_node(node: tuple[str, dict]) -> bool:
+    return str(node[1].get("graph_role", "")).lower() == "event"
+
+
+def _event_display_name(node: tuple[str, dict]) -> str:
+    description = str(node[1].get("description", ""))
+    title, _, _ = description.partition(": ")
+    return title.strip() or "event"
+
+
+def build_generation_view(item: dict[str, Any]) -> dict[str, Any]:
+    """Replace internal event IDs with readable names for generation only.
+
+    Canonical node/edge records remain unchanged for lineage. Duplicate event
+    titles receive deterministic suffixes; support names are mapped back and
+    validated against the original partition node IDs and edges.
+    """
+    original_nodes = [node for node in item.get("nodes", []) if isinstance(node, (list, tuple)) and len(node) == 2]
+    event_nodes = [node for node in original_nodes if _is_event_node(node)]
+    used_names = {
+        str(node[0]) for node in original_nodes if not _is_event_node(node)
+    }
+    display_by_id: dict[str, str] = {}
+    for node in sorted(event_nodes, key=lambda row: str(row[0])):
+        base = _event_display_name(node)
+        display = base
+        suffix = 2
+        while display in used_names:
+            display = f"{base} ({suffix})"
+            suffix += 1
+        used_names.add(display)
+        display_by_id[str(node[0])] = display
+
+    nodes = [
+        (display_by_id.get(str(node_id), node_id), data)
+        for node_id, data in original_nodes
+    ]
+    edges = []
+    canonical_edges = set()
+    support_aliases = {name: node_id for node_id, name in display_by_id.items()}
+    for edge in item.get("edges", []):
+        if not isinstance(edge, (list, tuple)) or len(edge) != 3:
+            continue
+        source, target, data = edge
+        display_source = display_by_id.get(str(source), source)
+        display_target = display_by_id.get(str(target), target)
+        edges.append((display_source, display_target, data))
+        canonical_edge = f"{source}{_EDGE_SEPARATOR}{target}"
+        canonical_edges.add(canonical_edge)
+        if str(source) in display_by_id or str(target) in display_by_id:
+            # the generator renders edges as "<src> -- <tgt>"; accept that
+            # form in both endpoint orders and map it back to canonical IDs
+            support_aliases[f"{display_source} -- {display_target}"] = canonical_edge
+            support_aliases[f"{display_source}{_EDGE_SEPARATOR}{display_target}"] = canonical_edge
+            support_aliases[f"{display_target} -- {display_source}"] = (
+                f"{target}{_EDGE_SEPARATOR}{source}"
+            )
+            support_aliases[f"{display_target}{_EDGE_SEPARATOR}{display_source}"] = (
+                f"{target}{_EDGE_SEPARATOR}{source}"
+            )
+
+    def resolve_support(support: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]:
+        if not support:
+            return None
+        resolved = []
+        for cited in support.get("cited", []):
+            if cited in support_aliases:
+                resolved.append(support_aliases[cited])
+            else:
+                resolved.append(cited)
+        return {"cited": resolved}
+
+    return {
+        "nodes": nodes,
+        "edges": edges,
+        "canonical_nodes": [str(node_id) for node_id, _ in original_nodes],
+        "canonical_edges": canonical_edges,
+        "resolve_support": resolve_support,
+    }
+
+
+def validate_support(
+    support: Optional[dict[str, Any]],
+    node_names: set[str],
+    edge_names: Optional[set[str]] = None,
+) -> Optional[dict[str, Any]]:
+    """Keep support only when citations reference known partition nodes or edges."""
+    if not support:
+        return None
+    cited = support.get("cited") or []
+    if not cited:
+        return None
+    if all(
+        citation in node_names
+        or (edge_names is not None and citation in edge_names)
+        for citation in cited
+    ):
+        return support
+    return None

graphgen/models/kg_builder/__init__.py

@@ -1,2 +1,3 @@
+from .event_entity_kg_builder import EventEntityKGBuilder
 from .light_rag_kg_builder import LightRAGKGBuilder
 from .mm_kg_builder import MMKGBuilder

graphgen/models/kg_builder/event_contract.py

@@ -0,0 +1,194 @@
+"""Pure stdlib contracts for profile-driven event extraction.
+
+All domain policy (entity types, anchor rules, number grounding) is supplied
+explicitly by a validated event extraction profile; no implicit defaults.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from collections import defaultdict
+from typing import Optional
+
+_ALLOWED_EVENT_KEYS = {"title", "content", "entities", "is_valid"}
+_ALLOWED_ENTITY_KEYS = {"type", "name", "description"}
+_NUMBER_RE = re.compile(
+    r"(?<![A-Za-z0-9_.])[-+]?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?%?(?![A-Za-z0-9_.])"
+)
+
+
+def numbers_in(text: str) -> set[str]:
+    """Return numeric strings using the deterministic grounding policy."""
+    return set(_NUMBER_RE.findall(text or ""))
+
+
+def event_anchor(content: str, pattern: Optional[str]) -> Optional[str]:
+    """Return the first configured anchor match, if any."""
+    if pattern is None:
+        return None
+    match = re.compile(pattern).search(content or "")
+    if match is None:
+        return None
+    if match.lastindex:
+        return match.group(1)
+    return match.group(0)
+
+
+def event_node_name(anchor: Optional[str], title: str, content: str, *, legacy: bool = False) -> str:
+    """Create a deterministic event node ID."""
+    if legacy:
+        digest = hashlib.sha256(f"{anchor}\n{title}\n{content}".encode("utf-8")).hexdigest()[:12]
+        return f"event:{anchor}:{digest}"
+    digest = hashlib.sha256(
+        json.dumps([anchor, title, content], ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+    ).hexdigest()[:20]
+    return f"event:{digest}"
+
+
+def _parse_entity(entity: dict, allowed_types: set[str], dropped: dict) -> Optional[dict]:
+    """Validate a single participant entity under the profile policy."""
+    if not isinstance(entity, dict):
+        dropped["entity_not_object"] += 1
+        return None
+    if set(entity) - _ALLOWED_ENTITY_KEYS:
+        dropped["entity_extra_fields"] += 1
+        return None
+    etype = str(entity.get("type") or "").strip().upper()
+    ename = str(entity.get("name") or "").strip()
+    edesc = str(entity.get("description") or "").strip()
+    if not etype or not ename or not edesc:
+        dropped["entity_blank"] += 1
+        return None
+    if etype not in allowed_types:
+        dropped["entity_type_unknown"] += 1
+        return None
+    return {"type": etype, "name": ename, "description": edesc}
+
+
+def _event_rejection_reason(
+    item: dict,
+    *,
+    evidence_numbers: set[str],
+    anchor_pattern: Optional[str],
+    anchor_required: bool,
+    number_grounding: bool,
+) -> Optional[str]:
+    """Return the first contract violation of an event item, if any."""
+    if not isinstance(item, dict):
+        return "event_not_object"
+    title = str(item.get("title") or "").strip()
+    content = str(item.get("content") or "").strip()
+    violations = (
+        ("event_extra_fields", bool(set(item) - _ALLOWED_EVENT_KEYS)),
+        ("event_blank", not title or not content),
+        ("event_self_invalid", item.get("is_valid") is False),
+        ("event_missing_anchor", anchor_required and event_anchor(content, anchor_pattern) is None),
+        (
+            "event_number_not_grounded",
+            number_grounding
+            and any(number not in evidence_numbers for number in numbers_in(f"{title} {content}")),
+        ),
+        ("event_entities_invalid", not isinstance(item.get("entities") or [], list)),
+    )
+    for reason, violated in violations:
+        if violated:
+            return reason
+    return None
+
+
+def _parse_event(
+    item: dict,
+    *,
+    evidence_numbers: set[str],
+    allowed_types: set[str],
+    anchor_pattern: Optional[str],
+    anchor_required: bool,
+    number_grounding: bool,
+    dropped: dict,
+) -> Optional[dict]:
+    """Validate a single extracted event under the profile policy."""
+    reason = _event_rejection_reason(
+        item,
+        evidence_numbers=evidence_numbers,
+        anchor_pattern=anchor_pattern,
+        anchor_required=anchor_required,
+        number_grounding=number_grounding,
+    )
+    if reason:
+        dropped[reason] += 1
+        return None
+    title = str(item.get("title") or "").strip()
+    content = str(item.get("content") or "").strip()
+    entities = [
+        parsed
+        for entity in item.get("entities") or []
+        if (parsed := _parse_entity(entity, allowed_types, dropped)) is not None
+    ]
+    return {
+        "title": title,
+        "content": content,
+        "anchor": event_anchor(content, anchor_pattern),
+        "entities": entities,
+    }
+
+
+def parse_extraction_response(
+    response: str,
+    *,
+    chunk_content: str,
+    entity_types: list[dict],
+    anchor_pattern: Optional[str] = None,
+    anchor_required: bool = False,
+    number_grounding: bool = True,
+) -> tuple[list[dict], dict[str, int]]:
+    """Parse the strict event response contract under an explicit profile policy.
+
+    :param response: raw LLM response
+    :param chunk_content: source chunk text used for number grounding
+    :param entity_types: allowed entity types from the event profile
+    :param anchor_pattern: optional regex applied to the event content
+    :param anchor_required: whether an event without an anchor is dropped
+    :param number_grounding: whether numbers must appear in the chunk content
+    :return: (valid_events, dropped_counters)
+    """
+    dropped: dict[str, int] = defaultdict(int)
+    text = (response or "").strip()
+    if text.startswith("```"):
+        text = text.strip("`")
+        if text.startswith("json"):
+            text = text[4:]
+        text = text.strip()
+    try:
+        payload = json.loads(text)
+    except json.JSONDecodeError:
+        dropped["response_not_json"] += 1
+        return [], dict(dropped)
+    if not isinstance(payload, dict) or set(payload) != {"type", "data"}:
+        dropped["response_contract_invalid"] += 1
+        return [], dict(dropped)
+    data = payload.get("data")
+    if not isinstance(data, dict) or set(data) != {"items"} or not isinstance(data["items"], list):
+        dropped["response_contract_invalid"] += 1
+        return [], dict(dropped)
+
+    allowed_types = {str(item["type"]).strip().upper() for item in entity_types}
+    evidence_numbers = numbers_in(chunk_content) if number_grounding else set()
+    valid_events = [
+        event
+        for item in data["items"]
+        if (
+            event := _parse_event(
+                item,
+                evidence_numbers=evidence_numbers,
+                allowed_types=allowed_types,
+                anchor_pattern=anchor_pattern,
+                anchor_required=anchor_required,
+                number_grounding=number_grounding,
+                dropped=dropped,
+            )
+        )
+        is not None
+    ]
+    return valid_events, dict(dropped)