Underground-Digital/Workflow-Engine
0
1import datetime2import uuid3 4import pandas as pd5from flask_login import current_user6from sqlalchemy import or_7from werkzeug.datastructures import FileStorage8from werkzeug.exceptions import NotFound9 10from extensions.ext_database import db11from extensions.ext_redis import redis_client12from models.model import App, AppAnnotationHitHistory, AppAnnotationSetting, Message, MessageAnnotation13from services.feature_service import FeatureService14from tasks.annotation.add_annotation_to_index_task import add_annotation_to_index_task15from tasks.annotation.batch_import_annotations_task import batch_import_annotations_task16from tasks.annotation.delete_annotation_index_task import delete_annotation_index_task17from tasks.annotation.disable_annotation_reply_task import disable_annotation_reply_task18from tasks.annotation.enable_annotation_reply_task import enable_annotation_reply_task19from tasks.annotation.update_annotation_to_index_task import update_annotation_to_index_task20 21 22class AppAnnotationService:23 @classmethod24 def up_insert_app_annotation_from_message(cls, args: dict, app_id: str) -> MessageAnnotation:25 # get app info26 app = (27 db.session.query(App)28 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")29 .first()30 )31 32 if not app:33 raise NotFound("App not found")34 if args.get("message_id"):35 message_id = str(args["message_id"])36 # get message info37 message = db.session.query(Message).filter(Message.id == message_id, Message.app_id == app.id).first()38 39 if not message:40 raise NotFound("Message Not Exists.")41 42 annotation = message.annotation43 # save the message annotation44 if annotation:45 annotation.content = args["answer"]46 annotation.question = args["question"]47 else:48 annotation = MessageAnnotation(49 app_id=app.id,50 conversation_id=message.conversation_id,51 message_id=message.id,52 content=args["answer"],53 question=args["question"],54 account_id=current_user.id,55 )56 else:57 annotation = MessageAnnotation(58 app_id=app.id, content=args["answer"], question=args["question"], account_id=current_user.id59 )60 db.session.add(annotation)61 db.session.commit()62 # if annotation reply is enabled , add annotation to index63 annotation_setting = (64 db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()65 )66 if annotation_setting:67 add_annotation_to_index_task.delay(68 annotation.id,69 args["question"],70 current_user.current_tenant_id,71 app_id,72 annotation_setting.collection_binding_id,73 )74 return annotation75 76 @classmethod77 def enable_app_annotation(cls, args: dict, app_id: str) -> dict:78 enable_app_annotation_key = "enable_app_annotation_{}".format(str(app_id))79 cache_result = redis_client.get(enable_app_annotation_key)80 if cache_result is not None:81 return {"job_id": cache_result, "job_status": "processing"}82 83 # async job84 job_id = str(uuid.uuid4())85 enable_app_annotation_job_key = "enable_app_annotation_job_{}".format(str(job_id))86 # send batch add segments task87 redis_client.setnx(enable_app_annotation_job_key, "waiting")88 enable_annotation_reply_task.delay(89 str(job_id),90 app_id,91 current_user.id,92 current_user.current_tenant_id,93 args["score_threshold"],94 args["embedding_provider_name"],95 args["embedding_model_name"],96 )97 return {"job_id": job_id, "job_status": "waiting"}98 99 @classmethod100 def disable_app_annotation(cls, app_id: str) -> dict:101 disable_app_annotation_key = "disable_app_annotation_{}".format(str(app_id))102 cache_result = redis_client.get(disable_app_annotation_key)103 if cache_result is not None:104 return {"job_id": cache_result, "job_status": "processing"}105 106 # async job107 job_id = str(uuid.uuid4())108 disable_app_annotation_job_key = "disable_app_annotation_job_{}".format(str(job_id))109 # send batch add segments task110 redis_client.setnx(disable_app_annotation_job_key, "waiting")111 disable_annotation_reply_task.delay(str(job_id), app_id, current_user.current_tenant_id)112 return {"job_id": job_id, "job_status": "waiting"}113 114 @classmethod115 def get_annotation_list_by_app_id(cls, app_id: str, page: int, limit: int, keyword: str):116 # get app info117 app = (118 db.session.query(App)119 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")120 .first()121 )122 123 if not app:124 raise NotFound("App not found")125 if keyword:126 annotations = (127 db.session.query(MessageAnnotation)128 .filter(MessageAnnotation.app_id == app_id)129 .filter(130 or_(131 MessageAnnotation.question.ilike("%{}%".format(keyword)),132 MessageAnnotation.content.ilike("%{}%".format(keyword)),133 )134 )135 .order_by(MessageAnnotation.created_at.desc(), MessageAnnotation.id.desc())136 .paginate(page=page, per_page=limit, max_per_page=100, error_out=False)137 )138 else:139 annotations = (140 db.session.query(MessageAnnotation)141 .filter(MessageAnnotation.app_id == app_id)142 .order_by(MessageAnnotation.created_at.desc(), MessageAnnotation.id.desc())143 .paginate(page=page, per_page=limit, max_per_page=100, error_out=False)144 )145 return annotations.items, annotations.total146 147 @classmethod148 def export_annotation_list_by_app_id(cls, app_id: str):149 # get app info150 app = (151 db.session.query(App)152 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")153 .first()154 )155 156 if not app:157 raise NotFound("App not found")158 annotations = (159 db.session.query(MessageAnnotation)160 .filter(MessageAnnotation.app_id == app_id)161 .order_by(MessageAnnotation.created_at.desc())162 .all()163 )164 return annotations165 166 @classmethod167 def insert_app_annotation_directly(cls, args: dict, app_id: str) -> MessageAnnotation:168 # get app info169 app = (170 db.session.query(App)171 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")172 .first()173 )174 175 if not app:176 raise NotFound("App not found")177 178 annotation = MessageAnnotation(179 app_id=app.id, content=args["answer"], question=args["question"], account_id=current_user.id180 )181 db.session.add(annotation)182 db.session.commit()183 # if annotation reply is enabled , add annotation to index184 annotation_setting = (185 db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()186 )187 if annotation_setting:188 add_annotation_to_index_task.delay(189 annotation.id,190 args["question"],191 current_user.current_tenant_id,192 app_id,193 annotation_setting.collection_binding_id,194 )195 return annotation196 197 @classmethod198 def update_app_annotation_directly(cls, args: dict, app_id: str, annotation_id: str):199 # get app info200 app = (201 db.session.query(App)202 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")203 .first()204 )205 206 if not app:207 raise NotFound("App not found")208 209 annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first()210 211 if not annotation:212 raise NotFound("Annotation not found")213 214 annotation.content = args["answer"]215 annotation.question = args["question"]216 217 db.session.commit()218 # if annotation reply is enabled , add annotation to index219 app_annotation_setting = (220 db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()221 )222 223 if app_annotation_setting:224 update_annotation_to_index_task.delay(225 annotation.id,226 annotation.question,227 current_user.current_tenant_id,228 app_id,229 app_annotation_setting.collection_binding_id,230 )231 232 return annotation233 234 @classmethod235 def delete_app_annotation(cls, app_id: str, annotation_id: str):236 # get app info237 app = (238 db.session.query(App)239 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")240 .first()241 )242 243 if not app:244 raise NotFound("App not found")245 246 annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first()247 248 if not annotation:249 raise NotFound("Annotation not found")250 251 db.session.delete(annotation)252 253 annotation_hit_histories = (254 db.session.query(AppAnnotationHitHistory)255 .filter(AppAnnotationHitHistory.annotation_id == annotation_id)256 .all()257 )258 if annotation_hit_histories:259 for annotation_hit_history in annotation_hit_histories:260 db.session.delete(annotation_hit_history)261 262 db.session.commit()263 # if annotation reply is enabled , delete annotation index264 app_annotation_setting = (265 db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()266 )267 268 if app_annotation_setting:269 delete_annotation_index_task.delay(270 annotation.id, app_id, current_user.current_tenant_id, app_annotation_setting.collection_binding_id271 )272 273 @classmethod274 def batch_import_app_annotations(cls, app_id, file: FileStorage) -> dict:275 # get app info276 app = (277 db.session.query(App)278 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")279 .first()280 )281 282 if not app:283 raise NotFound("App not found")284 285 try:286 # Skip the first row287 df = pd.read_csv(file)288 result = []289 for index, row in df.iterrows():290 content = {"question": row[0], "answer": row[1]}291 result.append(content)292 if len(result) == 0:293 raise ValueError("The CSV file is empty.")294 # check annotation limit295 features = FeatureService.get_features(current_user.current_tenant_id)296 if features.billing.enabled:297 annotation_quota_limit = features.annotation_quota_limit298 if annotation_quota_limit.limit < len(result) + annotation_quota_limit.size:299 raise ValueError("The number of annotations exceeds the limit of your subscription.")300 # async job301 job_id = str(uuid.uuid4())302 indexing_cache_key = "app_annotation_batch_import_{}".format(str(job_id))303 # send batch add segments task304 redis_client.setnx(indexing_cache_key, "waiting")305 batch_import_annotations_task.delay(306 str(job_id), result, app_id, current_user.current_tenant_id, current_user.id307 )308 except Exception as e:309 return {"error_msg": str(e)}310 return {"job_id": job_id, "job_status": "waiting"}311 312 @classmethod313 def get_annotation_hit_histories(cls, app_id: str, annotation_id: str, page, limit):314 # get app info315 app = (316 db.session.query(App)317 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")318 .first()319 )320 321 if not app:322 raise NotFound("App not found")323 324 annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first()325 326 if not annotation:327 raise NotFound("Annotation not found")328 329 annotation_hit_histories = (330 db.session.query(AppAnnotationHitHistory)331 .filter(332 AppAnnotationHitHistory.app_id == app_id,333 AppAnnotationHitHistory.annotation_id == annotation_id,334 )335 .order_by(AppAnnotationHitHistory.created_at.desc())336 .paginate(page=page, per_page=limit, max_per_page=100, error_out=False)337 )338 return annotation_hit_histories.items, annotation_hit_histories.total339 340 @classmethod341 def get_annotation_by_id(cls, annotation_id: str) -> MessageAnnotation | None:342 annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first()343 344 if not annotation:345 return None346 return annotation347 348 @classmethod349 def add_annotation_history(350 cls,351 annotation_id: str,352 app_id: str,353 annotation_question: str,354 annotation_content: str,355 query: str,356 user_id: str,357 message_id: str,358 from_source: str,359 score: float,360 ):361 # add hit count to annotation362 db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).update(363 {MessageAnnotation.hit_count: MessageAnnotation.hit_count + 1}, synchronize_session=False364 )365 366 annotation_hit_history = AppAnnotationHitHistory(367 annotation_id=annotation_id,368 app_id=app_id,369 account_id=user_id,370 question=query,371 source=from_source,372 score=score,373 message_id=message_id,374 annotation_question=annotation_question,375 annotation_content=annotation_content,376 )377 db.session.add(annotation_hit_history)378 db.session.commit()379 380 @classmethod381 def get_app_annotation_setting_by_app_id(cls, app_id: str):382 # get app info383 app = (384 db.session.query(App)385 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")386 .first()387 )388 389 if not app:390 raise NotFound("App not found")391 392 annotation_setting = (393 db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first()394 )395 if annotation_setting:396 collection_binding_detail = annotation_setting.collection_binding_detail397 return {398 "id": annotation_setting.id,399 "enabled": True,400 "score_threshold": annotation_setting.score_threshold,401 "embedding_model": {402 "embedding_provider_name": collection_binding_detail.provider_name,403 "embedding_model_name": collection_binding_detail.model_name,404 },405 }406 return {"enabled": False}407 408 @classmethod409 def update_app_annotation_setting(cls, app_id: str, annotation_setting_id: str, args: dict):410 # get app info411 app = (412 db.session.query(App)413 .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")414 .first()415 )416 417 if not app:418 raise NotFound("App not found")419 420 annotation_setting = (421 db.session.query(AppAnnotationSetting)422 .filter(423 AppAnnotationSetting.app_id == app_id,424 AppAnnotationSetting.id == annotation_setting_id,425 )426 .first()427 )428 if not annotation_setting:429 raise NotFound("App annotation not found")430 annotation_setting.score_threshold = args["score_threshold"]431 annotation_setting.updated_user_id = current_user.id432 annotation_setting.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)433 db.session.add(annotation_setting)434 db.session.commit()435 436 collection_binding_detail = annotation_setting.collection_binding_detail437 438 return {439 "id": annotation_setting.id,440 "enabled": True,441 "score_threshold": annotation_setting.score_threshold,442 "embedding_model": {443 "embedding_provider_name": collection_binding_detail.provider_name,444 "embedding_model_name": collection_binding_detail.model_name,445 },446 }447