main.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234
  1. import datetime
  2. import hashlib
  3. import uuid
  4. from concurrent.futures import ThreadPoolExecutor
  5. from io import BytesIO
  6. from typing import Optional, List
  7. from urllib.parse import quote
  8. import pymysql
  9. import pandas as pd
  10. import uvicorn
  11. import yaml
  12. from fastapi import FastAPI
  13. from fastapi.middleware.cors import CORSMiddleware
  14. from fastapi.responses import StreamingResponse
  15. from pydantic import BaseModel, Field
  16. from asr_client import send_asr_request, send_task_request
  17. from common_func import get_db_engine, mysql_replace_into
  18. from config.url import toutiao_static_video_url
  19. from database import insert, update, query, Task
  20. with open('/data/pythonProject/video_to_word/config/config.yaml', mode='r', encoding='utf-8') as f:
  21. config = yaml.load(f.read(), Loader=yaml.FullLoader)
  22. source_name_map = config['source_name_map']
  23. ai_word_engine = get_db_engine(config['ai_word_dev_db'])
  24. threadPool = ThreadPoolExecutor(max_workers=4)
  25. app = FastAPI()
  26. origins = [
  27. "http://192.168.1.34",
  28. "http://192.168.1.34:8000",
  29. "http://192.168.1.105",
  30. "http://192.168.1.105:3000",
  31. "http://111.206.86.186",
  32. "http://111.206.86.186:3000",
  33. "http://adsp.tjyourong.com.cn",
  34. "http://adsp.tjyourong.com.cn:3000",
  35. "http://adsp.c-top.com.cn",
  36. "http://adsp.c-top.com.cn:3000"
  37. ]
  38. app.add_middleware(
  39. CORSMiddleware,
  40. allow_origins=origins,
  41. allow_credentials=True,
  42. allow_methods=["*"],
  43. allow_headers=["*"],
  44. )
  45. class QueryItem():
  46. signature: Optional[str] = None
  47. url: Optional[str] = None
  48. @app.get('/')
  49. def index():
  50. return {'message': '你已经正确创建 FastApi 服务!'}
  51. @app.post('/asr/task/submit')
  52. def task_submit(signature: str, url: str):
  53. json = send_asr_request(url)
  54. task = Task(signature=signature, task_id=json.Data.TaskId, task_result=json.to_json_string(), task_status=1)
  55. insert(task)
  56. return {'code': 0, 'taskId': json.Data.TaskId}
  57. @app.post('/asr/task/result')
  58. def task_submit(task_id: int):
  59. json = send_task_request(task_id)
  60. task = query(None, None, task_id)[0]
  61. task.task_status = json.Data.Status
  62. task.task_result = json.to_json_string()
  63. try:
  64. if json.Data.Status == 2:
  65. task.word_text = json.Data.ResultDetail[0].FinalSentence
  66. task.word_split = json.Data.ResultDetail[0].SliceSentence
  67. task.word_text_md5 = hashlib.md5(task.word_text.encode('utf-8')).hexdigest()
  68. except:
  69. # 提取原始文本内容和分词内容发生异常,把 task_status 置为 -1
  70. task.task_status = -1
  71. update(task)
  72. return {'code': 0, 'status': json.Data.StatusStr}
  73. @app.post('/asr/task/list')
  74. def task_submit(task_status: int):
  75. task = query(None, task_status, None)
  76. return {'code': 0, 'data': task}
  77. class QueryWordItem(BaseModel):
  78. query_word: str = Field(..., description="查询词", min_length=1)
  79. stat_date: str = Field(..., description="日期", min_length=10, max_length=10)
  80. source: int = Field(..., description="来源,")
  81. @app.post('/export_excel/')
  82. def export_excel(item: List[QueryWordItem]):
  83. video_df = pd.DataFrame()
  84. if len(item) == 1:
  85. # 单个条目,直接导出
  86. pass
  87. else:
  88. # 1 从数据库获取视频数据
  89. # 多个条目,如果同一个素材有多个查询词,则合并打上这多个查询词
  90. for obj in item:
  91. query_word = obj.query_word
  92. stat_date = obj.stat_date
  93. source = obj.source
  94. sql = f"select signature, video_url, query_word, stat_date, {source} source from {source_name_map[source]['table']} " \
  95. f"where query_word = '{query_word}' " \
  96. f"and stat_date = '{stat_date}'"
  97. df = pd.read_sql(sql, ai_word_engine)
  98. video_df = video_df.append(df)
  99. # 按 'signature' + 'query_word' + 'stat_date' 进行去重
  100. video_df.drop_duplicates(['signature', 'query_word', 'stat_date', 'source'], keep='last', inplace=True)
  101. g = video_df.groupby('signature')
  102. query_word_lst_df = g.apply(lambda x: x['query_word'].unique())
  103. query_word_lst_df.name = 'query_word_lst'
  104. url_df = g.apply(lambda x: x['video_url'].values[0])
  105. url_df.name = 'video_url'
  106. source_df = g.apply(lambda x: x['source'].values[0])
  107. source_df.name = 'source'
  108. video_query_word_df = pd.concat([query_word_lst_df, url_df, source_df], axis=1)
  109. video_query_word_df.reset_index(inplace=True, drop=False)
  110. video_query_word_df['video_url'] = video_query_word_df.apply(
  111. lambda row: toutiao_static_video_url + row['signature'] if row.get('source') == 2 else row['video_url'], axis=1)
  112. # 2 根据第一步的视频数据获取脚本
  113. if not video_query_word_df.empty:
  114. sql = f"select signature, word_text from tb_asr_result where signature in " \
  115. f"{tuple(video_query_word_df.signature.values) if len(video_query_word_df.signature.values) > 1 else tuple(list(video_query_word_df.signature.values) * 2)} " \
  116. f"and task_status = 2"
  117. script_df = pd.read_sql(sql, ai_word_engine)
  118. out_df = video_query_word_df.merge(script_df, on='signature', how='inner')
  119. else:
  120. pass
  121. # 3 返回流数据
  122. if not out_df.empty:
  123. bio = BytesIO()
  124. writer = pd.ExcelWriter(bio, engine='xlsxwriter')
  125. out_df[['signature', 'query_word_lst', 'word_text', 'video_url']].to_excel(writer, index=False, encoding='utf8mb4')
  126. writer.save()
  127. bio.seek(0)
  128. # 组装header
  129. now_date = datetime.date.today().strftime('%Y-%m-%d')
  130. headers = {"content-type": "application/vnd.ms-excel",
  131. "content-disposition": f"attachment;filename={quote('优质素材脚本_')}{now_date}.xlsx"
  132. }
  133. return StreamingResponse(bio, media_type='xlsx', headers=headers)
  134. return None
  135. class ScriptConfig(BaseModel):
  136. query_word_lst: List = Field(..., description="关键词组")
  137. operator: str = Field(..., description="操作者")
  138. @app.post('/get_script_config_lst/')
  139. def get_script_config_lst():
  140. pass
  141. @app.post('/add_script_config/')
  142. def add_script_config(item: ScriptConfig):
  143. config_id = str(uuid.uuid4())
  144. config_lst = []
  145. for query_word in item.query_word_lst:
  146. sql = f"select * from ctop_ai_query_word where query_word = '{query_word}'"
  147. query_word_df = pd.read_sql(sql, ai_word_engine)
  148. if not query_word_df.empty:
  149. # 更新 ctop_ai_query_word
  150. query_word_id = query_word_df.query_word_id.values[0]
  151. script_config_conn_num = query_word_df.script_config_conn_num.values[0] + 1
  152. db_con = pymysql.connect(**config['ai_word_dev_db'])
  153. db_cur = db_con.cursor()
  154. sql = f"update ctop_ai_query_word set script_config_conn_num = {script_config_conn_num} where query_word_id = '{query_word_id}'"
  155. db_cur.execute(sql)
  156. db_con.commit()
  157. db_con.close()
  158. # update_query_word_df = pd.DataFrame([{"query_word_id": query_word_id,
  159. # "query_word": query_word,
  160. # "script_conn_num": script_conn_num}])
  161. # update_query_word_df.to_sql(name="ctop_ai_query_word",
  162. # con=ai_word_engine,
  163. # if_exists="append",
  164. # method=mysql_replace_into,
  165. # index=False)
  166. else:
  167. query_word_id = str(uuid.uuid4())
  168. new_query_word_df = pd.DataFrame([{"query_word_id": query_word_id, "query_word": query_word, "script_config_conn_num": 1}])
  169. new_query_word_df.to_sql(name="ctop_ai_query_word",
  170. con=ai_word_engine,
  171. if_exists="append",
  172. index=False)
  173. config_lst.append({"config_id": config_id, "query_word_id": query_word_id})
  174. # 新增配置记录插入到 ctop_ai_script_query_word_config
  175. config_df = pd.DataFrame(config_lst)
  176. config_df['operator'] = item.operator
  177. config_df['operate_type'] = 1
  178. config_df.to_sql(name="ctop_ai_script_query_word_config",
  179. con=ai_word_engine,
  180. if_exists='append',
  181. index=False)
  182. return {"code": 0, "message": "success"}
  183. if __name__ == '__main__':
  184. # 1 读取配置文件
  185. # test_items = [{'query_word': '红包', 'stat_date': '2021-10-28', 'source': 2},
  186. # {'query_word': '红包', 'stat_date': '2021-10-28', 'source': 3},
  187. # {'query_word': '赚钱', 'stat_date': '2021-10-28', 'source': 2},
  188. # {'query_word': '赚钱', 'stat_date': '2021-10-28', 'source': 3}]
  189. # export_excel(test_items)
  190. uvicorn.run(app='main:app', host="0.0.0.0", port=31013, reload=True, debug=True)
  191. # gunicorn main:app -w 4 -k uvicorn.workers.UvicornWorker #线上启动命令