dify/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py

from typing import Optional

from extensions.ext_database import db
from models.dataset import Dataset, Pipeline, PipelineBuiltInTemplate
from services.rag_pipeline.pipeline_template.pipeline_template_base import PipelineTemplateRetrievalBase
from services.rag_pipeline.pipeline_template.pipeline_template_type import PipelineTemplateType


class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
    """
    Retrieval pipeline   template from database
    """

    def get_pipeline_templates(self, language: str) -> dict:
        result = self.fetch_pipeline_templates_from_db(language)
        return result

    def get_pipeline_template_detail(self, pipeline_id: str):
        result = self.fetch_pipeline_template_detail_from_db(pipeline_id)
        return result

    def get_type(self) -> str:
        return PipelineTemplateType.DATABASE

    @classmethod
    def fetch_pipeline_templates_from_db(cls, language: str) -> dict:
        """
        Fetch pipeline templates from db.
        :param language: language
        :return:
        """

        pipeline_built_in_templates: list[PipelineBuiltInTemplate] = (
            db.session.query(PipelineBuiltInTemplate).filter(PipelineBuiltInTemplate.language == language).all()
        )

        recommended_pipelines_results = []
        for pipeline_built_in_template in pipeline_built_in_templates:
            pipeline_model: Pipeline = pipeline_built_in_template.pipeline

            recommended_pipeline_result = {
                "id": pipeline_built_in_template.id,
                "name": pipeline_built_in_template.name,
                "pipeline_id": pipeline_model.id,
                "description": pipeline_built_in_template.description,
                "icon": pipeline_built_in_template.icon,
                "copyright": pipeline_built_in_template.copyright,
                "privacy_policy": pipeline_built_in_template.privacy_policy,
                "position": pipeline_built_in_template.position,
            }
            dataset: Dataset = pipeline_model.dataset
            if dataset:
                recommended_pipeline_result["chunk_structure"] = dataset.chunk_structure
                recommended_pipelines_results.append(recommended_pipeline_result)

        return {"pipeline_templates": recommended_pipelines_results}

    @classmethod
    def fetch_pipeline_template_detail_from_db(cls, pipeline_id: str) -> Optional[dict]:
        """
        Fetch pipeline template detail from db.
        :param pipeline_id: Pipeline ID
        :return:
        """
        from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService

        # is in public recommended list
        pipeline_template = (
            db.session.query(PipelineBuiltInTemplate).filter(PipelineBuiltInTemplate.id == pipeline_id).first()
        )

        if not pipeline_template:
            return None

        # get app detail
        pipeline = db.session.query(Pipeline).filter(Pipeline.id == pipeline_template.pipeline_id).first()
        if not pipeline or not pipeline.is_public:
            return None

        return {
            "id": pipeline.id,
            "name": pipeline.name,
            "icon": pipeline.icon,
            "mode": pipeline.mode,
            "export_data": RagPipelineDslService.export_rag_pipeline_dsl(pipeline=pipeline),
        }