GX-gp-notify/app/services/pipeline.py

import logging
from typing import Any

from app.config import settings
from app.crawler.base import PipelineConfig, PipelineResult
from app.services.ai_state import is_ai_enabled
from app.services.filter_service import dedup_by_hash

logger = logging.getLogger(__name__)


class PostCrawlPipeline:
    def __init__(self, db_session, notification_service):
        self.db = db_session
        self.notify = notification_service

    async def process(self, announcements: list[dict[str, Any]],
                      config: PipelineConfig) -> PipelineResult:
        result = PipelineResult()

        if not announcements:
            return result

        # 1. Dedup
        if config.dedup_enabled:
            announcements = dedup_by_hash(announcements)

        # 2. 先标记 keyword_matched，再存库
        if config.filter_enabled and config.keywords:
            for a in announcements:
                a["keyword_matched"] = self._match_keywords(a, config.keywords)

        # 3. Store to database
        stored = await self._save_to_db(announcements)
        result.stored = stored

        to_notify = announcements

        # 4. Filter
        if config.filter_enabled and config.keywords:
            before = len(to_notify)
            to_notify = [a for a in to_notify if a.get("keyword_matched")]
            result.filtered = before - len(to_notify)

        # 5. Skip already-notified
        if to_notify:
            to_notify = await self._exclude_sent(to_notify)

        # 5.5 AI 分析（可选，支持运行时开关）
        if is_ai_enabled() and to_notify:
            await self._ai_analyze(to_notify)

        # 6. Notify
        if config.notify_mode == "all":
            result.notified = await self._send_notifications(to_notify)
        elif config.notify_mode == "filtered":
            if config.filter_enabled and config.keywords:
                result.notified = await self._send_notifications(to_notify)
            elif not config.filter_enabled:
                result.notified = await self._send_notifications(to_notify)

        # 7. Mark sent
        if config.mark_sent and result.notified > 0:
            await self._mark_sent(to_notify)

        return result

    async def _save_to_db(self, announcements: list[dict[str, Any]]) -> int:
        from sqlalchemy.dialects.postgresql import insert

        from app.models.announcement import Announcement

        if not announcements:
            return 0

        values = [{
            "title": a["title"],
            "publish_date": a["publish_date"],
            "purchase_name": a.get("purchase_name", ""),
            "content_url": a.get("content_url", ""),
            "source_code": a["source_code"],
            "source_name": a["source_name"],
            "announcement_type": a.get("announcement_type", "purchase"),
            "content_hash": a["content_hash"],
            "crawl_mode": a.get("crawl_mode", "auto"),
            "is_new": a.get("is_new", True),
            "is_sent": False,
            "keyword_matched": a.get("keyword_matched", False),
        } for a in announcements]

        stmt = insert(Announcement).values(values)
        stmt = stmt.on_conflict_do_nothing(index_elements=["content_hash"])

        result_proxy = await self.db.execute(stmt)
        await self.db.commit()
        return result_proxy.rowcount if result_proxy.rowcount >= 0 else len(values)

    async def _send_notifications(self, announcements: list[dict[str, Any]]) -> int:
        return await self.notify.send(announcements)

    async def _mark_sent(self, announcements: list[dict[str, Any]]) -> int:
        from sqlalchemy import update

        from app.models.announcement import Announcement

        hashes = [a["content_hash"] for a in announcements if a.get("content_hash")]
        if not hashes:
            return 0

        stmt = (
            update(Announcement)
            .where(Announcement.content_hash.in_(hashes))
            .values(is_sent=True)
        )
        result = await self.db.execute(stmt)
        await self.db.commit()
        return result.rowcount

    async def _exclude_sent(self, announcements: list[dict[str, Any]]) -> list[dict[str, Any]]:
        from sqlalchemy import select

        from app.models.announcement import Announcement

        hashes = [a["content_hash"] for a in announcements if a.get("content_hash")]
        if not hashes:
            return announcements

        stmt = select(Announcement.content_hash).where(
            Announcement.content_hash.in_(hashes),
            Announcement.is_sent == True,  # noqa: E712
        )
        result = await self.db.execute(stmt)
        sent_hashes = {row[0] for row in result.fetchall()}

        return [a for a in announcements if a.get("content_hash") not in sent_hashes]

    @staticmethod
    def _match_keywords(announcement: dict[str, Any], keywords: list[str]) -> bool:
        text = f"{announcement.get('title', '')} {announcement.get('purchase_name', '')}"
        return any(kw in text for kw in keywords)

    async def _ai_analyze(self, announcements: list[dict[str, Any]]) -> None:
        """对公告列表执行 AI 分析，将结果附加到每条公告的 ai_result 字段，并更新数据库"""
        from app.services.ai_analyzer import AiAnalyzer

        analyzer = AiAnalyzer()
        logger.info("AI 分析开始：共 %d 条公告", len(announcements))

        results = await analyzer.analyze_batch(announcements)

        ai_updates = []
        for ann, ai_result in zip(announcements, results):
            ann["ai_result"] = {
                "is_relevant": ai_result.is_relevant,
                "reason": ai_result.reason,
                "business_type": ai_result.business_type,
            }
            if ai_result.error:
                logger.warning("AI 分析失败 [%s]: %s", ann.get("title", "")[:30], ai_result.error)
            else:
                ai_updates.append({
                    "content_hash": ann["content_hash"],
                    "ai_relevant": ai_result.is_relevant,
                    "ai_analysis": ai_result.reason,
                })
                if ai_result.is_relevant:
                    logger.info("AI 标记可承接项目: %s (%s)", ann.get("title", "")[:40], ai_result.business_type)

        # 批量更新数据库中的 AI 分析结果
        if ai_updates:
            await self._update_ai_results(ai_updates)

    async def _update_ai_results(self, updates: list[dict[str, Any]]) -> None:
        """批量更新公告的 AI 分析结果到数据库"""
        from sqlalchemy import update

        from app.models.announcement import Announcement

        for u in updates:
            stmt = (
                update(Announcement)
                .where(Announcement.content_hash == u["content_hash"])
                .values(
                    ai_relevant=u["ai_relevant"],
                    ai_analysis=u["ai_analysis"],
                )
            )
            await self.db.execute(stmt)
        await self.db.commit()