gateway/modules/services/serviceAi/mainServiceAi.py

import logging
from typing import Dict, Any, List, Optional, Tuple, 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.datamodels.datamodelExtraction import ChunkResult, ContentExtracted
from modules.datamodels.datamodelWeb import (
    WebResearchRequest,
    WebResearchActionResult,
    WebResearchDocumentData,
    WebResearchActionDocument,
    WebSearchResultItem,
)
from modules.interfaces.interfaceAiObjects import AiObjects
from modules.services.serviceAi.subCoreAi import SubCoreAi
from modules.services.serviceAi.subDocumentProcessing import SubDocumentProcessing
from modules.services.serviceAi.subWebResearch import SubWebResearch
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
    - SubWebResearch: Web research and crawling functionality
    - 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._webResearch = 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 webResearchService(self):
        """Lazy initialization of web research service."""
        if self._webResearch is None:
            logger.info("Lazy initializing SubWebResearch...")
            self._webResearch = SubWebResearch(self.services, self.aiObjects)
        return self._webResearch

    @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)

    # Web Research
    async def webResearch(self, request: WebResearchRequest) -> WebResearchActionResult:
        """Perform web research using interface functions."""
        await self._ensureAiObjectsInitialized()
        return await self.webResearchService.webResearch(request)

    # Core AI Methods - Delegating to SubCoreAi
    async def callAiPlanning(
        self,
        prompt: str,
        placeholders: Optional[List[PromptPlaceholder]] = None,
        options: Optional[AiCallOptions] = None,
        loopInstructionFormat: Optional[str] = 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, options, "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")