script_config.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442
  1. import os
  2. import sys
  3. import traceback
  4. import uuid
  5. from datetime import date, datetime
  6. from datetime import timedelta
  7. from io import BytesIO
  8. from typing import Optional, List
  9. from urllib.parse import quote
  10. import pandas as pd
  11. import yaml
  12. from fastapi import APIRouter
  13. from fastapi.responses import StreamingResponse
  14. from loguru import logger
  15. from pangres import upsert
  16. from pydantic import BaseModel, Field
  17. curr_path = os.path.abspath(os.path.dirname(__file__))
  18. project_root_path = curr_path[:curr_path.find("video_to_word") + len("video_to_word")]
  19. sys.path.append(project_root_path)
  20. from config.url_and_db import toutiao_static_video_url, ai_word_engine
  21. from sqlalchemy import VARCHAR
  22. from time_task.get_material_and_script_by_query_word import get_material_and_script
  23. router = APIRouter(tags=['script_config_server'])
  24. with open('/data/pythonProject/video_to_word/config/config.yaml', mode='r', encoding='utf-8') as f:
  25. config = yaml.load(f.read(), Loader=yaml.FullLoader)
  26. source_name_map = config['source_name_map']
  27. class BaseResponse(BaseModel):
  28. message: str = Field(..., description='消息')
  29. success: bool = Field(True, description='true or false')
  30. code: int = Field(0, description='')
  31. class TaskDetail(BaseModel):
  32. source_name: str = Field('内部创意', description='数据来源名称')
  33. query_word: str = Field('红包', description='关键词')
  34. stat_date: str = Field('2021-11-11', description='日期')
  35. script_num: str = Field('', description='脚本数量')
  36. task_status: str = Field('执行成功', description='状态')
  37. number: int = Field(0, description='序号')
  38. class ConfigDetail(BaseModel):
  39. config_id: str = Field(..., description="脚本配置id")
  40. query_word: List[str] = Field(..., description="关键词")
  41. recommended_word: List[str] = Field([""], description="推荐词")
  42. create_time: str = Field(..., description="创建时间")
  43. operator: str = Field(..., description="创建人")
  44. number: int = Field(..., description="序号")
  45. user_id: str = Field(..., description="用户id")
  46. class TaskResponse(BaseResponse):
  47. total_num: int = Field(0, description="总个数")
  48. page_num: int = Field(1, description="第几页")
  49. page_size: int = Field(10, description="每页个数")
  50. config_id: str = Field('', description="脚本配置id")
  51. result: List[TaskDetail] = Field([], description="结果详情")
  52. class ConfigResponse(BaseResponse):
  53. total_num: int = Field(0, description="总个数")
  54. page_num: int = Field(1, description="第几页")
  55. page_size: int = Field(10, description="每页个数")
  56. result: List[ConfigDetail] = Field([], description="结果详情")
  57. class QueryWordItem(BaseModel):
  58. query_word: str = Field("红包", description="查询词", min_length=1)
  59. stat_date: str = Field("2021-11-16", description="日期", min_length=10, max_length=10)
  60. source_code: int = Field(2, description="数据来源编码{1:'内部创意', 2:'巨量创意', 3:'开眼快创'}")
  61. class ScriptConfigLst(BaseModel):
  62. start_date: Optional[date] = Field(date.today() + timedelta(days=-29), description="开始日期-用于查询")
  63. end_date: Optional[date] = Field(date.today(), description="结束日期-用于查询")
  64. search_word: Optional[str] = Field('', description="关键词/推荐词-用于查询")
  65. page_num: int = Field(1, description="第几页")
  66. page_size: int = Field(10, description="每页的大小")
  67. class QueryWordTaskInfoLst(BaseModel):
  68. start_date: Optional[date] = Field(date.today() + timedelta(days=-30), description="开始日期-用于查询")
  69. end_date: Optional[date] = Field(date.today(), description="结束日期-用于查询")
  70. search_word: Optional[str] = Field('', description="关键词/推荐词-用于查询")
  71. page_num: int = Field(1, description="第几页")
  72. page_size: int = Field(10, description="每页的大小")
  73. config_id: Optional[str] = Field('', description="脚本配置id")
  74. source_code: Optional[List[int]] = Field([], description="数据来源编码{1:'内部创意', 2:'巨量创意', 3:'开眼快创'}")
  75. class QueryWordAndRecommendedWordPair(BaseModel):
  76. query_word: str = Field(..., description='关键字')
  77. recommended_word: List[str] = Field([], description='推荐词')
  78. class AddScriptConfig(BaseModel):
  79. query_word_pair: List[QueryWordAndRecommendedWordPair] = Field(..., description="关键词-推荐词")
  80. operator: str = Field(..., description="操作者")
  81. user_id: str = Field(..., description="user_id")
  82. class Config:
  83. schema_extra = {
  84. "example": {
  85. "query_word_pair": [{"query_word": "水蜜桃", "recommended_word": ["我等你", "陌陌", "寻爱", "聊吧"]}],
  86. "operator": "管理员",
  87. "user_id": "e9ca23d68d884d4ebb19d07889727dae"
  88. }
  89. }
  90. class DeleteScriptConfig(BaseModel):
  91. config_id: str = Field(..., description="脚本配置id")
  92. operator: str = Field(..., description="操作者")
  93. user_id: str = Field(..., description="user_id")
  94. class Config:
  95. schema_extra = {
  96. "example": {
  97. "config_id": "71951bcb-0ef7-4ce0-9be5-c8aaf3128ab7",
  98. "operator": "管理员",
  99. "user_id": "e9ca23d68d884d4ebb19d07889727dae"
  100. }
  101. }
  102. @logger.catch
  103. @router.post('/export_script_file/',
  104. description="导出文件",
  105. summary='导出文件',
  106. response_model=BaseResponse
  107. )
  108. def export_script_file(item: List[QueryWordItem]):
  109. try:
  110. logger.info(f"request body: {item}")
  111. video_df = pd.DataFrame()
  112. out_df = pd.DataFrame()
  113. # 1 从数据库获取视频数据
  114. # 如果同一个素材有多个查询词,则合并打上这多个查询词
  115. for obj in item:
  116. query_word = obj.query_word
  117. stat_date = obj.stat_date
  118. source_code = obj.source_code
  119. sql = f"select signature, video_url, query_word, stat_date, {source_code} source_code from {source_name_map[source_code]['table']} " \
  120. f"where query_word = '{query_word}' " \
  121. f"and stat_date = '{stat_date}'"
  122. df = pd.read_sql(sql, ai_word_engine)
  123. video_df = video_df.append(df)
  124. if not video_df.empty:
  125. # 按 'signature' + 'query_word' + 'stat_date' 进行去重
  126. video_df.drop_duplicates(['signature', 'query_word', 'stat_date', 'source_code'], keep='last', inplace=True)
  127. video_query_word_df = video_df.groupby('signature').apply(lambda x: pd.Series({'query_word_lst': x['query_word'].unique(),
  128. 'video_url': x['video_url'].values[0],
  129. 'source_code': x['source_code'].values[0]}))
  130. video_query_word_df.reset_index(inplace=True, drop=False)
  131. # 如果来源==2 (头条巨量引擎),把视频链接替换为永久链接
  132. video_query_word_df['video_url'] = video_query_word_df.apply(
  133. lambda row: toutiao_static_video_url + row['signature'] if row.get('source_code') == 2 else row['video_url'], axis=1)
  134. # 2 根据第一步的视频数据获取脚本
  135. signature_lst = list(video_query_word_df.signature.values) if len(video_query_word_df.signature.values) > 1 \
  136. else list(video_query_word_df.signature.values) * 2
  137. sql = f"select signature, word_text from tb_asr_result where signature in {tuple(signature_lst)}" \
  138. f"and word_text is not null"
  139. script_df = pd.read_sql(sql, ai_word_engine)
  140. out_df = video_query_word_df.merge(script_df, on='signature', how='inner')
  141. # 手动数据库关闭连接
  142. ai_word_engine.dispose()
  143. # 3 返回流数据
  144. if not out_df.empty:
  145. bio = BytesIO()
  146. writer = pd.ExcelWriter(bio, engine='xlsxwriter')
  147. out_df[['signature', 'query_word_lst', 'word_text', 'video_url']].to_excel(writer, index=False, encoding='utf8mb4')
  148. writer.save()
  149. bio.seek(0)
  150. # 组装header
  151. now_date = date.today().strftime('%Y-%m-%d')
  152. headers = {"content-type": "application/vnd.ms-excel",
  153. "content-disposition": f"attachment;filename={quote('优质素材脚本_')}{now_date}.xlsx"
  154. }
  155. logger.info(f"request body: {item}, message: 数据导出成功")
  156. return StreamingResponse(bio, media_type='xlsx', headers=headers)
  157. else:
  158. logger.info(f"request body: {item}, message: 没有获取到对应的数据")
  159. return {"code": 0,
  160. "message": "没有获取到对应的数据",
  161. "success": True}
  162. except:
  163. logger.error(f"request body: {item}, message: {traceback.format_exc()}")
  164. return {"code": 0,
  165. "message": {traceback.format_exc()},
  166. "success": False}
  167. @logger.catch
  168. @router.post('/get_script_config_lst/', response_model=ConfigResponse,
  169. description="脚本配置列表",
  170. summary='脚本配置列表'
  171. )
  172. def get_script_config_lst(item: ScriptConfigLst):
  173. logger.info(f"request body: {item}")
  174. response = ConfigResponse(message="查询成功")
  175. try:
  176. end_date = item.end_date + timedelta(days=1)
  177. sql = f"select * from ctop_ai_script_query_word_config where config_id in " \
  178. f"(select distinct(config_id) config_id from ctop_ai_script_query_word_config " \
  179. f"where ('{item.search_word}' = '' or query_word like '%%{item.search_word}%%' or recommended_word like '%%{item.search_word}%%') )" \
  180. f"and operate_type = 1 and end_time ='9999-12-31' " \
  181. f"and start_time >= '{item.start_date}' and start_time < '{end_date}' "
  182. org_df = pd.read_sql(sql, ai_word_engine)
  183. if not org_df.empty:
  184. g_df = org_df.groupby('config_id').apply(lambda x: pd.Series({'query_word': list(x['query_word'].unique()),
  185. 'recommended_word': list(x['recommended_word'].unique()),
  186. 'operator': x['operator'].min(),
  187. 'create_time': str(x['start_time'].min()),
  188. 'user_id': x['user_id'].min()}))
  189. g_df.reset_index(drop=False, inplace=True)
  190. g_df.sort_values(by='create_time', ascending=False, inplace=True)
  191. g_df['number'] = list(range(1, len(g_df) + 1))
  192. total_num = g_df.shape[0]
  193. detail = g_df.iloc[(item.page_num - 1) * item.page_size: item.page_num * item.page_size].to_dict('records')
  194. response.result = detail
  195. response.total_num = total_num
  196. response.page_num = item.page_num
  197. response.page_size = item.page_size
  198. else:
  199. response.message = "没有符合条件的数据"
  200. response.page_num = item.page_num
  201. response.page_size = item.page_size
  202. logger.info(f"request body: {item}, response body: {response}")
  203. except:
  204. response.code = -1
  205. response.message = traceback.format_exc()
  206. response.success = False
  207. logger.error(f"request body: {item}, response body: {response}")
  208. return response
  209. @logger.catch
  210. @router.post('/get_query_word_task_info_lst/', response_model=TaskResponse,
  211. description="脚本数据导出列表",
  212. summary='脚本数据导出列表'
  213. )
  214. def get_query_word_task_info_lst(item: QueryWordTaskInfoLst):
  215. logger.info(f"request body: {item}")
  216. response = TaskResponse(code=0, message="查询成功", success=True)
  217. try:
  218. end_date = item.end_date + timedelta(days=1)
  219. source_code_lst = [-1, -2] if len(item.source_code) == 0 else (item.source_code * 2 if len(item.source_code) == 1 else item.source_code)
  220. df = pd.DataFrame()
  221. if item.config_id != '':
  222. sql = f"select query_word, recommended_word from ctop_ai_script_query_word_config where config_id = '{item.config_id}'"
  223. config_df = pd.read_sql(sql, ai_word_engine)
  224. query_word_set = set(config_df['query_word'].values)
  225. recommended_word_set = set(config_df['recommended_word'].values)
  226. word_lst = list(query_word_set.union(recommended_word_set))
  227. if len(word_lst) > 0:
  228. word_lst = word_lst * 2 if len(word_lst) == 1 else word_lst
  229. sql = f"select * from ctop_ai_query_word_task_record where query_word in {tuple(word_lst)}" \
  230. f"and stat_date >= '{item.start_date}' and stat_date < '{end_date}' " \
  231. f"and ('{item.source_code}' = '[]' or source_code in {tuple(source_code_lst)}) " \
  232. f"and ('{item.search_word}' = '' or query_word = '{item.search_word}')"
  233. df = pd.read_sql(sql, ai_word_engine)
  234. else:
  235. sql = f"select * from ctop_ai_query_word_task_record where " \
  236. f"stat_date >= '{item.start_date}' and stat_date < '{end_date}' " \
  237. f"and ('{item.source_code}' = '[]' or source_code in {tuple(source_code_lst)}) " \
  238. f"and ('{item.search_word}' = '' or query_word = '{item.search_word}')"
  239. df = pd.read_sql(sql, ai_word_engine)
  240. # 手动数据库关闭连接
  241. ai_word_engine.dispose()
  242. if not df.empty:
  243. df['source_name'] = df['source_code'].apply(lambda x: source_name_map[x]['name'])
  244. df = df[['source_name', 'query_word', 'stat_date', 'script_num', 'task_status']]
  245. df.sort_values(['stat_date', 'source_name', 'query_word'], ascending=False, inplace=True)
  246. df['number'] = list(range(1, len(df) + 1))
  247. # script_num 字段类型由 np.array 转化为 str 类型,解决返回 np.nan 时, responseModel 验证不通过
  248. df['script_num'] = df['script_num'].astype(pd.Int64Dtype())
  249. df['script_num'] = df['script_num'].astype(str)
  250. df.replace('<NA>', '', inplace=True)
  251. total_num = df.shape[0]
  252. detail = df.iloc[(item.page_num - 1) * item.page_size: item.page_num * item.page_size].to_dict('records')
  253. response.result = detail
  254. response.total_num = total_num
  255. response.page_num = item.page_num
  256. response.page_size = item.page_size
  257. response.config_id = item.config_id
  258. else:
  259. response.message = "没有符合条件的数据"
  260. response.page_num = item.page_num
  261. response.page_size = item.page_size
  262. response.config_id = item.config_id
  263. logger.info(f"request body: {item}, response body: {response}")
  264. except:
  265. response.message = traceback.format_exc()
  266. response.code = -1
  267. response.success = False
  268. logger.error(f"request body: {item}, response body: {response}")
  269. return response
  270. @logger.catch
  271. @router.post('/add_script_config/',
  272. description="新增脚本配置",
  273. summary='新增脚本配置',
  274. response_model=BaseResponse)
  275. def add_script_config(item: AddScriptConfig):
  276. try:
  277. logger.info(f"request body: {item}")
  278. multi_config_df = pd.DataFrame()
  279. for pair in item.query_word_pair:
  280. config_id = str(uuid.uuid4())
  281. config_df = pd.DataFrame({'query_word': pair.query_word,
  282. 'recommended_word': [""] if len(pair.recommended_word) == 0 else pair.recommended_word,
  283. 'config_id': config_id})
  284. multi_config_df = multi_config_df.append(config_df)
  285. multi_config_df['operator'] = item.operator
  286. multi_config_df['operate_type'] = 1
  287. multi_config_df['user_id'] = item.user_id
  288. # 新增配置记录插入到 ctop_ai_script_query_word_config
  289. multi_config_df.to_sql(name="ctop_ai_script_query_word_config",
  290. con=ai_word_engine,
  291. if_exists='append',
  292. index=False)
  293. # 手动数据库关闭连接
  294. ai_word_engine.dispose()
  295. logger.info(f"request body: {item}, code:0, message: add_script_config success")
  296. return {"code": 0,
  297. "message": "add success",
  298. "success": True}
  299. except:
  300. logger.error(f"request body: {item}, code:-1, message: add_script_config fail {traceback.format_exc()}")
  301. return {"code": -1,
  302. "message": traceback.format_exc(),
  303. "success": False}
  304. @logger.catch
  305. @router.post('/delete_script_config/',
  306. description="删除脚本配置",
  307. summary='删除脚本配置',
  308. response_model=BaseResponse)
  309. def delete_script_config(item: DeleteScriptConfig):
  310. response = BaseResponse(code=0, message='delete success', success=True)
  311. try:
  312. logger.info(f"request body: {item}")
  313. sql = f"select * from ctop_ai_script_query_word_config where config_id = '{item.config_id}'"
  314. config_df = pd.read_sql(sql, ai_word_engine)
  315. if not config_df.empty:
  316. column_type = {'config_id': VARCHAR(36),
  317. 'query_word': VARCHAR(36),
  318. 'recommended_word': VARCHAR(36)}
  319. update_config_df = config_df.copy(deep=True)
  320. update_config_df['end_time'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
  321. update_config_df.set_index(['config_id', 'query_word', 'recommended_word', 'operate_type'], drop=True, inplace=True)
  322. upsert(engine=ai_word_engine,
  323. df=update_config_df,
  324. table_name='ctop_ai_script_query_word_config',
  325. if_row_exists='update',
  326. dtype=column_type)
  327. add_config_df = config_df.copy(deep=True)
  328. add_config_df['operate_type'] = 3
  329. add_config_df['user_id'] = item.user_id
  330. add_config_df['operator'] = item.operator
  331. add_config_df['start_time'] = datetime.today().strftime('%Y-%m-%d %H:%M:%S')
  332. add_config_df.drop(labels='end_time', axis=1, inplace=True)
  333. add_config_df.set_index(['config_id', 'query_word', 'recommended_word', 'operate_type'], drop=True, inplace=True)
  334. upsert(engine=ai_word_engine,
  335. df=add_config_df,
  336. table_name='ctop_ai_script_query_word_config',
  337. if_row_exists='update',
  338. dtype=column_type)
  339. else:
  340. response.message = f'没有获取到 {item.config_id} 对应的数据'
  341. # 手动数据库关闭连接
  342. ai_word_engine.dispose()
  343. logger.info(f"request body: {item}, response: {response}")
  344. except:
  345. response.code = -1
  346. response.message = traceback.format_exc()
  347. response.success = False
  348. logger.error(f"request body: {item}, response: {response}")
  349. return response
  350. @logger.catch
  351. @router.post('/get_material_and_script_time_task/',
  352. response_model=BaseResponse,
  353. description="获取素材和脚本任务",
  354. summary='获取素材和脚本任务')
  355. def get_material_and_script_time_task(start_time: Optional[str] = ""):
  356. try:
  357. get_material_and_script(start_time)
  358. logger.info(f"{date.today().strftime('%Y-%m-%d')}, 获取素材和脚本任务执行完成.")
  359. return {"code": 0,
  360. "success": True,
  361. "message": f"{date.today().strftime('%Y-%m-%d')},获取素材和脚本任务执行完成."}
  362. except:
  363. logger.error(f"{date.today().strftime('%Y-%m-%d')}, 获取素材和脚本任务执行发生异常. {traceback.format_exc()}")
  364. return {"code": -1,
  365. "success": False,
  366. "message": f"{date.today().strftime('%Y-%m-%d')},获取素材和脚本任务执行发生异常 .{traceback.format_exc()}"}
  367. if __name__ == '__main__':
  368. req = DeleteScriptConfig(config_id='4113ec60-4b92-4a4a-9dc5-417a29df9b65',
  369. operator='隋炎均',
  370. user_id='f75b91a1a23946688ab1d93a65d0a435')
  371. delete_script_config(req)