From 304a03244c66affc575acadea340f93d9ae25b66 Mon Sep 17 00:00:00 2001
From: zhanglei <zhanglei@fuzhi.ai>
Date: Sat, 7 Oct 2023 16:56:19 +0800
Subject: [PATCH] =?UTF-8?q?add=EF=BC=9A=201.moderation=20tools=202.unittes?=
 =?UTF-8?q?t?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 metagpt/tools/moderation.py            | 40 ++++++++++++++++++++++++
 tests/metagpt/tools/test_moderation.py | 42 ++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)
 create mode 100644 metagpt/tools/moderation.py
 create mode 100644 tests/metagpt/tools/test_moderation.py

diff --git a/metagpt/tools/moderation.py b/metagpt/tools/moderation.py
new file mode 100644
index 000000000..c56a6afc4
--- /dev/null
+++ b/metagpt/tools/moderation.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/9/26 14:27
+@Author  : zhanglei
+@File    : moderation.py
+"""
+from typing import Union
+
+from metagpt.llm import LLM
+
+
+class Moderation:
+    def __init__(self):
+        self.llm = LLM()
+
+    def moderation(self, content: Union[str, list[str]]):
+        resp = []
+        if content:
+            moderation_results = self.llm.moderation(content=content)
+            results = moderation_results.results
+            for item in results:
+                resp.append(item.flagged)
+
+        return resp
+
+    async def amoderation(self, content: Union[str, list[str]]):
+        resp = []
+        if content:
+            moderation_results = await self.llm.amoderation(content=content)
+            results = moderation_results.results
+            for item in results:
+                resp.append(item.flagged)
+
+        return resp
+
+
+if __name__ == "__main__":
+    moderation = Moderation()
+    print(moderation.moderation(content=["I will kill you", "The weather is really nice today", "I want to hit you"]))
diff --git a/tests/metagpt/tools/test_moderation.py b/tests/metagpt/tools/test_moderation.py
new file mode 100644
index 000000000..225acff75
--- /dev/null
+++ b/tests/metagpt/tools/test_moderation.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/9/26 14:46
+@Author  : zhanglei
+@File    : test_translate.py
+"""
+
+import pytest
+
+from metagpt.tools.moderation import Moderation
+
+
+@pytest.mark.parametrize(
+    ("content",),
+    [
+        [
+            ["I will kill you", "The weather is really nice today", "I want to hit you"],
+        ]
+    ],
+)
+def test_moderation(content):
+    moderation = Moderation()
+    results = moderation.moderation(content=content)
+    assert isinstance(results, list)
+    assert len(results) == len(content)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("content",),
+    [
+        [
+            ["I will kill you", "The weather is really nice today", "I want to hit you"],
+        ]
+    ],
+)
+async def test_amoderation(content):
+    moderation = Moderation()
+    results = await moderation.amoderation(content=content)
+    assert isinstance(results, list)
+    assert len(results) == len(content)