Source code for operators.update_knowledge_base

"""Operator that syncs a data source into a Bedrock Knowledge Base or a Kendra index."""

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.exceptions import AirflowSkipException
import logging
from typing import Optional, Dict




class UpdateKnowledgeBaseOperator(BaseOperator):
    """
    Updates either a Bedrock Knowledge Base or Kendra Index with data from a specified source.

    This operator handles:
    - Kendra index synchronization with status monitoring
    - Bedrock knowledge base data ingestion
    - Concurrent update protection
    - Automatic retry logic for Kendra sync (up to 10 attempts)

    Args:
        data_source_id: ID of the data source to sync/ingest
        knowledge_base_id: Optional Bedrock knowledge base ID
        kendra_index_id: Optional Kendra index ID
    """

    ui_color = "#fcba03"
    logger = logging.getLogger(__name__)
    MAX_SYNC_ATTEMPTS = 10
    SYNC_POLL_INTERVAL = 60  # seconds

    @apply_defaults
    def __init__(
        self,
        data_source_id: str,
        knowledge_base_id: Optional[str] = None,
        kendra_index_id: Optional[str] = None,
        *args,
        **kwargs,
    ):
        super(UpdateKnowledgeBaseOperator, self).__init__(*args, **kwargs)
        self.data_source_id = data_source_id
        self.knowledge_base_id = knowledge_base_id
        self.kendra_index_id = kendra_index_id



    def update_kendra_index(self) -> Dict:
        """
        Syncs Kendra index with specified data source and monitors until completion.
        Retries up to 10 times with 60-second intervals between checks.
        """
        import boto3
        import time

        client = boto3.client("kendra")
        self.logger.info(
            f"Starting Kendra sync - Index: {self.kendra_index_id}, Source: {self.data_source_id}"
        )

        try:
            response = client.start_data_source_sync_job(
                Id=self.data_source_id,
                IndexId=self.kendra_index_id,
            )
            execution_id = response["ExecutionId"]
            self.logger.info(f"Started sync job: {execution_id}")

            for attempt in range(1, self.MAX_SYNC_ATTEMPTS + 1):
                self.logger.info(
                    f"Checking sync status - Attempt {attempt}/{self.MAX_SYNC_ATTEMPTS}"
                )
                jobs = client.list_data_source_sync_jobs(
                    Id=self.data_source_id, IndexId=self.kendra_index_id, MaxResults=1
                ).get("History", [])

                if jobs and jobs[0]["ExecutionId"] == execution_id:
                    status = jobs[0]["Status"]
                    self.logger.info(f"Current status: {status}")

                    if status == "FAILED":
                        error_msg = jobs[0].get("ErrorCode", "Unknown error")
                        self.logger.error(f"Sync failed: {error_msg}")
                        raise Exception(f"Sync job failed: {error_msg}")

                    if status in ["SUCCEEDED", "ABORTED"]:
                        break

                if attempt < self.MAX_SYNC_ATTEMPTS:
                    self.logger.info(
                        f"Waiting {self.SYNC_POLL_INTERVAL}s before next check"
                    )
                    time.sleep(self.SYNC_POLL_INTERVAL)

            self.logger.info(f"Kendra sync completed - Status: {status}")
            return response

        except client.exceptions.ConflictException:
            raise AirflowSkipException("Kendra sync already in progress - skipping")
        except Exception as e:
            self.logger.error(f"Kendra sync error: {str(e)}")
            raise




    def update_knowledge_base(self, description: str) -> Dict:
        """Updates Bedrock knowledge base with specified data source."""
        import boto3

        client = boto3.client("bedrock-agent")
        self.logger.info(f"Starting Bedrock ingestion - KB: {self.knowledge_base_id}")

        try:
            response = client.start_ingestion_job(
                dataSourceId=self.data_source_id,
                description=description,
                knowledgeBaseId=self.knowledge_base_id,
            )
            self.logger.info("Successfully started knowledge base ingestion")
            return response

        except client.exceptions.ConflictException:
            raise AirflowSkipException("Knowledge base ingestion already in progress")
        except Exception as e:
            self.logger.error(f"Knowledge base update error: {str(e)}")
            raise




    def execute(self, context: Dict) -> Dict:
        """Executes the appropriate update operation based on configuration."""
        self.logger.info(
            f"Starting execution - Source: {self.data_source_id}, "
            f"KB: {self.knowledge_base_id}, Kendra: {self.kendra_index_id}"
        )

        if not self.data_source_id:
            raise AirflowSkipException(
                "data_source_id and (knowledge_base_id or kendra_index_id) required"
            )

        if self.kendra_index_id:
            return self.update_kendra_index()

        if self.knowledge_base_id:
            return self.update_knowledge_base(
                description=f"Ingestion from DAG {context['dag'].dag_id} at {context['execution_date']}"
            )