161 lines
5.5 KiB
Python
161 lines
5.5 KiB
Python
"""公司数据补充服务（持续消费模型）
|
|
|
|
启动 N 个 worker 协程，每个 worker 循环：取一条 → 处理 → 取下一条。
|
|
没数据时短暂休眠后重试。
|
|
"""
|
|
|
|
import asyncio
|
|
from datetime import datetime
|
|
|
|
from sqlalchemy import text
|
|
|
|
from app.config import settings
|
|
from app.core.database import MysqlSession
|
|
from app.core.logger import log
|
|
from app.ai.model_config import CompanyCleanModel
|
|
from app.ai.prompts import COMPANY_ENRICH_SYSTEM
|
|
from app.services.ai_tool import ai_chat_json
|
|
from app.services.dict_cache_service import dict_cache
|
|
|
|
# 停止信号
|
|
_stop_event = asyncio.Event()
|
|
|
|
|
|
def stop_company_clean():
|
|
"""外部调用，通知所有 worker 停止"""
|
|
_stop_event.set()
|
|
|
|
|
|
async def run_company_clean() -> None:
|
|
"""启动 N 个 worker 协程持续消费"""
|
|
_stop_event.clear()
|
|
worker_count = settings.company_worker_count
|
|
log.info("公司补充：启动 {} 个 worker 协程", worker_count)
|
|
|
|
workers = [asyncio.create_task(_worker(i)) for i in range(worker_count)]
|
|
await asyncio.gather(*workers)
|
|
|
|
log.info("公司补充：所有 worker 已退出")
|
|
|
|
|
|
async def _worker(worker_id: int) -> None:
|
|
"""单个 worker：循环取一条、处理一条"""
|
|
while not _stop_event.is_set():
|
|
data = await _fetch_one()
|
|
if data is None:
|
|
await asyncio.sleep(settings.company_idle_sleep)
|
|
continue
|
|
|
|
try:
|
|
await _do_clean(data)
|
|
except Exception as e:
|
|
log.error("[worker-{}] 公司补充异常, id={}, shortName={}: {}",
|
|
worker_id, data["id"], data.get("short_name"), e)
|
|
|
|
|
|
async def _fetch_one() -> dict | None:
|
|
"""从 MySQL 锁定一条待完善公司并标记为 status=3"""
|
|
async with MysqlSession() as mysql:
|
|
# MySQL 不支持 UPDATE ... RETURNING，分两步
|
|
result = await mysql.execute(
|
|
text("""
|
|
SELECT * FROM bg_company
|
|
WHERE status = 0
|
|
LIMIT 1
|
|
FOR UPDATE SKIP LOCKED
|
|
"""),
|
|
)
|
|
row = result.mappings().first()
|
|
if not row:
|
|
return None
|
|
|
|
company_id = row["id"]
|
|
await mysql.execute(
|
|
text("UPDATE bg_company SET status = 3, update_time = NOW() WHERE id = :id"),
|
|
{"id": company_id},
|
|
)
|
|
await mysql.commit()
|
|
return dict(row)
|
|
|
|
|
|
async def _do_clean(company: dict) -> None:
|
|
"""公司补充逻辑"""
|
|
company_id = company["id"]
|
|
short_name = company.get("short_name", "")
|
|
|
|
user_msg = f"【公司简称】\n{short_name}\n\n【行业列表】\n{dict_cache.industry_text}"
|
|
result = await ai_chat_json(CompanyCleanModel.ENRICH, COMPANY_ENRICH_SYSTEM, user_msg, scene="company_enrich")
|
|
|
|
if result is None or not result.get("valid", False):
|
|
await _update_status(company_id, 4)
|
|
return
|
|
|
|
# 地区匹配
|
|
city = result.get("city")
|
|
region_code = dict_cache.match_region_code(city) if city else None
|
|
|
|
# 回填数据
|
|
now = datetime.now()
|
|
async with MysqlSession() as mysql:
|
|
await mysql.execute(
|
|
text("""
|
|
UPDATE bg_company SET
|
|
name = COALESCE(:name, name),
|
|
region_code = COALESCE(:region_code, region_code),
|
|
company_type = COALESCE(:company_type, company_type),
|
|
industry_id = :industry_id,
|
|
tags = :tags,
|
|
summary = COALESCE(:summary, summary),
|
|
description = COALESCE(:description, description),
|
|
founded_year = COALESCE(:founded_year, founded_year),
|
|
address = COALESCE(:address, address),
|
|
scale = COALESCE(:scale, scale),
|
|
website = COALESCE(:website, website),
|
|
financing_stage = COALESCE(:financing_stage, financing_stage),
|
|
latest_valuation = COALESCE(:latest_valuation, latest_valuation),
|
|
news = :news,
|
|
status = 1,
|
|
update_time = :now
|
|
WHERE id = :id
|
|
"""),
|
|
{
|
|
"name": result.get("name"),
|
|
"region_code": region_code,
|
|
"company_type": result.get("companyType"),
|
|
"industry_id": result.get("industryId"),
|
|
"tags": _to_json(result.get("tags")),
|
|
"summary": result.get("summary"),
|
|
"description": result.get("description"),
|
|
"founded_year": result.get("foundedYear"),
|
|
"address": result.get("address"),
|
|
"scale": result.get("scale"),
|
|
"website": result.get("website"),
|
|
"financing_stage": result.get("financingStage"),
|
|
"latest_valuation": result.get("latestValuation"),
|
|
"news": _to_json(result.get("news")),
|
|
"now": now,
|
|
"id": company_id,
|
|
},
|
|
)
|
|
await mysql.commit()
|
|
|
|
log.info("公司补充完成, id={}, shortName={}", company_id, short_name)
|
|
|
|
|
|
async def _update_status(company_id: int, status: int) -> None:
|
|
"""更新公司状态"""
|
|
async with MysqlSession() as mysql:
|
|
await mysql.execute(
|
|
text("UPDATE bg_company SET status = :s, update_time = NOW() WHERE id = :id"),
|
|
{"s": status, "id": company_id},
|
|
)
|
|
await mysql.commit()
|
|
|
|
|
|
def _to_json(value) -> str | None:
|
|
"""列表转 JSON 字符串"""
|
|
import json
|
|
if value and isinstance(value, list):
|
|
return json.dumps(value, ensure_ascii=False)
|
|
return None
|