diff --git a/MANIFEST.in b/MANIFEST.in
index d1a9003b..cb5039c3 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,3 +1,4 @@
include requirements.txt
include graphgen/configs/*.yaml
recursive-include graphgen/configs *.yaml
+recursive-include graphgen/templates/event_profiles *
diff --git a/graphgen/models/__init__.py b/graphgen/models/__init__.py
index 8bd2c9d8..8aefb510 100644
--- a/graphgen/models/__init__.py
+++ b/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",
diff --git a/graphgen/models/generator/support.py b/graphgen/models/generator/support.py
new file mode 100644
index 00000000..bb5292b6
--- /dev/null
+++ b/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"(.*?)", 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 " -- "; 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
diff --git a/graphgen/models/kg_builder/__init__.py b/graphgen/models/kg_builder/__init__.py
index 1e7e2c44..12fc3d3e 100644
--- a/graphgen/models/kg_builder/__init__.py
+++ b/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
diff --git a/graphgen/models/kg_builder/event_contract.py b/graphgen/models/kg_builder/event_contract.py
new file mode 100644
index 00000000..f4b52e2a
--- /dev/null
+++ b/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"(? 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)
diff --git a/graphgen/models/kg_builder/event_entity_kg_builder.py b/graphgen/models/kg_builder/event_entity_kg_builder.py
new file mode 100644
index 00000000..53c44344
--- /dev/null
+++ b/graphgen/models/kg_builder/event_entity_kg_builder.py
@@ -0,0 +1,236 @@
+"""Profile-driven event/entity star graph builder.
+
+Domain policy (prompt, entity types, anchor rules and grounding) comes from a
+validated event extraction profile. Generic graph construction is shared with
+the LightRAG builder: merging event nodes, participant nodes and membership
+edges is deterministic (no extra LLM summarization) so that cross-event
+context is never rewritten into synthesized facts.
+"""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from typing import Dict, List, Tuple, Union
+
+from graphgen.bases import BaseGraphStorage, BaseKGBuilder, BaseLLMWrapper, Chunk
+from graphgen.models.kg_builder.event_contract import event_node_name, parse_extraction_response
+from graphgen.templates.event_profile_loader import EventExtractionProfile, load_event_profile
+from graphgen.utils import logger
+
+EVENT_DESC_CAP = 4000
+ENTITY_DESC_CAP = 2000
+
+
+class EventEntityKGBuilder(BaseKGBuilder):
+ """Extract events and participants into a star-shaped event/entity graph."""
+
+ def __init__(self, llm_client: BaseLLMWrapper, profile: Union[str, EventExtractionProfile]):
+ super().__init__(llm_client)
+ # an event profile is always required; there is no implicit default
+ if profile is None:
+ raise ValueError("event_profile_required")
+ self.profile = profile if isinstance(profile, EventExtractionProfile) else load_event_profile(profile)
+ from graphgen.models import LightRAGKGBuilder
+
+ self._merger = LightRAGKGBuilder(llm_client=llm_client)
+ self.tokenizer = self._merger.tokenizer
+
+ async def extract(
+ self, chunk: Chunk
+ ) -> Tuple[Dict[str, List[dict]], Dict[Tuple[str, str], List[dict]]]:
+ """Extract events and participants from a single chunk.
+
+ :param chunk: input chunk
+ :return: (nodes_data, edges_data) in the same shape as LightRAGKGBuilder
+ """
+ content = chunk.content
+ entity_types_json = json.dumps(list(self.profile.entity_types), ensure_ascii=False)
+ request_json = json.dumps(
+ {
+ "type": "request",
+ "data": {
+ "items": [{"id": 1, "content": content}],
+ "meta": {"entity_types": "__ENTITY_TYPES__"},
+ },
+ },
+ ensure_ascii=False,
+ ).replace("__ENTITY_TYPES__", entity_types_json)
+ prompt = self.profile.prompt.replace("{input_text}", request_json)
+
+ response = await self.llm_client.generate_answer(prompt)
+ events, dropped = parse_extraction_response(
+ response,
+ chunk_content=content,
+ entity_types=list(self.profile.entity_types),
+ anchor_pattern=self.profile.anchor_pattern,
+ anchor_required=self.profile.anchor_required,
+ number_grounding=self.profile.number_grounding,
+ )
+ if sum(dropped.values()):
+ logger.warning("Event extraction dropped: %s (chunk=%s)", dropped, chunk.id)
+ if not events:
+ return {}, {}
+
+ nodes: Dict[str, List[dict]] = defaultdict(list)
+ edges: Dict[Tuple[str, str], List[dict]] = defaultdict(list)
+ for event in events:
+ anchor = event["anchor"]
+ event_name = event_node_name(
+ anchor, event["title"], event["content"], legacy=self.profile.legacy_event_id
+ )
+ description = f"{event['title']}: {event['content']}"
+ # The event node carries generic node attributes; the partitioner
+ # reads the anchor attribute and never parses node IDs.
+ nodes[event_name].append(
+ {
+ "entity_type": self.profile.event_entity_type,
+ "entity_name": event_name,
+ "graph_role": "event",
+ "description": description,
+ "anchor": anchor or "",
+ "source_id": chunk.id,
+ }
+ )
+ # Add the configured anchor entity only when the profile requests one.
+ seen: set[str] = set()
+ participants: list[dict] = []
+ for entity in event["entities"]:
+ if entity["name"] not in seen:
+ seen.add(entity["name"])
+ participants.append(entity)
+ if anchor is not None and self.profile.anchor_entity_type and anchor not in seen:
+ participants.append(
+ {
+ "type": self.profile.anchor_entity_type,
+ "name": anchor,
+ "description": f"Anchor {anchor} identifying this event",
+ }
+ )
+ seen.add(anchor)
+ for entity in participants:
+ nodes[entity["name"]].append(
+ {
+ "entity_type": entity["type"],
+ "entity_name": entity["name"],
+ "description": entity["description"],
+ "source_id": chunk.id,
+ }
+ )
+ edges[(event_name, entity["name"])].append(
+ {
+ "src_id": event_name,
+ "tgt_id": entity["name"],
+ "description": (
+ f"Event participant: {entity['name']} ({entity['type']})"
+ f" - {entity['description']}"
+ ),
+ "source_id": chunk.id,
+ "role": entity["type"],
+ "event_member": True,
+ }
+ )
+ return dict(nodes), dict(edges)
+
+ async def merge_nodes(self, node_data: tuple, kg_instance: BaseGraphStorage) -> dict:
+ entity_name, node_list = node_data
+ if node_list and str(node_list[0].get("entity_type", "")).upper() == self.profile.event_entity_type:
+ return await self._merge_event_node(entity_name, node_list, kg_instance)
+ # Star participants are a graph projection and need no LLM summary;
+ # merge deterministically to keep the LightRAG summarizer from
+ # expanding cross-event context into synthesized facts.
+ return await self._merge_entity_node(entity_name, node_list, kg_instance)
+
+ async def _merge_entity_node(
+ self, entity_name: str, node_list: list, kg_instance: BaseGraphStorage
+ ) -> dict:
+ existing = kg_instance.get_node(entity_name) or {}
+ all_rows = list(node_list)
+ if existing:
+ all_rows.append(existing)
+ types = [str(dp.get("entity_type") or "KEYWORD") for dp in all_rows]
+ entity_type = sorted(set(types), key=lambda t: (-types.count(t), t))[0]
+ descriptions = sorted({str(dp.get("description") or "") for dp in all_rows if dp.get("description")})
+ merged = "".join(descriptions)[:ENTITY_DESC_CAP]
+ source_ids = sorted({sid for dp in all_rows for sid in str(dp.get("source_id", "")).split("") if sid})
+ payload = {
+ "entity_type": entity_type,
+ "entity_name": entity_name,
+ "description": merged,
+ "source_id": "".join(source_ids),
+ "length": self.tokenizer.count_tokens(merged),
+ }
+ kg_instance.upsert_node(entity_name, node_data=payload)
+ return payload
+
+ async def _merge_event_node(
+ self, entity_name: str, node_list: list, kg_instance: BaseGraphStorage
+ ) -> dict:
+ """Merge event nodes deterministically (SEP-joined, capped).
+
+ The LightRAG summarizer would treat event IDs as entity names and leak
+ them into descriptions, so events never go through LLM summarization.
+ """
+ existing = kg_instance.get_node(entity_name) or {}
+ anchors = {str(dp.get("anchor") or "") for dp in node_list if dp.get("anchor")}
+ existing_anchor = str(existing.get("anchor") or "")
+ if existing_anchor:
+ anchors.add(existing_anchor)
+ merged = "".join(sorted({dp["description"] for dp in node_list}))
+ if len(merged) > EVENT_DESC_CAP:
+ merged = merged[:EVENT_DESC_CAP]
+ existing_source_ids = {
+ sid
+ for sid in str(existing.get("source_id", "")).split("")
+ if sid
+ }
+ node_data_dict = {
+ "entity_type": self.profile.event_entity_type,
+ "entity_name": entity_name,
+ "graph_role": "event",
+ "description": merged,
+ "anchor": "".join(sorted(anchors)),
+ "source_id": "".join(
+ sorted(
+ existing_source_ids
+ | {sid for dp in node_list for sid in str(dp.get("source_id", "")).split("") if sid}
+ )
+ ),
+ "length": self.tokenizer.count_tokens(merged),
+ }
+ kg_instance.upsert_node(entity_name, node_data=node_data_dict)
+ return node_data_dict
+
+ async def merge_edges(self, edges_data: tuple, kg_instance: BaseGraphStorage) -> dict:
+ (src_id, tgt_id), edge_list = edges_data
+ event_ids = {
+ str(dp["src_id"])
+ for dp in edge_list
+ if dp.get("event_member") and dp.get("src_id")
+ }
+ if src_id in event_ids or tgt_id in event_ids:
+ # Event membership edges are merged deterministically as well.
+ event_id = src_id if src_id in event_ids else tgt_id
+ participant_id = tgt_id if event_id == src_id else src_id
+ merged = "".join(sorted({dp["description"] for dp in edge_list}))
+ if len(merged) > EVENT_DESC_CAP:
+ merged = merged[:EVENT_DESC_CAP]
+ existing_edge = kg_instance.get_edge(src_id, tgt_id) or kg_instance.get_edge(tgt_id, src_id) or {}
+ source_ids = {
+ sid for dp in edge_list for sid in str(dp.get("source_id", "")).split("") if sid
+ } | {
+ sid for sid in str(existing_edge.get("source_id", "")).split("") if sid
+ }
+ role = sorted({str(dp["role"]) for dp in edge_list if dp.get("role")})
+ edge_data = {
+ "src_id": event_id,
+ "tgt_id": participant_id,
+ "description": merged,
+ "source_id": "".join(sorted(source_ids)),
+ "length": self.tokenizer.count_tokens(merged),
+ }
+ if role:
+ edge_data["role"] = "".join(role)
+ kg_instance.upsert_edge(event_id, participant_id, edge_data=edge_data)
+ return edge_data
+ return await self._merger.merge_edges(edges_data, kg_instance=kg_instance)
diff --git a/graphgen/models/partitioner/__init__.py b/graphgen/models/partitioner/__init__.py
index 4306f247..8d5b20bb 100644
--- a/graphgen/models/partitioner/__init__.py
+++ b/graphgen/models/partitioner/__init__.py
@@ -2,6 +2,7 @@
from .bfs_partitioner import BFSPartitioner
from .dfs_partitioner import DFSPartitioner
from .ece_partitioner import ECEPartitioner
+from .event_join_partitioner import EventJoinPartitioner
from .leiden_partitioner import LeidenPartitioner
from .quintuple_partitioner import QuintuplePartitioner
from .triple_partitioner import TriplePartitioner
diff --git a/graphgen/models/partitioner/event_join_partitioner.py b/graphgen/models/partitioner/event_join_partitioner.py
new file mode 100644
index 00000000..e94913d7
--- /dev/null
+++ b/graphgen/models/partitioner/event_join_partitioner.py
@@ -0,0 +1,160 @@
+"""Shared-participant event join partitioner.
+
+Groups EVENT star sub-graphs into communities joined on shared participant
+entities. Anchor isolation reads an explicit node attribute and never infers
+domain data from event ID strings.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Any, Iterable, List, Optional, Set, Tuple
+
+from graphgen.bases import BaseGraphStorage, BasePartitioner
+from graphgen.bases.datatypes import Community
+
+
+def _validate_params(
+ max_events_per_community: int,
+ max_tokens_per_community: int,
+ min_events_per_community: int,
+ event_entity_type: str,
+ anchor_attribute: str,
+ require_same_anchor: bool,
+) -> None:
+ if not isinstance(require_same_anchor, bool):
+ raise ValueError("event_join_require_same_anchor_must_be_bool")
+ if not isinstance(event_entity_type, str) or not event_entity_type.strip():
+ raise ValueError("event_join_event_entity_type_required")
+ if not isinstance(anchor_attribute, str) or not anchor_attribute.strip():
+ raise ValueError("event_join_anchor_attribute_required")
+ if max_events_per_community < 1 or min_events_per_community < 1 or max_tokens_per_community < 1:
+ raise ValueError("event_join_limits_must_be_positive")
+
+
+def _is_expandable(
+ other: str,
+ other_anchor: str,
+ *,
+ used: Set[str],
+ event_set: Set[str],
+ anchor: str,
+ require_same_anchor: bool,
+ token_sum: int,
+ max_tokens_per_community: int,
+ other_length: int,
+) -> bool:
+ """Decide whether a shared-participant event may join the community."""
+ if other in used or other in event_set:
+ return False
+ if require_same_anchor and other_anchor != anchor:
+ return False
+ return token_sum + other_length <= max_tokens_per_community
+
+
+class EventJoinPartitioner(BasePartitioner):
+ """Join event star sub-graphs that share participant entities."""
+
+ def partition(
+ self,
+ g: BaseGraphStorage,
+ max_events_per_community: int = 6,
+ max_tokens_per_community: int = 4096,
+ min_events_per_community: int = 1,
+ event_entity_type: str = "EVENT",
+ anchor_attribute: str = "anchor",
+ require_same_anchor: bool = True,
+ **kwargs: Any,
+ ) -> Iterable[Community]:
+ _validate_params(
+ max_events_per_community,
+ max_tokens_per_community,
+ min_events_per_community,
+ event_entity_type,
+ anchor_attribute,
+ require_same_anchor,
+ )
+ events: List[Tuple[str, dict]] = []
+ node_dict: dict[str, dict] = {}
+ for nid, data in g.get_all_nodes():
+ node_dict[nid] = data
+ if str(data.get("entity_type", "")).upper() == event_entity_type.upper():
+ events.append((nid, data))
+ events.sort(key=lambda item: item[0])
+
+ # participant -> events inverted index (shared entity as the join key)
+ events_by_entity: dict[str, List[str]] = {}
+ star_edges: dict[str, List[str]] = {}
+ for event_id, _ in events:
+ star_edges[event_id] = []
+ for neighbor in g.get_neighbors(event_id):
+ star_edges[event_id].append(neighbor)
+ events_by_entity.setdefault(neighbor, []).append(event_id)
+
+ used: Set[str] = set()
+
+ def _anchor(event_id: str) -> str:
+ return str(node_dict[event_id].get(anchor_attribute, ""))
+
+ def _grow(seed: Tuple[str, dict]) -> Optional[Community]:
+ member_events: List[str] = []
+ event_set: Set[str] = set()
+ member_entities: dict[str, dict] = {}
+ member_edges: Set[frozenset] = set()
+ token_sum = 0
+ anchor = _anchor(seed[0])
+ # skip events without the required anchor instead of failing
+ if require_same_anchor and not anchor:
+ return None
+ queue = deque([seed[0]])
+
+ while queue:
+ event_id = queue.popleft()
+ if event_id in event_set:
+ continue
+ event_data = node_dict[event_id]
+ if token_sum + int(event_data.get("length", 0)) > max_tokens_per_community:
+ continue
+ event_set.add(event_id)
+ member_events.append(event_id)
+ token_sum += int(event_data.get("length", 0))
+ for entity_id in star_edges.get(event_id, []):
+ if entity_id not in member_entities:
+ entity_data = node_dict.get(entity_id)
+ if entity_data is not None:
+ member_entities[entity_id] = entity_data
+ edge_key = frozenset((event_id, entity_id))
+ member_edges.add(edge_key)
+ # one-hop expansion over shared entities; events from
+ # another anchor are rejected under require_same_anchor
+ for other in events_by_entity.get(entity_id, []):
+ if _is_expandable(
+ other,
+ _anchor(other),
+ used=used,
+ event_set=event_set,
+ anchor=anchor,
+ require_same_anchor=require_same_anchor,
+ token_sum=token_sum,
+ max_tokens_per_community=max_tokens_per_community,
+ other_length=int(node_dict[other].get("length", 0)),
+ ):
+ queue.append(other)
+ if len(event_set) >= max_events_per_community:
+ break
+
+ if len(member_events) < min_events_per_community:
+ return None
+ used.update(event_set)
+ return Community(
+ id=seed[0],
+ nodes=member_events + sorted(member_entities),
+ edges=[tuple(sorted(edge)) for edge in sorted(member_edges, key=sorted)],
+ )
+
+ for event_id, data in events:
+ if event_id in used:
+ continue
+ community = _grow((event_id, data))
+ if community:
+ yield community
diff --git a/graphgen/operators/build_kg/build_event_kg.py b/graphgen/operators/build_kg/build_event_kg.py
new file mode 100644
index 00000000..3860ce72
--- /dev/null
+++ b/graphgen/operators/build_kg/build_event_kg.py
@@ -0,0 +1,61 @@
+"""Construct an event/entity star graph from text chunks."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import List, Tuple
+
+from graphgen.bases import BaseGraphStorage, BaseLLMWrapper
+from graphgen.bases.datatypes import Chunk
+from graphgen.models.kg_builder.event_entity_kg_builder import EventEntityKGBuilder
+from graphgen.utils import run_concurrent
+
+
+def build_event_entity_kg(
+ llm_client: BaseLLMWrapper,
+ kg_instance: BaseGraphStorage,
+ chunks: List[Chunk],
+ event_profile: str,
+) -> Tuple[list, list]:
+ """Build an event/entity star graph; merging reuses deterministic joins."""
+ kg_builder = EventEntityKGBuilder(llm_client=llm_client, profile=event_profile)
+
+ results = run_concurrent(
+ kg_builder.extract,
+ chunks,
+ desc="[2/4]Extracting events and entities from chunks",
+ unit="chunk",
+ )
+ results = [res for res in results if res]
+
+ nodes = defaultdict(list)
+ edges = defaultdict(list)
+ event_ids = {
+ event_id
+ for node_map, _ in results
+ for event_id, rows in node_map.items()
+ if rows and str(rows[0].get("entity_type", "")).upper() == kg_builder.profile.event_entity_type
+ }
+ for node_map, edge_map in results:
+ for key, value in node_map.items():
+ nodes[key].extend(value)
+ for key, value in edge_map.items():
+ normalized_key = tuple(sorted(key))
+ if any(endpoint in event_ids for endpoint in normalized_key):
+ for edge in value:
+ edge["event_member"] = True
+ edges[normalized_key].extend(value)
+
+ merged_nodes = run_concurrent(
+ lambda kv: kg_builder.merge_nodes(kv, kg_instance=kg_instance),
+ list(nodes.items()),
+ desc="Inserting events and entities into storage",
+ )
+
+ merged_edges = run_concurrent(
+ lambda kv: kg_builder.merge_edges(kv, kg_instance=kg_instance),
+ list(edges.items()),
+ desc="Inserting event membership into storage",
+ )
+
+ return merged_nodes, merged_edges
diff --git a/graphgen/operators/build_kg/build_kg_service.py b/graphgen/operators/build_kg/build_kg_service.py
index 19fa9e7a..bd334983 100644
--- a/graphgen/operators/build_kg/build_kg_service.py
+++ b/graphgen/operators/build_kg/build_kg_service.py
@@ -6,6 +6,7 @@
from graphgen.common.init_storage import init_storage
from graphgen.utils import logger
+from .build_event_kg import build_event_entity_kg
from .build_mm_kg import build_mm_kg
from .build_text_kg import build_text_kg
@@ -27,6 +28,14 @@ def __init__(
)
self.build_kwargs = build_kwargs
self.max_loop: int = int(self.build_kwargs.get("max_loop", 3))
+ # Opt-in event/entity star graph mode. Existing configs keep using
+ # LightRAG by default; an event profile is always required.
+ self.kg_method: str = str(self.build_kwargs.get("kg_method", "light_rag"))
+ self.event_profile = self.build_kwargs.get("event_profile")
+ if self.kg_method == "event_entity":
+ from graphgen.templates.event_profile_loader import load_event_profile
+
+ load_event_profile(self.event_profile)
def process(self, batch: list) -> Tuple[list, dict]:
"""
@@ -49,6 +58,16 @@ def process(self, batch: list) -> Tuple[list, dict]:
if len(text_chunks) == 0:
logger.info("All text chunks are already in the storage")
+ elif self.kg_method == "event_entity":
+ logger.info("[Text Event and Entity Extraction] processing ...")
+ text_nodes, text_edges = build_event_entity_kg(
+ llm_client=self.llm_client,
+ kg_instance=self.graph_storage,
+ chunks=text_chunks,
+ event_profile=self.event_profile,
+ )
+ nodes += text_nodes
+ edges += text_edges
else:
logger.info("[Text Entity and Relation Extraction] processing ...")
text_nodes, text_edges = build_text_kg(
diff --git a/graphgen/operators/partition/partition_service.py b/graphgen/operators/partition/partition_service.py
index b55c5159..cd1c1dca 100644
--- a/graphgen/operators/partition/partition_service.py
+++ b/graphgen/operators/partition/partition_service.py
@@ -29,6 +29,11 @@ def __init__(
self.tokenizer_instance: BaseTokenizer = Tokenizer(model_name=tokenizer_model)
method = partition_kwargs["method"]
self.method_params = partition_kwargs.get("method_params", {})
+ if method == "event_join":
+ self.method_params = dict(self.method_params)
+ self.method_params.setdefault("event_entity_type", partition_kwargs.get("event_entity_type", "EVENT"))
+ self.method_params.setdefault("anchor_attribute", partition_kwargs.get("anchor_attribute", "anchor"))
+ self.method_params.setdefault("require_same_anchor", partition_kwargs.get("require_same_anchor", True))
if method == "bfs":
from graphgen.models import BFSPartitioner
@@ -38,6 +43,10 @@ def __init__(
from graphgen.models import DFSPartitioner
self.partitioner = DFSPartitioner()
+ elif method == "event_join":
+ from graphgen.models import EventJoinPartitioner
+
+ self.partitioner = EventJoinPartitioner()
elif method == "ece":
# before ECE partitioning, we need to:
# 'quiz' and 'judge' to get the comprehension loss if unit_sampling is not random
diff --git a/graphgen/templates/event_profile_loader.py b/graphgen/templates/event_profile_loader.py
new file mode 100644
index 00000000..2b0324ee
--- /dev/null
+++ b/graphgen/templates/event_profile_loader.py
@@ -0,0 +1,161 @@
+"""Load and validate packaged event extraction profiles."""
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from dataclasses import dataclass
+from importlib.resources import files
+from pathlib import Path
+from typing import Any, Optional
+
+PROFILE_SCHEMA = "graphgen_event_profile.v1"
+PROFILE_ROOT = files("graphgen.templates").joinpath("event_profiles")
+_PROFILE_ID_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}")
+_ENTITY_TYPE_RE = re.compile(r"[A-Z][A-Z0-9_]{0,63}")
+
+
+@dataclass(frozen=True)
+class EventExtractionProfile:
+ profile_id: str
+ entity_types: tuple[dict[str, str], ...]
+ prompt: str
+ anchor_pattern: Optional[str]
+ anchor_required: bool
+ number_grounding: bool
+ event_entity_type: str
+ anchor_entity_type: Optional[str]
+ legacy_event_id: bool
+ manifest_sha256: str
+ prompt_sha256: str
+
+
+def _load_prompt_resource(profile_id: str, prompt_file: str) -> bytes:
+ if not prompt_file or Path(prompt_file).name != prompt_file:
+ raise ValueError("event_profile_prompt_file_invalid")
+ return PROFILE_ROOT.joinpath(profile_id, prompt_file).read_bytes()
+
+
+def _read_profile_bytes(profile_id: str) -> tuple[bytes, bytes]:
+ """Read manifest and prompt bytes for a packaged ID or a filesystem path."""
+ if "/" in profile_id or "\\" in profile_id:
+ # an explicit manifest path is preserved as given; its prompt file is
+ # resolved next to the manifest, never inside the package
+ manifest_path = Path(profile_id).resolve()
+ if not manifest_path.is_file():
+ raise ValueError(f"event_profile_not_found:{profile_id}")
+ manifest_bytes = manifest_path.read_bytes()
+ data = json.loads(manifest_bytes.decode("utf-8"))
+ prompt_file = str(data.get("prompt_file", ""))
+ if not prompt_file or Path(prompt_file).name != prompt_file:
+ raise ValueError("event_profile_prompt_file_invalid")
+ prompt_path = (manifest_path.parent / prompt_file).resolve()
+ if prompt_path.parent != manifest_path.parent:
+ raise ValueError("event_profile_prompt_file_invalid")
+ return manifest_bytes, prompt_path.read_bytes()
+
+ if not _PROFILE_ID_RE.fullmatch(profile_id):
+ raise ValueError("event_profile_id_invalid")
+ manifest_resource = PROFILE_ROOT.joinpath(profile_id, "profile.json")
+ if not manifest_resource.is_file():
+ raise ValueError(f"event_profile_not_found:{profile_id}")
+ manifest_bytes = manifest_resource.read_bytes()
+ data = json.loads(manifest_bytes.decode("utf-8"))
+ return manifest_bytes, _load_prompt_resource(profile_id, str(data.get("prompt_file", "")))
+
+
+def _parse_entity_types(raw: Any) -> tuple[list[dict[str, str]], set[str]]:
+ if not isinstance(raw, list) or not raw:
+ raise ValueError("event_profile_entity_types_required")
+ normalized: list[dict[str, str]] = []
+ seen: set[str] = set()
+ for item in raw:
+ if not isinstance(item, dict) or set(item) != {"type", "description"}:
+ raise ValueError("event_profile_entity_type_invalid")
+ name = str(item["type"]).strip().upper()
+ description = str(item["description"]).strip()
+ if not _ENTITY_TYPE_RE.fullmatch(name) or not description or name in seen:
+ raise ValueError("event_profile_entity_type_invalid")
+ seen.add(name)
+ normalized.append({"type": name, "description": description})
+ return normalized, seen
+
+
+def _parse_anchor(raw: Any, known_types: set[str]) -> tuple[bool, Optional[str], Optional[str]]:
+ if not isinstance(raw, dict) or set(raw) != {"required", "pattern", "entity_type"}:
+ raise ValueError("event_profile_anchor_invalid")
+ required, pattern, anchor_entity_type = raw["required"], raw["pattern"], raw["entity_type"]
+ if not isinstance(required, bool) or (pattern is not None and not isinstance(pattern, str)):
+ raise ValueError("event_profile_anchor_invalid")
+ if required and not pattern:
+ raise ValueError("event_profile_required_anchor_pattern_missing")
+ if pattern:
+ try:
+ re.compile(pattern)
+ except re.error as error:
+ raise ValueError("event_profile_anchor_pattern_invalid") from error
+ if anchor_entity_type is None:
+ return required, pattern, None
+ if not isinstance(anchor_entity_type, str) or anchor_entity_type.upper() not in known_types:
+ raise ValueError("event_profile_anchor_entity_type_unknown")
+ return required, pattern, anchor_entity_type.upper()
+
+
+def _parse_grounding(raw: Any) -> bool:
+ if not isinstance(raw, dict) or set(raw) != {"number_grounding"}:
+ raise ValueError("event_profile_grounding_invalid")
+ if not isinstance(raw["number_grounding"], bool):
+ raise ValueError("event_profile_grounding_invalid")
+ return raw["number_grounding"]
+
+
+def _parse_graph(raw: Any) -> tuple[str, bool]:
+ if not isinstance(raw, dict) or set(raw) != {"event_entity_type", "legacy_event_id"}:
+ raise ValueError("event_profile_graph_invalid")
+ event_entity_type = str(raw["event_entity_type"]).strip().upper()
+ legacy_event_id = raw["legacy_event_id"]
+ if not _ENTITY_TYPE_RE.fullmatch(event_entity_type) or not isinstance(legacy_event_id, bool):
+ raise ValueError("event_profile_event_entity_type_invalid")
+ return event_entity_type, legacy_event_id
+
+
+def load_event_profile(profile: str) -> EventExtractionProfile:
+ """Load a built-in profile ID or an explicit manifest path; fail closed.
+
+ An event profile is always required: there is no implicit default.
+ """
+ if not isinstance(profile, str) or not profile.strip():
+ raise ValueError("event_profile_must_be_non_empty_string")
+ profile_id = profile
+ manifest_bytes, prompt_bytes = _read_profile_bytes(profile_id)
+ data = json.loads(manifest_bytes.decode("utf-8"))
+
+ if not isinstance(data, dict) or data.get("schema_version") != PROFILE_SCHEMA:
+ raise ValueError("event_profile_schema_unsupported")
+ if not _PROFILE_ID_RE.fullmatch(str(data.get("profile_id", ""))):
+ raise ValueError("event_profile_id_invalid")
+ if data.get("profile_id") != profile_id and "/" not in profile_id and "\\" not in profile_id:
+ raise ValueError("event_profile_id_mismatch")
+
+ normalized, known_types = _parse_entity_types(data.get("entity_types"))
+ anchor_required, anchor_pattern, anchor_entity_type = _parse_anchor(data.get("anchor"), known_types)
+ number_grounding = _parse_grounding(data.get("grounding"))
+ event_entity_type, legacy_event_id = _parse_graph(data.get("graph"))
+
+ prompt = prompt_bytes.decode("utf-8")
+ if not prompt.strip() or prompt.count("{input_text}") != 1:
+ raise ValueError("event_profile_prompt_invalid")
+
+ return EventExtractionProfile(
+ profile_id=str(data["profile_id"]),
+ entity_types=tuple(normalized),
+ prompt=prompt,
+ anchor_pattern=anchor_pattern,
+ anchor_required=anchor_required,
+ number_grounding=number_grounding,
+ event_entity_type=event_entity_type,
+ anchor_entity_type=anchor_entity_type,
+ legacy_event_id=legacy_event_id,
+ manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(),
+ prompt_sha256=hashlib.sha256(prompt_bytes).hexdigest(),
+ )
diff --git a/graphgen/templates/event_profiles/generic_event.v1/profile.json b/graphgen/templates/event_profiles/generic_event.v1/profile.json
new file mode 100644
index 00000000..af572f6c
--- /dev/null
+++ b/graphgen/templates/event_profiles/generic_event.v1/profile.json
@@ -0,0 +1,12 @@
+{
+ "schema_version": "graphgen_event_profile.v1",
+ "profile_id": "generic_event.v1",
+ "prompt_file": "prompt.txt",
+ "entity_types": [
+ {"type": "ACTION", "description": "An action or occurrence"},
+ {"type": "ACTOR", "description": "A participant in the event"}
+ ],
+ "anchor": {"required": false, "pattern": null, "entity_type": null},
+ "grounding": {"number_grounding": true},
+ "graph": {"event_entity_type": "EVENT", "legacy_event_id": false}
+}
diff --git a/graphgen/templates/event_profiles/generic_event.v1/prompt.txt b/graphgen/templates/event_profiles/generic_event.v1/prompt.txt
new file mode 100644
index 00000000..b7d9b189
--- /dev/null
+++ b/graphgen/templates/event_profiles/generic_event.v1/prompt.txt
@@ -0,0 +1,27 @@
+## Role
+
+You extract events and participants from the supplied text. Report only facts supported by the input.
+
+## Rules
+
+- Write each event as a complete, independently readable factual statement.
+- Distinguish completed, ongoing, planned, and uncertain actions.
+- Use only entity types listed in `meta.entity_types`; do not invent types.
+- Describe each participant's role in the event and split distinct participants.
+- Preserve numbers and dates only when supported by the source. Keep relative dates as written.
+- Return no event when the source does not support a useful standalone fact.
+
+## Input
+
+The input is JSON with `type=request`, `data.items` and `data.meta.entity_types`.
+
+## Output
+
+Return only JSON matching this exact shape:
+{"type":"response","data":{"items":[{"title":"...","content":"...","entities":[{"type":"...","name":"...","description":"..."}],"is_valid":true}]}}
+For no events return `{"type":"response","data":{"items":[]}}`.
+
+Input:
+{input_text}
+
+Output:
diff --git a/tests/unit_tests/__init__.py b/tests/unit_tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/unit_tests/models/generator/test_support_edge_citations.py b/tests/unit_tests/models/generator/test_support_edge_citations.py
new file mode 100644
index 00000000..7492730e
--- /dev/null
+++ b/tests/unit_tests/models/generator/test_support_edge_citations.py
@@ -0,0 +1,61 @@
+"""Tests for generator support edge citations and closure validation."""
+
+from __future__ import annotations
+
+from graphgen.models.generator.support import build_generation_view, validate_support
+
+
+def _item():
+ event_id = "event:opaque-hash"
+ return {
+ "nodes": [
+ (event_id, {"entity_type": "EVENT", "graph_role": "event", "description": "Migration completed: fact"}),
+ ("Operator", {"entity_type": "ACTOR", "description": "responsible"}),
+ ],
+ "edges": [(event_id, "Operator", {"description": "participated"})],
+ }
+
+
+def test_rendered_edge_alias_resolves_and_only_present_canonical_edges_validate():
+ item = _item()
+ view = build_generation_view(item)
+ citation = f'{view["nodes"][0][0]} -- Operator'
+ resolved = view["resolve_support"]({"cited": [citation]})
+ known_edge = "event:opaque-hash - Operator"
+
+ assert resolved == {"cited": [known_edge]}
+ assert validate_support(
+ resolved,
+ set(view["canonical_nodes"]),
+ set(view["canonical_edges"]),
+ ) == resolved
+ assert validate_support(
+ {"cited": ["event:opaque-hash - Missing"]},
+ set(view["canonical_nodes"]),
+ set(view["canonical_edges"]),
+ ) is None
+
+
+def test_rendered_edge_alias_resolves_when_partition_reverses_edge_endpoints():
+ item = _item()
+ event_id = item["nodes"][0][0]
+ item["edges"] = [("Operator", event_id, {"description": "participated"})]
+ view = build_generation_view(item)
+ citation = f'Operator -- {view["nodes"][0][0]}'
+
+ resolved = view["resolve_support"]({"cited": [citation]})
+
+ assert resolved == {"cited": [f"Operator - {event_id}"]}
+ assert validate_support(
+ resolved,
+ set(view["canonical_nodes"]),
+ set(view["canonical_edges"]),
+ ) == resolved
+
+
+def test_parse_support_ignores_malformed_blocks():
+ from graphgen.models.generator.support import parse_support
+
+ assert parse_support("no block here") is None
+ assert parse_support("not json") is None
+ assert parse_support('{"cited": ["A -- B"]}') == {"cited": ["A -- B"]}
diff --git a/tests/unit_tests/models/kg_builder/test_event_entity_kg_builder.py b/tests/unit_tests/models/kg_builder/test_event_entity_kg_builder.py
new file mode 100644
index 00000000..3ff2da0d
--- /dev/null
+++ b/tests/unit_tests/models/kg_builder/test_event_entity_kg_builder.py
@@ -0,0 +1,126 @@
+"""Tests for the event/entity star graph builder."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from types import SimpleNamespace
+
+from graphgen.bases.datatypes import Chunk
+from graphgen.models.kg_builder.event_entity_kg_builder import EventEntityKGBuilder
+
+
+class FakeLLM:
+ def __init__(self, response: str = ""):
+ self.response = response
+ self.calls = []
+ self.tokenizer = SimpleNamespace(count_tokens=lambda text: len(text))
+
+ async def generate_answer(self, prompt, **kwargs):
+ self.calls.append(prompt)
+ return self.response
+
+
+class FakeStorage:
+ def __init__(self):
+ self.nodes = {}
+ self.edges = {}
+
+ def get_node(self, node_id):
+ return self.nodes.get(node_id)
+
+ def get_edge(self, src_id, tgt_id):
+ return self.edges.get((src_id, tgt_id)) or self.edges.get((tgt_id, src_id))
+
+ def upsert_node(self, node_id, node_data):
+ self.nodes[node_id] = node_data
+
+ def upsert_edge(self, src_id, tgt_id, edge_data):
+ self.edges[(src_id, tgt_id)] = edge_data
+
+
+def _response(events):
+ return json.dumps({"type": "response", "data": {"items": events}}, ensure_ascii=False)
+
+
+def test_extract_sets_event_member_marker_and_role():
+ llm = FakeLLM(
+ _response(
+ [
+ {
+ "title": "Server migration completed",
+ "content": "The team completed the server migration.",
+ "entities": [
+ {"type": "ACTOR", "name": "Ops Team", "description": "executed the migration"}
+ ],
+ "is_valid": True,
+ }
+ ]
+ )
+ )
+ builder = EventEntityKGBuilder(llm_client=llm, profile="generic_event.v1")
+ chunk = Chunk(id="chunk-1", content="The team completed the server migration.", type="text")
+
+ nodes, edges = asyncio.run(builder.extract(chunk))
+
+ event_ids = [nid for nid, rows in nodes.items() if rows[0].get("graph_role") == "event"]
+ assert len(event_ids) == 1
+ event_id = event_ids[0]
+ assert nodes[event_id][0]["entity_type"] == "EVENT"
+ assert nodes["Ops Team"][0]["entity_type"] == "ACTOR"
+
+ assert len(edges) == 1
+ (src, tgt), edge_rows = next(iter(edges.items()))
+ assert (src, tgt) == (event_id, "Ops Team")
+ assert edge_rows[0]["event_member"] is True
+ assert edge_rows[0]["role"] == "ACTOR"
+
+
+def test_merge_edges_persists_role_through_deterministic_merge():
+ builder = EventEntityKGBuilder(llm_client=FakeLLM(), profile="generic_event.v1")
+ storage = FakeStorage()
+ edge = {
+ "src_id": "event:case-1",
+ "tgt_id": "Ops Team",
+ "description": "Event participant: Ops Team (ACTOR) - executed the migration",
+ "source_id": "chunk-1",
+ "role": "ACTOR",
+ "event_member": True,
+ }
+
+ merged = asyncio.run(builder.merge_edges((("event:case-1", "Ops Team"), [edge]), storage))
+
+ assert merged["role"] == "ACTOR"
+ assert storage.edges[("event:case-1", "Ops Team")]["role"] == "ACTOR"
+ assert storage.edges[("event:case-1", "Ops Team")]["src_id"] == "event:case-1"
+ assert storage.edges[("event:case-1", "Ops Team")]["tgt_id"] == "Ops Team"
+
+
+def test_merge_nodes_preserves_event_role_and_anchor():
+ builder = EventEntityKGBuilder(llm_client=FakeLLM(), profile="generic_event.v1")
+ storage = FakeStorage()
+ node_rows = [
+ {
+ "entity_type": "EVENT",
+ "entity_name": "event:case-1",
+ "graph_role": "event",
+ "description": "Server migration completed: fact one",
+ "anchor": "CASE-1",
+ "source_id": "chunk-1",
+ },
+ {
+ "entity_type": "EVENT",
+ "entity_name": "event:case-1",
+ "graph_role": "event",
+ "description": "Server migration completed: fact two",
+ "anchor": "CASE-1",
+ "source_id": "chunk-2",
+ },
+ ]
+
+ merged = asyncio.run(builder.merge_nodes(("event:case-1", node_rows), storage))
+
+ assert merged["graph_role"] == "event"
+ assert merged["anchor"] == "CASE-1"
+ assert "fact one" in merged["description"] and "fact two" in merged["description"]
+ assert merged["source_id"] == "chunk-1chunk-2"
diff --git a/tests/unit_tests/models/partitioner/test_event_join_partitioner.py b/tests/unit_tests/models/partitioner/test_event_join_partitioner.py
new file mode 100644
index 00000000..56a5f602
--- /dev/null
+++ b/tests/unit_tests/models/partitioner/test_event_join_partitioner.py
@@ -0,0 +1,84 @@
+"""Tests for the event join partitioner."""
+
+from __future__ import annotations
+
+from graphgen.models.partitioner.event_join_partitioner import EventJoinPartitioner
+
+
+class FakeStorage:
+ def __init__(self, nodes, neighbors):
+ self._nodes = nodes
+ self._neighbors = neighbors
+
+ def get_all_nodes(self):
+ return list(self._nodes.items())
+
+ def get_neighbors(self, node_id):
+ return self._neighbors.get(node_id, [])
+
+
+def _nodes():
+ return {
+ "event:anchored-1": {"entity_type": "EVENT", "graph_role": "event", "anchor": "CASE-1", "length": 1},
+ "event:anchored-2": {"entity_type": "EVENT", "graph_role": "event", "anchor": "CASE-1", "length": 1},
+ "event:unanchored": {"entity_type": "EVENT", "graph_role": "event", "anchor": "", "length": 1},
+ "entity:operator": {"entity_type": "ACTOR", "length": 1},
+ }
+
+
+def test_event_join_skips_missing_anchor():
+ storage = FakeStorage(
+ _nodes(),
+ {
+ "event:anchored-1": ["entity:operator"],
+ "event:anchored-2": ["entity:operator"],
+ "event:unanchored": ["entity:operator"],
+ },
+ )
+
+ communities = list(EventJoinPartitioner().partition(storage, require_same_anchor=True))
+
+ assert len(communities) == 1
+ assert communities[0].id == "event:anchored-1"
+ assert "event:unanchored" not in communities[0].nodes
+ assert set(communities[0].nodes) == {"event:anchored-1", "event:anchored-2", "entity:operator"}
+
+
+def test_event_join_rejects_cross_anchor_events():
+ nodes = _nodes()
+ nodes["event:anchored-2"]["anchor"] = "CASE-2"
+ storage = FakeStorage(
+ nodes,
+ {
+ "event:anchored-1": ["entity:operator"],
+ "event:anchored-2": ["entity:operator"],
+ },
+ )
+
+ communities = list(EventJoinPartitioner().partition(storage, require_same_anchor=True))
+
+ assert len(communities) == 2
+ joined = {community.id: set(community.nodes) for community in communities}
+ assert joined["event:anchored-1"] == {"event:anchored-1", "entity:operator"}
+ assert joined["event:anchored-2"] == {"event:anchored-2", "entity:operator"}
+
+
+def test_event_join_keeps_unanchored_events_without_anchor_requirement():
+ storage = FakeStorage(
+ _nodes(),
+ {
+ "event:unanchored": ["entity:operator"],
+ "event:anchored-1": ["entity:operator"],
+ "event:anchored-2": ["entity:operator"],
+ },
+ )
+
+ communities = list(EventJoinPartitioner().partition(storage, require_same_anchor=False))
+
+ assert len(communities) == 1
+ assert set(communities[0].nodes) == {
+ "event:anchored-1",
+ "event:anchored-2",
+ "event:unanchored",
+ "entity:operator",
+ }
diff --git a/tests/unit_tests/templates/test_event_profile_loader.py b/tests/unit_tests/templates/test_event_profile_loader.py
new file mode 100644
index 00000000..90135615
--- /dev/null
+++ b/tests/unit_tests/templates/test_event_profile_loader.py
@@ -0,0 +1,112 @@
+"""Tests for the event profile loader and the opt-in service wiring."""
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+
+import pytest
+
+from graphgen.models.kg_builder.event_entity_kg_builder import EventEntityKGBuilder
+from graphgen.templates.event_profile_loader import load_event_profile
+from graphgen.utils import CURRENT_LOGGER_VAR
+
+
+class FakeLLM:
+ tokenizer = SimpleNamespace(count_tokens=lambda text: len(text))
+
+
+def _write_profile(root, profile_id: str = "custom_event.v1", prompt: str = "{input_text}") -> str:
+ manifest = {
+ "schema_version": "graphgen_event_profile.v1",
+ "profile_id": profile_id,
+ "prompt_file": "prompt.txt",
+ "entity_types": [
+ {"type": "ACTION", "description": "An action or occurrence"},
+ {"type": "ACTOR", "description": "A participant in the event"},
+ ],
+ "anchor": {"required": False, "pattern": None, "entity_type": None},
+ "grounding": {"number_grounding": False},
+ "graph": {"event_entity_type": "EVENT", "legacy_event_id": False},
+ }
+ root.mkdir(parents=True)
+ (root / "profile.json").write_text(json.dumps(manifest), encoding="utf-8")
+ (root / "prompt.txt").write_text(prompt, encoding="utf-8")
+ return str((root / "profile.json").resolve())
+
+
+def test_generic_builtin_profile_loads():
+ profile = load_event_profile("generic_event.v1")
+ assert profile.profile_id == "generic_event.v1"
+ assert profile.event_entity_type == "EVENT"
+ assert profile.anchor_required is False
+ assert profile.anchor_pattern is None
+ assert "{input_text}" in profile.prompt
+ assert [item["type"] for item in profile.entity_types] == ["ACTION", "ACTOR"]
+
+
+def test_event_profile_is_always_required():
+ with pytest.raises(ValueError, match="event_profile_must_be_non_empty_string"):
+ load_event_profile(None)
+ with pytest.raises(ValueError, match="event_profile_not_found"):
+ load_event_profile("pms_event.v1")
+
+
+def test_disk_profile_preserved_through_service_and_builder(tmp_path, monkeypatch):
+ manifest = _write_profile(tmp_path / "profiles" / "custom_event.v1")
+
+ import graphgen.common.init_storage as init_storage_module
+ import graphgen.operators.build_kg.build_kg_service as service_module
+
+ captured = {}
+
+ def fake_build_event_entity_kg(**kwargs):
+ captured.update(kwargs)
+ builder = EventEntityKGBuilder(llm_client=FakeLLM(), profile=kwargs["event_profile"])
+ assert builder.profile.prompt == "{input_text}"
+ return [], []
+
+ fake_storage = lambda **kwargs: SimpleNamespace(index_done_callback=lambda: None) # noqa: E731
+ monkeypatch.setattr(init_storage_module, "init_storage", fake_storage)
+ monkeypatch.setattr(service_module, "init_storage", fake_storage)
+ monkeypatch.setattr(service_module, "init_llm", lambda *args, **kwargs: FakeLLM())
+ monkeypatch.setattr(service_module, "build_event_entity_kg", fake_build_event_entity_kg)
+
+ service = service_module.BuildKGService(
+ working_dir=str(tmp_path / "cache"),
+ kv_backend="json_kv",
+ kg_method="event_entity",
+ event_profile=manifest,
+ )
+ assert service.event_profile == manifest
+
+ # BaseOperator.__call__ installs the context logger before process()
+ token = CURRENT_LOGGER_VAR.set(service.get_logger())
+ try:
+ service.process(
+ [{"_trace_id": "chunk-1", "type": "text", "content": "The team completed the migration."}]
+ )
+ finally:
+ CURRENT_LOGGER_VAR.reset(token)
+
+ assert captured["event_profile"] == manifest
+ builder = EventEntityKGBuilder(llm_client=FakeLLM(), profile=captured["event_profile"])
+ assert builder.profile.profile_id == "custom_event.v1"
+ assert builder.profile.prompt == "{input_text}"
+
+
+def test_event_entity_mode_without_profile_fails_closed(tmp_path, monkeypatch):
+ import graphgen.common.init_storage as init_storage_module
+ import graphgen.operators.build_kg.build_kg_service as service_module
+
+ fake_storage = lambda **kwargs: SimpleNamespace(index_done_callback=lambda: None) # noqa: E731
+ monkeypatch.setattr(init_storage_module, "init_storage", fake_storage)
+ monkeypatch.setattr(service_module, "init_storage", fake_storage)
+ monkeypatch.setattr(service_module, "init_llm", lambda *args, **kwargs: FakeLLM())
+
+ with pytest.raises(ValueError, match="event_profile_must_be_non_empty_string"):
+ service_module.BuildKGService(
+ working_dir=str(tmp_path / "cache"),
+ kv_backend="json_kv",
+ kg_method="event_entity",
+ )