conversation.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306
  1. from datetime import datetime, timezone
  2. import pytz
  3. from flask_login import current_user
  4. from flask_restful import Resource, marshal_with, reqparse
  5. from flask_restful.inputs import int_range
  6. from sqlalchemy import func, or_
  7. from sqlalchemy.orm import joinedload
  8. from werkzeug.exceptions import Forbidden, NotFound
  9. from controllers.console import api
  10. from controllers.console.app.wraps import get_app_model
  11. from controllers.console.setup import setup_required
  12. from controllers.console.wraps import account_initialization_required
  13. from core.app.entities.app_invoke_entities import InvokeFrom
  14. from extensions.ext_database import db
  15. from fields.conversation_fields import (
  16. conversation_detail_fields,
  17. conversation_message_detail_fields,
  18. conversation_pagination_fields,
  19. conversation_with_summary_pagination_fields,
  20. )
  21. from libs.helper import datetime_string
  22. from libs.login import login_required
  23. from models.model import AppMode, Conversation, EndUser, Message, MessageAnnotation
  24. class CompletionConversationApi(Resource):
  25. @setup_required
  26. @login_required
  27. @account_initialization_required
  28. @get_app_model(mode=AppMode.COMPLETION)
  29. @marshal_with(conversation_pagination_fields)
  30. def get(self, app_model):
  31. if not current_user.is_editor:
  32. raise Forbidden()
  33. parser = reqparse.RequestParser()
  34. parser.add_argument('keyword', type=str, location='args')
  35. parser.add_argument('start', type=datetime_string('%Y-%m-%d %H:%M'), location='args')
  36. parser.add_argument('end', type=datetime_string('%Y-%m-%d %H:%M'), location='args')
  37. parser.add_argument('annotation_status', type=str,
  38. choices=['annotated', 'not_annotated', 'all'], default='all', location='args')
  39. parser.add_argument('page', type=int_range(1, 99999), default=1, location='args')
  40. parser.add_argument('limit', type=int_range(1, 100), default=20, location='args')
  41. args = parser.parse_args()
  42. query = db.select(Conversation).where(Conversation.app_id == app_model.id, Conversation.mode == 'completion')
  43. if args['keyword']:
  44. query = query.join(
  45. Message, Message.conversation_id == Conversation.id
  46. ).filter(
  47. or_(
  48. Message.query.ilike('%{}%'.format(args['keyword'])),
  49. Message.answer.ilike('%{}%'.format(args['keyword']))
  50. )
  51. )
  52. account = current_user
  53. timezone = pytz.timezone(account.timezone)
  54. utc_timezone = pytz.utc
  55. if args['start']:
  56. start_datetime = datetime.strptime(args['start'], '%Y-%m-%d %H:%M')
  57. start_datetime = start_datetime.replace(second=0)
  58. start_datetime_timezone = timezone.localize(start_datetime)
  59. start_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
  60. query = query.where(Conversation.created_at >= start_datetime_utc)
  61. if args['end']:
  62. end_datetime = datetime.strptime(args['end'], '%Y-%m-%d %H:%M')
  63. end_datetime = end_datetime.replace(second=59)
  64. end_datetime_timezone = timezone.localize(end_datetime)
  65. end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
  66. query = query.where(Conversation.created_at < end_datetime_utc)
  67. if args['annotation_status'] == "annotated":
  68. query = query.options(joinedload(Conversation.message_annotations)).join(
  69. MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
  70. )
  71. elif args['annotation_status'] == "not_annotated":
  72. query = query.outerjoin(
  73. MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
  74. ).group_by(Conversation.id).having(func.count(MessageAnnotation.id) == 0)
  75. query = query.order_by(Conversation.created_at.desc())
  76. conversations = db.paginate(
  77. query,
  78. page=args['page'],
  79. per_page=args['limit'],
  80. error_out=False
  81. )
  82. return conversations
  83. class CompletionConversationDetailApi(Resource):
  84. @setup_required
  85. @login_required
  86. @account_initialization_required
  87. @get_app_model(mode=AppMode.COMPLETION)
  88. @marshal_with(conversation_message_detail_fields)
  89. def get(self, app_model, conversation_id):
  90. if not current_user.is_editor:
  91. raise Forbidden()
  92. conversation_id = str(conversation_id)
  93. return _get_conversation(app_model, conversation_id)
  94. @setup_required
  95. @login_required
  96. @account_initialization_required
  97. @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT])
  98. def delete(self, app_model, conversation_id):
  99. if not current_user.is_editor:
  100. raise Forbidden()
  101. conversation_id = str(conversation_id)
  102. conversation = db.session.query(Conversation) \
  103. .filter(Conversation.id == conversation_id, Conversation.app_id == app_model.id).first()
  104. if not conversation:
  105. raise NotFound("Conversation Not Exists.")
  106. conversation.is_deleted = True
  107. db.session.commit()
  108. return {'result': 'success'}, 204
  109. class ChatConversationApi(Resource):
  110. @setup_required
  111. @login_required
  112. @account_initialization_required
  113. @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT])
  114. @marshal_with(conversation_with_summary_pagination_fields)
  115. def get(self, app_model):
  116. if not current_user.is_editor:
  117. raise Forbidden()
  118. parser = reqparse.RequestParser()
  119. parser.add_argument('keyword', type=str, location='args')
  120. parser.add_argument('start', type=datetime_string('%Y-%m-%d %H:%M'), location='args')
  121. parser.add_argument('end', type=datetime_string('%Y-%m-%d %H:%M'), location='args')
  122. parser.add_argument('annotation_status', type=str,
  123. choices=['annotated', 'not_annotated', 'all'], default='all', location='args')
  124. parser.add_argument('message_count_gte', type=int_range(1, 99999), required=False, location='args')
  125. parser.add_argument('page', type=int_range(1, 99999), required=False, default=1, location='args')
  126. parser.add_argument('limit', type=int_range(1, 100), required=False, default=20, location='args')
  127. parser.add_argument('sort_by', type=str, choices=['created_at', '-created_at', 'updated_at', '-updated_at'],
  128. required=False, default='-updated_at', location='args')
  129. args = parser.parse_args()
  130. subquery = (
  131. db.session.query(
  132. Conversation.id.label('conversation_id'),
  133. EndUser.session_id.label('from_end_user_session_id')
  134. )
  135. .outerjoin(EndUser, Conversation.from_end_user_id == EndUser.id)
  136. .subquery()
  137. )
  138. query = db.select(Conversation).where(Conversation.app_id == app_model.id)
  139. if args['keyword']:
  140. keyword_filter = '%{}%'.format(args['keyword'])
  141. query = query.join(
  142. Message, Message.conversation_id == Conversation.id,
  143. ).join(
  144. subquery, subquery.c.conversation_id == Conversation.id
  145. ).filter(
  146. or_(
  147. Message.query.ilike(keyword_filter),
  148. Message.answer.ilike(keyword_filter),
  149. Conversation.name.ilike(keyword_filter),
  150. Conversation.introduction.ilike(keyword_filter),
  151. subquery.c.from_end_user_session_id.ilike(keyword_filter)
  152. ),
  153. )
  154. account = current_user
  155. timezone = pytz.timezone(account.timezone)
  156. utc_timezone = pytz.utc
  157. if args['start']:
  158. start_datetime = datetime.strptime(args['start'], '%Y-%m-%d %H:%M')
  159. start_datetime = start_datetime.replace(second=0)
  160. start_datetime_timezone = timezone.localize(start_datetime)
  161. start_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
  162. query = query.where(Conversation.created_at >= start_datetime_utc)
  163. if args['end']:
  164. end_datetime = datetime.strptime(args['end'], '%Y-%m-%d %H:%M')
  165. end_datetime = end_datetime.replace(second=59)
  166. end_datetime_timezone = timezone.localize(end_datetime)
  167. end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
  168. query = query.where(Conversation.created_at < end_datetime_utc)
  169. if args['annotation_status'] == "annotated":
  170. query = query.options(joinedload(Conversation.message_annotations)).join(
  171. MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
  172. )
  173. elif args['annotation_status'] == "not_annotated":
  174. query = query.outerjoin(
  175. MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
  176. ).group_by(Conversation.id).having(func.count(MessageAnnotation.id) == 0)
  177. if args['message_count_gte'] and args['message_count_gte'] >= 1:
  178. query = (
  179. query.options(joinedload(Conversation.messages))
  180. .join(Message, Message.conversation_id == Conversation.id)
  181. .group_by(Conversation.id)
  182. .having(func.count(Message.id) >= args['message_count_gte'])
  183. )
  184. if app_model.mode == AppMode.ADVANCED_CHAT.value:
  185. query = query.where(Conversation.invoke_from != InvokeFrom.DEBUGGER.value)
  186. match args['sort_by']:
  187. case 'created_at':
  188. query = query.order_by(Conversation.created_at.asc())
  189. case '-created_at':
  190. query = query.order_by(Conversation.created_at.desc())
  191. case 'updated_at':
  192. query = query.order_by(Conversation.updated_at.asc())
  193. case '-updated_at':
  194. query = query.order_by(Conversation.updated_at.desc())
  195. case _:
  196. query = query.order_by(Conversation.created_at.desc())
  197. conversations = db.paginate(
  198. query,
  199. page=args['page'],
  200. per_page=args['limit'],
  201. error_out=False
  202. )
  203. return conversations
  204. class ChatConversationDetailApi(Resource):
  205. @setup_required
  206. @login_required
  207. @account_initialization_required
  208. @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT])
  209. @marshal_with(conversation_detail_fields)
  210. def get(self, app_model, conversation_id):
  211. if not current_user.is_editor:
  212. raise Forbidden()
  213. conversation_id = str(conversation_id)
  214. return _get_conversation(app_model, conversation_id)
  215. @setup_required
  216. @login_required
  217. @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT])
  218. @account_initialization_required
  219. def delete(self, app_model, conversation_id):
  220. if not current_user.is_editor:
  221. raise Forbidden()
  222. conversation_id = str(conversation_id)
  223. conversation = db.session.query(Conversation) \
  224. .filter(Conversation.id == conversation_id, Conversation.app_id == app_model.id).first()
  225. if not conversation:
  226. raise NotFound("Conversation Not Exists.")
  227. conversation.is_deleted = True
  228. db.session.commit()
  229. return {'result': 'success'}, 204
  230. api.add_resource(CompletionConversationApi, '/apps/<uuid:app_id>/completion-conversations')
  231. api.add_resource(CompletionConversationDetailApi, '/apps/<uuid:app_id>/completion-conversations/<uuid:conversation_id>')
  232. api.add_resource(ChatConversationApi, '/apps/<uuid:app_id>/chat-conversations')
  233. api.add_resource(ChatConversationDetailApi, '/apps/<uuid:app_id>/chat-conversations/<uuid:conversation_id>')
  234. def _get_conversation(app_model, conversation_id):
  235. conversation = db.session.query(Conversation) \
  236. .filter(Conversation.id == conversation_id, Conversation.app_id == app_model.id).first()
  237. if not conversation:
  238. raise NotFound("Conversation Not Exists.")
  239. if not conversation.read_at:
  240. conversation.read_at = datetime.now(timezone.utc).replace(tzinfo=None)
  241. conversation.read_account_id = current_user.id
  242. db.session.commit()
  243. return conversation