229 lines
8.4 KiB
Python
229 lines
8.4 KiB
Python
"""
|
||
异步流式处理管线模块
|
||
实现三级并行：RSS抓取与文章处理流水线 / 网摘并行生成 / 分类介绍并行生成
|
||
"""
|
||
import asyncio
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from .config import AI_CONCURRENCY, MAX_THREADS
|
||
from .logger import get_logger
|
||
|
||
logger = get_logger()
|
||
|
||
_shared_executor = None
|
||
|
||
|
||
def _get_executor():
|
||
global _shared_executor
|
||
if _shared_executor is None:
|
||
_shared_executor = ThreadPoolExecutor(max_workers=max(MAX_THREADS, AI_CONCURRENCY) * 2)
|
||
return _shared_executor
|
||
|
||
|
||
def cleanup_executor():
|
||
global _shared_executor
|
||
if _shared_executor:
|
||
_shared_executor.shutdown(wait=True)
|
||
_shared_executor = None
|
||
|
||
|
||
async def fetch_all_feeds_async(feeds):
|
||
"""异步并行抓取所有RSS源，每个feed独立提交到线程池"""
|
||
from .rss_fetcher import _fetch_single_feed
|
||
|
||
loop = asyncio.get_running_loop()
|
||
executor = _get_executor()
|
||
|
||
async def fetch_one(feed):
|
||
return await loop.run_in_executor(executor, _fetch_single_feed, feed)
|
||
|
||
tasks = [asyncio.create_task(fetch_one(f)) for f in feeds]
|
||
results = await asyncio.gather(*tasks)
|
||
|
||
all_articles = []
|
||
for articles in results:
|
||
all_articles.extend(articles)
|
||
return all_articles
|
||
|
||
|
||
async def process_articles_streaming(feeds, cutoff, api_key, base_url, model, cache):
|
||
"""
|
||
流水线处理：抓取与AI处理并行执行
|
||
使用生产者-消费者模型：
|
||
- 上游：每个feed抓取完成后文章立即入队
|
||
- 下游：AI_CONCURRENCY个worker从队列取文章，执行翻译+分析
|
||
- 两级队列：标题翻译 → 收集过滤 → 正文翻译+AI分析
|
||
"""
|
||
from .rss_fetcher import _fetch_single_feed
|
||
from .translator import translate_title, translate_content
|
||
from .ai_analyzer import score_and_classify_article, calculate_weighted_score
|
||
from .keyword_filter import keyword_filter
|
||
|
||
loop = asyncio.get_running_loop()
|
||
executor = _get_executor()
|
||
|
||
fetch_queue = asyncio.Queue(maxsize=200)
|
||
|
||
# =========================================================================
|
||
# Phase 1: 流水线抓取 + 标题翻译
|
||
# =========================================================================
|
||
|
||
async def produce(feed):
|
||
articles = await loop.run_in_executor(executor, _fetch_single_feed, feed)
|
||
for a in articles:
|
||
if a['published'] >= cutoff:
|
||
await fetch_queue.put(a)
|
||
|
||
async def worker_title(worker_id, local_cached, local_new):
|
||
"""第一级worker：标题翻译 + 缓存检查"""
|
||
while True:
|
||
article = await fetch_queue.get()
|
||
if article is None:
|
||
fetch_queue.task_done()
|
||
break
|
||
|
||
cached = cache.get_article(article['id'], expire_hours=24)
|
||
if cached:
|
||
local_cached.append(cached)
|
||
fetch_queue.task_done()
|
||
continue
|
||
|
||
try:
|
||
article = await loop.run_in_executor(
|
||
executor, translate_title, article, api_key, base_url, model
|
||
)
|
||
local_new.append(article)
|
||
except Exception as e:
|
||
logger.warning("Worker%d 标题翻译失败: %s", worker_id, e)
|
||
fetch_queue.task_done()
|
||
|
||
producer_tasks = [asyncio.create_task(produce(f)) for f in feeds]
|
||
|
||
cached_all = []
|
||
title_translated = []
|
||
worker_tasks_title = [
|
||
asyncio.create_task(worker_title(i, cached_all, title_translated))
|
||
for i in range(AI_CONCURRENCY)
|
||
]
|
||
|
||
await asyncio.gather(*producer_tasks)
|
||
for _ in range(AI_CONCURRENCY):
|
||
await fetch_queue.put(None)
|
||
await asyncio.gather(*worker_tasks_title)
|
||
|
||
logger.info("流水线第一阶段完成：缓存命中 %d 篇，标题翻译 %d 篇",
|
||
len(cached_all), len(title_translated))
|
||
|
||
if not title_translated:
|
||
return cached_all, []
|
||
|
||
# =========================================================================
|
||
# Phase 2: 关键词筛选
|
||
# =========================================================================
|
||
process_queue = asyncio.Queue(maxsize=200)
|
||
|
||
if len(title_translated) > 15:
|
||
logger.info("文章数超过15篇，执行关键词筛选...")
|
||
filtered = keyword_filter(title_translated)
|
||
logger.info("关键词筛选: %d → %d 篇", len(title_translated), len(filtered))
|
||
else:
|
||
filtered = title_translated
|
||
|
||
if not filtered:
|
||
return cached_all, []
|
||
|
||
# =========================================================================
|
||
# Phase 3: 流水线正文翻译 + AI分析
|
||
# =========================================================================
|
||
for a in filtered:
|
||
await process_queue.put(a)
|
||
for _ in range(AI_CONCURRENCY):
|
||
await process_queue.put(None)
|
||
|
||
async def worker_full(worker_id, local_processed):
|
||
"""第二级worker：正文翻译 + AI评分分类 + 写缓存"""
|
||
while True:
|
||
article = await process_queue.get()
|
||
if article is None:
|
||
process_queue.task_done()
|
||
break
|
||
|
||
try:
|
||
article = await loop.run_in_executor(
|
||
executor, translate_content, article, api_key, base_url, model
|
||
)
|
||
article, ai_result = await loop.run_in_executor(
|
||
executor, score_and_classify_article, article, api_key, base_url, model
|
||
)
|
||
|
||
article['ai_result'] = ai_result
|
||
article['final_score'] = calculate_weighted_score(ai_result)
|
||
article['summary'] = ai_result.get('summary', article['translated_title'])
|
||
article['category'] = ai_result.get('category', '装备动态')
|
||
article['scores'] = ai_result.get('scores', {})
|
||
|
||
cache.save_article(article)
|
||
local_processed.append(article)
|
||
|
||
except Exception as e:
|
||
logger.warning("Worker%d 文章处理失败: %s", worker_id, e)
|
||
process_queue.task_done()
|
||
|
||
processed_all = []
|
||
worker_tasks_full = [
|
||
asyncio.create_task(worker_full(i, processed_all))
|
||
for i in range(AI_CONCURRENCY)
|
||
]
|
||
await asyncio.gather(*worker_tasks_full)
|
||
|
||
logger.info("流水线第二阶段完成：新处理 %d 篇文章", len(processed_all))
|
||
return cached_all, processed_all
|
||
|
||
|
||
async def generate_webzines_parallel(top3, api_key, base_url, webzine_model, cache, max_concurrent=1):
|
||
"""受控并行生成TOP3网摘（信号量限流，避免API 429）"""
|
||
from .ai_analyzer import generate_webzine_for_article
|
||
|
||
loop = asyncio.get_running_loop()
|
||
executor = _get_executor()
|
||
semaphore = asyncio.Semaphore(max_concurrent)
|
||
|
||
async def gen_one(index, article):
|
||
cached_wz = article.get('webzine_text', '')
|
||
if cached_wz:
|
||
logger.info("网摘 %d/3（缓存命中）: %s...", index, article['translated_title'][:35])
|
||
return cached_wz
|
||
async with semaphore:
|
||
logger.info("生成网摘 %d/3: %s...", index, article['translated_title'][:35])
|
||
webzine = await loop.run_in_executor(
|
||
executor, generate_webzine_for_article, article, api_key, base_url, webzine_model
|
||
)
|
||
cache.save_article_webzine(article['id'], webzine)
|
||
return webzine
|
||
|
||
tasks = [asyncio.create_task(gen_one(i + 1, a)) for i, a in enumerate(top3)]
|
||
return await asyncio.gather(*tasks)
|
||
|
||
|
||
async def generate_overviews_parallel(date_str, overview_tasks, cache, api_key, base_url, model, max_concurrent=2):
|
||
"""受控并行生成各分类总体介绍（信号量限流，避免API 429）"""
|
||
from .ai_analyzer import generate_category_overview
|
||
|
||
loop = asyncio.get_running_loop()
|
||
executor = _get_executor()
|
||
semaphore = asyncio.Semaphore(max_concurrent)
|
||
|
||
async def gen_one(cache_key, label, articles):
|
||
cached = cache.get_category_summary(date_str, cache_key)
|
||
if cached is not None:
|
||
logger.info("分类介绍 [%s]（缓存命中）", label)
|
||
return cached
|
||
async with semaphore:
|
||
overview = await loop.run_in_executor(
|
||
executor, generate_category_overview, label, articles, api_key, base_url, model
|
||
)
|
||
cache.save_category_summary(date_str, cache_key, overview)
|
||
return overview
|
||
|
||
tasks = [asyncio.create_task(gen_one(k, l, a)) for k, l, a in overview_tasks]
|
||
return await asyncio.gather(*tasks)
|