PMFlow/backend/app/services/semantic_fact_service/hooks.py

from __future__ import annotations

from sqlalchemy import and_, event, select
from sqlalchemy.orm import Session

from ...models.ingest_event import IngestEvent
from ...models.knowledge_base_doc import KnowledgeBaseDoc
from ...models.project import Project
from ...models.signal import Signal
from ...models.task import Task
from ...models.truth_node import TruthNode
from ...models.user import User
from .shared import (
    ENTITY_TYPE_TO_TRUTH_NODE_TYPE,
    SEMANTIC_FACT_EXTRACT_CHANNEL,
    TRUTH_NODE_TYPE_TO_ENTITY_TYPE,
    _normalized_entity_type,
    _status_is_inactive,
)

_SEMANTIC_FACT_INFO_KEY = "_pmflow_semantic_fact_events"
_SEMANTIC_FACT_APPLYING_KEY = "_pmflow_semantic_fact_applying"


def _truth_node_entity_type(node: TruthNode) -> str:
    return TRUTH_NODE_TYPE_TO_ENTITY_TYPE.get(str(getattr(node, "node_type", "") or "").strip().lower(), "")


def _build_payload(obj: object, *, deleted: bool) -> dict | None:
    entity_type = ""
    if isinstance(obj, Project):
        entity_type = "project"
    elif isinstance(obj, Task):
        entity_type = "task"
    elif isinstance(obj, Signal):
        entity_type = "signal"
    elif isinstance(obj, KnowledgeBaseDoc):
        entity_type = "document"
    elif isinstance(obj, User):
        entity_type = "user"
    elif isinstance(obj, TruthNode):
        entity_type = _truth_node_entity_type(obj)
    if not entity_type:
        return None

    entity_id = getattr(obj, "id", None)
    tenant_id = getattr(obj, "tenant_id", None)
    if entity_id is None or tenant_id is None:
        return None

    operation = "deactivate" if deleted else "extract"
    if not deleted:
        if isinstance(obj, User):
            operation = "extract" if bool(getattr(obj, "is_active", True)) and not _status_is_inactive(getattr(obj, "status", "")) else "deactivate"
        elif _status_is_inactive(getattr(obj, "status", "")):
            operation = "deactivate"

    updated_at = getattr(obj, "updated_at", None)
    return {
        "event_key": f"{_normalized_entity_type(entity_type)}:{int(entity_id)}",
        "payload": {
            "entity_type": _normalized_entity_type(entity_type),
            "entity_id": int(entity_id),
            "operation": operation,
            "trigger_type": "orm_commit",
            "tenant_id": int(tenant_id),
            "source_updated_at": updated_at.isoformat() if updated_at is not None else None,
        },
    }


def register_semantic_fact_hooks() -> None:
    if getattr(register_semantic_fact_hooks, "_registered", False):
        return

    @event.listens_for(Session, "after_flush")
    def _collect_semantic_fact_events(session: Session, _flush_context) -> None:  # pragma: no cover
        if bool(session.info.get(_SEMANTIC_FACT_APPLYING_KEY)):
            return
        payload_map = session.info.get(_SEMANTIC_FACT_INFO_KEY)
        if not isinstance(payload_map, dict):
            payload_map = {}
            session.info[_SEMANTIC_FACT_INFO_KEY] = payload_map
        objects = set(session.new).union(session.dirty).union(session.deleted)
        deleted_objects = set(session.deleted)
        for obj in objects:
            if isinstance(obj, IngestEvent):
                continue
            payload = _build_payload(obj, deleted=obj in deleted_objects)
            if not payload:
                continue
            payload_map[str(payload["event_key"])] = payload["payload"]

    @event.listens_for(Session, "after_flush_postexec")
    def _upsert_semantic_fact_events(session: Session, _flush_context) -> None:  # pragma: no cover
        if bool(session.info.get(_SEMANTIC_FACT_APPLYING_KEY)):
            return
        payload_map = session.info.get(_SEMANTIC_FACT_INFO_KEY)
        if not isinstance(payload_map, dict) or not payload_map:
            return
        session.info[_SEMANTIC_FACT_APPLYING_KEY] = True
        try:
            for event_key, payload in payload_map.items():
                tenant_id = int(payload.get("tenant_id") or 0)
                if tenant_id <= 0:
                    continue
                row = session.execute(
                    select(IngestEvent).where(
                        and_(
                            IngestEvent.tenant_id == tenant_id,
                            IngestEvent.channel == SEMANTIC_FACT_EXTRACT_CHANNEL,
                            IngestEvent.event_key == str(event_key),
                        )
                    )
                ).scalar_one_or_none()
                if row is None:
                    session.add(
                        IngestEvent(
                            tenant_id=tenant_id,
                            channel=SEMANTIC_FACT_EXTRACT_CHANNEL,
                            event_key=str(event_key),
                            status="pending",
                            payload=payload,
                            retry_count=0,
                            error_message=None,
                            processed_at=None,
                        )
                    )
                else:
                    row.status = "pending"
                    row.payload = payload
                    row.retry_count = 0
                    row.error_message = None
                    row.processed_at = None
            payload_map.clear()
        finally:
            session.info[_SEMANTIC_FACT_APPLYING_KEY] = False

    @event.listens_for(Session, "after_commit")
    def _clear_semantic_fact_events_after_commit(session: Session) -> None:  # pragma: no cover
        session.info.pop(_SEMANTIC_FACT_INFO_KEY, None)
        session.info.pop(_SEMANTIC_FACT_APPLYING_KEY, None)

    @event.listens_for(Session, "after_rollback")
    def _clear_semantic_fact_events_after_rollback(session: Session) -> None:  # pragma: no cover
        session.info.pop(_SEMANTIC_FACT_INFO_KEY, None)
        session.info.pop(_SEMANTIC_FACT_APPLYING_KEY, None)

    setattr(register_semantic_fact_hooks, "_registered", True)


__all__ = ["register_semantic_fact_hooks"]