From cad13905343abe816ad2cf69da934eed88f4586d Mon Sep 17 00:00:00 2001 From: v6ole Date: Sat, 9 May 2026 13:57:37 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20PostCrawlPipeline?= =?UTF-8?q?=20=E7=BB=9F=E4=B8=80=E7=AE=A1=E9=81=93=20+=20=E7=AD=9B?= =?UTF-8?q?=E9=80=89=E6=9C=8D=E5=8A=A1=20+=20=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/services/filter_service.py | 53 +++++++++++ app/services/pipeline.py | 102 +++++++++++++++++++++ tests/test_services/test_filter_service.py | 38 ++++++++ tests/test_services/test_pipeline.py | 93 +++++++++++++++++++ 4 files changed, 286 insertions(+) create mode 100644 app/services/filter_service.py create mode 100644 app/services/pipeline.py create mode 100644 tests/test_services/test_filter_service.py create mode 100644 tests/test_services/test_pipeline.py diff --git a/app/services/filter_service.py b/app/services/filter_service.py new file mode 100644 index 0000000..fe09fa1 --- /dev/null +++ b/app/services/filter_service.py @@ -0,0 +1,53 @@ +from datetime import date +from typing import Any, Dict, List, Optional + + +def filter_by_keywords(announcements: List[Dict[str, Any]], + keywords: List[str]) -> List[Dict[str, Any]]: + if not keywords: + return announcements + + filtered = [] + for ann in announcements: + search_text = f"{ann.get('title', '')} {ann.get('purchase_name', '')}" + if any(kw in search_text for kw in keywords): + ann["keyword_matched"] = True + filtered.append(ann) + + return filtered + + +def filter_by_date(announcements: List[Dict[str, Any]], + start_date: Optional[date] = None, + end_date: Optional[date] = None) -> List[Dict[str, Any]]: + if not start_date and not end_date: + return announcements + + filtered = [] + for ann in announcements: + pub_date = ann.get("publish_date") + if not pub_date: + continue + if isinstance(pub_date, date): + pub_date = pub_date + else: + pub_date = pub_date.date() if hasattr(pub_date, "date") else pub_date + + if start_date and pub_date < start_date: + continue + if end_date and pub_date > end_date: + continue + filtered.append(ann) + + return filtered + + +def dedup_by_hash(announcements: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + seen = set() + result = [] + for ann in announcements: + h = ann.get("content_hash") + if h and h not in seen: + seen.add(h) + result.append(ann) + return result diff --git a/app/services/pipeline.py b/app/services/pipeline.py new file mode 100644 index 0000000..ff3287f --- /dev/null +++ b/app/services/pipeline.py @@ -0,0 +1,102 @@ +from typing import Any, Dict, List +from app.crawler.base import PipelineConfig, PipelineResult +from app.services.filter_service import filter_by_keywords, dedup_by_hash + + +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. Store to database + stored = await self._save_to_db(announcements) + result.stored = stored + + to_notify = announcements + + # 3. Filter + if config.filter_enabled and config.keywords: + before = len(to_notify) + to_notify = [a for a in to_notify + if self._match_keywords(a, config.keywords)] + result.filtered = before - len(to_notify) + + # 4. 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) + + # 5. 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 or 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 app.models.announcement import Announcement + from sqlalchemy import update + + 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 + + @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) diff --git a/tests/test_services/test_filter_service.py b/tests/test_services/test_filter_service.py new file mode 100644 index 0000000..0c477f1 --- /dev/null +++ b/tests/test_services/test_filter_service.py @@ -0,0 +1,38 @@ +from datetime import datetime +from app.services.filter_service import filter_by_keywords, filter_by_date + + +def test_filter_by_keywords_match(): + announcements = [ + {"title": "大化县采购公告", "purchase_name": "大化县财政局", + "source_code": "test", "source_name": "test"}, + {"title": "南宁市采购公告", "purchase_name": "南宁市财政局", + "source_code": "test", "source_name": "test"}, + ] + result = filter_by_keywords(announcements, ["大化"]) + assert len(result) == 1 + assert result[0]["title"] == "大化县采购公告" + + +def test_filter_by_keywords_no_keywords(): + announcements = [ + {"title": "大化县采购公告", "purchase_name": "x", + "source_code": "test", "source_name": "test"}, + ] + result = filter_by_keywords(announcements, []) + assert len(result) == 1 + + +def test_filter_by_date_range(): + today = datetime(2026, 5, 9) + announcements = [ + {"title": "t1", "publish_date": datetime(2026, 5, 9), + "source_code": "test", "source_name": "test"}, + {"title": "t2", "publish_date": datetime(2026, 5, 1), + "source_code": "test", "source_name": "test"}, + {"title": "t3", "publish_date": datetime(2026, 4, 30), + "source_code": "test", "source_name": "test"}, + ] + result = filter_by_date(announcements, start_date=today, end_date=today) + assert len(result) == 1 + assert result[0]["title"] == "t1" diff --git a/tests/test_services/test_pipeline.py b/tests/test_services/test_pipeline.py new file mode 100644 index 0000000..402bf2c --- /dev/null +++ b/tests/test_services/test_pipeline.py @@ -0,0 +1,93 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock, patch +from datetime import datetime +from app.services.pipeline import PostCrawlPipeline +from app.crawler.base import CrawlResult, PipelineConfig + + +@pytest.mark.asyncio +async def test_pipeline_filtered_mode(): + config = PipelineConfig( + filter_enabled=True, + keywords=["大化"], + dedup_enabled=True, + notify_mode="filtered", + mark_sent=False, + ) + result = CrawlResult( + source_code="test", source_name="test", + total_count=3, new_count=3, + announcements=[ + { + "title": "大化县公告", "publish_date": datetime(2026, 5, 9), + "purchase_name": "大化县", "content_url": "https://1.com", + "source_code": "test", "source_name": "测试", + "announcement_type": "purchase", "crawl_mode": "auto", + "is_new": True, "is_today": True, + "content_hash": "abc123", + }, + { + "title": "南宁市公告", "publish_date": datetime(2026, 5, 9), + "purchase_name": "南宁市", "content_url": "https://2.com", + "source_code": "test", "source_name": "测试", + "announcement_type": "purchase", "crawl_mode": "auto", + "is_new": True, "is_today": True, + "content_hash": "def456", + }, + ], + ) + + mock_db = AsyncMock() + mock_notify = AsyncMock() + + pipeline = PostCrawlPipeline(db_session=mock_db, notification_service=mock_notify) + + with patch.object(pipeline, "_save_to_db", AsyncMock(return_value=2)): + with patch.object(pipeline, "_send_notifications", AsyncMock(return_value=1)): + pipe_result = await pipeline.process( + result.announcements, config + ) + assert pipe_result.stored == 2 + assert pipe_result.filtered == 1 + assert pipe_result.notified == 1 + + +@pytest.mark.asyncio +async def test_pipeline_all_mode(): + config = PipelineConfig( + filter_enabled=False, + keywords=[], + dedup_enabled=True, + notify_mode="all", + mark_sent=True, + ) + result = CrawlResult( + source_code="dahuagov", source_name="大化县政府网", + total_count=2, new_count=2, + announcements=[ + { + "title": f"公告{i}", "publish_date": datetime(2026, 5, 9), + "purchase_name": "大化县", "content_url": f"https://x.com/{i}", + "source_code": "dahuagov", "source_name": "大化县政府网采购公告", + "announcement_type": "purchase", "crawl_mode": "auto", + "is_new": True, "is_today": True, + "content_hash": f"hash{i}", + } + for i in range(2) + ], + ) + + mock_db = AsyncMock() + mock_notify = AsyncMock() + pipeline = PostCrawlPipeline(db_session=mock_db, notification_service=mock_notify) + + with patch.object(pipeline, "_save_to_db", AsyncMock(return_value=2)): + with patch.object(pipeline, "_send_notifications", AsyncMock(return_value=2)): + with patch.object(pipeline, "_mark_sent", AsyncMock(return_value=2)): + pipe_result = await pipeline.process( + result.announcements, config + ) + assert pipe_result.stored == 2 + assert pipe_result.filtered == 0 + assert pipe_result.notified == 2 + pipeline._mark_sent.assert_awaited_once()