api_server.py 28 KB

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