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

"""GraphRAGSearchProvider — registers GraphRAG as a search provider in unified_search."""

from __future__ import annotations

import logging
import uuid
from typing import Any

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from app.plugins.builtins.unified_search.base_provider import BaseSearchProvider

logger = logging.getLogger(__name__)


class GraphRAGSearchProvider(BaseSearchProvider):
    """Search provider for GraphRAG entity relationships.

    Enables full-text and semantic search over relationship metadata
    and entity types in the knowledge graph.
    """

    entity_type = "graph_relationship"

    async def _search_fts_filtered(
        self,
        db: AsyncSession,
        tsquery: str,
        tenant_id: uuid.UUID,
        limit: int,
        visible_ids: set[uuid.UUID] | None,
    ) -> list[dict[str, Any]]:
        """Full-text search on relationship metadata and types."""
        if visible_ids is not None:
            sql = text(
                """
                SELECT r.*, ts_rank(
                    to_tsvector('pg_catalog.german',
                        coalesce(r.relationship_type, '') || ' ' ||
                        coalesce(r.source_type, '') || ' ' ||
                        coalesce(r.target_type, '') || ' ' ||
                        coalesce(r.metadata::text, '')
                    ),
                    to_tsquery('pg_catalog.german', :q)
                ) AS rank
                FROM entity_relationships r
                WHERE r.tenant_id = :tid
                  AND r.deleted_at IS NULL
                  AND to_tsvector('pg_catalog.german',
                      coalesce(r.relationship_type, '') || ' ' ||
                      coalesce(r.source_type, '') || ' ' ||
                      coalesce(r.target_type, '') || ' ' ||
                      coalesce(r.metadata::text, '')
                  ) @@ to_tsquery('pg_catalog.german', :q)
                  AND r.id = ANY(:visible_ids)
                ORDER BY rank DESC
                LIMIT :lim
                """
            )
            result = await db.execute(
                sql,
                {
                    "q": tsquery,
                    "tid": tenant_id,
                    "lim": limit,
                    "visible_ids": list(visible_ids),
                },
            )
        else:
            sql = text(
                """
                SELECT r.*, ts_rank(
                    to_tsvector('pg_catalog.german',
                        coalesce(r.relationship_type, '') || ' ' ||
                        coalesce(r.source_type, '') || ' ' ||
                        coalesce(r.target_type, '') || ' ' ||
                        coalesce(r.metadata::text, '')
                    ),
                    to_tsquery('pg_catalog.german', :q)
                ) AS rank
                FROM entity_relationships r
                WHERE r.tenant_id = :tid
                  AND r.deleted_at IS NULL
                  AND to_tsvector('pg_catalog.german',
                      coalesce(r.relationship_type, '') || ' ' ||
                      coalesce(r.source_type, '') || ' ' ||
                      coalesce(r.target_type, '') || ' ' ||
                      coalesce(r.metadata::text, '')
                  ) @@ to_tsquery('pg_catalog.german', :q)
                ORDER BY rank DESC
                LIMIT :lim
                """
            )
            result = await db.execute(
                sql,
                {"q": tsquery, "tid": tenant_id, "lim": limit},
            )
        rows = result.mappings().all()
        return [dict(r) for r in rows]

    async def _search_vector_filtered(
        self,
        db: AsyncSession,
        embedding: list[float],
        tenant_id: uuid.UUID,
        limit: int,
        visible_ids: set[uuid.UUID] | None,
    ) -> list[dict[str, Any]]:
        """Semantic search is not yet supported for graph relationships.

        Returns empty list — relationships are searched via FTS on metadata.
        """
        return []

    async def get_embedding_text(
        self, db: AsyncSession, entity_id: uuid.UUID, tenant_id: uuid.UUID
    ) -> str:
        """Get text for embedding generation."""
        sql = text(
            """
            SELECT relationship_type, source_type, source_id, target_type, target_id, metadata
            FROM entity_relationships
            WHERE id = :eid AND tenant_id = :tid
            """
        )
        result = await db.execute(sql, {"eid": entity_id, "tid": tenant_id})
        row = result.mappings().first()
        if not row:
            return ""
        parts = [
            row.get("relationship_type", ""),
            row.get("source_type", ""),
            str(row.get("source_id", "")),
            row.get("target_type", ""),
            str(row.get("target_id", "")),
            str(row.get("metadata", {})),
        ]
        return " ".join(str(p) for p in parts if p)

    def to_search_result(self, entity: object) -> dict[str, Any]:
        """Convert relationship to search result dict."""
        if isinstance(entity, dict):
            return {
                "entity_type": self.entity_type,
                "entity_id": str(entity.get("id", "")),
                "title": f"{entity.get('source_type', '?')} --[{entity.get('relationship_type', '?')}]--> {entity.get('target_type', '?')}",
                "snippet": str(entity.get("metadata", {})),
                "score": float(entity.get("rank", 0.0)),
                "data": {
                    "source_type": entity.get("source_type"),
                    "source_id": str(entity.get("source_id", "")),
                    "target_type": entity.get("target_type"),
                    "target_id": str(entity.get("target_id", "")),
                    "relationship_type": entity.get("relationship_type"),
                },
            }
        return {
            "entity_type": self.entity_type,
            "entity_id": str(getattr(entity, "id", "")),
            "title": f"{getattr(entity, 'source_type', '?')} --[{getattr(entity, 'relationship_type', '?')}]--> {getattr(entity, 'target_type', '?')}",
            "snippet": str(getattr(entity, "metadata", {})),
            "score": 0.0,
            "data": {},
        }