Spaces:
Running
Running
| from __future__ import annotations | |
| from dataclasses import dataclass, field | |
| from typing import Any, Optional, Sequence | |
| from pydantic import BaseModel, Field | |
| from src.core.saving.grounding import ( | |
| align_span, | |
| endpoints_well_formed, | |
| find_span_offsets, | |
| realign_span_offsets, | |
| ) | |
| class BatchEndpoint(BaseModel): | |
| uuid: Optional[str] = None | |
| name: str | |
| type: str | |
| description: Optional[str] = None | |
| happened_at: Optional[str] = None | |
| polarity: Optional[str] = "neutral" | |
| properties: Optional[dict] = Field(default_factory=dict) | |
| class BatchRelationship(BaseModel): | |
| tail: BatchEndpoint | |
| tip: BatchEndpoint | |
| name: str | |
| description: Optional[str] = None | |
| amount: Optional[float] = None | |
| source_span: str = "" | |
| span_start: Optional[int] = None | |
| span_end: Optional[int] = None | |
| happened_at: Optional[str] = None | |
| properties: Optional[dict] = Field(default_factory=dict) | |
| class BatchExtractResponse(BaseModel): | |
| new_nodes: list[BatchEndpoint] = Field(default_factory=list) | |
| relationships: list[BatchRelationship] = Field(default_factory=list) | |
| class BatchValidationIssue: | |
| index: int | |
| reason: str | |
| relationship: Optional[BatchRelationship] = None | |
| class BatchValidationResult: | |
| accepted: list[BatchRelationship] = field(default_factory=list) | |
| rejected: list[BatchValidationIssue] = field(default_factory=list) | |
| new_nodes: list[BatchEndpoint] = field(default_factory=list) | |
| def ok(self) -> bool: | |
| return bool(self.accepted) and not self.rejected | |
| def usable(self) -> bool: | |
| """True when at least one relationship survived validation.""" | |
| return bool(self.accepted) | |
| def reasons(self) -> list[str]: | |
| return [f"[{i.index}] {i.reason}" for i in self.rejected] | |
| def _is_type_named_placeholder(endpoint: Any) -> bool: | |
| name = (getattr(endpoint, "name", None) or "").strip() | |
| typ = (getattr(endpoint, "type", None) or "").strip() | |
| if not name or not typ: | |
| return True | |
| return name.lower() == typ.lower() | |
| def _is_event_endpoint(endpoint: Any) -> bool: | |
| return (getattr(endpoint, "type", None) or "").strip().lower() == "event" | |
| def _endpoint_exists( | |
| endpoint: BatchEndpoint, | |
| *, | |
| known_uuids: set[str], | |
| known_keys: set[tuple[str, str]], | |
| ) -> bool: | |
| if endpoint.uuid and endpoint.uuid in known_uuids: | |
| return True | |
| key = ( | |
| (endpoint.name or "").strip().lower(), | |
| (endpoint.type or "").strip().lower(), | |
| ) | |
| return key in known_keys | |
| def _scout_lookup( | |
| scout_entities: Sequence[Any], | |
| ) -> tuple[dict[str, Any], dict[tuple[str, str], Any]]: | |
| by_uuid: dict[str, Any] = {} | |
| by_key: dict[tuple[str, str], Any] = {} | |
| for entity in scout_entities: | |
| uuid = getattr(entity, "uuid", None) or ( | |
| entity.get("uuid") if isinstance(entity, dict) else None | |
| ) | |
| name = getattr(entity, "name", None) or ( | |
| entity.get("name") if isinstance(entity, dict) else None | |
| ) | |
| typ = getattr(entity, "type", None) or ( | |
| entity.get("type") if isinstance(entity, dict) else None | |
| ) | |
| if uuid: | |
| by_uuid[str(uuid)] = entity | |
| if name and typ: | |
| by_key[(str(name).strip().lower(), str(typ).strip().lower())] = entity | |
| return by_uuid, by_key | |
| def _entity_description(entity: Any) -> str: | |
| if entity is None: | |
| return "" | |
| if isinstance(entity, dict): | |
| return (entity.get("description") or "").strip() | |
| return (getattr(entity, "description", None) or "").strip() | |
| def _fill_event_description( | |
| endpoint: BatchEndpoint, | |
| *, | |
| rel: BatchRelationship, | |
| scout_by_uuid: dict[str, Any], | |
| scout_by_key: dict[tuple[str, str], Any], | |
| ) -> bool: | |
| """ | |
| Ensure EVENT endpoint has a non-empty description. | |
| Prefers Scout description, then relationship description, then source_span, | |
| then a non-placeholder endpoint name. | |
| Returns True if description is non-empty after fill. | |
| """ | |
| if not _is_event_endpoint(endpoint): | |
| return True | |
| if (endpoint.description or "").strip(): | |
| return True | |
| scout = None | |
| if endpoint.uuid and endpoint.uuid in scout_by_uuid: | |
| scout = scout_by_uuid[endpoint.uuid] | |
| else: | |
| key = ( | |
| (endpoint.name or "").strip().lower(), | |
| (endpoint.type or "").strip().lower(), | |
| ) | |
| scout = scout_by_key.get(key) | |
| name = (endpoint.name or "").strip() | |
| for candidate in ( | |
| _entity_description(scout), | |
| (rel.description or "").strip(), | |
| (rel.source_span or "").strip(), | |
| name if name and not _is_type_named_placeholder(endpoint) else "", | |
| ): | |
| if candidate: | |
| endpoint.description = candidate | |
| return True | |
| return False | |
| def _ensure_endpoint_uuid( | |
| endpoint: BatchEndpoint, | |
| *, | |
| scout_by_uuid: dict[str, Any], | |
| scout_by_key: dict[tuple[str, str], Any], | |
| ) -> None: | |
| """Fill missing endpoint uuid from scout match or mint a fresh one.""" | |
| if endpoint.uuid: | |
| return | |
| key = ( | |
| (endpoint.name or "").strip().lower(), | |
| (endpoint.type or "").strip().lower(), | |
| ) | |
| scout = scout_by_key.get(key) | |
| if scout is not None: | |
| scout_uuid = getattr(scout, "uuid", None) or ( | |
| scout.get("uuid") if isinstance(scout, dict) else None | |
| ) | |
| if scout_uuid: | |
| endpoint.uuid = str(scout_uuid) | |
| return | |
| import uuid as _uuid | |
| endpoint.uuid = str(_uuid.uuid4()) | |
| def _admit_endpoint( | |
| endpoint: BatchEndpoint, | |
| *, | |
| known_uuids: set[str], | |
| known_keys: set[tuple[str, str]], | |
| admitted: list[BatchEndpoint], | |
| scout_by_uuid: dict[str, Any] | None = None, | |
| scout_by_key: dict[tuple[str, str], Any] | None = None, | |
| ) -> bool: | |
| """ | |
| Soft-register a well-formed unknown relationship endpoint so dense extracts | |
| stay on the schema path instead of all-reject → tooler escalate. | |
| """ | |
| if _is_type_named_placeholder(endpoint) or not (endpoint.name or "").strip(): | |
| return False | |
| typ = (endpoint.type or "").strip() | |
| if not typ: | |
| return False | |
| key = (endpoint.name.strip().lower(), typ.lower()) | |
| _ensure_endpoint_uuid( | |
| endpoint, | |
| scout_by_uuid=scout_by_uuid or {}, | |
| scout_by_key=scout_by_key or {}, | |
| ) | |
| if endpoint.uuid: | |
| known_uuids.add(endpoint.uuid) | |
| if key not in known_keys: | |
| known_keys.add(key) | |
| admitted.append(endpoint) | |
| return True | |
| def ensure_event_descriptions( | |
| rel: BatchRelationship, | |
| *, | |
| scout_by_uuid: dict[str, Any], | |
| scout_by_key: dict[tuple[str, str], Any], | |
| ) -> tuple[bool, str]: | |
| for endpoint in (rel.tail, rel.tip): | |
| if not _is_event_endpoint(endpoint): | |
| continue | |
| if not _fill_event_description( | |
| endpoint, | |
| rel=rel, | |
| scout_by_uuid=scout_by_uuid, | |
| scout_by_key=scout_by_key, | |
| ): | |
| return False, "missing_event_description" | |
| return True, "" | |
| def _span_in_bounds( | |
| rel: BatchRelationship, | |
| source_text: str, | |
| ) -> tuple[bool, str]: | |
| text = source_text or "" | |
| span = (rel.source_span or "").strip() | |
| if not span: | |
| return False, "missing_source_span" | |
| ok, new_start, new_end, reason = realign_span_offsets( | |
| span=span, | |
| source_text=text, | |
| span_start=rel.span_start, | |
| span_end=rel.span_end, | |
| ) | |
| if not ok: | |
| return False, reason or "source_span_not_grounded" | |
| rel.span_start = new_start | |
| rel.span_end = new_end | |
| # Prefer exact quote from source when offsets were repaired. | |
| if new_start is not None and new_end is not None: | |
| sliced = text[new_start:new_end] | |
| if sliced.strip(): | |
| rel.source_span = sliced | |
| return True, "" | |
| def normalize_loose(value: str) -> str: | |
| return " ".join((value or "").split()).strip().lower() | |
| def validate_batch_extract( | |
| payload: BatchExtractResponse | dict[str, Any], | |
| *, | |
| source_text: str, | |
| scout_entities: Sequence[Any], | |
| ) -> BatchValidationResult: | |
| if isinstance(payload, dict): | |
| payload = BatchExtractResponse.model_validate(payload) | |
| known_uuids: set[str] = set() | |
| known_keys: set[tuple[str, str]] = set() | |
| scout_by_uuid, scout_by_key = _scout_lookup(scout_entities) | |
| for entity in scout_entities: | |
| uuid = getattr(entity, "uuid", None) or (entity.get("uuid") if isinstance(entity, dict) else None) | |
| name = getattr(entity, "name", None) or (entity.get("name") if isinstance(entity, dict) else None) | |
| typ = getattr(entity, "type", None) or (entity.get("type") if isinstance(entity, dict) else None) | |
| if uuid: | |
| known_uuids.add(str(uuid)) | |
| if name and typ: | |
| known_keys.add((str(name).strip().lower(), str(typ).strip().lower())) | |
| kept_nodes: list[BatchEndpoint] = [] | |
| result = BatchValidationResult(new_nodes=[]) | |
| for node in payload.new_nodes or []: | |
| if _is_type_named_placeholder(node): | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=-1, | |
| reason=f"new_node_type_named_placeholder:{node.name}", | |
| ) | |
| ) | |
| continue | |
| if _is_event_endpoint(node) and not (node.description or "").strip(): | |
| scout = None | |
| if node.uuid and node.uuid in scout_by_uuid: | |
| scout = scout_by_uuid[node.uuid] | |
| else: | |
| scout = scout_by_key.get( | |
| ( | |
| (node.name or "").strip().lower(), | |
| (node.type or "").strip().lower(), | |
| ) | |
| ) | |
| filled = _entity_description(scout) | |
| if filled: | |
| node.description = filled | |
| else: | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=-1, | |
| reason=f"new_node_missing_event_description:{node.name}", | |
| ) | |
| ) | |
| continue | |
| if not node.uuid: | |
| import uuid as _uuid | |
| node.uuid = str(_uuid.uuid4()) | |
| kept_nodes.append(node) | |
| if node.uuid: | |
| known_uuids.add(node.uuid) | |
| known_keys.add( | |
| ((node.name or "").strip().lower(), (node.type or "").strip().lower()) | |
| ) | |
| result.new_nodes = kept_nodes | |
| for idx, rel in enumerate(payload.relationships or []): | |
| if _is_type_named_placeholder(rel.tail) or _is_type_named_placeholder(rel.tip): | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=idx, | |
| reason="type_named_placeholder_endpoint", | |
| relationship=rel, | |
| ) | |
| ) | |
| continue | |
| if not endpoints_well_formed(rel): | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=idx, | |
| reason="malformed_endpoints", | |
| relationship=rel, | |
| ) | |
| ) | |
| continue | |
| if not (rel.name or "").strip(): | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=idx, | |
| reason="missing_predicate", | |
| relationship=rel, | |
| ) | |
| ) | |
| continue | |
| if not _endpoint_exists(rel.tail, known_uuids=known_uuids, known_keys=known_keys): | |
| if not _admit_endpoint( | |
| rel.tail, | |
| known_uuids=known_uuids, | |
| known_keys=known_keys, | |
| admitted=kept_nodes, | |
| scout_by_uuid=scout_by_uuid, | |
| scout_by_key=scout_by_key, | |
| ): | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=idx, | |
| reason="unknown_tail_endpoint", | |
| relationship=rel, | |
| ) | |
| ) | |
| continue | |
| if not _endpoint_exists(rel.tip, known_uuids=known_uuids, known_keys=known_keys): | |
| if not _admit_endpoint( | |
| rel.tip, | |
| known_uuids=known_uuids, | |
| known_keys=known_keys, | |
| admitted=kept_nodes, | |
| scout_by_uuid=scout_by_uuid, | |
| scout_by_key=scout_by_key, | |
| ): | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=idx, | |
| reason="unknown_tip_endpoint", | |
| relationship=rel, | |
| ) | |
| ) | |
| continue | |
| ok_span, span_reason = _span_in_bounds(rel, source_text) | |
| if not ok_span: | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=idx, | |
| reason=span_reason, | |
| relationship=rel, | |
| ) | |
| ) | |
| continue | |
| ok_event, event_reason = ensure_event_descriptions( | |
| rel, | |
| scout_by_uuid=scout_by_uuid, | |
| scout_by_key=scout_by_key, | |
| ) | |
| if not ok_event: | |
| result.rejected.append( | |
| BatchValidationIssue( | |
| index=idx, | |
| reason=event_reason, | |
| relationship=rel, | |
| ) | |
| ) | |
| continue | |
| _ensure_endpoint_uuid( | |
| rel.tail, scout_by_uuid=scout_by_uuid, scout_by_key=scout_by_key | |
| ) | |
| _ensure_endpoint_uuid( | |
| rel.tip, scout_by_uuid=scout_by_uuid, scout_by_key=scout_by_key | |
| ) | |
| result.accepted.append(rel) | |
| return result | |
| def event_leg_incomplete( | |
| relationships: Sequence[BatchRelationship], | |
| scout_entities: Sequence[Any], | |
| ) -> bool: | |
| event_uuids = set() | |
| for entity in scout_entities: | |
| typ = getattr(entity, "type", None) or (entity.get("type") if isinstance(entity, dict) else None) | |
| uuid = getattr(entity, "uuid", None) or (entity.get("uuid") if isinstance(entity, dict) else None) | |
| if typ and str(typ).strip().lower() == "event" and uuid: | |
| event_uuids.add(str(uuid)) | |
| if not event_uuids: | |
| return False | |
| touched: set[str] = set() | |
| for rel in relationships: | |
| for endpoint in (rel.tail, rel.tip): | |
| if endpoint.uuid and endpoint.uuid in event_uuids: | |
| touched.add(endpoint.uuid) | |
| elif (endpoint.type or "").strip().lower() == "event": | |
| for entity in scout_entities: | |
| name = getattr(entity, "name", None) or ( | |
| entity.get("name") if isinstance(entity, dict) else None | |
| ) | |
| uuid = getattr(entity, "uuid", None) or ( | |
| entity.get("uuid") if isinstance(entity, dict) else None | |
| ) | |
| if ( | |
| uuid | |
| and name | |
| and normalize_loose(name) == normalize_loose(endpoint.name) | |
| ): | |
| touched.add(str(uuid)) | |
| return bool(event_uuids - touched) | |
| # Re-export for callers/tests that previously imported align helpers here. | |
| __all__ = [ | |
| "BatchEndpoint", | |
| "BatchRelationship", | |
| "BatchExtractResponse", | |
| "BatchValidationIssue", | |
| "BatchValidationResult", | |
| "validate_batch_extract", | |
| "event_leg_incomplete", | |
| "ensure_event_descriptions", | |
| "normalize_loose", | |
| "find_span_offsets", | |
| "align_span", | |
| ] | |