Merge branch 'main' into feat/mcp-06-18

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)

api/services/datasource_provider_service.py

@@ -3,7 +3,6 @@ import time
 from collections.abc import Mapping
 from typing import Any
 
-from flask_login import current_user
 from sqlalchemy.orm import Session
 
 from configs import dify_config
@@ -18,6 +17,7 @@ from core.tools.entities.tool_entities import CredentialType
 from core.tools.utils.encryption import ProviderConfigCache, ProviderConfigEncrypter, create_provider_encrypter
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
+from libs.login import current_account_with_tenant
 from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig, DatasourceProvider
 from models.provider_ids import DatasourceProviderID
 from services.plugin.plugin_service import PluginService
@@ -93,6 +93,8 @@ class DatasourceProviderService:
         """
         get credential by id
         """
+        current_user, _ = current_account_with_tenant()
+
         with Session(db.engine) as session:
             if credential_id:
                 datasource_provider = (
@@ -157,6 +159,8 @@ class DatasourceProviderService:
         """
         get all datasource credentials by provider
         """
+        current_user, _ = current_account_with_tenant()
+
         with Session(db.engine) as session:
             datasource_providers = (
                 session.query(DatasourceProvider)
@@ -604,6 +608,8 @@ class DatasourceProviderService:
         """
         provider_name = provider_id.provider_name
         plugin_id = provider_id.plugin_id
+        current_user, _ = current_account_with_tenant()
+
         with Session(db.engine) as session:
             lock = f"datasource_provider_create_lock:{tenant_id}_{provider_id}_{CredentialType.API_KEY}"
             with redis_client.lock(lock, timeout=20):
@@ -901,6 +907,8 @@ class DatasourceProviderService:
         """
         update datasource credentials.
         """
+        current_user, _ = current_account_with_tenant()
+
         with Session(db.engine) as session:
             datasource_provider = (
                 session.query(DatasourceProvider)