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

from __future__ import annotations

from .shared import *


class SemanticFactQueryMixin:
    async def upsert_fact(
        self,
        *,
        entity_type: str,
        entity_id: int,
        fact_namespace: str,
        fact_key: str,
        fact_value_text: str,
        fact_value_norm: Optional[str],
        value_type: str,
        confidence: float,
        evidence_source_type: Optional[str],
        evidence_source_id: Optional[int],
        evidence_excerpt: Optional[str],
        extract_version: Optional[str],
        freshness_at: Optional[str],
    ) -> SemanticFact:
        normalized_fact_value = _normalize_fact_value(fact_value_norm or fact_value_text)
        row = (
            await self.db.execute(
                select(SemanticFact).where(
                    and_(
                        SemanticFact.tenant_id == self.tenant_id,
                        SemanticFact.entity_type == str(entity_type or "").strip(),
                        SemanticFact.entity_id == int(entity_id),
                        SemanticFact.fact_key == str(fact_key or "").strip(),
                        SemanticFact.fact_value_norm == normalized_fact_value,
                    )
                )
            )
        ).scalar_one_or_none()

        if row is None:
            row = SemanticFact(
                tenant_id=self.tenant_id,
                entity_type=str(entity_type or "").strip(),
                entity_id=int(entity_id),
                fact_namespace=str(fact_namespace or "").strip() or "general",
                fact_key=str(fact_key or "").strip(),
                fact_value_text=str(fact_value_text or "").strip(),
                fact_value_norm=normalized_fact_value,
                value_type=str(value_type or "").strip() or "keyword",
                confidence=max(0.0, min(1.0, float(confidence or 0.0))),
                evidence_source_type=str(evidence_source_type or "").strip() or None,
                evidence_source_id=int(evidence_source_id) if evidence_source_id is not None else None,
                evidence_excerpt=str(evidence_excerpt or "").strip() or None,
                extract_version=str(extract_version or "").strip() or None,
                freshness_at=str(freshness_at or "").strip() or None,
                is_active=True,
            )
            self.db.add(row)
            await self.db.flush()
            return row

        row.fact_namespace = str(fact_namespace or "").strip() or row.fact_namespace
        row.fact_value_text = str(fact_value_text or "").strip() or row.fact_value_text
        row.value_type = str(value_type or "").strip() or row.value_type
        row.confidence = max(0.0, min(1.0, float(confidence or 0.0)))
        row.evidence_source_type = str(evidence_source_type or "").strip() or row.evidence_source_type
        row.evidence_source_id = int(evidence_source_id) if evidence_source_id is not None else row.evidence_source_id
        row.evidence_excerpt = str(evidence_excerpt or "").strip() or row.evidence_excerpt
        row.extract_version = str(extract_version or "").strip() or row.extract_version
        row.freshness_at = str(freshness_at or "").strip() or row.freshness_at
        row.is_active = True
        await self.db.flush()
        return row

    async def deactivate_entity_facts(
        self,
        *,
        entity_type: str,
        entity_id: int,
    ) -> int:
        normalized_entity_type = _normalized_entity_type(entity_type)
        rows = (
            await self.db.execute(
                select(SemanticFact).where(
                    and_(
                        SemanticFact.tenant_id == self.tenant_id,
                        SemanticFact.entity_type == normalized_entity_type,
                        SemanticFact.entity_id == int(entity_id),
                        SemanticFact.is_active.is_(True),
                    )
                )
            )
        ).scalars().all()
        for row in rows:
            row.is_active = False
        await self.db.flush()
        return len(rows)

    async def enqueue_entity_event(
        self,
        *,
        entity_type: str,
        entity_id: int,
        operation: str,
        trigger_type: str,
        extra_payload: Optional[dict[str, Any]] = None,
    ) -> IngestEvent:
        normalized_entity_type = _normalized_entity_type(entity_type)
        event_key = f"{normalized_entity_type}:{int(entity_id)}"
        payload = {
            "entity_type": normalized_entity_type,
            "entity_id": int(entity_id),
            "operation": str(operation or "extract").strip() or "extract",
            "trigger_type": str(trigger_type or "manual_retry").strip() or "manual_retry",
            "tenant_id": int(self.tenant_id),
        }
        if isinstance(extra_payload, dict):
            for key, value in extra_payload.items():
                payload[str(key)] = value

        row = (
            await self.db.execute(
                select(IngestEvent).where(
                    and_(
                        IngestEvent.tenant_id == self.tenant_id,
                        IngestEvent.channel == SEMANTIC_FACT_EXTRACT_CHANNEL,
                        IngestEvent.event_key == event_key,
                    )
                )
            )
        ).scalar_one_or_none()
        if row is None:
            row = IngestEvent(
                tenant_id=self.tenant_id,
                channel=SEMANTIC_FACT_EXTRACT_CHANNEL,
                event_key=event_key,
                status="pending",
                payload=payload,
                retry_count=0,
                error_message=None,
                processed_at=None,
            )
            self.db.add(row)
        else:
            row.status = "pending"
            row.payload = payload
            row.retry_count = 0
            row.error_message = None
            row.processed_at = None
        await self.db.flush()
        return row

    async def query_facts(
        self,
        *,
        entity_type: str,
        query_type: str,
        filters: Optional[list[dict[str, Any]]] = None,
        entity_id: Optional[int] = None,
        limit: int = 20,
    ) -> list[SemanticFact]:
        _ = query_type
        stmt = select(SemanticFact).where(
            and_(
                SemanticFact.tenant_id == self.tenant_id,
                SemanticFact.entity_type == _normalized_entity_type(entity_type),
                SemanticFact.is_active.is_(True),
            )
        )
        if entity_id is not None:
            stmt = stmt.where(SemanticFact.entity_id == int(entity_id))

        for item in list(filters or []):
            if not isinstance(item, dict):
                continue
            fact_key = str(item.get("fact_key") or item.get("field") or "").strip()
            if fact_key:
                stmt = stmt.where(SemanticFact.fact_key == fact_key)
            raw_value = str(item.get("value") or item.get("normalized_value") or "").strip()
            normalized_value = _normalize_fact_value(item.get("normalized_value") or raw_value)
            if normalized_value:
                stmt = stmt.where(
                    or_(
                        SemanticFact.fact_value_norm == normalized_value,
                        func.lower(SemanticFact.fact_value_text) == raw_value.lower(),
                    )
                )

        rows = (
            await self.db.execute(
                stmt.order_by(SemanticFact.confidence.desc(), SemanticFact.updated_at.desc(), SemanticFact.id.desc()).limit(
                    max(1, int(limit))
                )
            )
        ).scalars().all()
        return list(rows)


__all__ = ["SemanticFactQueryMixin"]