"""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.contracts 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.meta::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.meta::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.meta::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.meta::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": {}, }