codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4from typing import Optional, Type5 6from langchain_core.callbacks import CallbackManagerForToolRun7from pydantic import BaseModel, Field8 9from langchain_community.tools.edenai.edenai_base_tool import EdenaiTool10 11logger = logging.getLogger(__name__)12 13 14class TextModerationInput(BaseModel):15    query: str = Field(description="Text to moderate")16 17 18class EdenAiTextModerationTool(EdenaiTool):19    """Tool that queries the Eden AI Explicit text detection.20 21    for api reference check edenai documentation:22    https://docs.edenai.co/reference/image_explicit_content_create.23 24    To use, you should have25    the environment variable ``EDENAI_API_KEY`` set with your API token.26    You can find your token here: https://app.edenai.run/admin/account/settings27 28    """29 30    name: str = "edenai_explicit_content_detection_text"31    description: str = (32        "A wrapper around edenai Services explicit content detection for text. "33        """Useful for when you have to scan text for offensive, 34        sexually explicit or suggestive content,35        it checks also if there is any content of self-harm,36        violence, racist or hate speech."""37        """the structure of the output is : 38        'the type of the explicit content : the likelihood of it being explicit'39        the likelihood is a number 40        between 1 and 5, 1 being the lowest and 5 the highest.41        something is explicit if the likelihood is equal or higher than 3.42        for example : 43        nsfw_likelihood: 144        this is not explicit.45        for example : 46        nsfw_likelihood: 347        this is explicit.48        """49        "Input should be a string."50    )51    args_schema: Type[BaseModel] = TextModerationInput52 53    language: str54 55    feature: str = "text"56    subfeature: str = "moderation"57 58    def _parse_response(self, response: list) -> str:59        formatted_result = []60        for result in response:61            if "nsfw_likelihood" in result.keys():62                formatted_result.append(63                    "nsfw_likelihood: " + str(result["nsfw_likelihood"])64                )65 66            for label, likelihood in zip(result["label"], result["likelihood"]):67                formatted_result.append(f'"{label}": {str(likelihood)}')68 69        return "\n".join(formatted_result)70 71    def _run(72        self,73        query: str,74        run_manager: Optional[CallbackManagerForToolRun] = None,75    ) -> str:76        """Use the tool."""77        query_params = {"text": query, "language": self.language}78        return self._call_eden_ai(query_params)79