Initializer

class AugmentationInitializer(EmbeddingInitializer):
    """Initializer for the augmentation process.

    Extends the EmbeddingInitializer to set up the environment for augmentation tasks.
    This initializer is responsible for loading the required configuration and
    registering all necessary components with the dependency injection container.

    Multiple components are used in the embedding, augmentation and evaluation processes.
    To avoid code duplication, this initializer is used to bind the components to the injector.
    It is intended to be subclassed by the specific initializers for each process.
    """

    def __init__(
        self,
        configuration_class: Type[
            BaseConfiguration
        ] = AugmentationConfiguration,
        package_loader: BasePackageLoader = AugmentationPackageLoader(),
    ):
        """Initialize the AugmentationInitializer.

        Args:
            configuration_class: The configuration class to use for loading settings.
                                Defaults to AugmentationConfiguration.
            package_loader: Package loader instance responsible for loading required packages.
                           Defaults to a new AugmentationPackageLoader instance.
        """
        super().__init__(
            configuration_class=configuration_class,
            package_loader=package_loader,
        )
        self.scheduler = AugmentationScheduler(
            configuration=self.get_configuration(),
            logger=LoggerConfiguration.get_logger(__name__),
        )
        self._initialize_default_prompt()

    def get_scheduler(self) -> AugmentationScheduler:
        """Get the scheduler instance for managing scheduled tasks.

        Returns:
            AsyncIOScheduler: The scheduler instance used for scheduling jobs.
        """
        return self.scheduler

    def _initialize_default_prompt(self) -> None:
        """
        Initialize the default prompt templates for the augmentation process managed by Langfuse.
        """
        configuration = self.get_configuration()
        langfuse_prompt_service = LangfusePromptServiceFactory.create(
            configuration=configuration.augmentation.langfuse
        )

        # Use LlamaIndex's default condense prompt
        # Domain-specific prompts can be configured in Langfuse or loaded from prompts/ directory
        langfuse_prompt_service.create_prompt_if_not_exists(
            prompt_name="default_condense_prompt",
            prompt_template="""Given the following conversation between a user and an AI assistant and a follow up question from user,
rephrase the follow up question to be a standalone question.

Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:""",
        )

        langfuse_prompt_service.create_prompt_if_not_exists(
            prompt_name="default_context_prompt",
            prompt_template=DEFAULT_CONTEXT_PROMPT_TEMPLATE,
        )

        langfuse_prompt_service.create_prompt_if_not_exists(
            prompt_name="default_context_refine_prompt",
            prompt_template=DEFAULT_CONTEXT_REFINE_PROMPT_TEMPLATE,
        )

        # Generic system prompt for RAG assistant
        # Domain-specific prompts (e.g., Bundestag) should be configured in Langfuse
        # or loaded from prompts/ directory for specific deployments
        langfuse_prompt_service.create_prompt_if_not_exists(
            prompt_name="default_system_prompt",
            prompt_template="""You are a helpful AI assistant that answers questions based on the provided context documents.

CRITICAL: GROUNDING IN RETRIEVED DOCUMENTS
- Base your answers ONLY on the information in the provided context documents
- Do not use your training data or prior knowledge unless explicitly asked
- If the retrieved documents do not contain information about a topic, clearly state this
- Always cite or reference the source documents when possible

IMPORTANT GUIDELINES:
- Provide accurate, objective information based on the retrieved documents
- If information is ambiguous or conflicting, acknowledge this
- Be concise but thorough in your responses
- Maintain a helpful and professional tone""",
        )

        langfuse_prompt_service.create_prompt_if_not_exists(
            prompt_name="default_input_guardrail_prompt",
            prompt_template=DEFAULT_INPUT_GUARDRAIL_PROMPT_TEMPLATE,
        )

        langfuse_prompt_service.create_prompt_if_not_exists(
            prompt_name="default_output_guardrail_prompt",
            prompt_template=DEFAULT_OUTPUT_GUARDRAIL_PROMPT_TEMPLATE,
        )

class AugmentationPackageLoader(EmbeddingPackageLoader):
    """Package loader for augmentation components.

    Extends the EmbeddingPackageLoader to load additional packages required
    for the augmentation process, including LLMs, retrievers,
    postprocessors, and chat engines.
    """

    def __init__(
        self, logger: logging.Logger = LoggerConfiguration.get_logger(__name__)
    ):
        """Initialize the AugmentationPackageLoader.

        Args:
            logger: Logger instance for logging information. Defaults to a logger
                   configured with the current module name.
        """
        super().__init__(logger)

    def load_packages(self) -> None:
        """Load all required packages for augmentation.

        Calls the parent class's load_packages method first to load embedding packages,
        then loads additional packages specific to augmentation.
        """
        super().load_packages()
        self._load_packages(
            [
                "src.augmentation.components.guardrails",
                "src.augmentation.components.llms",
                "src.augmentation.components.retrievers",
                "src.augmentation.components.postprocessors",
                "src.augmentation.components.chat_engines",
            ]
        )

class AugmentationScheduler:

    def __init__(
        self,
        configuration: AugmentationConfiguration,
        logger: logging.Logger = LoggerConfiguration.get_logger(__name__),
    ):
        """
        Args:
            configuration (AugmentationConfiguration): The configuration object for the augmentation process.
            logger: Logger instance for logging information. Defaults to a logger
                   configured with the current module name.
        """
        self.logger = logger
        self.configuration = configuration

        self.scheduler = AsyncIOScheduler()

    def start(self) -> None:
        """Start the scheduler and schedule the daily queries retention job."""
        langfuse_configuration = self.configuration.augmentation.langfuse
        queries_retention_job = LangfuseRenetionJobFactory.create(
            langfuse_configuration
        )

        try:
            self.scheduler.add_job(
                queries_retention_job.run,
                CronTrigger.from_crontab(
                    langfuse_configuration.retention_job.crontab
                ),
                id=langfuse_configuration.retention_job.name,
                replace_existing=True,
            )
            self.logger.info(
                "Daily queries retention job scheduled successfully"
            )

            self.scheduler.start()
            self.logger.info("Scheduler started successfully")
        except Exception as e:
            self.logger.error(f"Failed to initialize scheduler: {e}")

    def stop(self) -> None:
        """Stop the scheduler if it is running."""
        if self.scheduler.running:
            self.scheduler.shutdown()
            self.logger.info("Scheduler stopped successfully")
        else:
            self.logger.warning("Scheduler was not running")