# ruff: noqa: E402 import asyncio import logging import re import sys import time from pathlib import Path from dotenv import load_dotenv PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) load_dotenv(Path(__file__).resolve().parents[1] / ".env") from datetime import datetime from app.db.models import AIDetection, GroupConfig, MediaAsset, MessageEvent, StatsMessagesDaily, StatsWordcloudTerm from app.db.session import SessionLocal from app.services.ai_client import detect_image, detect_text from app.services.queue import claim_job, complete_job, fail_job, get_pending_topics from app.services.telegram_moderator import delete_message, punish_user from app.services.telegram_sender import send_broadcast, send_post, send_scheduled_message logger = logging.getLogger(__name__) QUEUE_TOPICS = [ "ai_text_detect", "ai_image_detect", "moderation_action", "broadcast_send", "post_send", "schedule_send", "schedule_tick", "stats_rollup", ] async def _handle_ai_detection(message_id: str, kind: str) -> None: db = SessionLocal() try: message = db.get(MessageEvent, message_id) if not message: return config = db.query(GroupConfig).filter(GroupConfig.group_id == message.group_id).first() ad_cfg = (config.config_json.get("ad_filter") if config else {}) or {} if kind == "text" and not ad_cfg.get("text_ai_enabled", True): return if kind == "image" and not ad_cfg.get("image_ai_enabled", True): return threshold = float(ad_cfg.get("text_ai_threshold" if kind == "text" else "image_ai_threshold", 0.85)) score = 0.0 if kind == "text": score = await detect_text(message.text or "") else: file_id = "" if message.media_id: asset = db.get(MediaAsset, message.media_id) if asset: file_id = asset.telegram_file_id score = await detect_image(file_id) result = "pass" if score < threshold else "reject" db.add( AIDetection( group_id=message.group_id, message_id=message.id, type=kind, score=score, threshold=threshold, model_id=None, result=result, ) ) db.commit() if result == "reject" and message.chat_id: action = ad_cfg.get("action", "delete") duration = ad_cfg.get("action_duration") if action == "delete": await delete_message( chat_id=message.chat_id, message_id=message.telegram_message_id, group_id=str(message.group_id), user_id=message.telegram_user_id, reason="AI广告过滤", source="ai", ) else: await punish_user( chat_id=message.chat_id, user_id=message.telegram_user_id, group_id=str(message.group_id), action=action, duration=duration, reason="AI广告过滤", source="ai", ) finally: db.close() def _tokenize(text: str) -> list[str]: tokens = re.findall(r"[A-Za-z0-9]+|[\u4e00-\u9fff]+", text) return [token for token in tokens if len(token) >= 2] def _handle_stats_rollup(message_id: str) -> None: db = SessionLocal() try: message = db.get(MessageEvent, message_id) if not message: return day = (message.created_at or datetime.utcnow()).date() stats = ( db.query(StatsMessagesDaily) .filter(StatsMessagesDaily.group_id == message.group_id, StatsMessagesDaily.date == day) .first() ) if not stats: stats = StatsMessagesDaily(group_id=message.group_id, date=day, message_count=0, active_users=0) db.add(stats) stats.message_count += 1 active_users = ( db.query(MessageEvent.telegram_user_id) .filter(MessageEvent.group_id == message.group_id) .filter(MessageEvent.created_at >= datetime.combine(day, datetime.min.time())) .filter(MessageEvent.created_at <= datetime.combine(day, datetime.max.time())) .distinct() .count() ) stats.active_users = active_users if message.text: for token in _tokenize(message.text): term = ( db.query(StatsWordcloudTerm) .filter( StatsWordcloudTerm.group_id == message.group_id, StatsWordcloudTerm.date == day, StatsWordcloudTerm.term == token, ) .first() ) if not term: term = StatsWordcloudTerm(group_id=message.group_id, date=day, term=token, count=0) db.add(term) term.count += 1 db.commit() finally: db.close() async def _dispatch(topic: str, payload: dict) -> None: """Dispatch a single queue message by topic.""" if topic == "post_send": post_id = payload.get("post_id") if post_id: await send_post(post_id) elif topic == "broadcast_send": broadcast_id = payload.get("broadcast_id") if broadcast_id: await send_broadcast(broadcast_id) elif topic == "schedule_send": schedule_id = payload.get("schedule_id") if schedule_id: await send_scheduled_message(schedule_id) elif topic == "ai_text_detect": message_id = payload.get("message_id") if message_id: await _handle_ai_detection(message_id, "text") elif topic == "ai_image_detect": message_id = payload.get("message_id") if message_id: await _handle_ai_detection(message_id, "image") elif topic == "stats_rollup": message_id = payload.get("message_id") if message_id: # stats_rollup is sync (DB only), run in thread to avoid blocking await asyncio.to_thread(_handle_stats_rollup, message_id) else: logger.info("[worker] %s %s", topic, payload) async def worker_loop() -> None: logger.info("worker started (PostgreSQL queue)") while True: try: pending_topics = await asyncio.to_thread(get_pending_topics, QUEUE_TOPICS) if not pending_topics: await asyncio.sleep(0.5) continue for topic in pending_topics: job = await asyncio.to_thread(claim_job, topic) if not job: continue try: await _dispatch(job.topic, job.payload) await asyncio.to_thread(complete_job, job.id) except Exception: logger.exception("worker error processing %s", job.topic) await asyncio.to_thread(fail_job, job.id) except Exception: logger.exception("worker loop error") await asyncio.sleep(1) def main(): """Standalone entry point (backward compatible).""" logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") asyncio.run(worker_loop()) if __name__ == "__main__": main()