leocrm/app/plugins/builtins/unified_search/provider_registry.py

"""SearchProvider protocol and global registry."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Protocol, runtime_checkable

from sqlalchemy import select

if TYPE_CHECKING:
    import uuid
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)


@runtime_checkable
class SearchProvider(Protocol):
    """Protocol for entity-specific search providers."""

    entity_type: str
    supports_fts: bool
    supports_vector: bool
    supports_rag: bool
    supports_graph: bool

    async def search_fts(
        self,
        db: AsyncSession,
        tsquery: str,
        tenant_id: uuid.UUID,
        limit: int,
    ) -> list[dict]:
        """Full-text search using PostgreSQL tsquery."""
        ...

    async def search_vector(
        self,
        db: AsyncSession,
        embedding: list[float],
        tenant_id: uuid.UUID,
        limit: int,
    ) -> list[dict]:
        """Semantic vector search using pgvector."""
        ...

    async def get_embedding_text(
        self,
        db: AsyncSession,
        entity_id: uuid.UUID,
        tenant_id: uuid.UUID,
    ) -> str:
        """Get text representation for embedding generation."""
        ...

    def to_search_result(self, entity: object) -> dict:
        """Convert an ORM entity to a search result dict."""
        ...


class SearchProviderRegistry:
    """In-memory registry of search providers."""

    def __init__(self) -> None:
        self._providers: dict[str, SearchProvider] = {}

    def register(self, provider: SearchProvider) -> None:
        """Register a search provider."""
        self._providers[provider.entity_type] = provider
        logger.debug("Registered search provider: %s", provider.entity_type)

    def unregister(self, entity_type: str) -> None:
        """Unregister a search provider by entity type."""
        self._providers.pop(entity_type, None)

    def get(self, entity_type: str) -> SearchProvider | None:
        """Get a provider by entity type."""
        return self._providers.get(entity_type)

    def get_all(self) -> list[SearchProvider]:
        """Get all registered providers."""
        return list(self._providers.values())

    def get_entity_types(self) -> list[str]:
        """Get all registered entity types."""
        return list(self._providers.keys())

    def get_providers_by_capability(self, capability: str) -> list[SearchProvider]:
        """Get all providers that support a given capability (fts/vector/rag/graph)."""
        flag_attr = f"supports_{capability}"
        return [p for p in self._providers.values() if getattr(p, flag_attr, False)]

    def get_capabilities(self, entity_type: str) -> dict[str, bool]:
        """Get capability flags for a specific entity type."""
        provider = self.get(entity_type)
        if provider is None:
            return {"fts": False, "vector": False, "rag": False, "graph": False}
        return {
            "fts": getattr(provider, "supports_fts", False),
            "vector": getattr(provider, "supports_vector", False),
            "rag": getattr(provider, "supports_rag", False),
            "graph": getattr(provider, "supports_graph", False),
        }

    def clear(self) -> None:
        """Clear all registered providers."""
        self._providers.clear()


# Global singleton
_registry = SearchProviderRegistry()


def get_search_registry() -> SearchProviderRegistry:
    """Get the global search provider registry."""
    return _registry


async def auto_register_providers(db: AsyncSession) -> None:
    """Auto-register providers for active plugins.

    Checks which plugins are active and registers corresponding providers.
    """
    from app.plugins.builtins.unified_search.providers.contact_provider import (
        ContactSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.mail_provider import (
        MailSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.file_provider import (
        FileSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.event_provider import (
        EventSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.company_provider import (
        CompanySearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.task_provider import (
        TaskSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.contactperson_provider import (
        ContactPersonSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.tag_provider import (
        TagSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.conversation_provider import (
        ConversationSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.user_provider import (
        UserSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.agent_memory_provider import (
        AgentMemorySearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.ai_chat_provider import (
        AIChatSearchProvider,
    )
    from app.plugins.builtins.unified_search.providers.workflow_provider import (
        WorkflowSearchProvider,
    )
    from app.plugins.builtins.graph_rag.contracts import GraphRagContract
    GraphRAGSearchProvider = GraphRagContract.GraphRAGSearchProvider

    registry = get_search_registry()
    registry.clear()

    # Register all built-in providers
    for provider_cls in [
        ContactSearchProvider,
        MailSearchProvider,
        FileSearchProvider,
        EventSearchProvider,
        CompanySearchProvider,
        TaskSearchProvider,
        ContactPersonSearchProvider,
        TagSearchProvider,
        ConversationSearchProvider,
        UserSearchProvider,
        AgentMemorySearchProvider,
        AIChatSearchProvider,
        WorkflowSearchProvider,
        GraphRAGSearchProvider,
    ]:
        try:
            registry.register(provider_cls())
        except Exception:
            logger.exception("Failed to register %s", provider_cls.__name__)