mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 19:59:50 +08:00
Add AWS builtin Tools (#6721)
Co-authored-by: Yuanbo Li <ybalbert@amazon.com> Co-authored-by: crazywoola <427733928@qq.com>
This commit is contained in:
ybalbert001
GitHub
parent
8eb0d0fddd
commit
c23aa50bea
9
api/core/tools/provider/builtin/aws/_assets/icon.svg
Normal file
9
api/core/tools/provider/builtin/aws/_assets/icon.svg
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<!-- Uploaded to: SVG Repo, www.svgrepo.com, Generator: SVG Repo Mixer Tools -->
|
||||
<svg width="800px" height="800px" viewBox="0 0 16 16" xmlns="http://www.w3.org/2000/svg" fill="none">
|
||||
<path fill="#252F3E" d="M4.51 7.687c0 .197.02.357.058.475.042.117.096.245.17.384a.233.233 0 01.037.123c0 .053-.032.107-.1.16l-.336.224a.255.255 0 01-.138.048c-.054 0-.107-.026-.16-.074a1.652 1.652 0 01-.192-.251 4.137 4.137 0 01-.165-.315c-.415.491-.936.737-1.564.737-.447 0-.804-.129-1.064-.385-.261-.256-.394-.598-.394-1.025 0-.454.16-.822.484-1.1.325-.278.756-.416 1.304-.416.18 0 .367.016.564.042.197.027.4.07.612.118v-.39c0-.406-.085-.689-.25-.854-.17-.166-.458-.246-.868-.246-.186 0-.377.022-.574.07a4.23 4.23 0 00-.575.181 1.525 1.525 0 01-.186.07.326.326 0 01-.085.016c-.075 0-.112-.054-.112-.166v-.262c0-.085.01-.15.037-.186a.399.399 0 01.15-.113c.185-.096.409-.176.67-.24.26-.07.537-.101.83-.101.633 0 1.096.144 1.394.432.293.288.442.726.442 1.314v1.73h.01zm-2.161.811c.175 0 .356-.032.548-.096.191-.064.362-.182.505-.342a.848.848 0 00.181-.341c.032-.129.054-.283.054-.465V7.03a4.43 4.43 0 00-.49-.09 3.996 3.996 0 00-.5-.033c-.357 0-.618.07-.793.214-.176.144-.26.347-.26.614 0 .25.063.437.196.566.128.133.314.197.559.197zm4.273.577c-.096 0-.16-.016-.202-.054-.043-.032-.08-.106-.112-.208l-1.25-4.127a.938.938 0 01-.049-.214c0-.085.043-.133.128-.133h.522c.1 0 .17.016.207.053.043.032.075.107.107.208l.894 3.535.83-3.535c.026-.106.058-.176.1-.208a.365.365 0 01.214-.053h.425c.102 0 .17.016.213.053.043.032.08.107.101.208l.841 3.578.92-3.578a.458.458 0 01.107-.208.346.346 0 01.208-.053h.495c.085 0 .133.043.133.133 0 .027-.006.054-.01.086a.76.76 0 01-.038.133l-1.283 4.127c-.032.107-.069.177-.111.209a.34.34 0 01-.203.053h-.457c-.101 0-.17-.016-.213-.053-.043-.038-.08-.107-.101-.214L8.213 5.37l-.82 3.439c-.026.107-.058.176-.1.213-.043.038-.118.054-.213.054h-.458zm6.838.144a3.51 3.51 0 01-.82-.096c-.266-.064-.473-.134-.612-.214-.085-.048-.143-.101-.165-.15a.378.378 0 01-.031-.149v-.272c0-.112.042-.166.122-.166a.3.3 0 01.096.016c.032.011.08.032.133.054.18.08.378.144.585.187.213.042.42.064.633.064.336 0 .596-.059.777-.176a.575.575 0 00.277-.508.52.52 0 00-.144-.373c-.095-.102-.276-.193-.537-.278l-.772-.24c-.388-.123-.676-.305-.851-.545a1.275 1.275 0 01-.266-.774c0-.224.048-.422.143-.593.096-.17.224-.32.384-.438.16-.122.34-.213.553-.277.213-.064.436-.091.67-.091.118 0 .24.005.357.021.122.016.234.038.346.06.106.026.208.052.303.085.096.032.17.064.224.096a.46.46 0 01.16.133.289.289 0 01.047.176v.251c0 .112-.042.171-.122.171a.552.552 0 01-.202-.064 2.427 2.427 0 00-1.022-.208c-.303 0-.543.048-.708.15-.165.1-.25.256-.25.475 0 .149.053.277.16.379.106.101.303.202.585.293l.756.24c.383.123.66.294.825.513.165.219.244.47.244.748 0 .23-.047.437-.138.619a1.436 1.436 0 01-.388.47c-.165.133-.362.23-.591.299-.24.075-.49.112-.761.112z"/>
|
||||
<g fill="#F90" fill-rule="evenodd" clip-rule="evenodd">
|
||||
<path d="M14.465 11.813c-1.75 1.297-4.294 1.986-6.481 1.986-3.065 0-5.827-1.137-7.913-3.027-.165-.15-.016-.353.18-.235 2.257 1.313 5.04 2.109 7.92 2.109 1.941 0 4.075-.406 6.039-1.239.293-.133.543.192.255.406z"/>
|
||||
<path d="M15.194 10.98c-.223-.287-1.479-.138-2.048-.069-.17.022-.197-.128-.043-.24 1-.705 2.645-.502 2.836-.267.192.24-.053 1.89-.99 2.68-.143.123-.281.06-.218-.1.213-.53.687-1.72.463-2.003z"/>
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 3.3 KiB |
25
api/core/tools/provider/builtin/aws/aws.py
Normal file
25
api/core/tools/provider/builtin/aws/aws.py
Normal file
|
|
@ -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))
|
||||
|
||||
15
api/core/tools/provider/builtin/aws/aws.yaml
Normal file
15
api/core/tools/provider/builtin/aws/aws.yaml
Normal file
|
|
@ -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:
|
||||
83
api/core/tools/provider/builtin/aws/tools/apply_guardrail.py
Normal file
83
api/core/tools/provider/builtin/aws/tools/apply_guardrail.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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}')
|
||||
|
|
@ -0,0 +1,134 @@
|
|||
identity:
|
||||
name: lambda_translate_utils
|
||||
author: AWS
|
||||
label:
|
||||
en_US: TranslateTool
|
||||
zh_Hans: 翻译工具
|
||||
pt_BR: TranslateTool
|
||||
icon: icon.svg
|
||||
description:
|
||||
human:
|
||||
en_US: A util tools for LLM translation, extra deployment is needed on AWS. Please refer Github Repo - https://github.com/ybalbert001/dynamodb-rag
|
||||
zh_Hans: 大语言模型翻译工具(专词映射获取)，需要在AWS上进行额外部署，可参考Github Repo - https://github.com/ybalbert001/dynamodb-rag
|
||||
pt_BR: A util tools for LLM translation, specfic Lambda Function deployment is needed on AWS. Please refer Github Repo - https://github.com/ybalbert001/dynamodb-rag
|
||||
llm: A util tools for translation.
|
||||
parameters:
|
||||
- name: text_content
|
||||
type: string
|
||||
required: true
|
||||
label:
|
||||
en_US: source content for translation
|
||||
zh_Hans: 待翻译原文
|
||||
pt_BR: source content for translation
|
||||
human_description:
|
||||
en_US: source content for translation
|
||||
zh_Hans: 待翻译原文
|
||||
pt_BR: source content for translation
|
||||
llm_description: source content for translation
|
||||
form: llm
|
||||
- name: src_lang
|
||||
type: string
|
||||
required: true
|
||||
label:
|
||||
en_US: source language code
|
||||
zh_Hans: 原文语言代号
|
||||
pt_BR: source language code
|
||||
human_description:
|
||||
en_US: source language code
|
||||
zh_Hans: 原文语言代号
|
||||
pt_BR: source language code
|
||||
llm_description: source language code
|
||||
form: llm
|
||||
- name: dest_lang
|
||||
type: string
|
||||
required: true
|
||||
label:
|
||||
en_US: target language code
|
||||
zh_Hans: 目标语言代号
|
||||
pt_BR: target language code
|
||||
human_description:
|
||||
en_US: target language code
|
||||
zh_Hans: 目标语言代号
|
||||
pt_BR: target language code
|
||||
llm_description: target language code
|
||||
form: llm
|
||||
- name: aws_region
|
||||
type: string
|
||||
required: false
|
||||
label:
|
||||
en_US: region of Lambda
|
||||
zh_Hans: Lambda 所在的region
|
||||
pt_BR: region of Lambda
|
||||
human_description:
|
||||
en_US: region of Lambda
|
||||
zh_Hans: Lambda 所在的region
|
||||
pt_BR: region of Lambda
|
||||
llm_description: region of Lambda
|
||||
form: form
|
||||
- name: model_id
|
||||
type: string
|
||||
required: false
|
||||
default: anthropic.claude-3-sonnet-20240229-v1:0
|
||||
label:
|
||||
en_US: LLM model_id in bedrock
|
||||
zh_Hans: bedrock上的大语言模型model_id
|
||||
pt_BR: LLM model_id in bedrock
|
||||
human_description:
|
||||
en_US: LLM model_id in bedrock
|
||||
zh_Hans: bedrock上的大语言模型model_id
|
||||
pt_BR: LLM model_id in bedrock
|
||||
llm_description: LLM model_id in bedrock
|
||||
form: form
|
||||
- name: dictionary_name
|
||||
type: string
|
||||
required: false
|
||||
label:
|
||||
en_US: dictionary name for term mapping
|
||||
zh_Hans: 专词映射表名称
|
||||
pt_BR: dictionary name for term mapping
|
||||
human_description:
|
||||
en_US: dictionary name for term mapping
|
||||
zh_Hans: 专词映射表名称
|
||||
pt_BR: dictionary name for term mapping
|
||||
llm_description: dictionary name for term mapping
|
||||
form: form
|
||||
- name: request_type
|
||||
type: select
|
||||
required: false
|
||||
label:
|
||||
en_US: request type
|
||||
zh_Hans: 请求类型
|
||||
pt_BR: request type
|
||||
human_description:
|
||||
en_US: request type
|
||||
zh_Hans: 请求类型
|
||||
pt_BR: request type
|
||||
default: term_mapping
|
||||
options:
|
||||
- value: term_mapping
|
||||
label:
|
||||
en_US: term_mapping
|
||||
zh_Hans: 专词映射
|
||||
- value: segment_only
|
||||
label:
|
||||
en_US: segment_only
|
||||
zh_Hans: 仅切词
|
||||
- value: translate
|
||||
label:
|
||||
en_US: translate
|
||||
zh_Hans: 翻译内容
|
||||
form: form
|
||||
- name: lambda_name
|
||||
type: string
|
||||
default: "translate_tool"
|
||||
required: true
|
||||
label:
|
||||
en_US: AWS Lambda for term mapping retrieval
|
||||
zh_Hans: 专词召回映射 - AWS Lambda
|
||||
pt_BR: lambda name for term mapping retrieval
|
||||
human_description:
|
||||
en_US: AWS Lambda for term mapping retrieval
|
||||
zh_Hans: 专词召回映射 - AWS Lambda
|
||||
pt_BR: AWS Lambda for term mapping retrieval
|
||||
llm_description: AWS Lambda for term mapping retrieval
|
||||
form: form
|
||||
|
|
@ -0,0 +1,86 @@
|
|||
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 SageMakerReRankTool(BuiltinTool):
|
||||
sagemaker_client: Any = None
|
||||
sagemaker_endpoint:str = None
|
||||
topk:int = None
|
||||
|
||||
def _sagemaker_rerank(self, query_input: str, docs: list[str], rerank_endpoint:str):
|
||||
inputs = [query_input]*len(docs)
|
||||
response_model = self.sagemaker_client.invoke_endpoint(
|
||||
EndpointName=rerank_endpoint,
|
||||
Body=json.dumps(
|
||||
{
|
||||
"inputs": inputs,
|
||||
"docs": docs
|
||||
}
|
||||
),
|
||||
ContentType="application/json",
|
||||
)
|
||||
json_str = response_model['Body'].read().decode('utf8')
|
||||
json_obj = json.loads(json_str)
|
||||
scores = json_obj['scores']
|
||||
return scores if isinstance(scores, list) else [scores]
|
||||
|
||||
def _invoke(self,
|
||||
user_id: str,
|
||||
tool_parameters: dict[str, Any],
|
||||
) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
|
||||
"""
|
||||
invoke tools
|
||||
"""
|
||||
line = 0
|
||||
try:
|
||||
if not self.sagemaker_client:
|
||||
aws_region = tool_parameters.get('aws_region')
|
||||
if aws_region:
|
||||
self.sagemaker_client = boto3.client("sagemaker-runtime", region_name=aws_region)
|
||||
else:
|
||||
self.sagemaker_client = boto3.client("sagemaker-runtime")
|
||||
|
||||
line = 1
|
||||
if not self.sagemaker_endpoint:
|
||||
self.sagemaker_endpoint = tool_parameters.get('sagemaker_endpoint')
|
||||
|
||||
line = 2
|
||||
if not self.topk:
|
||||
self.topk = tool_parameters.get('topk', 5)
|
||||
|
||||
line = 3
|
||||
query = tool_parameters.get('query', '')
|
||||
if not query:
|
||||
return self.create_text_message('Please input query')
|
||||
|
||||
line = 4
|
||||
candidate_texts = tool_parameters.get('candidate_texts')
|
||||
if not candidate_texts:
|
||||
return self.create_text_message('Please input candidate_texts')
|
||||
|
||||
line = 5
|
||||
candidate_docs = json.loads(candidate_texts)
|
||||
docs = [ item.get('content') for item in candidate_docs ]
|
||||
|
||||
line = 6
|
||||
scores = self._sagemaker_rerank(query_input=query, docs=docs, rerank_endpoint=self.sagemaker_endpoint)
|
||||
|
||||
line = 7
|
||||
for idx in range(len(candidate_docs)):
|
||||
candidate_docs[idx]["score"] = scores[idx]
|
||||
|
||||
line = 8
|
||||
sorted_candidate_docs = sorted(candidate_docs, key=lambda x: x['score'], reverse=True)
|
||||
|
||||
line = 9
|
||||
results_str = json.dumps(sorted_candidate_docs[:self.topk], ensure_ascii=False)
|
||||
return self.create_text_message(text=results_str)
|
||||
|
||||
except Exception as e:
|
||||
return self.create_text_message(f'Exception {str(e)}, line : {line}')
|
||||
|
||||
|
|
@ -0,0 +1,82 @@
|
|||
identity:
|
||||
name: sagemaker_text_rerank
|
||||
author: AWS
|
||||
label:
|
||||
en_US: SagemakerRerank
|
||||
zh_Hans: Sagemaker重排序
|
||||
pt_BR: SagemakerRerank
|
||||
icon: icon.svg
|
||||
description:
|
||||
human:
|
||||
en_US: A tool for performing text similarity ranking. You can find deploy notebook on Github Repo - https://github.com/aws-samples/dify-aws-tool
|
||||
zh_Hans: Sagemaker重排序工具, 请参考 Github Repo - https://github.com/aws-samples/dify-aws-tool上的部署脚本
|
||||
pt_BR: A tool for performing text similarity ranking.
|
||||
llm: A tool for performing text similarity ranking. You can find deploy notebook on Github Repo - https://github.com/aws-samples/dify-aws-tool
|
||||
parameters:
|
||||
- name: sagemaker_endpoint
|
||||
type: string
|
||||
required: true
|
||||
label:
|
||||
en_US: sagemaker endpoint for reranking
|
||||
zh_Hans: 重排序的SageMaker 端点
|
||||
pt_BR: sagemaker endpoint for reranking
|
||||
human_description:
|
||||
en_US: sagemaker endpoint for reranking
|
||||
zh_Hans: 重排序的SageMaker 端点
|
||||
pt_BR: sagemaker endpoint for reranking
|
||||
llm_description: sagemaker endpoint for reranking
|
||||
form: form
|
||||
- name: query
|
||||
type: string
|
||||
required: true
|
||||
label:
|
||||
en_US: Query string
|
||||
zh_Hans: 查询语句
|
||||
pt_BR: Query string
|
||||
human_description:
|
||||
en_US: key words for searching
|
||||
zh_Hans: 查询关键词
|
||||
pt_BR: key words for searching
|
||||
llm_description: key words for searching
|
||||
form: llm
|
||||
- name: candidate_texts
|
||||
type: string
|
||||
required: true
|
||||
label:
|
||||
en_US: text candidates
|
||||
zh_Hans: 候选文本
|
||||
pt_BR: text candidates
|
||||
human_description:
|
||||
en_US: searched candidates by query
|
||||
zh_Hans: 查询文本搜到候选文本
|
||||
pt_BR: searched candidates by query
|
||||
llm_description: searched candidates by query
|
||||
form: llm
|
||||
- name: topk
|
||||
type: number
|
||||
required: false
|
||||
form: form
|
||||
label:
|
||||
en_US: Limit for results count
|
||||
zh_Hans: 返回个数限制
|
||||
pt_BR: Limit for results count
|
||||
human_description:
|
||||
en_US: Limit for results count
|
||||
zh_Hans: 返回个数限制
|
||||
pt_BR: Limit for results count
|
||||
min: 1
|
||||
max: 10
|
||||
default: 5
|
||||
- name: aws_region
|
||||
type: string
|
||||
required: false
|
||||
label:
|
||||
en_US: region of sagemaker endpoint
|
||||
zh_Hans: SageMaker 端点所在的region
|
||||
pt_BR: region of sagemaker endpoint
|
||||
human_description:
|
||||
en_US: region of sagemaker endpoint
|
||||
zh_Hans: SageMaker 端点所在的region
|
||||
pt_BR: region of sagemaker endpoint
|
||||
llm_description: region of sagemaker endpoint
|
||||
form: form
|
||||
Loading…
Reference in New Issue
Block a user