eveAI/common/services/user/tenant_services.py

from typing import Dict, List

from flask import session, current_app
from sqlalchemy import desc
from sqlalchemy.exc import SQLAlchemyError

from common.extensions import db, cache_manager
from common.models.user import Partner, PartnerTenant, PartnerService, Tenant, TenantConsent, ConsentStatus, \
    ConsentVersion
from common.utils.eveai_exceptions import EveAINoManagementPartnerService
from common.utils.model_logging_utils import set_logging_information
from datetime import datetime as dt, timezone as tz



class TenantServices:
    @staticmethod
    def associate_tenant_with_partner(tenant_id):
        """Associate a tenant with a partner"""
        try:
            partner_id = session['partner']['id']
            # Get partner service (MANAGEMENT_SERVICE type)
            partner = Partner.query.get(partner_id)
            if not partner:
                return

            # Find a management service for this partner
            management_service = next((service for service in session['partner']['services']
                                       if service.get('type') == 'MANAGEMENT_SERVICE'), None)

            if not management_service:
                current_app.logger.error(f"No Management Service defined for partner {partner_id} "
                                         f"while associating tenant {tenant_id} with partner.")
                raise EveAINoManagementPartnerService()

            # Create the association
            tenant_partner = PartnerTenant(
                partner_service_id=management_service['id'],
                tenant_id=tenant_id,
            )
            set_logging_information(tenant_partner, dt.now(tz.utc))

            db.session.add(tenant_partner)
            db.session.commit()

        except SQLAlchemyError as e:
            db.session.rollback()
            current_app.logger.error(f"Error associating tenant {tenant_id} with partner: {str(e)}")
            raise e

    @staticmethod
    def get_available_types_for_tenant(tenant_id: int, config_type: str) -> Dict[str, Dict[str, str]]:
        """
        Get available configuration types for a tenant based on partner relationships

        Args:
            tenant_id: The tenant ID
            config_type: The configuration type ('specialists', 'agents', 'tasks', etc.)

        Returns:
            Dictionary of available types for the tenant
        """
        # Get the appropriate cache handler based on config_type
        cache_handler = None
        if config_type == 'specialists':
            cache_handler = cache_manager.specialists_types_cache
        elif config_type == 'agents':
            cache_handler = cache_manager.agents_types_cache
        elif config_type == 'tasks':
            cache_handler = cache_manager.tasks_types_cache
        elif config_type == 'tools':
            cache_handler = cache_manager.tools_types_cache
        elif config_type == 'catalogs':
            cache_handler = cache_manager.catalogs_types_cache
        elif config_type == 'retrievers':
            cache_handler = cache_manager.retrievers_types_cache
        else:
            raise ValueError(f"Unsupported config type: {config_type}")

        # Get all types with their metadata (including partner info)
        all_types = cache_handler.get_types()

        # Filter to include:
        # 1. Types with no partner (global)
        # 2. Types with partners that have a SPECIALIST_SERVICE relationship with this tenant
        available_partners = TenantServices.get_tenant_partner_specialist_denominators(tenant_id)

        available_types = {
            type_id: info for type_id, info in all_types.items()
            if info.get('partner') is None or info.get('partner') in available_partners
        }

        return available_types

    @staticmethod
    def get_tenant_partner_specialist_denominators(tenant_id: int) -> List[str]:
        """
        Get names of partners that have a SPECIALIST_SERVICE relationship with this tenant, that can be used for
        filtering configurations.

        Args:
            tenant_id: The tenant ID

        Returns:
            List of partner names (tenant names)
        """
        # Find all PartnerTenant relationships for this tenant
        partner_service_denominators = []
        try:
            # Get all partner services of type SPECIALIST_SERVICE
            specialist_services = (
                PartnerService.query
                .filter_by(type='SPECIALIST_SERVICE')
                .all()
            )

            if not specialist_services:
                return []

            # Find tenant relationships with these services
            partner_tenants = (
                PartnerTenant.query
                .filter_by(tenant_id=tenant_id)
                .filter(PartnerTenant.partner_service_id.in_([svc.id for svc in specialist_services]))
                .all()
            )

            # Get the partner names (their tenant names)
            for pt in partner_tenants:
                partner_service = (
                    PartnerService.query
                    .filter_by(id=pt.partner_service_id)
                    .first()
                )

                if partner_service:
                    partner_service_denominators.append(partner_service.configuration.get("specialist_denominator", ""))

        except SQLAlchemyError as e:
            current_app.logger.error(f"Database error retrieving partner names: {str(e)}")

        return partner_service_denominators

    @staticmethod
    def can_use_specialist_type(tenant_id: int, specialist_type: str) -> bool:
            """
            Check if a tenant can use a specific specialist type

            Args:
                tenant_id: The tenant ID
                specialist_type: The specialist type ID

            Returns:
                True if the tenant can use the specialist type, False otherwise
            """
            # Get the specialist type definition
            try:
                specialist_types = cache_manager.specialists_types_cache.get_types()
                specialist_def = specialist_types.get(specialist_type)

                if not specialist_def:
                    return False

                # If it's a global specialist, anyone can use it
                if specialist_def.get('partner') is None:
                    return True

                # If it's a partner-specific specialist, check if tenant has access
                partner_name = specialist_def.get('partner')
                available_partners = TenantServices.get_tenant_partner_specialist_denominators(tenant_id)

                return partner_name in available_partners

            except Exception as e:
                current_app.logger.error(f"Error checking specialist type access: {str(e)}")
                return False

    @staticmethod
    def get_consent_status(tenant_id: int) -> ConsentStatus:
        # Delegate to centralized ConsentService to ensure consistent logic
        from common.services.user.consent_service import ConsentService
        return ConsentService.get_consent_status(tenant_id)