bedrock_agent.py

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from aws_lambda_powertools.utilities.parser.envelopes.base import BaseEnvelope
from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel, BedrockAgentFunctionEventModel

if TYPE_CHECKING:
    from aws_lambda_powertools.utilities.parser.types import T

logger = logging.getLogger(__name__)


class BedrockAgentEnvelope(BaseEnvelope):
    """Bedrock Agent envelope to extract data within input_text key"""

    def parse(self, data: dict[str, Any] | Any | None, model: type[T]) -> T | None:
        """Parses data found with model provided

        Parameters
        ----------
        data : dict
            Lambda event to be parsed
        model : type[T]
            Data model provided to parse after extracting data using envelope

        Returns
        -------
        T | None
            Parsed detail payload with model provided
        """
        logger.debug(f"Parsing incoming data with Bedrock Agent model {BedrockAgentEventModel}")
        parsed_envelope: BedrockAgentEventModel = BedrockAgentEventModel.model_validate(data)
        logger.debug(f"Parsing event payload in `input_text` with {model}")
        return self._parse(data=parsed_envelope.input_text, model=model)


class BedrockAgentFunctionEnvelope(BaseEnvelope):
    """Bedrock Agent Function envelope to extract data within input_text key"""

    def parse(self, data: dict[str, Any] | Any | None, model: type[T]) -> T | None:
        """Parses data found with model provided

        Parameters
        ----------
        data : dict
            Lambda event to be parsed
        model : type[T]
            Data model provided to parse after extracting data using envelope

        Returns
        -------
        T | None
            Parsed detail payload with model provided
        """
        logger.debug(f"Parsing incoming data with Bedrock Agent Function model {BedrockAgentFunctionEventModel}")
        parsed_envelope: BedrockAgentFunctionEventModel = BedrockAgentFunctionEventModel.model_validate(data)
        logger.debug(f"Parsing event payload in `input_text` with {model}")
        return self._parse(data=parsed_envelope.input_text, model=model)