gateway/modules/services/serviceAi/mainServiceAi.py

import logging
from typing import Dict, Any, List, Optional, Union
from modules.datamodels.datamodelChat import PromptPlaceholder, ChatDocument
from modules.services.serviceExtraction.mainServiceExtraction import ExtractionService
from modules.datamodels.datamodelAi import AiCallRequest, AiCallOptions, OperationTypeEnum, PriorityEnum
from modules.aicore.aicorePluginTavily import WebResearchRequest, WebResearchResult
from modules.interfaces.interfaceAiObjects import AiObjects
from modules.services.serviceAi.subCoreAi import SubCoreAi
from modules.services.serviceAi.subDocumentProcessing import SubDocumentProcessing
from modules.services.serviceAi.subDocumentGeneration import SubDocumentGeneration
from modules.services.serviceAi.subSharedAiUtils import sanitizePromptContent


logger = logging.getLogger(__name__)

class AiService:
    """Lightweight AI service orchestrator that delegates to specialized sub-modules.
    
    Manager delegates to specialized sub-modules:
    - SubCoreAi: Core AI operations (readImage, generateImage, callAi, planning, text calls)
    - SubDocumentProcessing: Document chunking, processing, and merging logic
    - SubDocumentGeneration: Single-file and multi-file document generation
    
    The main service  acts as a coordinator:
    1. Manages lazy initialization of sub-modules
    2. Delegates operations to appropriate sub-modules
    3. Maintains the same public API for backward compatibility
    """

    def __init__(self, serviceCenter=None) -> None:
        """Initialize AI service with service center access.
        
        Args:
            serviceCenter: Service center instance for accessing other services
        """
        self.services = serviceCenter
        # Only depend on interfaces
        self.aiObjects = None  # Will be initialized in create()
        self._extractionService = None  # Lazy initialization
        self._coreAi = None  # Lazy initialization
        self._documentProcessor = None  # Lazy initialization
        self._documentGenerator = None  # Lazy initialization

    @property
    def extractionService(self):
        """Lazy initialization of extraction service."""
        if self._extractionService is None:
            logger.info("Lazy initializing ExtractionService...")
            self._extractionService = ExtractionService(self.services)
        return self._extractionService

    @property
    def coreAi(self):
        """Lazy initialization of core AI service."""
        if self._coreAi is None:
            if self.aiObjects is None:
                raise RuntimeError("AiService.aiObjects must be initialized before accessing coreAi. Use await AiService.create() or await service._ensureAiObjectsInitialized()")
            logger.info("Lazy initializing SubCoreAi...")
            self._coreAi = SubCoreAi(self.services, self.aiObjects)
        return self._coreAi

    @property
    def documentProcessor(self):
        """Lazy initialization of document processing service."""
        if self._documentProcessor is None:
            logger.info("Lazy initializing SubDocumentProcessing...")
            self._documentProcessor = SubDocumentProcessing(self.services, self.aiObjects)
        return self._documentProcessor


    @property
    def documentGenerator(self):
        """Lazy initialization of document generation service."""
        if self._documentGenerator is None:
            logger.info("Lazy initializing SubDocumentGeneration...")
            self._documentGenerator = SubDocumentGeneration(self.services, self.aiObjects, self.documentProcessor)
        return self._documentGenerator

    async def _ensureAiObjectsInitialized(self):
        """Ensure aiObjects is initialized."""
        if self.aiObjects is None:
            logger.info("Lazy initializing AiObjects...")
            self.aiObjects = await AiObjects.create()
            logger.info("AiObjects initialization completed")

    @classmethod
    async def create(cls, serviceCenter=None) -> "AiService":
        """Create AiService instance with all connectors initialized."""
        logger.info("AiService.create() called")
        instance = cls(serviceCenter)
        logger.info("AiService created, about to call AiObjects.create()...")
        instance.aiObjects = await AiObjects.create()
        logger.info("AiObjects.create() completed")
        return instance

    # AI Image Analysis
    async def readImage(
        self,
        prompt: str,
        imageData: Union[str, bytes],
        mimeType: str = None,
        options: Optional[AiCallOptions] = None,
        ) -> str:
        """Call AI for image analysis using interface.callImage()."""
        await self._ensureAiObjectsInitialized()
        return await self.coreAi.readImage(prompt, imageData, mimeType, options)

    # AI Image Generation
    async def generateImage(
        self,
        prompt: str,
        size: str = "1024x1024",
        quality: str = "standard",
        style: str = "vivid",
        options: Optional[AiCallOptions] = None,
        ) -> Dict[str, Any]:
        """Generate an image using AI using interface.generateImage()."""
        await self._ensureAiObjectsInitialized()
        return await self.coreAi.generateImage(prompt, size, quality, style, options)


    # Core AI Methods - Delegating to SubCoreAi
    async def callAiPlanning(
        self,
        prompt: str,
        placeholders: Optional[List[PromptPlaceholder]] = None
        ) -> str:
        """Planning AI call for task planning, action planning, action selection, etc."""
        await self._ensureAiObjectsInitialized()
        # Always use "json" for planning calls since they return JSON
        return await self.coreAi.callAiPlanning(prompt, placeholders, "json")

    async def callAiDocuments(
        self,
        prompt: str,
        documents: Optional[List[ChatDocument]] = None,
        options: Optional[AiCallOptions] = None,
        outputFormat: Optional[str] = None,
        title: Optional[str] = None
        ) -> Union[str, Dict[str, Any]]:
        """Document generation AI call for all non-planning calls."""
        await self._ensureAiObjectsInitialized()
        # Use "json" for document generation calls since they return JSON
        return await self.coreAi.callAiDocuments(prompt, documents, options, outputFormat, title, "json")

    def sanitizePromptContent(self, content: str, contentType: str = "text") -> str:
        """Sanitize prompt content to prevent injection attacks and ensure safe presentation."""
        return sanitizePromptContent(content, contentType)