add dataset service api enable

api/core/app/apps/pipeline/pipeline_generator.py

@@ -25,6 +25,7 @@ from core.app.apps.pipeline.pipeline_runner import PipelineRunner
 from core.app.apps.workflow.generate_response_converter import WorkflowAppGenerateResponseConverter
 from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline
 from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerateEntity
+from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity
 from core.app.entities.task_entities import WorkflowAppBlockingResponse, WorkflowAppStreamResponse
 from core.datasource.entities.datasource_entities import (
     DatasourceProviderType,
@@ -41,6 +42,7 @@ from core.workflow.repositories.workflow_execution_repository import WorkflowExe
 from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository
 from core.workflow.variable_loader import DUMMY_VARIABLE_LOADER, VariableLoader
 from extensions.ext_database import db
+from extensions.ext_redis import redis_client
 from libs.flask_utils import preserve_flask_contexts
 from models import Account, EndUser, Workflow, WorkflowNodeExecutionTriggeredFrom
 from models.dataset import Document, DocumentPipelineExecutionLog, Pipeline
@@ -48,7 +50,10 @@ from models.enums import WorkflowRunTriggeredFrom
 from models.model import AppMode
 from services.dataset_service import DocumentService
 from services.datasource_provider_service import DatasourceProviderService
+from services.feature_service import FeatureService
+from services.file_service import FileService
 from services.workflow_draft_variable_service import DraftVarLoader, WorkflowDraftVariableService
+from tasks.rag_pipeline.priority_rag_pipeline_run_task import priority_rag_pipeline_run_task
 from tasks.rag_pipeline.rag_pipeline_run_task import rag_pipeline_run_task
 
 logger = logging.getLogger(__name__)
@@ -147,6 +152,7 @@ class PipelineGenerator(BaseAppGenerator):
             db.session.commit()
 
         # run in child thread
+        rag_pipeline_invoke_entities = []
         for i, datasource_info in enumerate(datasource_info_list):
             workflow_run_id = str(uuid.uuid4())
             document_id = None
@@ -223,7 +229,7 @@ class PipelineGenerator(BaseAppGenerator):
                     workflow_thread_pool_id=workflow_thread_pool_id,
                 )
             else:
-                rag_pipeline_run_task.delay(  # type: ignore
+                rag_pipeline_invoke_entities.append(RagPipelineInvokeEntity(
                     pipeline_id=pipeline.id,
                     user_id=user.id,
                     tenant_id=pipeline.tenant_id,
@@ -232,7 +238,36 @@ class PipelineGenerator(BaseAppGenerator):
                     workflow_execution_id=workflow_run_id,
                     workflow_thread_pool_id=workflow_thread_pool_id,
                     application_generate_entity=application_generate_entity.model_dump(),
+                ))
+
+        if rag_pipeline_invoke_entities:
+            # store the rag_pipeline_invoke_entities to object storage
+            text = [item.model_dump() for item in rag_pipeline_invoke_entities]
+            name = "rag_pipeline_invoke_entities.json"
+            # Convert list to proper JSON string
+            json_text = json.dumps(text)
+            upload_file = FileService(db.engine).upload_text(json_text, name)
+            features = FeatureService.get_features(dataset.tenant_id)
+            if features.billing.subscription.plan == "sandbox":
+                tenant_pipeline_task_key = f"tenant_pipeline_task:{dataset.tenant_id}"
+                tenant_self_pipeline_task_queue = f"tenant_self_pipeline_task_queue:{dataset.tenant_id}"
+
+                if redis_client.get(tenant_pipeline_task_key):
+                    # Add to waiting queue using List operations (lpush)
+                    redis_client.lpush(tenant_self_pipeline_task_queue, upload_file.id)
+                else:
+                    # Set flag and execute task
+                    redis_client.set(tenant_pipeline_task_key, 1, ex=60 * 60)
+                    rag_pipeline_run_task.delay(  # type: ignore
+                        rag_pipeline_invoke_entities_file_id=upload_file.id,
+                        tenant_id=dataset.tenant_id,
+                    )
+
+            else:
+                priority_rag_pipeline_run_task.delay(  # type: ignore
+                    rag_pipeline_invoke_entities_file_id=upload_file.id,
                 )
+
         # return batch, dataset, documents
         return {
             "batch": batch,

api/core/app/entities/rag_pipeline_invoke_entities.py

@@ -0,0 +1,14 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class RagPipelineInvokeEntity(BaseModel):
+    pipeline_id: str
+    application_generate_entity: dict[str, Any]
+    user_id: str
+    tenant_id: str
+    workflow_id: str
+    streaming: bool
+    workflow_execution_id: str | None = None
+    workflow_thread_pool_id: str | None = None