leocrm/app/plugins/builtins/graph_rag/services.py

"""GraphRAG services — relationship management and BFS graph traversal."""

from __future__ import annotations

import uuid
from collections import deque
from typing import Any

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

from app.plugins.builtins.graph_rag.models import EntityRelationship


async def create_relationship(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    source_type: str,
    source_id: uuid.UUID,
    target_type: str,
    target_id: uuid.UUID,
    relationship_type: str,
    metadata: dict[str, Any] | None = None,
    owner_id: uuid.UUID | None = None,
) -> dict[str, Any]:
    """Create a new relationship between two entities.

    Args:
        db: Database session.
        tenant_id: Tenant UUID.
        source_type: Entity type of the source.
        source_id: UUID of the source entity.
        target_type: Entity type of the target.
        target_id: UUID of the target entity.
        relationship_type: Type of relationship.
        metadata: Optional metadata dict.
        owner_id: Optional owner user UUID.

    Returns:
        Dict with the created relationship data.
    """
    # Check for duplicate
    existing = await db.execute(
        select(EntityRelationship).where(
            EntityRelationship.tenant_id == tenant_id,
            EntityRelationship.source_type == source_type,
            EntityRelationship.source_id == source_id,
            EntityRelationship.target_type == target_type,
            EntityRelationship.target_id == target_id,
            EntityRelationship.relationship_type == relationship_type,
            EntityRelationship.deleted_at.is_(None),
        )
    )
    if existing.scalar_one_or_none() is not None:
        return {"error": "Relationship already exists", "code": "duplicate"}

    rel = EntityRelationship(
        tenant_id=tenant_id,
        source_type=source_type,
        source_id=source_id,
        target_type=target_type,
        target_id=target_id,
        relationship_type=relationship_type,
        meta=metadata or {},
        owner_id=owner_id,
    )
    db.add(rel)
    await db.flush()
    await db.refresh(rel)

    return {
        "id": str(rel.id),
        "source_type": rel.source_type,
        "source_id": str(rel.source_id),
        "target_type": rel.target_type,
        "target_id": str(rel.target_id),
        "relationship_type": rel.relationship_type,
        "metadata": rel.meta,
        "owner_id": str(rel.owner_id) if rel.owner_id else None,
        "created_at": rel.created_at.isoformat() if rel.created_at else None,
    }


async def traverse_graph(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    source_type: str,
    source_id: uuid.UUID,
    max_hops: int = 3,
    relationship_types: list[str] | None = None,
) -> dict[str, Any]:
    """BFS traversal of the knowledge graph from a starting node.

    Uses Breadth-First Search to find all entities reachable within
    `max_hops` steps, following directed relationships.

    Args:
        db: Database session.
        tenant_id: Tenant UUID.
        source_type: Entity type of the starting node.
        source_id: UUID of the starting node.
        max_hops: Maximum traversal depth (1-10).
        relationship_types: Optional filter to follow only specific relationship types.

    Returns:
        Dict with nodes and edges discovered during traversal.
    """
    visited: set[tuple[str, str]] = set()
    nodes: list[dict[str, Any]] = []
    edges: list[dict[str, Any]] = []

    # BFS queue: (source_type, source_id, depth, path)
    queue: deque[tuple[str, str, int, list[str]]] = deque()
    start_key = (source_type, str(source_id))
    visited.add(start_key)
    nodes.append({
        "entity_type": source_type,
        "entity_id": str(source_id),
        "depth": 0,
        "path": [],
    })
    queue.append((source_type, str(source_id), 0, []))

    while queue:
        current_type, current_id, depth, path = queue.popleft()

        if depth >= max_hops:
            continue

        # Build query for outgoing relationships
        stmt = select(EntityRelationship).where(
            EntityRelationship.tenant_id == tenant_id,
            EntityRelationship.source_type == current_type,
            EntityRelationship.source_id == uuid.UUID(current_id),
            EntityRelationship.deleted_at.is_(None),
        )
        if relationship_types:
            stmt = stmt.where(EntityRelationship.relationship_type.in_(relationship_types))

        result = await db.execute(stmt)
        relationships = result.scalars().all()

        for rel in relationships:
            target_key = (rel.target_type, str(rel.target_id))
            new_path = path + [rel.relationship_type]

            edges.append({
                "source_type": rel.source_type,
                "source_id": str(rel.source_id),
                "target_type": rel.target_type,
                "target_id": str(rel.target_id),
                "relationship_type": rel.relationship_type,
                "metadata": rel.meta,
            })

            if target_key not in visited:
                visited.add(target_key)
                nodes.append({
                    "entity_type": rel.target_type,
                    "entity_id": str(rel.target_id),
                    "depth": depth + 1,
                    "path": new_path,
                })
                queue.append((rel.target_type, str(rel.target_id), depth + 1, new_path))

        # Also traverse incoming relationships (bidirectional graph)
        stmt_in = select(EntityRelationship).where(
            EntityRelationship.tenant_id == tenant_id,
            EntityRelationship.target_type == current_type,
            EntityRelationship.target_id == uuid.UUID(current_id),
            EntityRelationship.deleted_at.is_(None),
        )
        if relationship_types:
            stmt_in = stmt_in.where(EntityRelationship.relationship_type.in_(relationship_types))

        result_in = await db.execute(stmt_in)
        relationships_in = result_in.scalars().all()

        for rel in relationships_in:
            source_key = (rel.source_type, str(rel.source_id))
            new_path = path + [f"inverse_{rel.relationship_type}"]

            edges.append({
                "source_type": rel.source_type,
                "source_id": str(rel.source_id),
                "target_type": rel.target_type,
                "target_id": str(rel.target_id),
                "relationship_type": rel.relationship_type,
                "metadata": rel.meta,
            })

            if source_key not in visited:
                visited.add(source_key)
                nodes.append({
                    "entity_type": rel.source_type,
                    "entity_id": str(rel.source_id),
                    "depth": depth + 1,
                    "path": new_path,
                })
                queue.append((rel.source_type, str(rel.source_id), depth + 1, new_path))

    return {
        "nodes": nodes,
        "edges": edges,
        "total_nodes": len(nodes),
        "total_edges": len(edges),
    }


async def delete_relationship(
    db: AsyncSession,
    tenant_id: uuid.UUID,
    relationship_id: uuid.UUID,
) -> bool:
    """Delete a relationship by ID."""
    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)
    return True