106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
|
import requests
|
|||
|
|
import json
|
|||
|
|
import math
|
|||
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|||
|
|
|
|||
|
|
def parse_urls(map_json: list[dict]):
|
|||
|
|
"""
|
|||
|
|
解析 Firecrawl 返回的数据
|
|||
|
|
"""
|
|||
|
|
if not map_json:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
map_obj = map_json[0]
|
|||
|
|
# 稍微做个容错，防止有些时候结构不一样
|
|||
|
|
if not map_obj.get("success", False):
|
|||
|
|
# 如果不是必须抛异常，可以打印日志并返回空
|
|||
|
|
print(f"Firecrawl Map节点返回失败或无数据：{map_obj}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
urls = map_obj.get("links", [])
|
|||
|
|
return urls
|
|||
|
|
|
|||
|
|
def send_batch_request(urls_batch: list[str], task_id: int, BASE_URL: str):
|
|||
|
|
"""
|
|||
|
|
发送单个批次的请求
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
# 设置 timeout 是好习惯，防止卡死
|
|||
|
|
# 因为我们不关心返回值，只要发出去就行，timeout设置短一点
|
|||
|
|
res = requests.post(
|
|||
|
|
f"{BASE_URL}/add_urls",
|
|||
|
|
json={
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"urls": urls_batch
|
|||
|
|
},
|
|||
|
|
timeout=10 # 10秒没发完就拉倒，防止拖累主进程
|
|||
|
|
)
|
|||
|
|
if res.status_code == 200:
|
|||
|
|
return True
|
|||
|
|
else:
|
|||
|
|
print(f"Batch failed with status {res.status_code}: {res.text[:100]}")
|
|||
|
|
return False
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Batch request error: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def main(map_json: list[dict], BASE_URL: str, task_id: float):
|
|||
|
|
|
|||
|
|
# 1. 解析 URL
|
|||
|
|
all_urls = parse_urls(map_json)
|
|||
|
|
total_count = len(all_urls)
|
|||
|
|
|
|||
|
|
if total_count == 0:
|
|||
|
|
return {"msg": "没有解析到URL"}
|
|||
|
|
|
|||
|
|
# ================= 配置区 =================
|
|||
|
|
BATCH_SIZE = 50 # 每一批发送 50 个 URL (根据你后端性能调整)
|
|||
|
|
MAX_WORKERS = 10 # 同时开 10 个线程并发发送
|
|||
|
|
# ==========================================
|
|||
|
|
|
|||
|
|
# 2. 将 URL 切片 (分批)
|
|||
|
|
# 比如 1000 个 URL，切成 20 个 batch
|
|||
|
|
batches = [all_urls[i:i + BATCH_SIZE] for i in range(0, total_count, BATCH_SIZE)]
|
|||
|
|
|
|||
|
|
print(f"总共 {total_count} 个URL，分为 {len(batches)} 批发送，并发数: {MAX_WORKERS}")
|
|||
|
|
|
|||
|
|
# 3. 多线程并发发送
|
|||
|
|
# 在 Dify/Lambda 等环境中，必须等待线程池执行完才能退出 main，
|
|||
|
|
# 否则请求还没发出去进程就被杀掉了。但因为是并发，速度会非常快。
|
|||
|
|
success_batches = 0
|
|||
|
|
|
|||
|
|
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
|||
|
|
# 提交所有任务
|
|||
|
|
futures = [
|
|||
|
|
executor.submit(send_batch_request, batch, int(task_id), BASE_URL)
|
|||
|
|
for batch in batches
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# 等待完成 (as_completed)
|
|||
|
|
for future in as_completed(futures):
|
|||
|
|
if future.result():
|
|||
|
|
success_batches += 1
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"status": "done",
|
|||
|
|
"total_urls": total_count,
|
|||
|
|
"batches_sent": len(batches),
|
|||
|
|
"success_batches": success_batches,
|
|||
|
|
"msg": "已使用多线程并发发送数据，忽略详细返回值"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def test():
|
|||
|
|
import json
|
|||
|
|
from time import time
|
|||
|
|
with open("nodes\parse_and_add_urls.json", "r", encoding="utf-8") as f:
|
|||
|
|
data = json.load(f)
|
|||
|
|
map_json = data["json"]
|
|||
|
|
BASE_URL = "http://47.122.127.178"
|
|||
|
|
task_id = 6
|
|||
|
|
start_time = time()
|
|||
|
|
res = main(map_json, BASE_URL, task_id)
|
|||
|
|
end_time = time()
|
|||
|
|
print(f"添加URL耗时：{end_time - start_time}秒")
|
|||
|
|
print(res)
|
|||
|
|
|
|||
|
|
test()
|