Merge pull request #182 from hermitgreen/master

新增文生图敏感词风控判断逻辑
baidubce · Mar 13, 2024 · fad4cd4 · fad4cd4
2 parents 6e3cbcc + ba31333
commit fad4cd4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/appbuilder/core/_exception.py b/appbuilder/core/_exception.py
@@ -89,3 +89,9 @@ class InvalidRequestArgumentError(BaseRPCException):
     r"""InvalidRequestArgumentError invalid request param
     """
     pass
+
+
+class RiskInputException(BaseRPCException):
+    r"""RiskInputException
+    """
+    pass
diff --git a/appbuilder/core/components/text_to_image/component.py b/appbuilder/core/components/text_to_image/component.py
@@ -20,7 +20,7 @@
 from appbuilder.core.component import Component
 from appbuilder.core.message import Message
 from appbuilder.core._client import HTTPClient
-from appbuilder.core._exception import AppBuilderServerException
+from appbuilder.core._exception import AppBuilderServerException, RiskInputException
 from appbuilder.core.components.text_to_image.model import Text2ImageSubmitRequest, Text2ImageQueryRequest, \
     Text2ImageQueryResponse, Text2ImageSubmitResponse, Text2ImageOutMessage, Text2ImageInMessage
 
@@ -197,13 +197,17 @@ def tool_eval(
             query = origin_query
         try:
             result = self.run(Message({"prompt": query}))
-            result = {
-                'event': 'image',
-                'type': 'image',
-                'text': result.content['img_urls'][0],
-            }
         except Exception as e:
-            raise AppBuilderServerException(f'绘图时发生错误：{e}')
+            raise AppBuilderServerException(f'绘图服务发生错误：{e}')
+
+        if len(result.content['img_urls']) == 0:
+            raise RiskInputException(f'query：{query} 中可能存在敏感词')
+
+        result = {
+            'event': 'image',
+            'type': 'image',
+            'text': result.content['img_urls'][0],
+        }
         if streaming:
             yield result
         else: