Skip to content

Commit

Permalink
Merge pull request geekan#399 from Justin-ZL/main
Browse files Browse the repository at this point in the history
add:
  • Loading branch information
stellaHSR authored Oct 7, 2023
2 parents e822093 + 304a032 commit 9c0f47c
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 0 deletions.
40 changes: 40 additions & 0 deletions metagpt/tools/moderation.py
Original file line number Diff line number Diff line change
@@ -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"]))
42 changes: 42 additions & 0 deletions tests/metagpt/tools/test_moderation.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 9c0f47c

Please sign in to comment.