戒酒的李白
371 lines
12 KiB
Python
371 lines
12 KiB
Python
from flask import Blueprint, jsonify, request, render_template
|
|
import json
|
|
import os
|
|
from datetime import datetime
|
|
import threading
|
|
from queue import Queue
|
|
import asyncio
|
|
import websockets
|
|
import logging
|
|
from spider.spiderData import SpiderData
|
|
from openai import OpenAI
|
|
from anthropic import Anthropic
|
|
import aiohttp
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from ratelimit import limits, sleep_and_retry
|
|
from tenacity import retry, stop_after_attempt, wait_exponential
|
|
|
|
# 创建蓝图
|
|
spider_bp = Blueprint('spider', __name__)
|
|
|
|
# 创建日志记录器
|
|
logger = logging.getLogger('spider_control')
|
|
logger.setLevel(logging.INFO)
|
|
|
|
# 存储WebSocket连接的集合
|
|
websocket_connections = set()
|
|
|
|
# 创建消息队列
|
|
message_queue = Queue()
|
|
|
|
# 创建线程池
|
|
thread_pool = ThreadPoolExecutor(max_workers=3)
|
|
|
|
# 创建异步事件循环
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
# 默认配置
|
|
DEFAULT_CONFIG = {
|
|
'crawlDepth': 3,
|
|
'interval': 5,
|
|
'maxRetries': 3,
|
|
'timeout': 30,
|
|
'maxConcurrent': 2
|
|
}
|
|
|
|
# 限流装饰器
|
|
@sleep_and_retry
|
|
@limits(calls=100, period=60) # 每分钟最多100个请求
|
|
def rate_limited_request():
|
|
pass
|
|
|
|
class SpiderWorker:
|
|
def __init__(self, topics, parameters):
|
|
self.topics = topics
|
|
self.parameters = parameters
|
|
self.total_topics = len(topics)
|
|
self.completed_topics = 0
|
|
self.spider = SpiderData()
|
|
self.message_buffer = []
|
|
self.message_buffer_size = 10
|
|
self.semaphore = asyncio.Semaphore(parameters.get('maxConcurrent', DEFAULT_CONFIG['maxConcurrent']))
|
|
|
|
async def send_message(self, message):
|
|
"""异步发送消息，使用缓冲区优化"""
|
|
self.message_buffer.append(message)
|
|
if len(self.message_buffer) >= self.message_buffer_size:
|
|
await self.flush_messages()
|
|
|
|
async def flush_messages(self):
|
|
"""刷新消息缓冲区"""
|
|
if not self.message_buffer:
|
|
return
|
|
|
|
try:
|
|
await broadcast_message(self.message_buffer)
|
|
self.message_buffer.clear()
|
|
except Exception as e:
|
|
logger.error(f"发送消息失败: {e}")
|
|
|
|
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
|
|
async def crawl_single_topic(self, topic):
|
|
"""爬取单个话题"""
|
|
try:
|
|
rate_limited_request()
|
|
|
|
await self.send_message({
|
|
'type': 'log',
|
|
'message': f'开始爬取话题: {topic}'
|
|
})
|
|
|
|
async with self.semaphore:
|
|
await asyncio.get_event_loop().run_in_executor(
|
|
thread_pool,
|
|
self.spider.crawl_topic,
|
|
topic,
|
|
self.parameters['crawlDepth'],
|
|
self.parameters['interval'],
|
|
self.parameters['maxRetries'],
|
|
self.parameters['timeout']
|
|
)
|
|
|
|
self.completed_topics += 1
|
|
progress = int((self.completed_topics / self.total_topics) * 100)
|
|
|
|
await self.send_message({
|
|
'type': 'progress',
|
|
'value': progress
|
|
})
|
|
|
|
await self.send_message({
|
|
'type': 'log',
|
|
'message': f'话题 {topic} 爬取完成'
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"爬取话题 {topic} 失败: {e}")
|
|
await self.send_message({
|
|
'type': 'log',
|
|
'message': f'爬取话题 {topic} 时出错: {str(e)}'
|
|
})
|
|
raise
|
|
|
|
async def run(self):
|
|
"""运行爬虫任务"""
|
|
try:
|
|
tasks = [self.crawl_single_topic(topic) for topic in self.topics]
|
|
await asyncio.gather(*tasks)
|
|
await self.flush_messages()
|
|
|
|
await self.send_message({
|
|
'type': 'log',
|
|
'message': '所有话题爬取完成'
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"爬虫任务执行出错: {e}")
|
|
await self.send_message({
|
|
'type': 'log',
|
|
'message': f'爬虫任务执行出错: {str(e)}'
|
|
})
|
|
finally:
|
|
await self.flush_messages()
|
|
|
|
async def broadcast_message(messages):
|
|
"""广播消息到所有WebSocket连接"""
|
|
if not websocket_connections:
|
|
return
|
|
|
|
for websocket in websocket_connections.copy():
|
|
try:
|
|
if isinstance(messages, list):
|
|
for message in messages:
|
|
await websocket.send(json.dumps(message))
|
|
else:
|
|
await websocket.send(json.dumps(messages))
|
|
except websockets.exceptions.ConnectionClosed:
|
|
websocket_connections.remove(websocket)
|
|
except Exception as e:
|
|
logger.error(f"发送WebSocket消息失败: {e}")
|
|
websocket_connections.remove(websocket)
|
|
|
|
@spider_bp.route('/spider/control')
|
|
def spider_control():
|
|
"""渲染爬虫控制页面"""
|
|
return render_template('spider_control.html')
|
|
|
|
@spider_bp.route('/api/spider/start', methods=['POST'])
|
|
async def start_spider():
|
|
"""启动爬虫任务"""
|
|
try:
|
|
data = request.get_json()
|
|
topics = data.get('topics', [])
|
|
parameters = {**DEFAULT_CONFIG, **data.get('parameters', {})}
|
|
|
|
if not topics:
|
|
return jsonify({
|
|
'success': False,
|
|
'message': '请选择至少一个话题'
|
|
})
|
|
|
|
# 创建爬虫工作器
|
|
worker = SpiderWorker(topics, parameters)
|
|
|
|
# 在事件循环中运行爬虫任务
|
|
asyncio.create_task(worker.run())
|
|
|
|
return jsonify({
|
|
'success': True,
|
|
'message': '爬虫任务已启动'
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"启动爬虫任务失败: {e}")
|
|
return jsonify({
|
|
'success': False,
|
|
'message': str(e)
|
|
})
|
|
|
|
@spider_bp.route('/api/spider/save-config', methods=['POST'])
|
|
def save_spider_config():
|
|
"""保存爬虫配置"""
|
|
try:
|
|
config = request.get_json()
|
|
if save_config(config):
|
|
return jsonify({
|
|
'success': True,
|
|
'message': '配置保存成功'
|
|
})
|
|
else:
|
|
return jsonify({
|
|
'success': False,
|
|
'message': '配置保存失败'
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"保存配置失败: {e}")
|
|
return jsonify({
|
|
'success': False,
|
|
'message': str(e)
|
|
})
|
|
|
|
@spider_bp.websocket('/ws/spider-status')
|
|
async def spider_status_socket(websocket):
|
|
"""WebSocket连接处理"""
|
|
try:
|
|
websocket_connections.add(websocket)
|
|
logging.info("新的WebSocket连接已建立")
|
|
|
|
try:
|
|
while True:
|
|
# 等待消息，保持连接活跃
|
|
message = await websocket.receive()
|
|
if message is None:
|
|
break
|
|
except websockets.exceptions.ConnectionClosed:
|
|
logging.info("WebSocket连接已关闭")
|
|
finally:
|
|
websocket_connections.remove(websocket)
|
|
logging.info("WebSocket连接已移除")
|
|
except Exception as e:
|
|
logger.error(f"WebSocket连接处理失败: {e}")
|
|
if websocket in websocket_connections:
|
|
websocket_connections.remove(websocket)
|
|
|
|
def get_ai_client():
|
|
"""获取可用的AI客户端"""
|
|
# 按优先级尝试不同的AI服务
|
|
if os.getenv('ANTHROPIC_API_KEY'):
|
|
return {
|
|
'type': 'anthropic',
|
|
'client': Anthropic(api_key=os.getenv('ANTHROPIC_API_KEY'))
|
|
}
|
|
elif os.getenv('OPENAI_API_KEY'):
|
|
return {
|
|
'type': 'openai',
|
|
'client': OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
|
|
}
|
|
else:
|
|
raise ValueError("未找到可用的AI API密钥")
|
|
|
|
def parse_ai_response(response_text):
|
|
"""解析AI响应中的JSON配置"""
|
|
try:
|
|
# 查找JSON内容
|
|
start = response_text.find('{')
|
|
end = response_text.rfind('}') + 1
|
|
if start == -1 or end == 0:
|
|
raise ValueError("未找到有效的JSON配置")
|
|
|
|
json_str = response_text[start:end]
|
|
config = json.loads(json_str)
|
|
|
|
# 验证配置格式
|
|
if not isinstance(config.get('topics'), list):
|
|
raise ValueError("配置必须包含话题列表")
|
|
|
|
parameters = config.get('parameters', {})
|
|
if not all(key in parameters for key in ['crawlDepth', 'interval', 'maxRetries', 'timeout']):
|
|
raise ValueError("配置缺少必要的参数")
|
|
|
|
# 提取建议文本（JSON之前的部分）
|
|
suggestion = response_text[:start].strip()
|
|
|
|
return config, suggestion
|
|
except Exception as e:
|
|
raise ValueError(f"解析AI响应失败: {str(e)}")
|
|
|
|
@spider_bp.route('/api/spider/ai-config', methods=['POST'])
|
|
def generate_ai_config():
|
|
"""使用AI生成爬虫配置"""
|
|
try:
|
|
prompt = request.json.get('prompt', '')
|
|
if not prompt:
|
|
return jsonify({
|
|
'success': False,
|
|
'message': '请提供爬虫需求描述'
|
|
})
|
|
|
|
# 构建AI提示
|
|
system_prompt = """你是一个专业的爬虫配置助手。请根据用户的自然语言描述，生成合适的微博爬虫配置。
|
|
配置应包含以下内容：
|
|
1. 要爬取的话题列表
|
|
2. 爬虫参数（爬取深度、间隔时间、重试次数、超时时间）
|
|
|
|
请先用通俗易懂的语言解释你的配置建议，然后在最后提供一个JSON格式的具体配置。
|
|
注意：
|
|
- 爬取深度(crawlDepth)范围：1-10页
|
|
- 间隔时间(interval)范围：3-30秒
|
|
- 重试次数(maxRetries)范围：1-5次
|
|
- 超时时间(timeout)范围：10-60秒
|
|
- 所有参数都必须是整数
|
|
|
|
示例输出格式：
|
|
根据您的需求，我建议...
|
|
|
|
{
|
|
"topics": ["话题1", "话题2"],
|
|
"parameters": {
|
|
"crawlDepth": 5,
|
|
"interval": 5,
|
|
"maxRetries": 3,
|
|
"timeout": 30
|
|
}
|
|
}"""
|
|
|
|
# 获取AI客户端
|
|
ai = get_ai_client()
|
|
|
|
try:
|
|
if ai['type'] == 'anthropic':
|
|
response = ai['client'].messages.create(
|
|
model="claude-3-sonnet-20240229",
|
|
max_tokens=1000,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt}
|
|
]
|
|
)
|
|
response_text = response.content[0].text
|
|
else: # OpenAI
|
|
response = ai['client'].chat.completions.create(
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt}
|
|
]
|
|
)
|
|
response_text = response.choices[0].message.content
|
|
|
|
# 解析AI响应
|
|
config, suggestion = parse_ai_response(response_text)
|
|
|
|
return jsonify({
|
|
'success': True,
|
|
'config': config,
|
|
'suggestion': suggestion
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"AI服务调用失败: {e}")
|
|
return jsonify({
|
|
'success': False,
|
|
'message': f"AI配置生成失败: {str(e)}"
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"生成配置失败: {e}")
|
|
return jsonify({
|
|
'success': False,
|
|
'message': str(e)
|
|
}) |