Skip to content

A
aidso-data
Commit 94534a08 authored by Yaowentong's avatar Yaowentong
Browse files
Options

修复链接池 

修改为队列
修改百度ai默认返回值
parent 9af6ae57
Show whitespace changes
Inline Side-by-side
Showing with 286 additions and 204 deletions
aidso_geo/config/base_config.py
View file @ 94534a08
......@@ -35,13 +35,13 @@ def init_redis8():
port=6379,
db=8,
password='aiyingli@@123',
socket_timeout=5,
socket_timeout=120,
socket_connect_timeout=5,
decode_responses=True,
retry_on_timeout=True,
health_check_interval=30,
socket_keepalive=True,
max_connections=20,
max_connections=100,
)
_redis8_client = redis.Redis(connection_pool=_redis8_pool)
......
aidso_geo/core/commit_process.py
View file @ 94534a08
......@@ -5,11 +5,12 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
from loguru import logger
from apscheduler.schedulers.blocking import BlockingScheduler
from aidso_geo.config.base_config import init_redis
from aidso_geo.config.base_config import init_redis, init_redis8
from aidso_geo.models.process import main_process
redis_client = init_redis()
# redis_client = init_redis8()
QUEUE_KEY = "geo:task_commit:list"
......@@ -66,6 +67,7 @@ def handle_one_task(task):
prompt = task.get("prompt", "")
try:
main_process(task)
return True
......@@ -140,6 +142,5 @@ if __name__ == "__main__":
f"max_instances=5"
)
consume_task_queue()
scheduler.start()
aidso_geo/core/routes/interface.py
View file @ 94534a08
......@@ -12,7 +12,7 @@ from aidso_geo.config.base_config import init_redis, init_redis8
from aidso_geo.utils import bh_utils, tos_utils, url_utils
line_app = Blueprint("line_app", __name__)
redis_client = init_redis()
redis_client8 = init_redis8()
MAX_WORKERS = 30
MAX_PENDING_TASKS = 200
......@@ -167,7 +167,10 @@ def task_commit():
"reqId": req_id
})
logger.success(f"{data['reqId']}--{platform}--{prompt}--任务提交--{type}")
# ok = submit_background_task(main_process, data)
ret = redis_client.lpush("geo:task_commit:list",json.dumps(data, ensure_ascii=False))
# ret = redis_client8.lpush("geo:task_commit:list",json.dumps(data, ensure_ascii=False))
if ret and ret > 0:
resp_cache_key = f"geo:task_check:resp:{req_id}"
resp = {
......@@ -215,6 +218,7 @@ def data_call_back():
platform = task_data.get('platform')
logger.success(f"{req_id}--{platform}-------CALL_BACK")
ok = submit_background_task(process_call_back, task_data, result)
# ret = redis_client8.lpush("geo:call_back:list", json.dumps(data, ensure_ascii=False))
if not ok:
return jsonify({
"code": 503,
......
aidso_geo/core/task_process.py
View file @ 94534a08
import json
import time
import threading
import signal
from loguru import logger
from concurrent.futures import ThreadPoolExecutor, as_completed
from aidso_geo.config.base_config import init_redis8
from loguru import logger
from apscheduler.schedulers.blocking import BlockingScheduler
from aidso_geo.config.base_config import init_redis, init_redis8
from aidso_geo.models.process import main_process, process_call_back
from aidso_geo.utils import bh_utils
redis_client8 = init_redis8()
RUNNING = True
def stop_handler(signum, frame):
global RUNNING
logger.success("开始退出")
RUNNING = False
def handle_item(data):
if data.get('comWordsMap'):
data['comWordsMap'] = json.loads(data.get('comWordsMap'))
if data.get('brandWords'):
data['brandWords'] = json.loads(data.get('brandWords'))
if data.get('comWords'):
data['comWords'] = json.loads(data.get('comWords'))
if data.get('keywords'):
data['keywords'] = json.loads(data.get('keywords'))
if data.get('productWordsMap'):
data['productWordsMap'] = json.loads(data.get('productWordsMap'))
return main_process(data)
def handle_one_task(task):
# redis_client = init_redis()
redis_client = init_redis8()
QUEUE_KEY = "geo:call_back:list"
# 每轮最多拉取任务数量
BATCH_SIZE = 50
# 每 30 秒执行一次
INTERVAL_SECONDS = 30
# 每轮内部并发数
CONCURRENT_WORKERS = 50
def parse_task(raw):
try:
task_data = task.get("task_data")
result = task.get("result")
process_call_back(task_data, result)
if isinstance(raw, bytes):
raw = raw.decode("utf-8")
return json.loads(raw)
except Exception as e:
print(f"回调任务处理失败: {e}, task={task}")
return None
def run_stream_task():
global RUNNING
while RUNNING:
status = 'ING'
stream_task = bh_utils.query_data(
f"select * from geo_commit_task where type = 'stream' and status = '{status}' limit 20"
)
if not stream_task:
time.sleep(5)
continue
if stream_task:
with ThreadPoolExecutor(max_workers=20) as executor:
futures = [executor.submit(handle_item, item) for item in stream_task]
def safe_json_loads(value, default=None):
if value is None:
return default
# 已经是 list/dict 了，直接返回
if isinstance(value, (list, dict)):
return value
# Redis 读出来有时是 bytes
if isinstance(value, bytes):
value = value.decode("utf-8")
# 字符串才尝试 json.loads
if isinstance(value, str):
value = value.strip()
if value == "":
return default
for future in as_completed(futures):
try:
future.result()
return json.loads(value)
except Exception as e:
print(f"任务执行失败: {e}")
return default
return default
def run_callback_task():
global RUNNING
REDIS_KEY = "geo:call_back:list"
def pull_tasks():
"""
每次最多从 Redis 队列拉取 BATCH_SIZE 条任务
Redis 连接异常时，只停止本轮拉取，不影响下次定时任务继续跑
"""
task_list = []
while RUNNING:
batch = []
for _ in range(BATCH_SIZE):
try:
raw = redis_client.rpop(QUEUE_KEY)
# 一次从 redis 取 20 个
for _ in range(20):
result = redis_client8.rpop(REDIS_KEY)
if not result:
if not raw:
break
batch.append(json.loads(result))
# 没数据就休眠，避免死循环空转
if not batch:
time.sleep(1)
continue
data = parse_task(raw)
if data:
task_list.append(data)
# 20 个并发处理
with ThreadPoolExecutor(max_workers=20) as executor:
futures = [executor.submit(handle_one_task, task) for task in batch]
except (ConnectionError, TimeoutError) as e:
logger.warning(f"Redis连接异常，本轮停止拉取: {e}")
try:
redis_client.connection_pool.disconnect()
except Exception:
pass
break
except Exception as e:
logger.exception(f"拉取队列任务异常: {e}")
break
return task_list
def handle_one_task(task):
"""
单条任务处理。
成功：正常结束
失败：重新放回 Redis 队列
"""
task_data = task.get("task_data")
result = task.get("result")
req_id = task_data.get('reqId')
platform = task_data.get('platform')
task_data["comWordsMap"] = safe_json_loads(task_data.get("comWordsMap"), {})
task_data["brandWords"] = safe_json_loads(task_data.get("brandWords"), [])
task_data["comWords"] = safe_json_loads(task_data.get("comWords"), [])
task_data["keywords"] = safe_json_loads(task_data.get("keywords"), [])
task_data["productWordsMap"] = safe_json_loads(task_data.get("productWordsMap"), {})
try:
process_call_back(task_data,result)
return True
except Exception as e:
logger.exception(f"{req_id}--{platform}--处理失败，重新放回队列: {e}")
redis_client.lpush(
QUEUE_KEY,
json.dumps(task, ensure_ascii=False)
)
return False
def consume_call_back():
"""
定时任务：
1. 拉取 Redis 队列
2. 使用线程池并发执行 main_process
"""
start_time = time.time()
try:
tasks = pull_tasks()
if not tasks:
logger.info("本轮没有任务")
return
with ThreadPoolExecutor(max_workers=CONCURRENT_WORKERS) as executor:
futures = [
executor.submit(handle_one_task, task)
for task in tasks
]
for future in as_completed(futures):
try:
future.result()
ok = future.result()
except Exception as e:
print(f"线程执行异常: {e}")
logger.exception(f"线程任务异常: {e}")
cost = round(time.time() - start_time, 2)
if __name__ == '__main__':
# signal.signal(signal.SIGINT, stop_handler)
# signal.signal(signal.SIGTERM, stop_handler)
logger.success(
f"本轮处理完成，总数: {len(tasks)}, "
f"耗时: {cost}s"
)
t1 = threading.Thread(target=run_stream_task, name="stream-thread")
# t2 = threading.Thread(target=run_callback_task, name="callback-thread")
t1.start()
# t2.start()
# logger.success("stream-thread 线程已启动")
# logger.success("callback-thread 线程已启动")
t1.join()
# t2.join()
# logger.success("所有线程已退出")
except Exception as e:
logger.exception(f"定时消费任务异常: {e}")
if __name__ == "__main__":
scheduler = BlockingScheduler(timezone="Asia/Shanghai")
scheduler.add_job(
consume_call_back,
trigger="interval",
seconds=INTERVAL_SECONDS,
id="consume_call_back",
max_instances=2, # 允许最多 5 个调度批次同时跑
coalesce=True,
misfire_grace_time=300
)
logger.success(
f"geo task commit consumer 启动，每 {INTERVAL_SECONDS}s 执行一次，"
f"QUEUE_KEY={QUEUE_KEY}, "
f"BATCH_SIZE={BATCH_SIZE}, "
f"CONCURRENT_WORKERS={CONCURRENT_WORKERS}, "
f"max_instances=2"
)
scheduler.start()
\ No newline at end of file
aidso_geo/models/baidu_baikan_data_process.py
View file @ 94534a08
......@@ -28,7 +28,7 @@ def baidu_baikan_data_process_original_data(data):
citation_list_data = json_content.get('citationListData','')
err_code = json_content.get('there_is_no_answer_to_this_prompt','')
if err_code =='99':
response_content = "There is no answer to this prompt"
response_content = "很抱歉，百度AI暂未对该问题生成回答。当前请求已成功提交至百度AI，未返回内容可能与其内容审核策略或服务状态有关。您可以稍后重试，或换一种表述方式再次提问"
spider_save_tos.process_and_save_files(file_path, search_keyword, url_list, think_content, response_content,
suggestions)
return (file_path, search_keyword, url_list, think_content, response_content, suggestions)
......
aidso_geo/models/process.py
View file @ 94534a08
......@@ -18,8 +18,6 @@ from aidso_geo.utils.tos_utils import put_string_to_tos, check_file_in_tos, get_
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(BASE_DIR)
# BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# sys.path.append(BASE_DIR)
from aidso_geo.config.base_config import init_redis
from aidso_geo.utils.ai_utils import ai_get_brand_sentiment_and_mentions, \
ai_get_brand_sentiment_and_mentions_by_articles, ai_get_product_list, ai_get_product_sentiment_and_mentions, \
......@@ -1290,6 +1288,7 @@ def reset_rank(all_vos, brand_vos, competitor_vos):
item["rank"] = new_rank
return all_vos, brand_vos, competitor_vos
def result_v2(response_content, data):
data = process_com_map(data)
reqId = data.get('reqId')
......@@ -1719,7 +1718,7 @@ def main_process(data):
bool_result = spider_interface.get_platform_response(data)
if bool_result:
size = get_tos_file_size(f"geo/{task_id}/{platform}/original.text")
if size is None or size < 1024:
if (size is None or size < 1024) and platform!='BDAI':
scheduler(data)
else:
logger.success(f"{req_id}--{platform}--{prompt}--ING_IN_TOS")
......@@ -1822,7 +1821,7 @@ def process_call_back(task_data, result):
req_status = req_check[0].get("status", "")
if req_check and req_status != 'SUCCESS':
size = get_tos_file_size(f"geo/{task_id}/{platform}/original.text")
if size is None or size < 1024:
if (size is None or size < 1024) and platform!='BDAI':
scheduler(task_data)
else:
platform_process(task_data)
......@@ -1955,12 +1954,9 @@ def run_data(PAGE_SIZE,MAX_WORKERS):
if __name__ == '__main__':
data_list = bh_utils.query_data(f"select * from geo_commit_task where status ='ING'")
# data_list = bh_utils.query_data(query_sql)
# print(data_list)
# # #
# # #
# # # # #
#
data_list = bh_utils.query_data(f"select * from geo_commit_task where status !='SUCCESS'")
def handle_item(i):
if i.get('comWordsMap'):
i['comWordsMap'] = json.loads(i.get('comWordsMap'))
......@@ -1989,5 +1985,3 @@ if __name__ == '__main__':
except Exception as e:
logger.exception(f"platform_process 执行异常: {e}")
# keywords = []
# get_keyword_ranks()
aidso_geo/utils/bh_utils.py
View file @ 94534a08
import time
import uuid
from loguru import logger
from pymysql.cursors import DictCursor
import json
import pymysql
import os,sys
from dbutils.pooled_db import PooledDB
DB_CONFIG = {
'host': "tenant-2101894307-cn-beijing-public.bytehouse.volces.com",
'host': "tenant-2101894307-cn-beijing.bytehouse.ivolces.com",
'port': 3306,
'user': "bytehouse",
'password': "JkoyoRV3PH:7YaWLXRwAO",
......@@ -20,11 +18,11 @@ DB_CONFIG = {
POOL = PooledDB(
creator=pymysql,
maxconnections=50,
mincached=5,
maxcached=20,
maxconnections=30,
mincached=1,
maxcached=10,
blocking=True,
ping=4,
ping=7,
cursorclass=DictCursor,
autocommit=True,
connect_timeout=5,
......@@ -33,35 +31,73 @@ POOL = PooledDB(
**DB_CONFIG
)
def is_connection_error(e):
if isinstance(e, pymysql.err.OperationalError):
code = e.args[0] if e.args else None
return code in (0, 2006, 2013)
return False
def query_data(sql, params=None, size=None):
def get_conn(retry_times=2):
for attempt in range(retry_times + 1):
conn = None
try:
conn = POOL.connection()
with conn.cursor() as cursor:
conn.ping()
return conn
except Exception as e:
try:
if conn is not None:
conn.close()
except Exception:
pass
if attempt < retry_times:
time.sleep(0.2 * (attempt + 1))
continue
raise
def query_data(sql, params=None, size=None, retry_times=2):
for attempt in range(retry_times + 1):
conn = None
try:
conn = get_conn()
with conn.cursor() as cursor:
cursor.execute(sql, params)
if size is None:
return cursor.fetchall()
return cursor.fetchmany(size)
except Exception as e:
print(f"查询数据失败: {e}")
logger.warning(
f"查询数据失败 attempt={attempt + 1}/{retry_times + 1}: {e}"
)
if attempt < retry_times and is_connection_error(e):
time.sleep(0.2 * (attempt + 1))
continue
return None
finally:
try:
if conn is not None:
conn.close()
except Exception:
pass
def insert_data(table_name, items) -> bool:
conn = None
try:
# 检查输入是否为空
def insert_data(table_name, items, retry_times=2) -> bool:
if not items:
return False
conn = POOL.connection()
with conn.cursor() as cursor:
for attempt in range(retry_times + 1):
conn = None
try:
conn = get_conn()
with conn.cursor() as cursor:
cols = items[0].keys()
cols_str = ", ".join(cols)
placeholders = ", ".join(["%s"] * len(cols))
......@@ -71,7 +107,7 @@ def insert_data(table_name, items) -> bool:
row = []
for key in cols:
value = item[key]
if isinstance(value, list):
if isinstance(value, (list, dict)):
row.append(json.dumps(value, ensure_ascii=False))
else:
row.append(value)
......@@ -81,55 +117,24 @@ def insert_data(table_name, items) -> bool:
cursor.executemany(insert_sql, processed_data)
return True
except Exception as e:
print(f"插入数据失败: {str(e)}")
logger.error(
f"插入数据失败 attempt={attempt + 1}/{retry_times + 1}, "
f"table={table_name}: {e}"
)
if attempt < retry_times and is_connection_error(e):
time.sleep(0.2 * (attempt + 1))
continue
return False
finally:
try:
if conn is not None:
conn.close()
except Exception:
pass
# 示例调用
if __name__ == "__main__":
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
TASK_ID = "6751e32e-d87a-4ed4-a914-5b9893c45371"
PLATFORM = "KIMI"
SQL = f"""
INSERT INTO geo_commit_task (brandWords, comWords, platform, prompt, reqId, searchEnabled, taskId, thinkingEnabled, type, thinking_enabled, search_enabled, insertime, status) VALUES ('[\"爱搜\", \"AIDSO\"]', '[]', 'DOUBA', '视频内容优化服务适用于哪些具体场景？', '51f7307e004449e9a30acd87394fca1b', 1, '875eee49002a4182b4c8b1ce7a4d09b1', 1, 'stream_batch', '1', '1', 1774323442, 'ING')
"""
import uuid
def run_query():
random_uuid = str(uuid.uuid4())
data = [
{
"brandWords": '[\"爱搜\", \"AIDSO\"]',
"comWords": "[]",
"platform": 'DOUBA',
"prompt": '视频内容优化服务适用于哪些具体场景',
"reqId": random_uuid,
"searchEnabled": ["程序员", "Go"],
"taskId": random_uuid,
"thinkingEnabled": 3,
"search_enabled":3,
"insertime": 1774323442,
"status": 'SUCCESS',
}]
return insert_data("geo_commit_task",data)
while True:
with ThreadPoolExecutor(max_workers=200) as executor:
futures = [executor.submit(run_query) for _ in range(50)]
for future in as_completed(futures):
try:
result = future.result()
print(result)
except Exception as e:
print("查询异常:", e)
\ No newline at end of file
aidso_geo/utils/tos_utils.py
View file @ 94534a08
......@@ -134,9 +134,6 @@ def save_image_to_tos(image_url, tos_path=None):
buffer.seek(0)
tos_path = str(tos_path)
if not tos_path.endswith(".jpg"):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment