leocrm/app/ai/knowledge_lifecycle.py

"""Knowledge lifecycle — event-driven extraction, derived-data lifecycle,
retention policy, ask-knowledge, and review queue (H-EVT, H-DATA-LIFE, H-RET, H-ASK, H-REV).

Event-driven extraction: new mail/dokument/message → ARQ-Job → extraction.
Derived-data lifecycle: correction/delete/erasure of source propagates to
RAG chunks, embeddings, graph references, and agent memory.
Retention: configurable per-source retention policy, ARQ cleans up.
Ask Knowledge: RAG queries with evidence cards via workstream.
Review queue: low-confidence extracted relationships pending review.
"""

from __future__ import annotations

import logging
import uuid
from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy import select, delete
from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)


# ─── H-EVT: Event-Driven Extraction ──────────────────────────────────────────

# Events that trigger knowledge extraction
EXTRACTION_TRIGGERS = {
    "mail.received": {"source_type": "mail", "text_field": "body_text"},
    "dms.file_uploaded": {"source_type": "dms", "text_field": "extracted_text"},
    "wiki.article_published": {"source_type": "wiki", "text_field": "content"},
    "communication.message_created": {"source_type": "communication", "text_field": "content"},
}


def should_extract(event_name: str) -> bool:
    """Check if an event should trigger knowledge extraction (H-EVT)."""
    return event_name in EXTRACTION_TRIGGERS


async def handle_extraction_event(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    event_name: str,
    payload: dict[str, Any],
) -> dict[str, Any] | None:
    """Handle an event that may trigger knowledge extraction (H-EVT).

    Called by the event bus. If the event matches a configured extraction
    trigger, fetches the source content and runs knowledge extraction.
    """
    if not should_extract(event_name):
        return None

    trigger_config = EXTRACTION_TRIGGERS[event_name]
    source_type = trigger_config["source_type"]
    entity_id_str = payload.get("entity_id") or payload.get("file_id") or payload.get("message_id")
    if not entity_id_str:
        return None

    try:
        entity_id = uuid.UUID(str(entity_id_str))
    except (ValueError, TypeError):
        return None

    # Fetch source content
    from app.ai.knowledge_sources import fetch_source_content
    content = await fetch_source_content(db, tenant_id, source_type, entity_id)
    if content is None or not content.get("text"):
        return None

    # Run extraction
    from app.ai.knowledge_extraction import extract_knowledge, auto_create_relationships
    extraction = await extract_knowledge(
        text=content["text"],
        tenant_id=tenant_id,
        source_type=source_type,
        source_id=str(entity_id),
    )

    # Auto-create high-confidence relationships
    result = await auto_create_relationships(db, tenant_id, extraction)

    return {
        "event": event_name,
        "source_type": source_type,
        "source_id": str(entity_id),
        "entities_found": len(extraction.entities),
        "relationships_found": len(extraction.relationships),
        **result,
    }


# ─── H-DATA-LIFE: Derived-Data Lifecycle ─────────────────────────────────────

async def propagate_source_deletion(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    source_type: str,
    source_id: uuid.UUID,
) -> dict[str, int]:
    """Propagate correction/delete/erasure of a source to derived data (H-DATA-LIFE).

    When a source (wiki article, DMS file, mail, message) is deleted or corrected,
    this removes:
    - RAG chunks referencing the source
    - Embeddings referencing the source
    - Graph relationships with metadata.source_id matching
    - Agent memory entries referencing the source

    Returns counts of what was removed.
    """
    removed = {"graph_relationships": 0, "agent_memory": 0}

    # Remove graph relationships that were auto-extracted from this source
    try:
        from app.plugins.builtins.graph_rag.models import EntityRelationship

        result = await db.execute(
            select(EntityRelationship).where(
                EntityRelationship.tenant_id == tenant_id,
                EntityRelationship.metadata["source_id"].astext == str(source_id),
                EntityRelationship.metadata["source_type"].astext == source_type,
                EntityRelationship.metadata["auto_extracted"].astext == "true",
            )
        )
        rels = result.scalars().all()
        for rel in rels:
            await db.delete(rel)
        removed["graph_relationships"] = len(rels)
    except Exception as e:
        logger.warning("Failed to remove graph relationships for %s/%s: %s", source_type, source_id, e)

    # Remove agent memory entries referencing this source
    try:
        from app.ai.agent_memory import AgentMemory

        result = await db.execute(
            select(AgentMemory).where(
                AgentMemory.tenant_id == tenant_id,
                AgentMemory.metadata["source_type"].astext == source_type,
                AgentMemory.metadata["source_id"].astext == str(source_id),
            )
        )
        memories = result.scalars().all()
        for mem in memories:
            await db.delete(mem)
        removed["agent_memory"] = len(memories)
    except Exception as e:
        logger.warning("Failed to remove agent memory for %s/%s: %s", source_type, source_id, e)

    await db.flush()
    return removed


async def propagate_source_correction(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    source_type: str,
    source_id: uuid.UUID,
) -> dict[str, Any]:
    """Propagate source correction — re-extract knowledge from updated content (H-DATA-LIFE).

    Removes old derived data and re-runs extraction on the updated source.
    """
    # First remove old derived data
    removed = await propagate_source_deletion(db, tenant_id, source_type, source_id)

    # Then re-extract from updated content
    from app.ai.knowledge_sources import fetch_source_content
    content = await fetch_source_content(db, tenant_id, source_type, source_id)
    if content is None:
        return {"removed": removed, "re_extracted": False}

    from app.ai.knowledge_extraction import extract_knowledge, auto_create_relationships
    extraction = await extract_knowledge(
        text=content["text"],
        tenant_id=tenant_id,
        source_type=source_type,
        source_id=str(source_id),
    )
    created = await auto_create_relationships(db, tenant_id, extraction)

    return {
        "removed": removed,
        "re_extracted": True,
        "entities_found": len(extraction.entities),
        "relationships_found": len(extraction.relationships),
        **created,
    }


# ─── H-RET: Knowledge/Memory Retention ───────────────────────────────────────

# Default retention per source type (days). 0 = no retention limit.
DEFAULT_RETENTION_DAYS = {
    "wiki": 0,        # No limit — wiki articles are persistent knowledge
    "dms": 365,       # 1 year for document-derived knowledge
    "mail": 180,      # 6 months for mail-derived knowledge
    "communication": 90,  # 3 months for communication-derived knowledge
}


def get_retention_days(source_type: str) -> int:
    """Get retention period for a knowledge source type (H-RET)."""
    return DEFAULT_RETENTION_DAYS.get(source_type, 180)


async def cleanup_expired_knowledge(
    db: AsyncSession,
    tenant_id: uuid.UUID,
) -> dict[str, int]:
    """Clean up expired knowledge based on retention policy (H-RET).

    Called by ARQ cron job. Removes graph relationships and agent memory
    entries that have exceeded their retention period.
    """
    cleaned = {"graph_relationships": 0, "agent_memory": 0}
    now = datetime.now(UTC)

    # Clean up expired graph relationships
    try:
        from app.plugins.builtins.graph_rag.models import EntityRelationship

        for source_type, retention_days in DEFAULT_RETENTION_DAYS.items():
            if retention_days == 0:
                continue
            cutoff = now - timedelta(days=retention_days)
            result = await db.execute(
                select(EntityRelationship).where(
                    EntityRelationship.tenant_id == tenant_id,
                    EntityRelationship.metadata["source_type"].astext == source_type,
                    EntityRelationship.metadata["auto_extracted"].astext == "true",
                    EntityRelationship.created_at < cutoff,
                )
            )
            rels = result.scalars().all()
            for rel in rels:
                await db.delete(rel)
            cleaned["graph_relationships"] += len(rels)
    except Exception as e:
        logger.warning("Failed to cleanup expired graph relationships: %s", e)

    # Clean up expired agent memory
    try:
        from app.ai.agent_memory import AgentMemory

        for source_type, retention_days in DEFAULT_RETENTION_DAYS.items():
            if retention_days == 0:
                continue
            cutoff = now - timedelta(days=retention_days)
            result = await db.execute(
                select(AgentMemory).where(
                    AgentMemory.tenant_id == tenant_id,
                    AgentMemory.metadata["source_type"].astext == source_type,
                    AgentMemory.created_at < cutoff,
                )
            )
            memories = result.scalars().all()
            for mem in memories:
                await db.delete(mem)
            cleaned["agent_memory"] += len(memories)
    except Exception as e:
        logger.warning("Failed to cleanup expired agent memory: %s", e)

    await db.flush()
    return cleaned


# ─── H-ASK: Ask Knowledge in Workstream ──────────────────────────────────────

async def ask_knowledge(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
    query: str,
    *,
    source_types: list[str] | None = None,
    max_results: int = 5,
) -> dict[str, Any]:
    """Ask Knowledge — RAG query with evidence cards (H-ASK).

    Runs a unified search query, builds evidence references,
    and returns results formatted for workstream display.
    """
    if not query:
        return {"answer": "", "evidence": [], "query": ""}

    try:
        from app.plugins.builtins.unified_search.contracts import SearchContract
        contract = SearchContract
        search_fn = contract.get_function("unified_search")
        if search_fn is None:
            return {"answer": "Search not available", "evidence": [], "query": query}

        results = await search_fn(
            db=db,
            tenant_id=tenant_id,
            query=query,
            entity_type=None,
            limit=max_results * 2,
        )

        # Filter by source types if specified
        if source_types and results:
            results = [r for r in results if r.get("source_type") in source_types]

        # Build evidence references
        from app.ai.knowledge_sources import build_evidence_references
        refs = build_evidence_references(results, max_results=max_results)

        # Build answer from top results
        if not refs:
            return {"answer": "No relevant knowledge found.", "evidence": [], "query": query}

        # Summarize top results
        snippets = [f"- {r.title}: {r.snippet[:150]}" for r in refs[:3]]
        answer = f"Based on {len(refs)} source(s):\n\n" + "\n".join(snippets)

        return {
            "answer": answer,
            "evidence": [r.to_dict() for r in refs],
            "workstream_blocks": [r.to_workstream_block() for r in refs],
            "query": query,
        }
    except Exception as e:
        logger.warning("Ask knowledge failed: %s", e)
        return {"answer": f"Knowledge query failed: {e}", "evidence": [], "query": query}


# ─── H-REV: Review Queue for extracted relationships ─────────────────────────

async def get_review_queue(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    *,
    page: int = 1,
    page_size: int = 20,
) -> dict[str, Any]:
    """Get low-confidence extracted relationships pending review (H-REV)."""
    from app.plugins.builtins.graph_rag.models import EntityRelationship
    from sqlalchemy import func

    query = select(EntityRelationship).where(
        EntityRelationship.tenant_id == tenant_id,
        EntityRelationship.confidence < 0.6,
        EntityRelationship.metadata["auto_extracted"].astext == "true",
        EntityRelationship.metadata["reviewed"].astext != "true",
    )

    count_q = select(func.count()).select_from(query.subquery())
    total = (await db.execute(count_q)).scalar() or 0

    query = query.order_by(EntityRelationship.confidence.asc()).offset((page - 1) * page_size).limit(page_size)
    result = await db.execute(query)
    items = result.scalars().all()

    return {
        "items": [
            {
                "id": str(r.id),
                "source_type": r.source_type,
                "source_id": str(r.source_id),
                "target_type": r.target_type,
                "target_id": str(r.target_id),
                "relationship_type": r.relationship_type,
                "confidence": r.confidence,
                "evidence": (r.metadata or {}).get("evidence", ""),
                "source": (r.metadata or {}).get("source_type", ""),
            }
            for r in items
        ],
        "total": total,
        "page": page,
        "page_size": page_size,
    }


async def approve_relationship(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    relationship_id: uuid.UUID,
    user_id: uuid.UUID,
) -> bool:
    """Approve a low-confidence relationship (H-REV).

    Marks the relationship as reviewed and boosts its confidence.
    """
    from app.plugins.builtins.graph_rag.models import EntityRelationship

    result = await db.execute(
        select(EntityRelationship).where(
            EntityRelationship.id == relationship_id,
            EntityRelationship.tenant_id == tenant_id,
        )
    )
    rel = result.scalar_one_or_none()
    if rel is None:
        return False

    meta = dict(rel.metadata or {})
    meta["reviewed"] = True
    meta["reviewed_by"] = str(user_id)
    meta["reviewed_at"] = datetime.now(UTC).isoformat()
    rel.metadata = meta
    rel.confidence = max(rel.confidence, 0.8)  # Boost confidence after review
    await db.flush()
    return True


async def reject_relationship(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    relationship_id: uuid.UUID,
) -> bool:
    """Reject a low-confidence relationship — delete it (H-REV)."""
    from app.plugins.builtins.graph_rag.models import EntityRelationship

    result = await db.execute(
        select(EntityRelationship).where(
            EntityRelationship.id == relationship_id,
            EntityRelationship.tenant_id == tenant_id,
        )
    )
    rel = result.scalar_one_or_none()
    if rel is None:
        return False

    await db.delete(rel)
    await db.flush()
    return True


__all__ = [
    "should_extract",
    "handle_extraction_event",
    "propagate_source_deletion",
    "propagate_source_correction",
    "get_retention_days",
    "cleanup_expired_knowledge",
    "ask_knowledge",
    "get_review_queue",
    "approve_relationship",
    "reject_relationship",
    "EXTRACTION_TRIGGERS",
    "DEFAULT_RETENTION_DAYS",
]