diff --git a/src/dependencies/providers/model.py b/src/dependencies/providers/model.py index 239b4bc..2857247 100644 --- a/src/dependencies/providers/model.py +++ b/src/dependencies/providers/model.py @@ -12,12 +12,22 @@ from pydantic_ai.providers.google import GoogleProvider from utils.db.models import DynamicConfig, RespondSession from utils.env import env +from utils.logging import console from ..types.model import RespondModel, ReviewModel GOOGLE_SAFETY_SETTINGS = [ - SafetySettingDict(category=category, threshold=HarmBlockThreshold.OFF) - for category in HarmCategory + { + "category": HarmCategory.HARM_CATEGORY_HATE_SPEECH, + "threshold": HarmBlockThreshold.OFF, + } + for category in [ + HarmCategory.HARM_CATEGORY_HATE_SPEECH, + HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, + HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, + HarmCategory.HARM_CATEGORY_HARASSMENT, + HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY, + ] ]