from typing import Any, Dict, List, Optional from langchain.chains.base import Chain from langchain_core.callbacks.manager import CallbackManagerForChainRun from langchain_experimental.comprehend_moderation.base_moderation import BaseModeration from langchain_experimental.comprehend_moderation.base_moderation_callbacks import ( BaseModerationCallbackHandler, ) from langchain_experimental.comprehend_moderation.base_moderation_config import ( BaseModerationConfig, ) from langchain_experimental.pydantic_v1 import root_validator class AmazonComprehendModerationChain(Chain): """Moderation Chain, based on `Amazon Comprehend` service. See more at https://aws.amazon.com/comprehend/ """ output_key: str = "output" #: :meta private: """Key used to fetch/store the output in data containers. Defaults to `output`""" input_key: str = "input" #: :meta private: """Key used to fetch/store the input in data containers. Defaults to `input`""" moderation_config: BaseModerationConfig = BaseModerationConfig() """ Configuration settings for moderation, defaults to BaseModerationConfig with default values """ client: Optional[Any] = None """boto3 client object for connection to Amazon Comprehend""" region_name: Optional[str] = None """The aws region e.g., `us-west-2`. Fallsback to AWS_DEFAULT_REGION env variable or region specified in ~/.aws/config in case it is not provided here. """ credentials_profile_name: Optional[str] = None """The name of the profile in the ~/.aws/credentials or ~/.aws/config files, which has either access keys or role information specified. If not specified, the default credential profile or, if on an EC2 instance, credentials from IMDS will be used. See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html """ moderation_callback: Optional[BaseModerationCallbackHandler] = None """Callback handler for moderation, this is different from regular callbacks which can be used in addition to this.""" unique_id: Optional[str] = None """A unique id that can be used to identify or group a user or session""" @root_validator(pre=True) def create_client(cls, values: Dict[str, Any]) -> Dict[str, Any]: """ Creates an Amazon Comprehend client. Args: values (Dict[str, Any]): A dictionary containing configuration values. Returns: Dict[str, Any]: A dictionary with the updated configuration values, including the Amazon Comprehend client. Raises: ModuleNotFoundError: If the 'boto3' package is not installed. ValueError: If there is an issue importing 'boto3' or loading AWS credentials. Example: .. code-block:: python config = { "credentials_profile_name": "my-profile", "region_name": "us-west-2" } updated_config = create_client(config) comprehend_client = updated_config["client"] """ if values.get("client") is not None: return values try: import boto3 if values.get("credentials_profile_name"): session = boto3.Session(profile_name=values["credentials_profile_name"]) else: # use default credentials session = boto3.Session() client_params = {} if values.get("region_name"): client_params["region_name"] = values["region_name"] values["client"] = session.client("comprehend", **client_params) return values except ImportError: raise ModuleNotFoundError( "Could not import boto3 python package. " "Please install it with `pip install boto3`." ) except Exception as e: raise ValueError( "Could not load credentials to authenticate with AWS client. " "Please check that credentials in the specified " f"profile name are valid. {e}" ) from e @property def output_keys(self) -> List[str]: """ Returns a list of output keys. This method defines the output keys that will be used to access the output values produced by the chain or function. It ensures that the specified keys are available to access the outputs. Returns: List[str]: A list of output keys. Note: This method is considered private and may not be intended for direct external use. """ return [self.output_key] @property def input_keys(self) -> List[str]: """ Returns a list of input keys expected by the prompt. This method defines the input keys that the prompt expects in order to perform its processing. It ensures that the specified keys are available for providing input to the prompt. Returns: List[str]: A list of input keys. Note: This method is considered private and may not be intended for direct external use. """ return [self.input_key] def _call( self, inputs: Dict[str, Any], run_manager: Optional[CallbackManagerForChainRun] = None, ) -> Dict[str, str]: """ Executes the moderation process on the input text and returns the processed output. This internal method performs the moderation process on the input text. It converts the input prompt value to plain text, applies the specified filters, and then converts the filtered output back to a suitable prompt value object. Additionally, it provides the option to log information about the run using the provided `run_manager`. Args: inputs: A dictionary containing input values run_manager: A run manager to handle run-related events. Default is None Returns: Dict[str, str]: A dictionary containing the processed output of the moderation process. Raises: ValueError: If there is an error during the moderation process """ if run_manager: run_manager.on_text("Running AmazonComprehendModerationChain...\n") moderation = BaseModeration( client=self.client, config=self.moderation_config, moderation_callback=self.moderation_callback, unique_id=self.unique_id, run_manager=run_manager, ) response = moderation.moderate(prompt=inputs[self.input_keys[0]]) return {self.output_key: response}