"""SearchProvider protocol and global registry.""" from __future__ import annotations import logging from typing import TYPE_CHECKING, Protocol, runtime_checkable 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.graph_rag.contracts import GraphRagContract 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.company_provider import ( CompanySearchProvider, ) from app.plugins.builtins.unified_search.providers.contact_provider import ( ContactSearchProvider, ) from app.plugins.builtins.unified_search.providers.contactperson_provider import ( ContactPersonSearchProvider, ) from app.plugins.builtins.unified_search.providers.conversation_provider import ( ConversationSearchProvider, ) from app.plugins.builtins.unified_search.providers.event_provider import ( EventSearchProvider, ) from app.plugins.builtins.unified_search.providers.file_provider import ( FileSearchProvider, ) from app.plugins.builtins.unified_search.providers.mail_provider import ( MailSearchProvider, ) from app.plugins.builtins.unified_search.providers.tag_provider import ( TagSearchProvider, ) from app.plugins.builtins.unified_search.providers.task_provider import ( TaskSearchProvider, ) from app.plugins.builtins.unified_search.providers.user_provider import ( UserSearchProvider, ) from app.plugins.builtins.unified_search.providers.workflow_provider import ( WorkflowSearchProvider, ) graph_rag_search_provider = 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, graph_rag_search_provider, ]: try: registry.register(provider_cls()) except Exception: logger.exception("Failed to register %s", provider_cls.__name__)