Add AI filtering

3 files changed, 59 insertions, 2 deletions

diff --git a/requirements.txt b/requirements.txt
index 94e1553..9202358 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,5 +5,6 @@ slack-sdk
 aiosqlite
 python-dotenv
 httpx
-openrouter
+# openrouter
+google-genai
 aiohttp
diff --git a/src/botnet.py b/src/botnet.py
index 46b59d5..fbf8efe 100644
--- a/src/botnet.py
+++ b/src/botnet.py
@@ -2,11 +2,60 @@ from .db import create_deployment, get_daily_user_deployment_count, get_deployme
 from slack_sdk.web.async_client import AsyncWebClient
 import os
 from dotenv import load_dotenv
+from google import genai
+from google.genai import types
+import json
 
 DAILY_LIMIT = 2
 
 load_dotenv()
 selfbot_token = os.environ.get("SLACK_SELFBOT_OAUTH_XOXP")
+gemini_key = os.environ.get("GEMINI_API_KEY")
+
+system_prompt = """
+You are a strict content moderator for a professional Slack workspace called Hack Club.
+
+Your job is to analyze messages and determine if they are SAFE or UNSAFE to be posted by a bot on behalf of multiple users.
+
+A message is UNSAFE if it contains ANY of the following:
+- Swear words, profanity, or offensive language
+- Harassment, bullying, or personal attacks
+- Hate speech or discrimination (race, gender, sexuality, religion, etc.)
+- NSFW, sexual, or inappropriate content
+- Spam, excessive gibberish, or repetitive messages
+- Phishing, scams, or malicious links
+- Mentions of illegal activities (hacking, drugs, violence, etc.)
+- Anything that would make Slack admins (FD/Fire Dept.) upset or require them to intervene
+
+A message is SAFE if it is:
+- Friendly, positive, or neutral
+- On-topic for a tech/developer community
+- A harmless joke or meme (as long as it's not offensive)
+- Helpful, constructive, or collaborative
+- Literally just "chicken nugget" or similar harmless nonsense
+
+Respond with JSON only. Do not add any other text.
+Output format:
+{'safe': true/false, 'reasoning': '<short reasoning here>'}
+"""
+
+async def check_ai_safety(message: str):
+    try:
+        client = genai.Client(api_key=gemini_key)
+        response = await client.aio.models.generate_content(
+                model="gemini-3.1-flash-lite",
+                contents=f"Message to analyse: {message}",
+                config=types.GenerateContentConfig(
+                    system_instruction=system_prompt,
+                    response_mime_type="application/json"
+                    )
+                )
+        
+        resp_json = json.loads(response.text)
+        return resp_json.get("safe", False), resp_json.get("reasoning", "No reasoning provided")
+    except Exception as e:
+        print(f"ai safety check might be cooked :fear: {e}")
+        return False, f"AI safety check failed, please try again later.\nError: {e}"
 
 async def deploy_message(user_id: str, message: str):
     user = await get_user(user_id)
@@ -23,6 +72,13 @@ async def deploy_message(user_id: str, message: str):
             "error": f"You've reached your daily limit of {DAILY_LIMIT} messages!"
         }
 
+    safe, reasoning = await check_ai_safety(message)
+    if not safe:
+        return {
+                "success": False,
+                "error": f"Message flagged by AI: {reasoning}"
+                }
+
     deployment_id = await create_deployment(user_id, message)
 
     print(f"DEBUG: Executing the deployment {deployment_id} on arrival - REMOVE THIS LATER!")
diff --git a/src/main.py b/src/main.py
index 02375cc..1c43240 100644
--- a/src/main.py
+++ b/src/main.py
@@ -26,7 +26,7 @@ async def handle_messages(message, say):
         if message.get("user") == selfbot_user_id:
             return
 
-        if message.get("user") not in ["U0AKCBZHHMH", "U0BG82V0CP6"]:
+        if message.get("user") not in ["U0AKCBZHHMH"]:
             await say("Botnet is temporarily disabled for users other than <@U0AKCBZHHMH> for debugging and the safety of all OAuth-authorized users. Please DM <@U0AKCBZHHMH> if this is not resolved within a few days.")
             return