def build_safety_settings()

in src/dfcx_scrapi/core/scrapi_base.py [0:0]
54 lines of code
5 McCabe index (conditional complexity)

    def build_safety_settings(
            self,
            safety_config: Optional[Dict[str, str]] = None
            ):
        """Build a SafetyConfig payload for Gemini SDK calls.

        If a safety_config dict is not provided, we'll set the defaults to OFF.

        Args:
            safety_config: (Optional) A key / value pair dictionary containing
            the category names and thresholds to set for each category. If not
            provided, the default threshold will be used as defined below.
        """
        # https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters#harm_categories
        CATEGORY_MAP = {
            "hate_speech": HarmCategory.HARM_CATEGORY_HATE_SPEECH,
            "harassment": HarmCategory.HARM_CATEGORY_HARASSMENT,
            "sexually_explicit": HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
            "dangerous_content": HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT
        }

        # https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters#how_to_configure_safety_filters
        THRESHOLD_MAP = {
            "low": HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
            "medium": HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            "high": HarmBlockThreshold.BLOCK_ONLY_HIGH,
            "off": HarmBlockThreshold.OFF,
            "none": HarmBlockThreshold.BLOCK_NONE
            }

        if not safety_config:
            safety_list = [
                SafetySetting(
                    category=HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
                    threshold=HarmBlockThreshold.OFF
                ),
                SafetySetting(
                    category=HarmCategory.HARM_CATEGORY_HATE_SPEECH,
                    threshold=HarmBlockThreshold.OFF
                ),
                SafetySetting(
                    category=HarmCategory.HARM_CATEGORY_HARASSMENT,
                    threshold=HarmBlockThreshold.OFF
                ),
                SafetySetting(
                    category=HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
                    threshold=HarmBlockThreshold.OFF
                )
            ]

        elif safety_config:
            safety_list: List[SafetySetting] = []
            for category, threshold in safety_config.items():
                try:
                    safety_list.append(
                        SafetySetting(
                            category=CATEGORY_MAP[category],
                            threshold=THRESHOLD_MAP[threshold]
                            )
                        )
                except KeyError as ke:
                    available_categories = ", ".join(CATEGORY_MAP.keys())
                    available_thresholds = ", ".join(THRESHOLD_MAP.keys())
                    raise KeyError(
                        f"Invalid key provided: {ke}. Available categories: "\
                        f"[{available_categories}]. Available thresholds: "\
                        f"[{available_thresholds}]") from ke

        return safety_list