api_server.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843
  1. #! /usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. # vim:fenc=utf-8
  4. import json
  5. import logging
  6. from argparse import ArgumentParser
  7. import werkzeug.exceptions
  8. from flask import Flask, request, jsonify
  9. from sqlalchemy.orm import sessionmaker
  10. import pqai_agent_server.utils
  11. from pqai_agent import chat_service, prompt_templates
  12. from pqai_agent import configs
  13. from pqai_agent.chat_service import OpenAICompatible
  14. from pqai_agent.data_models.agent_configuration import AgentConfiguration
  15. from pqai_agent.data_models.service_module import ServiceModule
  16. from pqai_agent.history_dialogue_service import HistoryDialogueService
  17. from pqai_agent.logging import logger, setup_root_logger
  18. from pqai_agent.toolkit import global_tool_map
  19. from pqai_agent.user_manager import MySQLUserManager, MySQLUserRelationManager
  20. from pqai_agent.utils.db_utils import create_ai_agent_db_engine
  21. from pqai_agent.utils.prompt_utils import format_agent_profile, format_user_profile
  22. from pqai_agent_server.const import AgentApiConst
  23. from pqai_agent_server.const.status_enum import TestTaskStatus
  24. from pqai_agent_server.const.type_enum import EvaluateType
  25. from pqai_agent_server.dataset_service import DatasetService
  26. from pqai_agent_server.models import MySQLSessionManager
  27. from pqai_agent_server.task_server import TaskManager
  28. from pqai_agent_server.utils import (
  29. run_extractor_prompt,
  30. run_chat_prompt,
  31. run_response_type_prompt,
  32. )
  33. from pqai_agent_server.utils import wrap_response
  34. app = Flask('agent_api_server')
  35. const = AgentApiConst()
  36. @app.route('/api/listStaffs', methods=['GET'])
  37. def list_staffs():
  38. staff_data = app.user_relation_manager.list_staffs()
  39. return wrap_response(200, data=staff_data)
  40. @app.route('/api/getStaffProfile', methods=['GET'])
  41. def get_staff_profile():
  42. staff_id = request.args['staff_id']
  43. profile = app.user_manager.get_staff_profile(staff_id)
  44. if not profile:
  45. return wrap_response(404, msg='staff not found')
  46. else:
  47. return wrap_response(200, data=profile)
  48. @app.route('/api/getUserProfile', methods=['GET'])
  49. def get_user_profile():
  50. user_id = request.args['user_id']
  51. profile = app.user_manager.get_user_profile(user_id)
  52. if not profile:
  53. resp = {
  54. 'code': 404,
  55. 'msg': 'user not found'
  56. }
  57. else:
  58. resp = {
  59. 'code': 200,
  60. 'msg': 'success',
  61. 'data': profile
  62. }
  63. return jsonify(resp)
  64. @app.route('/api/listUsers', methods=['GET'])
  65. def list_users():
  66. user_name = request.args.get('user_name', None)
  67. user_union_id = request.args.get('user_union_id', None)
  68. if not user_name and not user_union_id:
  69. resp = {
  70. 'code': 400,
  71. 'msg': 'user_name or user_union_id is required'
  72. }
  73. return jsonify(resp)
  74. data = app.user_manager.list_users(user_name=user_name, user_union_id=user_union_id)
  75. return jsonify({'code': 200, 'data': data})
  76. @app.route('/api/getDialogueHistory', methods=['GET'])
  77. def get_dialogue_history():
  78. staff_id = request.args['staff_id']
  79. user_id = request.args['user_id']
  80. recent_minutes = int(request.args.get('recent_minutes', 1440))
  81. dialogue_history = app.history_dialogue_service.get_dialogue_history(staff_id, user_id, recent_minutes)
  82. return jsonify({'code': 200, 'data': dialogue_history})
  83. @app.route('/api/listModels', methods=['GET'])
  84. def list_models():
  85. models = {
  86. "deepseek-chat": chat_service.VOLCENGINE_MODEL_DEEPSEEK_V3,
  87. "gpt-4o": chat_service.OPENAI_MODEL_GPT_4o,
  88. "gpt-4o-mini": chat_service.OPENAI_MODEL_GPT_4o_mini,
  89. "doubao-pro-32k": chat_service.VOLCENGINE_MODEL_DOUBAO_PRO_32K,
  90. "doubao-pro-1.5": chat_service.VOLCENGINE_MODEL_DOUBAO_PRO_1_5_32K,
  91. "doubao-1.5-vision-pro": chat_service.VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO,
  92. "openrouter-gemini-2.5-pro": chat_service.OPENROUTER_MODEL_GEMINI_2_5_PRO,
  93. }
  94. ret_data = [
  95. {
  96. 'model_type': 'openai_compatible',
  97. 'model_name': model_name,
  98. 'display_name': f"{model_display_name} ({OpenAICompatible.get_price(model_name).get_cny_brief()})"
  99. }
  100. for model_display_name, model_name in models.items()
  101. ]
  102. return wrap_response(200, data=ret_data)
  103. @app.route('/api/listScenes', methods=['GET'])
  104. def list_scenes():
  105. scenes = [
  106. {'scene': 'greeting', 'display_name': '问候'},
  107. {'scene': 'chitchat', 'display_name': '闲聊'},
  108. {'scene': 'profile_extractor', 'display_name': '画像提取'},
  109. {'scene': 'response_type_detector', 'display_name': '回复模态判断'},
  110. {'scene': 'custom_debugging', 'display_name': '自定义调试场景'}
  111. ]
  112. return wrap_response(200, data=scenes)
  113. @app.route('/api/getBasePrompt', methods=['GET'])
  114. def get_base_prompt():
  115. scene = request.args['scene']
  116. prompt_map = {
  117. 'greeting': prompt_templates.GENERAL_GREETING_PROMPT,
  118. 'chitchat': prompt_templates.CHITCHAT_PROMPT_COZE,
  119. 'profile_extractor': prompt_templates.USER_PROFILE_EXTRACT_PROMPT_V2,
  120. 'response_type_detector': prompt_templates.RESPONSE_TYPE_DETECT_PROMPT,
  121. 'custom_debugging': '',
  122. }
  123. model_map = {
  124. 'greeting': chat_service.VOLCENGINE_MODEL_DOUBAO_PRO_32K,
  125. 'chitchat': chat_service.VOLCENGINE_MODEL_DOUBAO_PRO_32K,
  126. 'profile_extractor': chat_service.VOLCENGINE_MODEL_DOUBAO_PRO_1_5_32K,
  127. 'response_type_detector': chat_service.VOLCENGINE_MODEL_DOUBAO_PRO_1_5_32K,
  128. 'custom_debugging': chat_service.VOLCENGINE_BOT_DEEPSEEK_V3_SEARCH
  129. }
  130. if scene not in prompt_map:
  131. return wrap_response(404, msg='scene not found')
  132. data = {
  133. 'model_name': model_map[scene],
  134. 'content': prompt_map[scene]
  135. }
  136. return wrap_response(200, data=data)
  137. @app.route('/api/runPrompt', methods=['POST'])
  138. def run_prompt():
  139. try:
  140. req_data = request.json
  141. logger.debug(req_data)
  142. scene = req_data['scene']
  143. if scene == 'profile_extractor':
  144. response = run_extractor_prompt(req_data)
  145. return wrap_response(200, data=response)
  146. elif scene == 'response_type_detector':
  147. response = run_response_type_prompt(req_data)
  148. return wrap_response(200, data=response.choices[0].message.content)
  149. else:
  150. response = run_chat_prompt(req_data)
  151. return wrap_response(200, data=response.choices[0].message.content)
  152. except Exception as e:
  153. logger.error(e)
  154. return wrap_response(500, msg='Error: {}'.format(e))
  155. @app.route('/api/formatForPrompt', methods=['POST'])
  156. def format_data_for_prompt():
  157. try:
  158. req_data = request.json
  159. content = req_data['content']
  160. format_type = req_data['format_type']
  161. if format_type == 'staff_profile':
  162. if not isinstance(content, dict):
  163. return wrap_response(400, msg='staff_profile should be a dict')
  164. response = format_agent_profile(content)
  165. elif format_type == 'user_profile':
  166. if not isinstance(content, dict):
  167. return wrap_response(400, msg='user_profile should be a dict')
  168. response = format_user_profile(content)
  169. elif format_type == 'dialogue':
  170. if not isinstance(content, list):
  171. return wrap_response(400, msg='dialogue should be a list')
  172. from pqai_agent_server.utils.prompt_util import format_dialogue_history
  173. response = format_dialogue_history(content)
  174. else:
  175. return wrap_response(400, msg='Invalid format_type')
  176. return wrap_response(200, data=response)
  177. except Exception as e:
  178. logger.error(e)
  179. return wrap_response(500, msg='Error: {}'.format(e))
  180. @app.route("/api/healthCheck", methods=["GET"])
  181. def health_check():
  182. return wrap_response(200, msg="OK")
  183. @app.route("/api/getStaffSessionSummary", methods=["GET"])
  184. def get_staff_session_summary():
  185. staff_id = request.args.get("staff_id")
  186. status = request.args.get("status", const.DEFAULT_STAFF_STATUS)
  187. page_id = request.args.get("page_id", const.DEFAULT_PAGE_ID)
  188. page_size = request.args.get("page_size", const.DEFAULT_PAGE_SIZE)
  189. # check params
  190. try:
  191. page_id = int(page_id)
  192. page_size = int(page_size)
  193. status = int(status)
  194. except Exception as e:
  195. return wrap_response(404, msg="Invalid parameter: {}".format(e))
  196. staff_session_summary = app.session_manager.get_staff_sessions_summary(
  197. staff_id, page_id, page_size, status
  198. )
  199. if not staff_session_summary:
  200. return wrap_response(404, msg="staff not found")
  201. else:
  202. return wrap_response(200, data=staff_session_summary)
  203. @app.route("/api/getStaffSessionList", methods=["GET"])
  204. def get_staff_session_list():
  205. staff_id = request.args.get("staff_id")
  206. if not staff_id:
  207. return wrap_response(404, msg="staff_id is required")
  208. page_size = request.args.get("page_size", const.DEFAULT_PAGE_SIZE)
  209. page_id = request.args.get("page_id", const.DEFAULT_PAGE_ID)
  210. # check params
  211. try:
  212. page_id = int(page_id)
  213. page_size = int(page_size)
  214. except Exception as e:
  215. return wrap_response(404, msg="Invalid parameter: {}".format(e))
  216. staff_session_list = app.session_manager.get_staff_session_list(staff_id, page_id, page_size)
  217. if not staff_session_list:
  218. return wrap_response(404, msg="staff not found")
  219. return wrap_response(200, data=staff_session_list)
  220. @app.route("/api/getStaffList", methods=["GET"])
  221. def get_staff_list():
  222. page_size = request.args.get("page_size", const.DEFAULT_PAGE_SIZE)
  223. page_id = request.args.get("page_id", const.DEFAULT_PAGE_ID)
  224. # check params
  225. try:
  226. page_id = int(page_id)
  227. page_size = int(page_size)
  228. except Exception as e:
  229. return wrap_response(404, msg="Invalid parameter: {}".format(e))
  230. staff_list = app.user_manager.get_staff_list(page_id, page_size)
  231. if not staff_list:
  232. return wrap_response(404, msg="staff not found")
  233. return wrap_response(200, data=staff_list)
  234. @app.route("/api/getConversationList", methods=["GET"])
  235. def get_conversation_list():
  236. """
  237. 获取staff && user 私聊对话列表
  238. :return:
  239. """
  240. staff_id = request.args.get("staff_id")
  241. user_id = request.args.get("user_id")
  242. if not staff_id or not user_id:
  243. return wrap_response(404, msg="staff_id and user_id are required")
  244. page = request.args.get("page")
  245. response = app.session_manager.get_conversation_list(staff_id, user_id, page, const.DEFAULT_CONVERSATION_SIZE)
  246. return wrap_response(200, data=response)
  247. @app.route("/api/sendMessage", methods=["POST"])
  248. def send_message():
  249. return wrap_response(200, msg="暂不实现功能")
  250. @app.route("/api/quitHumanIntervention", methods=["POST"])
  251. def quit_human_intervention():
  252. """
  253. 退出人工介入状态
  254. :return:
  255. """
  256. req_data = request.json
  257. staff_id = req_data["staff_id"]
  258. user_id = req_data["user_id"]
  259. if not user_id or not staff_id:
  260. return wrap_response(404, msg="user_id and staff_id are required")
  261. if pqai_agent_server.utils.common.quit_human_intervention(user_id, staff_id):
  262. return wrap_response(200, msg="success")
  263. else:
  264. return wrap_response(500, msg="error")
  265. @app.route("/api/enterHumanIntervention", methods=["POST"])
  266. def enter_human_intervention():
  267. """
  268. 进入人工介入状态
  269. :return:
  270. """
  271. req_data = request.json
  272. staff_id = req_data["staff_id"]
  273. user_id = req_data["user_id"]
  274. if not user_id or not staff_id:
  275. return wrap_response(404, msg="user_id and staff_id are required")
  276. if pqai_agent_server.utils.common.enter_human_intervention(user_id, staff_id):
  277. return wrap_response(200, msg="success")
  278. else:
  279. return wrap_response(500, msg="error")
  280. ## Agent管理接口
  281. @app.route("/api/getNativeAgentList", methods=["GET"])
  282. def get_native_agent_list():
  283. """
  284. 获取所有的Agent列表
  285. :return:
  286. """
  287. page = request.args.get('page', 1)
  288. page_size = request.args.get('page_size', 50)
  289. create_user = request.args.get('create_user', None)
  290. update_user = request.args.get('update_user', None)
  291. offset = (int(page) - 1) * int(page_size)
  292. with app.session_maker() as session:
  293. query = session.query(AgentConfiguration) \
  294. .filter(AgentConfiguration.is_delete == 0)
  295. if create_user:
  296. query = query.filter(AgentConfiguration.create_user == create_user)
  297. if update_user:
  298. query = query.filter(AgentConfiguration.update_user == update_user)
  299. total = query.count()
  300. query = query.offset(offset).limit(int(page_size))
  301. data = query.all()
  302. ret_data = {
  303. 'total': total,
  304. 'agent_list': [
  305. {
  306. 'id': agent.id,
  307. 'name': agent.name,
  308. 'display_name': agent.display_name,
  309. 'type': agent.type,
  310. 'execution_model': agent.execution_model,
  311. 'create_user': agent.create_user,
  312. 'update_user': agent.update_user,
  313. 'create_time': agent.create_time.strftime('%Y-%m-%d %H:%M:%S'),
  314. 'update_time': agent.update_time.strftime('%Y-%m-%d %H:%M:%S')
  315. }
  316. for agent in data
  317. ]
  318. }
  319. return wrap_response(200, data=ret_data)
  320. @app.route("/api/getNativeAgentConfiguration", methods=["GET"])
  321. def get_native_agent_configuration():
  322. """
  323. 获取指定Agent的配置
  324. :return:
  325. """
  326. agent_id = request.args.get('agent_id')
  327. if not agent_id:
  328. return wrap_response(404, msg='agent_id is required')
  329. with app.session_maker() as session:
  330. agent = session.query(AgentConfiguration).filter(AgentConfiguration.id == agent_id).first()
  331. if not agent:
  332. return wrap_response(404, msg='Agent not found')
  333. data = {
  334. 'id': agent.id,
  335. 'name': agent.name,
  336. 'display_name': agent.display_name,
  337. 'type': agent.type,
  338. 'execution_model': agent.execution_model,
  339. 'system_prompt': agent.system_prompt,
  340. 'task_prompt': agent.task_prompt,
  341. 'tools': json.loads(agent.tools),
  342. 'sub_agents': json.loads(agent.sub_agents),
  343. 'extra_params': json.loads(agent.extra_params),
  344. 'create_time': agent.create_time.strftime('%Y-%m-%d %H:%M:%S'),
  345. 'update_time': agent.update_time.strftime('%Y-%m-%d %H:%M:%S')
  346. }
  347. return wrap_response(200, data=data)
  348. @app.route("/api/saveNativeAgentConfiguration", methods=["POST"])
  349. def save_native_agent_configuration():
  350. """
  351. 保存Agent配置
  352. :return:
  353. """
  354. req_data = request.json
  355. agent_id = req_data.get('agent_id', None)
  356. name = req_data.get('name')
  357. display_name = req_data.get('display_name', None)
  358. type_ = req_data.get('type', 0)
  359. execution_model = req_data.get('execution_model', None)
  360. system_prompt = req_data.get('system_prompt', None)
  361. task_prompt = req_data.get('task_prompt', None)
  362. tools = json.dumps(req_data.get('tools', []))
  363. sub_agents = json.dumps(req_data.get('sub_agents', []))
  364. extra_params = req_data.get('extra_params', {})
  365. operate_user = req_data.get('user', None)
  366. if isinstance(extra_params, dict):
  367. extra_params = json.dumps(extra_params)
  368. elif isinstance(extra_params, str):
  369. try:
  370. json.loads(extra_params)
  371. except json.JSONDecodeError:
  372. return wrap_response(400, msg='extra_params should be a valid JSON object or string')
  373. if not extra_params:
  374. extra_params = '{}'
  375. if not name:
  376. return wrap_response(400, msg='name is required')
  377. with app.session_maker() as session:
  378. if agent_id:
  379. agent_id = int(agent_id)
  380. agent = session.query(AgentConfiguration).filter(AgentConfiguration.id == agent_id).first()
  381. if not agent:
  382. return wrap_response(404, msg='Agent not found')
  383. agent.name = name
  384. agent.display_name = display_name
  385. agent.type = type_
  386. agent.execution_model = execution_model
  387. agent.system_prompt = system_prompt
  388. agent.task_prompt = task_prompt
  389. agent.tools = tools
  390. agent.sub_agents = sub_agents
  391. agent.extra_params = extra_params
  392. agent.update_user = operate_user
  393. else:
  394. agent = AgentConfiguration(
  395. name=name,
  396. display_name=display_name,
  397. type=type_,
  398. execution_model=execution_model,
  399. system_prompt=system_prompt,
  400. task_prompt=task_prompt,
  401. tools=tools,
  402. sub_agents=sub_agents,
  403. extra_params=extra_params,
  404. create_user=operate_user,
  405. update_user=operate_user
  406. )
  407. session.add(agent)
  408. session.commit()
  409. return wrap_response(200, msg='Agent configuration saved successfully', data={'id': agent.id})
  410. @app.route("/api/deleteNativeAgentConfiguration", methods=["POST"])
  411. def delete_native_agent_configuration():
  412. """
  413. 删除指定Agent配置(软删除,设置is_delete=1)
  414. :return:
  415. """
  416. req_data = request.json
  417. agent_id = req_data.get('agent_id', None)
  418. if not agent_id:
  419. return wrap_response(400, msg='agent_id is required')
  420. try:
  421. agent_id = int(agent_id)
  422. except ValueError:
  423. return wrap_response(400, msg='agent_id must be an integer')
  424. with app.session_maker() as session:
  425. agent = session.query(AgentConfiguration).filter(
  426. AgentConfiguration.id == agent_id,
  427. AgentConfiguration.is_delete == 0
  428. ).first()
  429. if not agent:
  430. return wrap_response(404, msg='Agent not found')
  431. agent.is_delete = 1
  432. session.commit()
  433. return wrap_response(200, msg='Agent configuration deleted successfully')
  434. @app.route("/api/getModuleList", methods=["GET"])
  435. def get_module_list():
  436. """
  437. 获取所有的模块列表,支持分页查询
  438. :return:
  439. """
  440. page = request.args.get('page', 1)
  441. page_size = request.args.get('page_size', 50)
  442. try:
  443. page = int(page)
  444. page_size = int(page_size)
  445. except Exception as e:
  446. return wrap_response(400, msg="Invalid parameter: {}".format(e))
  447. offset = (page - 1) * page_size
  448. with app.session_maker() as session:
  449. query = session.query(
  450. ServiceModule,
  451. AgentConfiguration.name.label("default_agent_name")
  452. ).outerjoin(
  453. AgentConfiguration,
  454. ServiceModule.default_agent_id == AgentConfiguration.id
  455. ).filter(ServiceModule.is_delete == 0)
  456. total = query.count()
  457. modules = query.offset(offset).limit(page_size).all()
  458. ret_data = {
  459. 'total': total,
  460. 'module_list': [
  461. {
  462. 'id': module.id,
  463. 'name': module.name,
  464. 'display_name': module.display_name,
  465. 'default_agent_type': module.default_agent_type,
  466. 'default_agent_id': module.default_agent_id,
  467. 'default_agent_name': default_agent_name,
  468. 'create_time': module.create_time.strftime('%Y-%m-%d %H:%M:%S'),
  469. 'update_time': module.update_time.strftime('%Y-%m-%d %H:%M:%S')
  470. }
  471. for module, default_agent_name in modules
  472. ]
  473. }
  474. return wrap_response(200, data=ret_data)
  475. @app.route("/api/getModuleConfiguration", methods=["GET"])
  476. def get_module_configuration():
  477. """
  478. 获取指定模块的配置
  479. :return:
  480. """
  481. module_id = request.args.get('module_id')
  482. if not module_id:
  483. return wrap_response(404, msg='module_id is required')
  484. with app.session_maker() as session:
  485. module = session.query(ServiceModule).filter(ServiceModule.id == module_id).first()
  486. if not module:
  487. return wrap_response(404, msg='Module not found')
  488. data = {
  489. 'id': module.id,
  490. 'name': module.name,
  491. 'display_name': module.display_name,
  492. 'default_agent_type': module.default_agent_type,
  493. 'default_agent_id': module.default_agent_id,
  494. 'create_time': module.create_time.strftime('%Y-%m-%d %H:%M:%S'),
  495. 'updated_time': module.updated_time.strftime('%Y-%m-%d %H:%M:%S')
  496. }
  497. return wrap_response(200, data=data)
  498. @app.route("/api/saveModuleConfiguration", methods=["POST"])
  499. def save_module_configuration():
  500. """
  501. 保存模块配置
  502. :return:
  503. """
  504. req_data = request.json
  505. module_id = req_data.get('module_id', None)
  506. name = req_data.get('name')
  507. display_name = req_data.get('display_name', None)
  508. default_agent_type = req_data.get('default_agent_type', 0)
  509. default_agent_id = req_data.get('default_agent_id', None)
  510. if not name:
  511. return wrap_response(400, msg='name is required')
  512. with app.session_maker() as session:
  513. if module_id:
  514. module_id = int(module_id)
  515. module = session.query(ServiceModule).filter(ServiceModule.id == module_id).first()
  516. if not module:
  517. return wrap_response(404, msg='Module not found')
  518. module.name = name
  519. module.display_name = display_name
  520. module.default_agent_type = default_agent_type
  521. module.default_agent_id = default_agent_id
  522. else:
  523. module = ServiceModule(
  524. name=name,
  525. display_name=display_name,
  526. default_agent_type=default_agent_type,
  527. default_agent_id=default_agent_id
  528. )
  529. session.add(module)
  530. session.commit()
  531. return wrap_response(200, msg='Module configuration saved successfully', data={'id': module.id})
  532. @app.route("/api/getToolList", methods=["GET"])
  533. def get_tool_list():
  534. """
  535. 获取所有的工具列表
  536. :return:
  537. """
  538. tools = []
  539. for tool_name, tool in global_tool_map.items():
  540. tools.append({
  541. 'name': tool_name,
  542. 'description': tool.get_function_description(),
  543. 'parameters': tool.parameters if hasattr(tool, 'parameters') else {}
  544. })
  545. return wrap_response(200, data=tools)
  546. @app.route("/api/getModuleAgentTypes", methods=["GET"])
  547. def get_agent_types():
  548. """
  549. 获取所有的Agent类型
  550. :return:
  551. """
  552. agent_types = [
  553. {'type': 0, 'display_name': '原生'},
  554. {'type': 1, 'display_name': 'Coze'}
  555. ]
  556. return wrap_response(200, data=agent_types)
  557. @app.route("/api/getTestTaskList", methods=["GET"])
  558. def get_test_task_list():
  559. """
  560. 获取单元测试任务列表
  561. :return:
  562. """
  563. page_num = request.args.get("pageNum", const.DEFAULT_PAGE_ID)
  564. page_size = request.args.get("pageSize", const.DEFAULT_PAGE_SIZE)
  565. try:
  566. page_num = int(page_num)
  567. page_size = int(page_size)
  568. except Exception as e:
  569. return wrap_response(404, msg="Invalid parameter: {}".format(e))
  570. response = app.task_manager.get_test_task_list(page_num, page_size)
  571. return wrap_response(200, data=response)
  572. @app.route("/api/getTestTaskConversations", methods=["GET"])
  573. def get_test_task_conversations():
  574. """
  575. 获取单元测试对话任务列表
  576. :return:
  577. """
  578. task_id = request.args.get("taskId", None)
  579. if not task_id:
  580. return wrap_response(404, msg='task_id is required')
  581. page_num = request.args.get("pageNum", const.DEFAULT_PAGE_ID)
  582. page_size = request.args.get("pageSize", const.DEFAULT_PAGE_SIZE)
  583. try:
  584. page_num = int(page_num)
  585. page_size = int(page_size)
  586. except Exception as e:
  587. return wrap_response(404, msg="Invalid parameter: {}".format(e))
  588. response = app.task_manager.get_test_task_conversations(int(task_id), page_num, page_size)
  589. return wrap_response(200, data=response)
  590. @app.route("/api/createTestTask", methods=["POST"])
  591. def create_test_task():
  592. """
  593. 创建单元测试任务
  594. :return:
  595. """
  596. req_data = request.json
  597. agent_id = req_data.get('agentId', None)
  598. module_id = req_data.get('moduleId', None)
  599. evaluate_type = req_data.get('evaluateType', None)
  600. if not agent_id:
  601. return wrap_response(404, msg='agent id is required')
  602. if not module_id:
  603. return wrap_response(404, msg='module id is required')
  604. if not evaluate_type:
  605. return wrap_response(404, msg='evaluate_type id is required')
  606. app.task_manager.create_task(agent_id, module_id, evaluate_type)
  607. return wrap_response(200)
  608. @app.route("/api/stopTestTask", methods=["POST"])
  609. def stop_test_task():
  610. """
  611. 停止单元测试任务
  612. :return:
  613. """
  614. req_data = request.json
  615. task_id = req_data.get('taskId', None)
  616. if not task_id:
  617. return wrap_response(400, msg='task id is required')
  618. task = app.task_manager.get_task(task_id)
  619. if task.status not in (TestTaskStatus.NOT_STARTED.value, TestTaskStatus.IN_PROGRESS.value):
  620. return wrap_response(400, msg='task status is invalid')
  621. app.task_manager.cancel_task(task_id)
  622. return wrap_response(200)
  623. @app.route("/api/resumeTestTask", methods=["POST"])
  624. def resume_test_task():
  625. """
  626. 恢复停止的单元测试任务
  627. :return:
  628. """
  629. req_data = request.json
  630. task_id = req_data.get('taskId', None)
  631. if not task_id:
  632. return wrap_response(400, msg='task id is required')
  633. task = app.task_manager.get_task(task_id)
  634. if task.status != TestTaskStatus.CANCELLED.value:
  635. return wrap_response(400, msg='task status is invalid')
  636. app.task_manager.resume_task(task_id)
  637. return wrap_response(200)
  638. @app.route("/api/getEvaluateType", methods=["GET"])
  639. def get_evaluate_type():
  640. """
  641. 获取评估类型
  642. :return:
  643. """
  644. name_desc_list = [
  645. {
  646. "type": item.value,
  647. "desc": item.description
  648. }
  649. for item in EvaluateType]
  650. return wrap_response(code=200, data=name_desc_list)
  651. @app.route("/api/getDatasetList", methods=["GET"])
  652. def get_dataset_list():
  653. """
  654. 获取数据集列表
  655. :return:
  656. """
  657. page_num = request.args.get("pageNum", const.DEFAULT_PAGE_ID)
  658. page_size = request.args.get("pageSize", const.DEFAULT_PAGE_SIZE)
  659. try:
  660. page_num = int(page_num)
  661. page_size = int(page_size)
  662. except Exception as e:
  663. return wrap_response(404, msg="Invalid parameter: {}".format(e))
  664. response = app.dataset_service.get_dataset_list(page_num, page_size)
  665. return wrap_response(200, data=response)
  666. @app.route("/api/getConversationDataList", methods=["GET"])
  667. def get_conversation_data_list():
  668. """
  669. 获取对话列表
  670. :return:
  671. """
  672. dataset_id = request.args.get("datasetId", None)
  673. if not dataset_id:
  674. return wrap_response(404, msg='dataset_id is required')
  675. page_num = request.args.get("pageNum", const.DEFAULT_PAGE_ID)
  676. page_size = request.args.get("pageSize", const.DEFAULT_PAGE_SIZE)
  677. try:
  678. page_num = int(page_num)
  679. page_size = int(page_size)
  680. except Exception as e:
  681. return wrap_response(404, msg="Invalid parameter: {}".format(e))
  682. response = app.dataset_service.get_conversation_data_list(int(dataset_id), page_num, page_size)
  683. return wrap_response(200, data=response)
  684. @app.errorhandler(werkzeug.exceptions.BadRequest)
  685. def handle_bad_request(e):
  686. logger.error(e)
  687. return wrap_response(400, msg='Bad Request: {}'.format(e.description))
  688. if __name__ == '__main__':
  689. parser = ArgumentParser()
  690. parser.add_argument('--prod', action='store_true')
  691. parser.add_argument('--host', default='127.0.0.1')
  692. parser.add_argument('--port', type=int, default=8083)
  693. parser.add_argument('--log-level', default='INFO')
  694. args = parser.parse_args()
  695. config = configs.get()
  696. logging_level = logging.getLevelName(args.log_level)
  697. setup_root_logger(level=logging_level, logfile_name='agent_api_server.log')
  698. # set db config
  699. agent_db_config = config['database']['ai_agent']
  700. growth_db_config = config['database']['growth']
  701. user_db_config = config['storage']['user']
  702. staff_db_config = config['storage']['staff']
  703. agent_state_db_config = config['storage']['agent_state']
  704. chat_history_db_config = config['storage']['chat_history']
  705. # init user manager
  706. user_manager = MySQLUserManager(agent_db_config, user_db_config['table'], staff_db_config['table'])
  707. app.user_manager = user_manager
  708. # init session manager
  709. session_manager = MySQLSessionManager(
  710. db_config=agent_db_config,
  711. staff_table=staff_db_config['table'],
  712. user_table=user_db_config['table'],
  713. agent_state_table=agent_state_db_config['table'],
  714. chat_history_table=chat_history_db_config['table']
  715. )
  716. app.session_manager = session_manager
  717. agent_db_engine = create_ai_agent_db_engine()
  718. app.session_maker = sessionmaker(bind=agent_db_engine)
  719. dataset_service = DatasetService(session_maker=sessionmaker(bind=agent_db_engine))
  720. app.dataset_service = dataset_service
  721. task_manager = TaskManager(session_maker=sessionmaker(bind=agent_db_engine), dataset_service=dataset_service)
  722. app.task_manager = task_manager
  723. task_manager.recover_tasks()
  724. wecom_db_config = config['storage']['user_relation']
  725. user_relation_manager = MySQLUserRelationManager(
  726. agent_db_config, growth_db_config,
  727. config['storage']['staff']['table'],
  728. user_db_config['table'],
  729. wecom_db_config['table']['staff'],
  730. wecom_db_config['table']['relation'],
  731. wecom_db_config['table']['user']
  732. )
  733. app.user_relation_manager = user_relation_manager
  734. app.history_dialogue_service = HistoryDialogueService(
  735. config['storage']['history_dialogue']['api_base_url']
  736. )
  737. app.run(debug=not args.prod, host=args.host, port=args.port)