Use hook to get userid (#26839)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

api/services/annotation_service.py

@@ -8,8 +8,7 @@ from werkzeug.exceptions import NotFound
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.datetime_utils import naive_utc_now
-from libs.login import current_user
-from models.account import Account
+from libs.login import current_account_with_tenant
 from models.model import App, AppAnnotationHitHistory, AppAnnotationSetting, Message, MessageAnnotation
 from services.feature_service import FeatureService
 from tasks.annotation.add_annotation_to_index_task import add_annotation_to_index_task
@@ -24,10 +23,10 @@ class AppAnnotationService:
     @classmethod
     def up_insert_app_annotation_from_message(cls, args: dict, app_id: str) -> MessageAnnotation:
         # get app info
-        assert isinstance(current_user, Account)
+        current_user, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -63,12 +62,12 @@ class AppAnnotationService:
         db.session.commit()
         # if annotation reply is enabled , add annotation to index
         annotation_setting = db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).first()
-        assert current_user.current_tenant_id is not None
+        assert current_tenant_id is not None
         if annotation_setting:
             add_annotation_to_index_task.delay(
                 annotation.id,
                 args["question"],
-                current_user.current_tenant_id,
+                current_tenant_id,
                 app_id,
                 annotation_setting.collection_binding_id,
             )
@@ -86,13 +85,12 @@ class AppAnnotationService:
         enable_app_annotation_job_key = f"enable_app_annotation_job_{str(job_id)}"
         # send batch add segments task
         redis_client.setnx(enable_app_annotation_job_key, "waiting")
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        current_user, current_tenant_id = current_account_with_tenant()
         enable_annotation_reply_task.delay(
             str(job_id),
             app_id,
             current_user.id,
-            current_user.current_tenant_id,
+            current_tenant_id,
             args["score_threshold"],
             args["embedding_provider_name"],
             args["embedding_model_name"],
@@ -101,8 +99,7 @@ class AppAnnotationService:
 
     @classmethod
     def disable_app_annotation(cls, app_id: str):
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         disable_app_annotation_key = f"disable_app_annotation_{str(app_id)}"
         cache_result = redis_client.get(disable_app_annotation_key)
         if cache_result is not None:
@@ -113,17 +110,16 @@ class AppAnnotationService:
         disable_app_annotation_job_key = f"disable_app_annotation_job_{str(job_id)}"
         # send batch add segments task
         redis_client.setnx(disable_app_annotation_job_key, "waiting")
-        disable_annotation_reply_task.delay(str(job_id), app_id, current_user.current_tenant_id)
+        disable_annotation_reply_task.delay(str(job_id), app_id, current_tenant_id)
         return {"job_id": job_id, "job_status": "waiting"}
 
     @classmethod
     def get_annotation_list_by_app_id(cls, app_id: str, page: int, limit: int, keyword: str):
         # get app info
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -153,11 +149,10 @@ class AppAnnotationService:
     @classmethod
     def export_annotation_list_by_app_id(cls, app_id: str):
         # get app info
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -174,11 +169,10 @@ class AppAnnotationService:
     @classmethod
     def insert_app_annotation_directly(cls, args: dict, app_id: str) -> MessageAnnotation:
         # get app info
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        current_user, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -196,7 +190,7 @@ class AppAnnotationService:
             add_annotation_to_index_task.delay(
                 annotation.id,
                 args["question"],
-                current_user.current_tenant_id,
+                current_tenant_id,
                 app_id,
                 annotation_setting.collection_binding_id,
             )
@@ -205,11 +199,10 @@ class AppAnnotationService:
     @classmethod
     def update_app_annotation_directly(cls, args: dict, app_id: str, annotation_id: str):
         # get app info
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -234,7 +227,7 @@ class AppAnnotationService:
             update_annotation_to_index_task.delay(
                 annotation.id,
                 annotation.question,
-                current_user.current_tenant_id,
+                current_tenant_id,
                 app_id,
                 app_annotation_setting.collection_binding_id,
             )
@@ -244,11 +237,10 @@ class AppAnnotationService:
     @classmethod
     def delete_app_annotation(cls, app_id: str, annotation_id: str):
         # get app info
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -277,17 +269,16 @@ class AppAnnotationService:
 
         if app_annotation_setting:
             delete_annotation_index_task.delay(
-                annotation.id, app_id, current_user.current_tenant_id, app_annotation_setting.collection_binding_id
+                annotation.id, app_id, current_tenant_id, app_annotation_setting.collection_binding_id
             )
 
     @classmethod
     def delete_app_annotations_in_batch(cls, app_id: str, annotation_ids: list[str]):
         # get app info
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -317,7 +308,7 @@ class AppAnnotationService:
         for annotation, annotation_setting in annotations_to_delete:
             if annotation_setting:
                 delete_annotation_index_task.delay(
-                    annotation.id, app_id, current_user.current_tenant_id, annotation_setting.collection_binding_id
+                    annotation.id, app_id, current_tenant_id, annotation_setting.collection_binding_id
                 )
 
         # Step 4: Bulk delete annotations in a single query
@@ -333,11 +324,10 @@ class AppAnnotationService:
     @classmethod
     def batch_import_app_annotations(cls, app_id, file: FileStorage):
         # get app info
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        current_user, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -354,7 +344,7 @@ class AppAnnotationService:
             if len(result) == 0:
                 raise ValueError("The CSV file is empty.")
             # check annotation limit
-            features = FeatureService.get_features(current_user.current_tenant_id)
+            features = FeatureService.get_features(current_tenant_id)
             if features.billing.enabled:
                 annotation_quota_limit = features.annotation_quota_limit
                 if annotation_quota_limit.limit < len(result) + annotation_quota_limit.size:
@@ -364,21 +354,18 @@ class AppAnnotationService:
             indexing_cache_key = f"app_annotation_batch_import_{str(job_id)}"
             # send batch add segments task
             redis_client.setnx(indexing_cache_key, "waiting")
-            batch_import_annotations_task.delay(
-                str(job_id), result, app_id, current_user.current_tenant_id, current_user.id
-            )
+            batch_import_annotations_task.delay(str(job_id), result, app_id, current_tenant_id, current_user.id)
         except Exception as e:
             return {"error_msg": str(e)}
         return {"job_id": job_id, "job_status": "waiting"}
 
     @classmethod
     def get_annotation_hit_histories(cls, app_id: str, annotation_id: str, page, limit):
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         # get app info
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -445,12 +432,11 @@ class AppAnnotationService:
 
     @classmethod
     def get_app_annotation_setting_by_app_id(cls, app_id: str):
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         # get app info
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -481,12 +467,11 @@ class AppAnnotationService:
 
     @classmethod
     def update_app_annotation_setting(cls, app_id: str, annotation_setting_id: str, args: dict):
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        current_user, current_tenant_id = current_account_with_tenant()
         # get app info
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -531,11 +516,10 @@ class AppAnnotationService:
 
     @classmethod
     def clear_all_annotations(cls, app_id: str):
-        assert isinstance(current_user, Account)
-        assert current_user.current_tenant_id is not None
+        _, current_tenant_id = current_account_with_tenant()
         app = (
             db.session.query(App)
-            .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal")
+            .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
             .first()
         )
 
@@ -558,7 +542,7 @@ class AppAnnotationService:
             # if annotation reply is enabled, delete annotation index
             if app_annotation_setting:
                 delete_annotation_index_task.delay(
-                    annotation.id, app_id, current_user.current_tenant_id, app_annotation_setting.collection_binding_id
+                    annotation.id, app_id, current_tenant_id, app_annotation_setting.collection_binding_id
                 )
 
             db.session.delete(annotation)