add app and dataset tag support

api/controllers/console/tag/tags.py

+from flask import request
+from flask_login import current_user
+from flask_restful import Resource, marshal_with, reqparse
+from werkzeug.exceptions import Forbidden
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from fields.tag_fields import tag_fields
+from libs.login import login_required
+from models.model import Tag
+from services.tag_service import TagService
+
+
+def _validate_name(name):
+    if not name or len(name) < 1 or len(name) > 40:
+        raise ValueError('Name must be between 1 to 50 characters.')
+    return name
+
+
+class TagListApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(tag_fields)
+    def get(self):
+        tag_type = request.args.get('type', type=str)
+        keyword = request.args.get('keyword', default=None, type=str)
+        tags = TagService.get_tags(tag_type, current_user.current_tenant_id, keyword)
+
+        return tags, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        # The role of the current user in the ta table must be admin or owner
+        if not current_user.is_admin_or_owner:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', nullable=False, required=True,
+                            help='Name must be between 1 to 50 characters.',
+                            type=_validate_name)
+        parser.add_argument('type', type=str, location='json',
+                            choices=Tag.TAG_TYPE_LIST,
+                            nullable=True,
+                            help='Invalid tag type.')
+        args = parser.parse_args()
+        tag = TagService.save_tags(args)
+
+        response = {
+            'id': tag.id,
+            'name': tag.name,
+            'type': tag.type,
+            'binding_count': 0
+        }
+
+        return response, 200
+
+
+class TagUpdateDeleteApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def patch(self, tag_id):
+        tag_id = str(tag_id)
+        # The role of the current user in the ta table must be admin or owner
+        if not current_user.is_admin_or_owner:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', nullable=False, required=True,
+                            help='Name must be between 1 to 50 characters.',
+                            type=_validate_name)
+        args = parser.parse_args()
+        tag = TagService.update_tags(args, tag_id)
+
+        binding_count = TagService.get_tag_binding_count(tag_id)
+
+        response = {
+            'id': tag.id,
+            'name': tag.name,
+            'type': tag.type,
+            'binding_count': binding_count
+        }
+
+        return response, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, tag_id):
+        tag_id = str(tag_id)
+        # The role of the current user in the ta table must be admin or owner
+        if not current_user.is_admin_or_owner:
+            raise Forbidden()
+
+        TagService.delete_tag(tag_id)
+
+        return 200
+
+
+class TagBindingCreateApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        # The role of the current user in the ta table must be admin or owner
+        if not current_user.is_admin_or_owner:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('tag_ids', type=list, nullable=False, required=True,
+                            help='Tag IDs is required.')
+        parser.add_argument('target_id', type=str, nullable=False, required=True,
+                            help='Target ID is required.')
+        parser.add_argument('type', type=str, location='json',
+                            choices=Tag.TAG_TYPE_LIST,
+                            nullable=True,
+                            help='Invalid tag type.')
+        args = parser.parse_args()
+        TagService.save_tag_binding(args)
+
+        return 200
+
+
+class TagBindingDeleteApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        # The role of the current user in the ta table must be admin or owner
+        if not current_user.is_admin_or_owner:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('tag_id', type=list, nullable=False, required=True,
+                            help='Tag ID is required.')
+        parser.add_argument('target_id', type=str, nullable=False, required=True,
+                            help='Target ID is required.')
+        parser.add_argument('type', type=str, location='json',
+                            choices=Tag.TAG_TYPE_LIST,
+                            nullable=True,
+                            help='Invalid tag type.')
+        args = parser.parse_args()
+        TagService.delete_tag_binding(args)
+
+        return 200
+
+
+api.add_resource(TagListApi, '/tags')
+api.add_resource(TagUpdateDeleteApi, '/tags/<uuid:tag_id>')
+api.add_resource(TagBindingCreateApi, '/tag-bindings/create')
+api.add_resource(TagBindingDeleteApi, '/tag-bindings/remove')

api/fields/tag_fields.py

+from flask_restful import fields
+
+tag_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'type': fields.String,
+    'binding_count': fields.String
+}
\ No newline at end of file

api/models/model.py

@@ -69,7 +69,7 @@ class App(db.Model):
     def tenant(self):
         tenant = db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first()
         return tenant
-    
+
     @property
     def is_agent(self) -> bool:
         app_model_config = self.app_model_config
@@ -78,10 +78,10 @@ class App(db.Model):
         if not app_model_config.agent_mode:
             return False
         if self.app_model_config.agent_mode_dict.get('enabled', False) \
-            and self.app_model_config.agent_mode_dict.get('strategy', '') in ['function_call', 'react']:
+                and self.app_model_config.agent_mode_dict.get('strategy', '') in ['function_call', 'react']:
             return True
         return False
-    
+
     @property
     def deleted_tools(self) -> list:
         # get agent mode tools
@@ -92,7 +92,7 @@ class App(db.Model):
             return []
         agent_mode = app_model_config.agent_mode_dict
         tools = agent_mode.get('tools', [])
-        
+
         provider_ids = []
 
         for tool in tools:
@@ -129,6 +129,7 @@ class App(db.Model):
 
         return deleted_tools
 
+
 class AppModelConfig(db.Model):
     __tablename__ = 'app_model_configs'
     __table_args__ = (
@@ -235,7 +236,8 @@ class AppModelConfig(db.Model):
 
     @property
     def agent_mode_dict(self) -> dict:
-        return json.loads(self.agent_mode) if self.agent_mode else {"enabled": False, "strategy": None, "tools": [], "prompt": None}
+        return json.loads(self.agent_mode) if self.agent_mode else {"enabled": False, "strategy": None, "tools": [],
+                                                                    "prompt": None}
 
     @property
     def chat_prompt_config_dict(self) -> dict:
@@ -415,6 +417,7 @@ class InstalledApp(db.Model):
             return False
         return app.is_agent
 
+
 class Conversation(db.Model):
     __tablename__ = 'conversations'
     __table_args__ = (
@@ -726,6 +729,7 @@ class MessageFile(db.Model):
     created_by = db.Column(UUID, nullable=False)
     created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
 
+
 class MessageAnnotation(db.Model):
     __tablename__ = 'message_annotations'
     __table_args__ = (
@@ -1032,7 +1036,7 @@ class MessageAgentThought(db.Model):
             return json.loads(self.message_files)
         else:
             return []
-        
+
     @property
     def tool_labels(self) -> dict:
         try:
@@ -1043,6 +1047,7 @@ class MessageAgentThought(db.Model):
         except Exception as e:
             return {}
 
+
 class DatasetRetrieverResource(db.Model):
     __tablename__ = 'dataset_retriever_resources'
     __table_args__ = (
@@ -1068,3 +1073,37 @@ class DatasetRetrieverResource(db.Model):
     retriever_from = db.Column(db.Text, nullable=False)
     created_by = db.Column(UUID, nullable=False)
     created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.current_timestamp())
+
+
+class Tag(db.Model):
+    __tablename__ = 'tags'
+    __table_args__ = (
+        db.PrimaryKeyConstraint('id', name='tag_pkey'),
+        db.Index('tag_type_idx', 'type'),
+        db.Index('tag_name_idx', 'name'),
+    )
+
+    TAG_TYPE_LIST = ['knowledge', 'app']
+
+    id = db.Column(UUID, server_default=db.text('uuid_generate_v4()'))
+    tenant_id = db.Column(UUID, nullable=True)
+    type = db.Column(db.String(16), nullable=False)
+    name = db.Column(db.String(255), nullable=False)
+    created_by = db.Column(UUID, nullable=False)
+    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
+
+
+class TagBinding(db.Model):
+    __tablename__ = 'tag_bindings'
+    __table_args__ = (
+        db.PrimaryKeyConstraint('id', name='tag_binding_pkey'),
+        db.Index('tag_bind_target_id_idx', 'target_id'),
+        db.Index('tag_bind_tag_id_idx', 'tag_id'),
+    )
+
+    id = db.Column(UUID, server_default=db.text('uuid_generate_v4()'))
+    tenant_id = db.Column(UUID, nullable=True)
+    tag_id = db.Column(UUID, nullable=True)
+    target_id = db.Column(UUID, nullable=True)
+    created_by = db.Column(UUID, nullable=False)
+    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))

api/services/dataset_service.py

@@ -44,10 +44,11 @@ from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task
 from tasks.delete_segment_from_index_task import delete_segment_from_index_task
 from tasks.document_indexing_task import document_indexing_task
 from tasks.document_indexing_update_task import document_indexing_update_task
-from tasks.recover_document_indexing_task import recover_document_indexing_task
 from tasks.duplicate_document_indexing_task import duplicate_document_indexing_task
+from tasks.recover_document_indexing_task import recover_document_indexing_task
 from tasks.retry_document_indexing_task import retry_document_indexing_task
 
+
 class DatasetService:
 
     @staticmethod
@@ -63,11 +64,11 @@ class DatasetService:
         if search:
             query = query.filter(db.and_(Dataset.name.ilike(f'%{search}%')))
         datasets = query.paginate(
-                    page=page,
-                    per_page=per_page,
-                    max_per_page=100,
-                    error_out=False
-                )
+            page=page,
+            per_page=per_page,
+            max_per_page=100,
+            error_out=False
+        )
 
         return datasets.items, datasets.total
 
@@ -667,7 +668,8 @@ class DocumentService:
     def check_documents_upload_quota(count: int, features: FeatureModel):
         can_upload_size = features.documents_upload_quota.limit - features.documents_upload_quota.size
         if count > can_upload_size:
-            raise ValueError(f'You have reached the limit of your subscription. Only {can_upload_size} documents can be uploaded.')
+            raise ValueError(
+                f'You have reached the limit of your subscription. Only {can_upload_size} documents can be uploaded.')
 
     @staticmethod
     def build_document(dataset: Dataset, process_rule_id: str, data_source_type: str, document_form: str,

api/services/tag_service.py

+import uuid
+
+from flask_login import current_user
+from sqlalchemy import func
+from werkzeug.exceptions import NotFound
+
+from extensions.ext_database import db
+from models.dataset import Dataset
+from models.model import App, Tag, TagBinding
+
+
+class TagService:
+    @staticmethod
+    def get_tags(tag_type: str, current_tenant_id: str, keyword: str = None) -> list:
+        query = db.session.query(
+            Tag.id, Tag.type, Tag.name, func.count(TagBinding.id).label('binding_count')
+        ).outerjoin(
+            TagBinding, Tag.id == TagBinding.tag_id
+        ).filter(
+            Tag.type == tag_type,
+            Tag.tenant_id == current_tenant_id
+        )
+        if keyword:
+            query = query.filter(db.and_(Tag.name.ilike(f'%{keyword}%')))
+        query = query.group_by(
+            Tag.id
+        )
+        results = query.all()
+        return results
+
+    @staticmethod
+    def save_tags(args: dict) -> Tag:
+        tag = Tag(
+            id=str(uuid.uuid4()),
+            name=args['name'],
+            type=args['type'],
+            created_by=current_user.id,
+            tenant_id=current_user.current_tenant_id
+        )
+        db.session.add(tag)
+        db.session.commit()
+        return tag
+
+    @staticmethod
+    def update_tags(args: dict, tag_id: str) -> Tag:
+        tag = db.session.query(Tag).filter(Tag.id == tag_id).first()
+        if not tag:
+            raise NotFound("Tag not found")
+        tag.name = args['name']
+        db.session.commit()
+        return tag
+
+    @staticmethod
+    def get_tag_binding_count(tag_id: str) -> int:
+        count = db.session.query(TagBinding).filter(TagBinding.tag_id == tag_id).count()
+        return count
+
+    @staticmethod
+    def delete_tag(tag_id: str):
+        tag = db.session.query(Tag).filter(Tag.id == tag_id).first()
+        if not tag:
+            raise NotFound("Tag not found")
+        db.session.delete(tag)
+        # delete tag binding
+        tag_bindings = db.session.query(TagBinding).filter(TagBinding.tag_id == tag_id).all()
+        if tag_bindings:
+            for tag_binding in tag_bindings:
+                db.session.delete(tag_binding)
+        db.session.commit()
+
+    @staticmethod
+    def save_tag_binding(args):
+        # check if target exists
+        TagService.check_target_exists(args['type'], args['target_id'])
+        # save tag binding
+        for tag_id in args['tag_ids']:
+            tag_binding = db.session.query(TagBinding).filter(
+                TagBinding.tag_id == tag_id,
+                TagBinding.target_id == args['target_id']
+            ).first()
+            if tag_binding:
+                continue
+            new_tag_binding = TagBinding(
+                id=str(uuid.uuid4()),
+                tag_id=tag_id,
+                target_id=args['target_id'],
+                tenant_id=current_user.current_tenant_id
+            )
+            db.session.add(new_tag_binding)
+        db.session.commit()
+
+    @staticmethod
+    def delete_tag_binding(args):
+        # check if target exists
+        TagService.check_target_exists(args['type'], args['target_id'])
+        # delete tag binding
+        tag_bindings = db.session.query(TagBinding).filter(
+            TagBinding.target_id == args['target_id'],
+            TagBinding.tag_id == (args['tag_id'])
+        ).first()
+        if tag_bindings:
+            db.session.delete(tag_bindings)
+            db.session.commit()
+
+
+
+    @staticmethod
+    def check_target_exists(type: str, target_id: str):
+        if type == 'knowledge':
+            dataset = db.session.query(Dataset).filter(
+                Dataset.tenant_id == current_user.current_tenant_id,
+                Dataset.id == target_id
+            ).first()
+            if not dataset:
+                raise NotFound("Dataset not found")
+        elif type == 'app':
+            app = db.session.query(App).filter(
+                App.tenant_id == current_user.current_tenant_id,
+                App.id == target_id
+            ).first()
+            if not app:
+                raise NotFound("App not found")
+        else:
+            raise NotFound("Invalid binding type")
+

api/tasks/retry_document_indexing_task.py

@@ -4,9 +4,8 @@ import time
 
 import click
 from celery import shared_task
-from flask import current_app
 
-from core.indexing_runner import DocumentIsPausedException, IndexingRunner
+from core.indexing_runner import IndexingRunner
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client