From dd0904f95c97dab8404f71e06212ee182db9c554 Mon Sep 17 00:00:00 2001
From: jiandanfeng <chenjh3@wangsu.com>
Date: Wed, 22 Jan 2025 19:26:25 +0800
Subject: [PATCH] feat: add giteeAI risk control identification. (#12946)

---
 .../builtin/gitee_ai/tools/risk-control.py    | 26 +++++++++++++++
 .../builtin/gitee_ai/tools/risk-control.yaml  | 32 +++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 api/core/tools/provider/builtin/gitee_ai/tools/risk-control.py
 create mode 100644 api/core/tools/provider/builtin/gitee_ai/tools/risk-control.yaml

diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.py b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.py
new file mode 100644
index 0000000000..e3558ce699
--- /dev/null
+++ b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.py
@@ -0,0 +1,26 @@
+from typing import Any, Union
+
+import requests
+
+from core.tools.entities.tool_entities import ToolInvokeMessage
+from core.tools.tool.builtin_tool import BuiltinTool
+
+
+class GiteeAIToolRiskControl(BuiltinTool):
+    def _invoke(
+        self, user_id: str, tool_parameters: dict[str, Any]
+    ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+        headers = {
+            "content-type": "application/json",
+            "authorization": f"Bearer {self.runtime.credentials['api_key']}",
+        }
+
+        inputs = [{"type": "text", "text": tool_parameters.get("input-text")}]
+        model = tool_parameters.get("model", "Security-semantic-filtering")
+        payload = {"model": model, "input": inputs}
+        url = "https://ai.gitee.com/v1/moderations"
+        response = requests.post(url, json=payload, headers=headers)
+        if response.status_code != 200:
+            return self.create_text_message(f"Got Error Response:{response.text}")
+
+        return [self.create_text_message(response.content.decode("utf-8"))]
diff --git a/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.yaml b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.yaml
new file mode 100644
index 0000000000..6e7229dc1c
--- /dev/null
+++ b/api/core/tools/provider/builtin/gitee_ai/tools/risk-control.yaml
@@ -0,0 +1,32 @@
+identity:
+  name: risk control
+  author: gitee_ai
+  label:
+    en_US: risk control identification
+    zh_Hans: 风控识别
+  icon: icon.svg
+description:
+  human:
+    en_US: Ensuring the protection and compliance of sensitive information through the filtering and analysis of data semantics
+    zh_Hans: 通过对数据语义的过滤和分析，确保敏感信息的保护和合规性
+  llm: This tool is used to risk control identification.
+parameters:
+  - name: model
+    type: string
+    required: true
+    default: Security-semantic-filtering
+    label:
+      en_US: Service Model
+      zh_Hans: 服务模型
+    form: form
+  - name: input-text
+    type: string
+    required: true
+    label:
+      en_US: Input Text
+      zh_Hans: 输入文本
+    human_description:
+      en_US: The text input for filtering and analysis.
+      zh_Hans: 用于分析过滤的文本
+    llm_description: The text input for filtering and analysis.
+    form: llm