Add AWS builtin Tools (langgenius#6721)

Co-authored-by: Yuanbo Li <[email protected]>
Co-authored-by: crazywoola <[email protected]>

api/core/tools/provider/builtin/aws/aws.py

@@ -0,0 +1,25 @@
+from core.tools.errors import ToolProviderCredentialValidationError
+from core.tools.provider.builtin.aws.tools.sagemaker_text_rerank import SageMakerReRankTool
+from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
+
+
+class SageMakerProvider(BuiltinToolProviderController):
+    def _validate_credentials(self, credentials: dict) -> None:
+        try:
+            SageMakerReRankTool().fork_tool_runtime(
+                runtime={
+                    "credentials": credentials,
+                }
+            ).invoke(
+                user_id='',
+                tool_parameters={
+                    "sagemaker_endpoint" : "",
+                    "query": "misaka mikoto",
+                    "candidate_texts" : "hello$$$hello world",
+                    "topk" : 5,
+                    "aws_region" : ""
+                },
+            )
+        except Exception as e:
+            raise ToolProviderCredentialValidationError(str(e))
+

api/core/tools/provider/builtin/aws/aws.yaml

@@ -0,0 +1,15 @@
+identity:
+  author: AWS
+  name: aws
+  label:
+    en_US: AWS
+    zh_Hans: 亚马逊云科技
+    pt_BR: AWS
+  description:
+    en_US: Services on AWS.
+    zh_Hans: 亚马逊云科技的各类服务
+    pt_BR: Services on AWS.
+  icon: icon.svg
+  tags:
+    - search
+credentials_for_provider:

api/core/tools/provider/builtin/aws/tools/apply_guardrail.py

@@ -0,0 +1,83 @@
+import json
+import logging
+from typing import Any, Union
+
+import boto3
+from pydantic import BaseModel, Field
+
+from core.tools.entities.tool_entities import ToolInvokeMessage
+from core.tools.tool.builtin_tool import BuiltinTool
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class GuardrailParameters(BaseModel):
+    guardrail_id: str = Field(..., description="The identifier of the guardrail")
+    guardrail_version: str = Field(..., description="The version of the guardrail")
+    source: str = Field(..., description="The source of the content")
+    text: str = Field(..., description="The text to apply the guardrail to")
+    aws_region: str = Field(default="us-east-1", description="AWS region for the Bedrock client")
+
+class ApplyGuardrailTool(BuiltinTool):
+    def _invoke(self,
+                user_id: str,
+                tool_parameters: dict[str, Any]
+                ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+        """
+        Invoke the ApplyGuardrail tool
+        """
+        try:
+            # Validate and parse input parameters
+            params = GuardrailParameters(**tool_parameters)
+
+            # Initialize AWS client
+            bedrock_client = boto3.client('bedrock-runtime', region_name=params.aws_region)
+
+            # Apply guardrail
+            response = bedrock_client.apply_guardrail(
+                guardrailIdentifier=params.guardrail_id,
+                guardrailVersion=params.guardrail_version,
+                source=params.source,
+                content=[{"text": {"text": params.text}}]
+            )
+
+            # Check for empty response
+            if not response:
+                return self.create_text_message(text="Received empty response from AWS Bedrock.")
+
+            # Process the result
+            action = response.get("action", "No action specified")
+            outputs = response.get("outputs", [])
+            output = outputs[0].get("text", "No output received") if outputs else "No output received"
+            assessments = response.get("assessments", [])
+
+            # Format assessments
+            formatted_assessments = []
+            for assessment in assessments:
+                for policy_type, policy_data in assessment.items():
+                    if isinstance(policy_data, dict) and 'topics' in policy_data:
+                        for topic in policy_data['topics']:
+                            formatted_assessments.append(f"Policy: {policy_type}, Topic: {topic['name']}, Type: {topic['type']}, Action: {topic['action']}")
+                    else:
+                        formatted_assessments.append(f"Policy: {policy_type}, Data: {policy_data}")
+
+            result = f"Action: {action}\n "
+            result += f"Output: {output}\n "
+            if formatted_assessments:
+                result += "Assessments:\n " + "\n ".join(formatted_assessments) + "\n "
+#           result += f"Full response: {json.dumps(response, indent=2, ensure_ascii=False)}"
+
+            return self.create_text_message(text=result)
+
+        except boto3.exceptions.BotoCoreError as e:
+            error_message = f'AWS service error: {str(e)}'
+            logger.error(error_message, exc_info=True)
+            return self.create_text_message(text=error_message)
+        except json.JSONDecodeError as e:
+            error_message = f'JSON parsing error: {str(e)}'
+            logger.error(error_message, exc_info=True)
+            return self.create_text_message(text=error_message)
+        except Exception as e:
+            error_message = f'An unexpected error occurred: {str(e)}'
+            logger.error(error_message, exc_info=True)
+            return self.create_text_message(text=error_message)

api/core/tools/provider/builtin/aws/tools/apply_guardrail.yaml

@@ -0,0 +1,56 @@
+identity:
+  name: apply_guardrail
+  author: AWS
+  label:
+    en_US: Content Moderation Guardrails
+    zh_Hans: 内容审查护栏
+description:
+  human:
+    en_US: Content Moderation Guardrails utilizes the ApplyGuardrail API, a feature of Guardrails for Amazon Bedrock. This API is capable of evaluating input prompts and model responses for all Foundation Models (FMs), including those on Amazon Bedrock, custom FMs, and third-party FMs. By implementing this functionality, organizations can achieve centralized governance across all their generative AI applications, thereby enhancing control and consistency in content moderation.
+    zh_Hans: 内容审查护栏采用 Guardrails for Amazon Bedrock 功能中的 ApplyGuardrail API 。ApplyGuardrail 可以评估所有基础模型(FMs)的输入提示和模型响应，包括 Amazon Bedrock 上的 FMs、自定义 FMs 和第三方 FMs。通过实施这一功能, 组织可以在所有生成式 AI 应用程序中实现集中化的治理，从而增强内容审核的控制力和一致性。
+  llm: Content Moderation Guardrails utilizes the ApplyGuardrail API, a feature of Guardrails for Amazon Bedrock. This API is capable of evaluating input prompts and model responses for all Foundation Models (FMs), including those on Amazon Bedrock, custom FMs, and third-party FMs. By implementing this functionality, organizations can achieve centralized governance across all their generative AI applications, thereby enhancing control and consistency in content moderation.
+parameters:
+  - name: guardrail_id
+    type: string
+    required: true
+    label:
+      en_US: Guardrail ID
+      zh_Hans: Guardrail ID
+    human_description:
+      en_US: Please enter the ID of the Guardrail that has already been created on Amazon Bedrock, for example 'qk5nk0e4b77b'.
+      zh_Hans: 请输入已经在 Amazon Bedrock 上创建好的 Guardrail ID, 例如 'qk5nk0e4b77b'.
+    llm_description: Please enter the ID of the Guardrail that has already been created on Amazon Bedrock, for example 'qk5nk0e4b77b'.
+    form: form
+  - name: guardrail_version
+    type: string
+    required: true
+    label:
+      en_US: Guardrail Version Number
+      zh_Hans: Guardrail 版本号码
+    human_description:
+      en_US: Please enter the published version of the Guardrail ID that has already been created on Amazon Bedrock. This is typically a version number, such as 2.
+      zh_Hans: 请输入已经在Amazon Bedrock 上创建好的Guardrail ID发布的版本, 通常使用版本号, 例如2.
+    llm_description: Please enter the published version of the Guardrail ID that has already been created on Amazon Bedrock. This is typically a version number, such as 2.
+    form: form
+  - name: source
+    type: string
+    required: true
+    label:
+      en_US: Content Source (INPUT or OUTPUT)
+      zh_Hans: 内容来源 (INPUT or OUTPUT)
+    human_description:
+      en_US: The source of data used in the request to apply the guardrail. Valid Values "INPUT | OUTPUT"
+      zh_Hans: 用于应用护栏的请求中所使用的数据来源。有效值为 "INPUT | OUTPUT"
+    llm_description: The source of data used in the request to apply the guardrail. Valid Values "INPUT | OUTPUT"
+    form: form
+  - name: text
+    type: string
+    required: true
+    label:
+      en_US: Content to be reviewed
+      zh_Hans: 待审查内容
+    human_description:
+      en_US: The content used for requesting guardrail review, which can be either user input or LLM output.
+      zh_Hans: 用于请求护栏审查的内容，可以是用户输入或 LLM 输出。
+    llm_description: The content used for requesting guardrail review, which can be either user input or LLM output.
+    form: llm

api/core/tools/provider/builtin/aws/tools/lambda_translate_utils.py

@@ -0,0 +1,88 @@
+import json
+from typing import Any, Union
+
+import boto3
+
+from core.tools.entities.tool_entities import ToolInvokeMessage
+from core.tools.tool.builtin_tool import BuiltinTool
+
+
+class LambdaTranslateUtilsTool(BuiltinTool):
+    lambda_client: Any = None
+
+    def _invoke_lambda(self, text_content, src_lang, dest_lang, model_id, dictionary_name, request_type, lambda_name):
+        msg = { 
+            "src_content":text_content, 
+            "src_lang": src_lang, 
+            "dest_lang":dest_lang, 
+            "dictionary_id": dictionary_name,
+            "request_type" : request_type, 
+            "model_id" : model_id
+        }
+
+        invoke_response = self.lambda_client.invoke(FunctionName=lambda_name,
+                                               InvocationType='RequestResponse',
+                                               Payload=json.dumps(msg))
+        response_body = invoke_response['Payload']
+
+        response_str = response_body.read().decode("unicode_escape")
+
+        return response_str
+
+    def _invoke(self, 
+                user_id: str, 
+               tool_parameters: dict[str, Any], 
+        ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+        """
+            invoke tools
+        """
+        line = 0
+        try:
+            if not self.lambda_client:
+                aws_region = tool_parameters.get('aws_region')
+                if aws_region:
+                    self.lambda_client = boto3.client("lambda", region_name=aws_region)
+                else:
+                    self.lambda_client = boto3.client("lambda")
+
+            line = 1
+            text_content = tool_parameters.get('text_content', '')
+            if not text_content:
+                return self.create_text_message('Please input text_content')
+
+            line = 2
+            src_lang = tool_parameters.get('src_lang', '')
+            if not src_lang:
+                return self.create_text_message('Please input src_lang')
+
+            line = 3
+            dest_lang = tool_parameters.get('dest_lang', '')
+            if not dest_lang:
+                return self.create_text_message('Please input dest_lang')
+
+            line = 4
+            lambda_name = tool_parameters.get('lambda_name', '')
+            if not lambda_name:
+                return self.create_text_message('Please input lambda_name')
+
+            line = 5
+            request_type = tool_parameters.get('request_type', '')
+            if not request_type:
+                return self.create_text_message('Please input request_type')
+
+            line = 6
+            model_id = tool_parameters.get('model_id', '')
+            if not model_id:
+                return self.create_text_message('Please input model_id')
+
+            line = 7
+            dictionary_name = tool_parameters.get('dictionary_name', '')
+            if not dictionary_name:
+                return self.create_text_message('Please input dictionary_name')
+
+            result = self._invoke_lambda(text_content, src_lang, dest_lang, model_id, dictionary_name, request_type, lambda_name)
+
+            return self.create_text_message(text=result)
+
+        except Exception as e:
+            return self.create_text_message(f'Exception {str(e)}, line : {line}')