main.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367
  1. import hashlib
  2. import traceback
  3. import uuid
  4. from concurrent.futures import ThreadPoolExecutor
  5. from datetime import date
  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 uvicorn
  12. import yaml
  13. from fastapi import FastAPI
  14. from fastapi.middleware.cors import CORSMiddleware
  15. from fastapi.responses import StreamingResponse
  16. from loguru import logger
  17. from pydantic import BaseModel, Field
  18. from asr_client import send_asr_request, send_task_request
  19. from common_func import get_db_engine
  20. from config.url import toutiao_static_video_url
  21. from database import insert, update, query, Task
  22. logger.add("logs/loguru.{time:YYYY-MM-DD}.log",
  23. rotation="00:00",
  24. format="{time:YYYY-MM-DD HH:mm:ss,SSS} [{process}] [{thread}] {level} {file} {line} - {message}",
  25. level="INFO")
  26. with open('/data/pythonProject/video_to_word/config/config.yaml', mode='r', encoding='utf-8') as f:
  27. config = yaml.load(f.read(), Loader=yaml.FullLoader)
  28. source_name_map = config['source_name_map']
  29. # 数据库连接引擎,依据开发、测试环境/生产环境 进行切换
  30. mac = uuid.UUID(int=uuid.getnode()).hex[-12:]
  31. if mac in ['5254003fa716', '52540003f5dd']:
  32. ai_word_engine = get_db_engine(config['ai_word_dev_db'])
  33. else:
  34. ai_word_engine = get_db_engine(config['ai_word_product_db'])
  35. threadPool = ThreadPoolExecutor(max_workers=4)
  36. app = FastAPI()
  37. origins = [
  38. "http://192.168.1.34",
  39. "http://192.168.1.34:8000",
  40. "http://192.168.1.105",
  41. "http://192.168.1.105:3000",
  42. "http://111.206.86.186",
  43. "http://111.206.86.186:3000",
  44. "http://adsp.tjyourong.com.cn",
  45. "http://adsp.tjyourong.com.cn:3000",
  46. "http://adsp.c-top.com.cn",
  47. "http://adsp.c-top.com.cn:3000"
  48. ]
  49. app.add_middleware(
  50. CORSMiddleware,
  51. allow_origins=origins,
  52. allow_credentials=True,
  53. allow_methods=["*"],
  54. allow_headers=["*"],
  55. )
  56. class QueryItem():
  57. signature: Optional[str] = None
  58. url: Optional[str] = None
  59. @app.get('/', tags=['back-end task'])
  60. def index():
  61. return {'message': '你已经正确创建 FastApi 服务!'}
  62. @app.post('/asr/task/submit', tags=['back-end task'])
  63. def task_submit(signature: str, url: str):
  64. json = send_asr_request(url)
  65. task = Task(signature=signature, task_id=json.Data.TaskId, task_result=json.to_json_string(), task_status=1)
  66. insert(task)
  67. return {'code': 0, 'taskId': json.Data.TaskId}
  68. @app.post('/asr/task/result', tags=['back-end task'])
  69. def task_submit(task_id: int):
  70. json = send_task_request(task_id)
  71. task = query(None, None, task_id)[0]
  72. task.task_status = json.Data.Status
  73. task.task_result = json.to_json_string()
  74. try:
  75. if json.Data.Status == 2:
  76. task.word_text = json.Data.ResultDetail[0].FinalSentence
  77. task.word_split = json.Data.ResultDetail[0].SliceSentence
  78. task.word_text_md5 = hashlib.md5(task.word_text.encode('utf-8')).hexdigest()
  79. except:
  80. # 提取原始文本内容和分词内容发生异常,把 task_status 置为 -1
  81. task.task_status = -1
  82. update(task)
  83. return {'code': 0, 'status': json.Data.StatusStr}
  84. @app.post('/asr/task/list', tags=['back-end task'])
  85. def task_submit(task_status: int):
  86. task = query(None, task_status, None)
  87. return {'code': 0, 'data': task}
  88. class BaseResponse(BaseModel):
  89. message: str = Field(..., description='消息')
  90. success: bool = Field(..., description='true or false')
  91. code: int = Field(..., description='')
  92. class TaskDetail(BaseModel):
  93. source_name: str = Field('内部创意', description='数据来源名称')
  94. query_word: str = Field('红包', description='关键词')
  95. stat_date: str = Field('2021-11-11', description='日期')
  96. script_num: int = Field(180, description='脚本数量')
  97. task_status: str = Field('执行成功', description='状态')
  98. number: int = Field(0, description='序号')
  99. class ConfigDetail(BaseModel):
  100. config_id: str = Field(..., description="脚本配置id")
  101. query_word_lst: List[str] = Field(..., description="关键词")
  102. create_time: str = Field(..., description="创建时间")
  103. operator: str = Field(..., description="创建人")
  104. number: int = Field(..., description="序号")
  105. class TaskResponse(BaseResponse):
  106. total_num: int = Field(0, description="总个数")
  107. page_num: int = Field(1, description="第几页")
  108. page_size: int = Field(10, description="每页个数")
  109. config_id: str = Field('', description="脚本配置id")
  110. result: List[TaskDetail] = Field(..., description="结果详情")
  111. class ConfigResponse(BaseResponse):
  112. total_num: int = Field(0, description="总个数")
  113. page_num: int = Field(1, description="第几页")
  114. page_size: int = Field(10, description="每页个数")
  115. result: List[ConfigDetail] = Field(..., description="结果详情")
  116. class QueryWordItem(BaseModel):
  117. query_word: str = Field("红包", description="查询词", min_length=1)
  118. stat_date: str = Field("2021-11-16", description="日期", min_length=10, max_length=10)
  119. source_code: int = Field(2, description="数据来源编码{1:'内部创意', 2:'巨量创意', 3:'开眼快创'}")
  120. class ScriptConfigLst(BaseModel):
  121. start_date: Optional[date] = Field(date.today() + timedelta(days=-6), description="开始日期-用于查询")
  122. end_date: Optional[date] = Field(date.today(), description="结束日期-用于查询")
  123. search_word: Optional[str] = Field('', description="关键词-用于查询")
  124. page_num: int = Field(1, description="第几页")
  125. page_size: int = Field(10, description="每页的大小")
  126. class QueryWordTaskInfoLst(BaseModel):
  127. start_date: Optional[date] = Field(date.today() + timedelta(days=-30), description="开始日期-用于查询")
  128. end_date: Optional[date] = Field(date.today(), description="结束日期-用于查询")
  129. search_word: Optional[str] = Field('', description="关键词-用于查询")
  130. page_num: int = Field(1, description="第几页")
  131. page_size: int = Field(10, description="每页的大小")
  132. config_id: Optional[str] = Field('', description="脚本配置id")
  133. source_code: Optional[List[int]] = Field([0], description="数据来源编码{1:'内部创意', 2:'巨量创意', 3:'开眼快创'}")
  134. class AddScriptConfig(BaseModel):
  135. query_word_lst: List[str] = Field(..., description="关键词组")
  136. operator: str = Field(..., description="操作者")
  137. class Config:
  138. schema_extra = {
  139. "example": {
  140. "query_word_lst": ["红包", "淘特"],
  141. "operator": "龙猫"
  142. }
  143. }
  144. @logger.catch
  145. @app.post('/export_script_file/', tags=['front-end interactive'],
  146. description="导出文件",
  147. summary='导出文件'
  148. )
  149. def export_script_file(item: List[QueryWordItem]):
  150. video_df = pd.DataFrame()
  151. # 1 从数据库获取视频数据
  152. # 如果同一个素材有多个查询词,则合并打上这多个查询词
  153. for obj in item:
  154. query_word = obj.query_word
  155. stat_date = obj.stat_date
  156. source_code = obj.source_code
  157. sql = f"select signature, video_url, query_word, stat_date, {source_code} source_code from {source_name_map[source_code]['table']} " \
  158. f"where query_word = '{query_word}' " \
  159. f"and stat_date = '{stat_date}'"
  160. df = pd.read_sql(sql, ai_word_engine)
  161. video_df = video_df.append(df)
  162. # 按 'signature' + 'query_word' + 'stat_date' 进行去重
  163. video_df.drop_duplicates(['signature', 'query_word', 'stat_date', 'source_code'], keep='last', inplace=True)
  164. video_query_word_df = video_df.groupby('signature').apply(lambda x: pd.Series({'query_word_lst': x['query_word'].unique(),
  165. 'video_url': x['video_url'].values[0],
  166. 'source_code': x['source_code'].values[0]}))
  167. video_query_word_df.reset_index(inplace=True, drop=False)
  168. # 如果来源==2 (头条巨量引擎),把视频链接替换为永久链接
  169. video_query_word_df['video_url'] = video_query_word_df.apply(
  170. lambda row: toutiao_static_video_url + row['signature'] if row.get('source_code') == 2 else row['video_url'], axis=1)
  171. # 2 根据第一步的视频数据获取脚本
  172. if not video_query_word_df.empty:
  173. signature_lst = list(video_query_word_df.signature.values) if len(video_query_word_df.signature.values) > 1 \
  174. else list(video_query_word_df.signature.values) * 2
  175. sql = f"select signature, word_text from tb_asr_result where signature in {tuple(signature_lst)}" \
  176. f"and word_text is not null"
  177. script_df = pd.read_sql(sql, ai_word_engine)
  178. out_df = video_query_word_df.merge(script_df, on='signature', how='inner')
  179. # 3 返回流数据
  180. if not out_df.empty:
  181. bio = BytesIO()
  182. writer = pd.ExcelWriter(bio, engine='xlsxwriter')
  183. out_df[['signature', 'query_word_lst', 'word_text', 'video_url']].to_excel(writer, index=False, encoding='utf8mb4')
  184. writer.save()
  185. bio.seek(0)
  186. # 组装header
  187. now_date = date.today().strftime('%Y-%m-%d')
  188. headers = {"content-type": "application/vnd.ms-excel",
  189. "content-disposition": f"attachment;filename={quote('优质素材脚本_')}{now_date}.xlsx"
  190. }
  191. return StreamingResponse(bio, media_type='xlsx', headers=headers)
  192. return None
  193. @logger.catch
  194. @app.post('/get_script_config_lst/', tags=['front-end interactive'], response_model=ConfigResponse,
  195. description="脚本配置列表",
  196. summary='脚本配置列表'
  197. )
  198. def get_script_config_lst(item: ScriptConfigLst):
  199. try:
  200. end_date = item.end_date + timedelta(days=1)
  201. sql = f"select * from ctop_ai_script_query_word_config " \
  202. f"where start_time >= '{item.start_date}' " \
  203. f"and start_time < '{end_date}' " \
  204. f"and ('{item.search_word}' = '' or query_word like '%%{item.search_word}%%')"
  205. org_df = pd.read_sql(sql, ai_word_engine)
  206. g_df = org_df.groupby('config_id').apply(lambda x: pd.Series({'query_word_lst': list(x['query_word'].unique()),
  207. 'operator': x['operator'].min(),
  208. 'create_time': str(x['start_time'].min())}))
  209. g_df.reset_index(drop=False, inplace=True)
  210. g_df.sort_values(by='create_time', ascending=False, inplace=True)
  211. g_df['number'] = list(range(1, len(g_df) + 1))
  212. total_num = g_df.shape[0]
  213. detail = g_df.iloc[(item.page_num - 1) * item.page_size: item.page_num * item.page_size].to_dict('records')
  214. response = {'code': 0,
  215. "message": "查询成功",
  216. "success": True,
  217. "result": detail,
  218. "total_num": total_num,
  219. "page_num": item.page_num,
  220. "page_size": item.page_size}
  221. logger.info(f"request body: {item}, response body: {response}")
  222. return response
  223. except:
  224. response = {"code": -1,
  225. "message": traceback.format_exc(),
  226. "success": False,
  227. "result": None}
  228. logger.error(f"request body: {item}, response body: {response}")
  229. return response
  230. @logger.catch
  231. @app.post('/get_query_word_task_info_lst/', tags=['front-end interactive'], response_model=TaskResponse,
  232. description="脚本数据导出列表",
  233. summary='脚本数据导出列表'
  234. )
  235. def get_query_word_task_info_lst(item: QueryWordTaskInfoLst):
  236. try:
  237. end_date = item.end_date + timedelta(days=1)
  238. source_code_lst = item.source_code * 2 if len(item.source_code) == 1 else item.source_code
  239. if item.config_id != '':
  240. sql = f"select distinct(query_word) query_word from ctop_ai_script_query_word_config where config_id = {item.config_id}"
  241. query_word_lst = list(pd.read_sql(sql, ai_word_engine).query_word.values)
  242. if len(query_word_lst) > 0:
  243. query_word_lst = query_word_lst * 2 if len(query_word_lst) == 1 else query_word_lst
  244. sql = f"select * from ctop_ai_query_word_task_record where query_word in {tuple(query_word_lst)}" \
  245. f"and stat_date >= '{item.start_date}' and stat_date < '{end_date}' " \
  246. f"and ('{item.source_code}' = '[0]' or source_code in {tuple(source_code_lst)}) " \
  247. f"and ('{item.search_word}' = '' or query_word = '{item.search_word}')"
  248. df = pd.read_sql(sql, ai_word_engine)
  249. else:
  250. sql = f"select * from ctop_ai_query_word_task_record where " \
  251. f"stat_date >= '{item.start_date}' and stat_date < '{end_date}' " \
  252. f"and ('{item.source_code}' = '[0]' or source_code in {tuple(source_code_lst)}) " \
  253. f"and ('{item.search_word}' = '' or query_word = '{item.search_word}')"
  254. df = pd.read_sql(sql, ai_word_engine)
  255. df['source_name'] = df['source_code'].apply(lambda x: source_name_map[x]['name'])
  256. df = df[['source_name', 'query_word', 'stat_date', 'script_num', 'task_status']]
  257. df.sort_values(['stat_date', 'source_name', 'query_word'], ascending=False, inplace=True)
  258. df['number'] = list(range(1, len(df) + 1))
  259. total_num = df.shape[0]
  260. detail = df.iloc[(item.page_num - 1) * item.page_size: item.page_num * item.page_size].to_dict('records')
  261. response = {'code': 0,
  262. "message": "查询成功",
  263. "success": True,
  264. "result": detail,
  265. "total_num": total_num,
  266. "page_num": item.page_num,
  267. "page_size": item.page_size,
  268. "config_id": item.config_id}
  269. logger.info(f"request body: {item}, response body: {response}")
  270. return response
  271. except:
  272. response = {"code": -1,
  273. "message": traceback.format_exc(),
  274. "success": False,
  275. "result": None}
  276. logger.error(f"request body: {item}, response body: {response}")
  277. return response
  278. @logger.catch
  279. @app.post('/add_script_config/',
  280. tags=['front-end interactive'],
  281. description="新增脚本配置",
  282. summary='新增脚本配置',
  283. response_model=BaseResponse)
  284. def add_script_config(item: AddScriptConfig):
  285. try:
  286. # 按查询词拆分配置记录
  287. config_id = str(uuid.uuid4())
  288. config_df = pd.DataFrame(data=item.query_word_lst, columns=['query_word'])
  289. config_df['config_id'] = config_id
  290. config_df['operator'] = item.operator
  291. config_df['operate_type'] = 1
  292. # 新增配置记录插入到 ctop_ai_script_query_word_config
  293. config_df.to_sql(name="ctop_ai_script_query_word_config",
  294. con=ai_word_engine,
  295. if_exists='append',
  296. index=False)
  297. logger.info(f"request body: {item}, code:0, message: add_script_config success")
  298. return {"code": 0,
  299. "message": "add success",
  300. "success": True}
  301. except:
  302. logger.error(f"request body: {item}, code:-1, message: add_script_config fail {traceback.format_exc()}")
  303. return {"code": -1,
  304. "message": traceback.format_exc(),
  305. "success": False}
  306. if __name__ == '__main__':
  307. uvicorn.run(app='main:app', host="0.0.0.0", port=31013, reload=True, debug=True)
  308. # gunicorn main:app -w 4 -k uvicorn.workers.UvicornWorker #线上启动命令