Upload 1150 files

.gitattributes

@@ -6,3 +6,9 @@
 
 *.sh      text eol=lf
 api/tests/integration_tests/model_runtime/assets/audio.mp3 filter=lfs diff=lfs merge=lfs -text
+api/core/tools/docs/images/index/image-1.png filter=lfs diff=lfs merge=lfs -text
+api/core/tools/docs/images/index/image-2.png filter=lfs diff=lfs merge=lfs -text
+api/core/tools/docs/images/index/image.png filter=lfs diff=lfs merge=lfs -text
+api/core/tools/provider/builtin/comfyui/_assets/icon.png filter=lfs diff=lfs merge=lfs -text
+api/core/tools/provider/builtin/dalle/_assets/icon.png filter=lfs diff=lfs merge=lfs -text
+api/core/tools/provider/builtin/wecom/_assets/icon.png filter=lfs diff=lfs merge=lfs -text

api/core/__init__.py

@@ -0,0 +1 @@
+import core.moderation.base

api/core/hosting_configuration.py

@@ -0,0 +1,255 @@
+from typing import Optional
+
+from flask import Flask
+from pydantic import BaseModel
+
+from configs import dify_config
+from core.entities.provider_entities import QuotaUnit, RestrictModel
+from core.model_runtime.entities.model_entities import ModelType
+from models.provider import ProviderQuotaType
+
+
+class HostingQuota(BaseModel):
+    quota_type: ProviderQuotaType
+    restrict_models: list[RestrictModel] = []
+
+
+class TrialHostingQuota(HostingQuota):
+    quota_type: ProviderQuotaType = ProviderQuotaType.TRIAL
+    quota_limit: int = 0
+    """Quota limit for the hosting provider models. -1 means unlimited."""
+
+
+class PaidHostingQuota(HostingQuota):
+    quota_type: ProviderQuotaType = ProviderQuotaType.PAID
+
+
+class FreeHostingQuota(HostingQuota):
+    quota_type: ProviderQuotaType = ProviderQuotaType.FREE
+
+
+class HostingProvider(BaseModel):
+    enabled: bool = False
+    credentials: Optional[dict] = None
+    quota_unit: Optional[QuotaUnit] = None
+    quotas: list[HostingQuota] = []
+
+
+class HostedModerationConfig(BaseModel):
+    enabled: bool = False
+    providers: list[str] = []
+
+
+class HostingConfiguration:
+    provider_map: dict[str, HostingProvider] = {}
+    moderation_config: Optional[HostedModerationConfig] = None
+
+    def init_app(self, app: Flask) -> None:
+        if dify_config.EDITION != "CLOUD":
+            return
+
+        self.provider_map["azure_openai"] = self.init_azure_openai()
+        self.provider_map["openai"] = self.init_openai()
+        self.provider_map["anthropic"] = self.init_anthropic()
+        self.provider_map["minimax"] = self.init_minimax()
+        self.provider_map["spark"] = self.init_spark()
+        self.provider_map["zhipuai"] = self.init_zhipuai()
+
+        self.moderation_config = self.init_moderation_config()
+
+    @staticmethod
+    def init_azure_openai() -> HostingProvider:
+        quota_unit = QuotaUnit.TIMES
+        if dify_config.HOSTED_AZURE_OPENAI_ENABLED:
+            credentials = {
+                "openai_api_key": dify_config.HOSTED_AZURE_OPENAI_API_KEY,
+                "openai_api_base": dify_config.HOSTED_AZURE_OPENAI_API_BASE,
+                "base_model_name": "gpt-35-turbo",
+            }
+
+            quotas: list[HostingQuota] = []
+            hosted_quota_limit = dify_config.HOSTED_AZURE_OPENAI_QUOTA_LIMIT
+            trial_quota = TrialHostingQuota(
+                quota_limit=hosted_quota_limit,
+                restrict_models=[
+                    RestrictModel(model="gpt-4", base_model_name="gpt-4", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-4o", base_model_name="gpt-4o", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-4o-mini", base_model_name="gpt-4o-mini", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-4-32k", base_model_name="gpt-4-32k", model_type=ModelType.LLM),
+                    RestrictModel(
+                        model="gpt-4-1106-preview", base_model_name="gpt-4-1106-preview", model_type=ModelType.LLM
+                    ),
+                    RestrictModel(
+                        model="gpt-4-vision-preview", base_model_name="gpt-4-vision-preview", model_type=ModelType.LLM
+                    ),
+                    RestrictModel(model="gpt-35-turbo", base_model_name="gpt-35-turbo", model_type=ModelType.LLM),
+                    RestrictModel(
+                        model="gpt-35-turbo-1106", base_model_name="gpt-35-turbo-1106", model_type=ModelType.LLM
+                    ),
+                    RestrictModel(
+                        model="gpt-35-turbo-instruct", base_model_name="gpt-35-turbo-instruct", model_type=ModelType.LLM
+                    ),
+                    RestrictModel(
+                        model="gpt-35-turbo-16k", base_model_name="gpt-35-turbo-16k", model_type=ModelType.LLM
+                    ),
+                    RestrictModel(
+                        model="text-davinci-003", base_model_name="text-davinci-003", model_type=ModelType.LLM
+                    ),
+                    RestrictModel(
+                        model="text-embedding-ada-002",
+                        base_model_name="text-embedding-ada-002",
+                        model_type=ModelType.TEXT_EMBEDDING,
+                    ),
+                    RestrictModel(
+                        model="text-embedding-3-small",
+                        base_model_name="text-embedding-3-small",
+                        model_type=ModelType.TEXT_EMBEDDING,
+                    ),
+                    RestrictModel(
+                        model="text-embedding-3-large",
+                        base_model_name="text-embedding-3-large",
+                        model_type=ModelType.TEXT_EMBEDDING,
+                    ),
+                ],
+            )
+            quotas.append(trial_quota)
+
+            return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
+
+        return HostingProvider(
+            enabled=False,
+            quota_unit=quota_unit,
+        )
+
+    def init_openai(self) -> HostingProvider:
+        quota_unit = QuotaUnit.CREDITS
+        quotas: list[HostingQuota] = []
+
+        if dify_config.HOSTED_OPENAI_TRIAL_ENABLED:
+            hosted_quota_limit = dify_config.HOSTED_OPENAI_QUOTA_LIMIT
+            trial_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_TRIAL_MODELS")
+            trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models)
+            quotas.append(trial_quota)
+
+        if dify_config.HOSTED_OPENAI_PAID_ENABLED:
+            paid_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_PAID_MODELS")
+            paid_quota = PaidHostingQuota(restrict_models=paid_models)
+            quotas.append(paid_quota)
+
+        if len(quotas) > 0:
+            credentials = {
+                "openai_api_key": dify_config.HOSTED_OPENAI_API_KEY,
+            }
+
+            if dify_config.HOSTED_OPENAI_API_BASE:
+                credentials["openai_api_base"] = dify_config.HOSTED_OPENAI_API_BASE
+
+            if dify_config.HOSTED_OPENAI_API_ORGANIZATION:
+                credentials["openai_organization"] = dify_config.HOSTED_OPENAI_API_ORGANIZATION
+
+            return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
+
+        return HostingProvider(
+            enabled=False,
+            quota_unit=quota_unit,
+        )
+
+    @staticmethod
+    def init_anthropic() -> HostingProvider:
+        quota_unit = QuotaUnit.TOKENS
+        quotas: list[HostingQuota] = []
+
+        if dify_config.HOSTED_ANTHROPIC_TRIAL_ENABLED:
+            hosted_quota_limit = dify_config.HOSTED_ANTHROPIC_QUOTA_LIMIT
+            trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit)
+            quotas.append(trial_quota)
+
+        if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED:
+            paid_quota = PaidHostingQuota()
+            quotas.append(paid_quota)
+
+        if len(quotas) > 0:
+            credentials = {
+                "anthropic_api_key": dify_config.HOSTED_ANTHROPIC_API_KEY,
+            }
+
+            if dify_config.HOSTED_ANTHROPIC_API_BASE:
+                credentials["anthropic_api_url"] = dify_config.HOSTED_ANTHROPIC_API_BASE
+
+            return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
+
+        return HostingProvider(
+            enabled=False,
+            quota_unit=quota_unit,
+        )
+
+    @staticmethod
+    def init_minimax() -> HostingProvider:
+        quota_unit = QuotaUnit.TOKENS
+        if dify_config.HOSTED_MINIMAX_ENABLED:
+            quotas: list[HostingQuota] = [FreeHostingQuota()]
+
+            return HostingProvider(
+                enabled=True,
+                credentials=None,  # use credentials from the provider
+                quota_unit=quota_unit,
+                quotas=quotas,
+            )
+
+        return HostingProvider(
+            enabled=False,
+            quota_unit=quota_unit,
+        )
+
+    @staticmethod
+    def init_spark() -> HostingProvider:
+        quota_unit = QuotaUnit.TOKENS
+        if dify_config.HOSTED_SPARK_ENABLED:
+            quotas: list[HostingQuota] = [FreeHostingQuota()]
+
+            return HostingProvider(
+                enabled=True,
+                credentials=None,  # use credentials from the provider
+                quota_unit=quota_unit,
+                quotas=quotas,
+            )
+
+        return HostingProvider(
+            enabled=False,
+            quota_unit=quota_unit,
+        )
+
+    @staticmethod
+    def init_zhipuai() -> HostingProvider:
+        quota_unit = QuotaUnit.TOKENS
+        if dify_config.HOSTED_ZHIPUAI_ENABLED:
+            quotas: list[HostingQuota] = [FreeHostingQuota()]
+
+            return HostingProvider(
+                enabled=True,
+                credentials=None,  # use credentials from the provider
+                quota_unit=quota_unit,
+                quotas=quotas,
+            )
+
+        return HostingProvider(
+            enabled=False,
+            quota_unit=quota_unit,
+        )
+
+    @staticmethod
+    def init_moderation_config() -> HostedModerationConfig:
+        if dify_config.HOSTED_MODERATION_ENABLED and dify_config.HOSTED_MODERATION_PROVIDERS:
+            return HostedModerationConfig(enabled=True, providers=dify_config.HOSTED_MODERATION_PROVIDERS.split(","))
+
+        return HostedModerationConfig(enabled=False)
+
+    @staticmethod
+    def parse_restrict_models_from_env(env_var: str) -> list[RestrictModel]:
+        models_str = dify_config.model_dump().get(env_var)
+        models_list = models_str.split(",") if models_str else []
+        return [
+            RestrictModel(model=model_name.strip(), model_type=ModelType.LLM)
+            for model_name in models_list
+            if model_name.strip()
+        ]

api/core/indexing_runner.py

@@ -0,0 +1,754 @@
+import concurrent.futures
+import datetime
+import json
+import logging
+import re
+import threading
+import time
+import uuid
+from typing import Any, Optional, cast
+
+from flask import current_app
+from flask_login import current_user  # type: ignore
+from sqlalchemy.orm.exc import ObjectDeletedError
+
+from configs import dify_config
+from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail
+from core.errors.error import ProviderTokenNotInitError
+from core.model_manager import ModelInstance, ModelManager
+from core.model_runtime.entities.model_entities import ModelType
+from core.rag.cleaner.clean_processor import CleanProcessor
+from core.rag.datasource.keyword.keyword_factory import Keyword
+from core.rag.docstore.dataset_docstore import DatasetDocumentStore
+from core.rag.extractor.entity.extract_setting import ExtractSetting
+from core.rag.index_processor.constant.index_type import IndexType
+from core.rag.index_processor.index_processor_base import BaseIndexProcessor
+from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
+from core.rag.models.document import ChildDocument, Document
+from core.rag.splitter.fixed_text_splitter import (
+    EnhanceRecursiveCharacterTextSplitter,
+    FixedRecursiveCharacterTextSplitter,
+)
+from core.rag.splitter.text_splitter import TextSplitter
+from core.tools.utils.web_reader_tool import get_image_upload_file_ids
+from extensions.ext_database import db
+from extensions.ext_redis import redis_client
+from extensions.ext_storage import storage
+from libs import helper
+from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment
+from models.dataset import Document as DatasetDocument
+from models.model import UploadFile
+from services.feature_service import FeatureService
+
+
+class IndexingRunner:
+    def __init__(self):
+        self.storage = storage
+        self.model_manager = ModelManager()
+
+    def run(self, dataset_documents: list[DatasetDocument]):
+        """Run the indexing process."""
+        for dataset_document in dataset_documents:
+            try:
+                # get dataset
+                dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first()
+
+                if not dataset:
+                    raise ValueError("no dataset found")
+
+                # get the process rule
+                processing_rule = (
+                    db.session.query(DatasetProcessRule)
+                    .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
+                    .first()
+                )
+                if not processing_rule:
+                    raise ValueError("no process rule found")
+                index_type = dataset_document.doc_form
+                index_processor = IndexProcessorFactory(index_type).init_index_processor()
+                # extract
+                text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict())
+
+                # transform
+                documents = self._transform(
+                    index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()
+                )
+                # save segment
+                self._load_segments(dataset, dataset_document, documents)
+
+                # load
+                self._load(
+                    index_processor=index_processor,
+                    dataset=dataset,
+                    dataset_document=dataset_document,
+                    documents=documents,
+                )
+            except DocumentIsPausedError:
+                raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id))
+            except ProviderTokenNotInitError as e:
+                dataset_document.indexing_status = "error"
+                dataset_document.error = str(e.description)
+                dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+                db.session.commit()
+            except ObjectDeletedError:
+                logging.warning("Document deleted, document id: {}".format(dataset_document.id))
+            except Exception as e:
+                logging.exception("consume document failed")
+                dataset_document.indexing_status = "error"
+                dataset_document.error = str(e)
+                dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+                db.session.commit()
+
+    def run_in_splitting_status(self, dataset_document: DatasetDocument):
+        """Run the indexing process when the index_status is splitting."""
+        try:
+            # get dataset
+            dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first()
+
+            if not dataset:
+                raise ValueError("no dataset found")
+
+            # get exist document_segment list and delete
+            document_segments = DocumentSegment.query.filter_by(
+                dataset_id=dataset.id, document_id=dataset_document.id
+            ).all()
+
+            for document_segment in document_segments:
+                db.session.delete(document_segment)
+                if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
+                    # delete child chunks
+                    db.session.query(ChildChunk).filter(ChildChunk.segment_id == document_segment.id).delete()
+            db.session.commit()
+            # get the process rule
+            processing_rule = (
+                db.session.query(DatasetProcessRule)
+                .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
+                .first()
+            )
+            if not processing_rule:
+                raise ValueError("no process rule found")
+
+            index_type = dataset_document.doc_form
+            index_processor = IndexProcessorFactory(index_type).init_index_processor()
+            # extract
+            text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict())
+
+            # transform
+            documents = self._transform(
+                index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()
+            )
+            # save segment
+            self._load_segments(dataset, dataset_document, documents)
+
+            # load
+            self._load(
+                index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents
+            )
+        except DocumentIsPausedError:
+            raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id))
+        except ProviderTokenNotInitError as e:
+            dataset_document.indexing_status = "error"
+            dataset_document.error = str(e.description)
+            dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+            db.session.commit()
+        except Exception as e:
+            logging.exception("consume document failed")
+            dataset_document.indexing_status = "error"
+            dataset_document.error = str(e)
+            dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+            db.session.commit()
+
+    def run_in_indexing_status(self, dataset_document: DatasetDocument):
+        """Run the indexing process when the index_status is indexing."""
+        try:
+            # get dataset
+            dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first()
+
+            if not dataset:
+                raise ValueError("no dataset found")
+
+            # get exist document_segment list and delete
+            document_segments = DocumentSegment.query.filter_by(
+                dataset_id=dataset.id, document_id=dataset_document.id
+            ).all()
+
+            documents = []
+            if document_segments:
+                for document_segment in document_segments:
+                    # transform segment to node
+                    if document_segment.status != "completed":
+                        document = Document(
+                            page_content=document_segment.content,
+                            metadata={
+                                "doc_id": document_segment.index_node_id,
+                                "doc_hash": document_segment.index_node_hash,
+                                "document_id": document_segment.document_id,
+                                "dataset_id": document_segment.dataset_id,
+                            },
+                        )
+                        if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
+                            child_chunks = document_segment.child_chunks
+                            if child_chunks:
+                                child_documents = []
+                                for child_chunk in child_chunks:
+                                    child_document = ChildDocument(
+                                        page_content=child_chunk.content,
+                                        metadata={
+                                            "doc_id": child_chunk.index_node_id,
+                                            "doc_hash": child_chunk.index_node_hash,
+                                            "document_id": document_segment.document_id,
+                                            "dataset_id": document_segment.dataset_id,
+                                        },
+                                    )
+                                    child_documents.append(child_document)
+                                document.children = child_documents
+                        documents.append(document)
+
+            # build index
+            # get the process rule
+            processing_rule = (
+                db.session.query(DatasetProcessRule)
+                .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
+                .first()
+            )
+
+            index_type = dataset_document.doc_form
+            index_processor = IndexProcessorFactory(index_type).init_index_processor()
+            self._load(
+                index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents
+            )
+        except DocumentIsPausedError:
+            raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id))
+        except ProviderTokenNotInitError as e:
+            dataset_document.indexing_status = "error"
+            dataset_document.error = str(e.description)
+            dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+            db.session.commit()
+        except Exception as e:
+            logging.exception("consume document failed")
+            dataset_document.indexing_status = "error"
+            dataset_document.error = str(e)
+            dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+            db.session.commit()
+
+    def indexing_estimate(
+        self,
+        tenant_id: str,
+        extract_settings: list[ExtractSetting],
+        tmp_processing_rule: dict,
+        doc_form: Optional[str] = None,
+        doc_language: str = "English",
+        dataset_id: Optional[str] = None,
+        indexing_technique: str = "economy",
+    ) -> IndexingEstimate:
+        """
+        Estimate the indexing for the document.
+        """
+        # check document limit
+        features = FeatureService.get_features(tenant_id)
+        if features.billing.enabled:
+            count = len(extract_settings)
+            batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT
+            if count > batch_upload_limit:
+                raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
+
+        embedding_model_instance = None
+        if dataset_id:
+            dataset = Dataset.query.filter_by(id=dataset_id).first()
+            if not dataset:
+                raise ValueError("Dataset not found.")
+            if dataset.indexing_technique == "high_quality" or indexing_technique == "high_quality":
+                if dataset.embedding_model_provider:
+                    embedding_model_instance = self.model_manager.get_model_instance(
+                        tenant_id=tenant_id,
+                        provider=dataset.embedding_model_provider,
+                        model_type=ModelType.TEXT_EMBEDDING,
+                        model=dataset.embedding_model,
+                    )
+                else:
+                    embedding_model_instance = self.model_manager.get_default_model_instance(
+                        tenant_id=tenant_id,
+                        model_type=ModelType.TEXT_EMBEDDING,
+                    )
+        else:
+            if indexing_technique == "high_quality":
+                embedding_model_instance = self.model_manager.get_default_model_instance(
+                    tenant_id=tenant_id,
+                    model_type=ModelType.TEXT_EMBEDDING,
+                )
+        preview_texts = []  # type: ignore
+
+        total_segments = 0
+        index_type = doc_form
+        index_processor = IndexProcessorFactory(index_type).init_index_processor()
+        for extract_setting in extract_settings:
+            # extract
+            processing_rule = DatasetProcessRule(
+                mode=tmp_processing_rule["mode"], rules=json.dumps(tmp_processing_rule["rules"])
+            )
+            text_docs = index_processor.extract(extract_setting, process_rule_mode=tmp_processing_rule["mode"])
+            documents = index_processor.transform(
+                text_docs,
+                embedding_model_instance=embedding_model_instance,
+                process_rule=processing_rule.to_dict(),
+                tenant_id=current_user.current_tenant_id,
+                doc_language=doc_language,
+                preview=True,
+            )
+            total_segments += len(documents)
+            for document in documents:
+                if len(preview_texts) < 10:
+                    if doc_form and doc_form == "qa_model":
+                        preview_detail = QAPreviewDetail(
+                            question=document.page_content, answer=document.metadata.get("answer") or ""
+                        )
+                        preview_texts.append(preview_detail)
+                    else:
+                        preview_detail = PreviewDetail(content=document.page_content)  # type: ignore
+                        if document.children:
+                            preview_detail.child_chunks = [child.page_content for child in document.children]  # type: ignore
+                        preview_texts.append(preview_detail)
+
+                # delete image files and related db records
+                image_upload_file_ids = get_image_upload_file_ids(document.page_content)
+                for upload_file_id in image_upload_file_ids:
+                    image_file = db.session.query(UploadFile).filter(UploadFile.id == upload_file_id).first()
+                    try:
+                        if image_file:
+                            storage.delete(image_file.key)
+                    except Exception:
+                        logging.exception(
+                            "Delete image_files failed while indexing_estimate, \
+                                          image_upload_file_is: {}".format(upload_file_id)
+                        )
+                    db.session.delete(image_file)
+
+        if doc_form and doc_form == "qa_model":
+            return IndexingEstimate(total_segments=total_segments * 20, qa_preview=preview_texts, preview=[])
+        return IndexingEstimate(total_segments=total_segments, preview=preview_texts)  # type: ignore
+
+    def _extract(
+        self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict
+    ) -> list[Document]:
+        # load file
+        if dataset_document.data_source_type not in {"upload_file", "notion_import", "website_crawl"}:
+            return []
+
+        data_source_info = dataset_document.data_source_info_dict
+        text_docs = []
+        if dataset_document.data_source_type == "upload_file":
+            if not data_source_info or "upload_file_id" not in data_source_info:
+                raise ValueError("no upload file found")
+
+            file_detail = (
+                db.session.query(UploadFile).filter(UploadFile.id == data_source_info["upload_file_id"]).one_or_none()
+            )
+
+            if file_detail:
+                extract_setting = ExtractSetting(
+                    datasource_type="upload_file", upload_file=file_detail, document_model=dataset_document.doc_form
+                )
+                text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
+        elif dataset_document.data_source_type == "notion_import":
+            if (
+                not data_source_info
+                or "notion_workspace_id" not in data_source_info
+                or "notion_page_id" not in data_source_info
+            ):
+                raise ValueError("no notion import info found")
+            extract_setting = ExtractSetting(
+                datasource_type="notion_import",
+                notion_info={
+                    "notion_workspace_id": data_source_info["notion_workspace_id"],
+                    "notion_obj_id": data_source_info["notion_page_id"],
+                    "notion_page_type": data_source_info["type"],
+                    "document": dataset_document,
+                    "tenant_id": dataset_document.tenant_id,
+                },
+                document_model=dataset_document.doc_form,
+            )
+            text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
+        elif dataset_document.data_source_type == "website_crawl":
+            if (
+                not data_source_info
+                or "provider" not in data_source_info
+                or "url" not in data_source_info
+                or "job_id" not in data_source_info
+            ):
+                raise ValueError("no website import info found")
+            extract_setting = ExtractSetting(
+                datasource_type="website_crawl",
+                website_info={
+                    "provider": data_source_info["provider"],
+                    "job_id": data_source_info["job_id"],
+                    "tenant_id": dataset_document.tenant_id,
+                    "url": data_source_info["url"],
+                    "mode": data_source_info["mode"],
+                    "only_main_content": data_source_info["only_main_content"],
+                },
+                document_model=dataset_document.doc_form,
+            )
+            text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
+        # update document status to splitting
+        self._update_document_index_status(
+            document_id=dataset_document.id,
+            after_indexing_status="splitting",
+            extra_update_params={
+                DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs),
+                DatasetDocument.parsing_completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+            },
+        )
+
+        # replace doc id to document model id
+        text_docs = cast(list[Document], text_docs)
+        for text_doc in text_docs:
+            if text_doc.metadata is not None:
+                text_doc.metadata["document_id"] = dataset_document.id
+                text_doc.metadata["dataset_id"] = dataset_document.dataset_id
+
+        return text_docs
+
+    @staticmethod
+    def filter_string(text):
+        text = re.sub(r"<\|", "<", text)
+        text = re.sub(r"\|>", ">", text)
+        text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\xEF\xBF\xBE]", "", text)
+        # Unicode  U+FFFE
+        text = re.sub("\ufffe", "", text)
+        return text
+
+    @staticmethod
+    def _get_splitter(
+        processing_rule_mode: str,
+        max_tokens: int,
+        chunk_overlap: int,
+        separator: str,
+        embedding_model_instance: Optional[ModelInstance],
+    ) -> TextSplitter:
+        """
+        Get the NodeParser object according to the processing rule.
+        """
+        if processing_rule_mode in ["custom", "hierarchical"]:
+            # The user-defined segmentation rule
+            max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH
+            if max_tokens < 50 or max_tokens > max_segmentation_tokens_length:
+                raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.")
+
+            if separator:
+                separator = separator.replace("\\n", "\n")
+
+            character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder(
+                chunk_size=max_tokens,
+                chunk_overlap=chunk_overlap,
+                fixed_separator=separator,
+                separators=["\n\n", "。", ". ", " ", ""],
+                embedding_model_instance=embedding_model_instance,
+            )
+        else:
+            # Automatic segmentation
+            automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"])
+            character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder(
+                chunk_size=automatic_rules["max_tokens"],
+                chunk_overlap=automatic_rules["chunk_overlap"],
+                separators=["\n\n", "。", ". ", " ", ""],
+                embedding_model_instance=embedding_model_instance,
+            )
+
+        return character_splitter  # type: ignore
+
+    def _split_to_documents_for_estimate(
+        self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule
+    ) -> list[Document]:
+        """
+        Split the text documents into nodes.
+        """
+        all_documents: list[Document] = []
+        for text_doc in text_docs:
+            # document clean
+            document_text = self._document_clean(text_doc.page_content, processing_rule)
+            text_doc.page_content = document_text
+
+            # parse document to nodes
+            documents = splitter.split_documents([text_doc])
+
+            split_documents = []
+            for document in documents:
+                if document.page_content is None or not document.page_content.strip():
+                    continue
+                if document.metadata is not None:
+                    doc_id = str(uuid.uuid4())
+                    hash = helper.generate_text_hash(document.page_content)
+                    document.metadata["doc_id"] = doc_id
+                    document.metadata["doc_hash"] = hash
+
+                split_documents.append(document)
+
+            all_documents.extend(split_documents)
+
+        return all_documents
+
+    @staticmethod
+    def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str:
+        """
+        Clean the document text according to the processing rules.
+        """
+        if processing_rule.mode == "automatic":
+            rules = DatasetProcessRule.AUTOMATIC_RULES
+        else:
+            rules = json.loads(processing_rule.rules) if processing_rule.rules else {}
+        document_text = CleanProcessor.clean(text, {"rules": rules})
+
+        return document_text
+
+    @staticmethod
+    def format_split_text(text: str) -> list[QAPreviewDetail]:
+        regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)"
+        matches = re.findall(regex, text, re.UNICODE)
+
+        return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a]
+
+    def _load(
+        self,
+        index_processor: BaseIndexProcessor,
+        dataset: Dataset,
+        dataset_document: DatasetDocument,
+        documents: list[Document],
+    ) -> None:
+        """
+        insert index and update document/segment status to completed
+        """
+
+        embedding_model_instance = None
+        if dataset.indexing_technique == "high_quality":
+            embedding_model_instance = self.model_manager.get_model_instance(
+                tenant_id=dataset.tenant_id,
+                provider=dataset.embedding_model_provider,
+                model_type=ModelType.TEXT_EMBEDDING,
+                model=dataset.embedding_model,
+            )
+
+        # chunk nodes by chunk size
+        indexing_start_at = time.perf_counter()
+        tokens = 0
+        if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX:
+            # create keyword index
+            create_keyword_thread = threading.Thread(
+                target=self._process_keyword_index,
+                args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents),  # type: ignore
+            )
+            create_keyword_thread.start()
+
+        max_workers = 10
+        if dataset.indexing_technique == "high_quality":
+            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+                futures = []
+
+                # Distribute documents into multiple groups based on the hash values of page_content
+                # This is done to prevent multiple threads from processing the same document,
+                # Thereby avoiding potential database insertion deadlocks
+                document_groups: list[list[Document]] = [[] for _ in range(max_workers)]
+                for document in documents:
+                    hash = helper.generate_text_hash(document.page_content)
+                    group_index = int(hash, 16) % max_workers
+                    document_groups[group_index].append(document)
+                for chunk_documents in document_groups:
+                    if len(chunk_documents) == 0:
+                        continue
+                    futures.append(
+                        executor.submit(
+                            self._process_chunk,
+                            current_app._get_current_object(),  # type: ignore
+                            index_processor,
+                            chunk_documents,
+                            dataset,
+                            dataset_document,
+                            embedding_model_instance,
+                        )
+                    )
+
+                for future in futures:
+                    tokens += future.result()
+        if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX:
+            create_keyword_thread.join()
+        indexing_end_at = time.perf_counter()
+
+        # update document status to completed
+        self._update_document_index_status(
+            document_id=dataset_document.id,
+            after_indexing_status="completed",
+            extra_update_params={
+                DatasetDocument.tokens: tokens,
+                DatasetDocument.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+                DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
+                DatasetDocument.error: None,
+            },
+        )
+
+    @staticmethod
+    def _process_keyword_index(flask_app, dataset_id, document_id, documents):
+        with flask_app.app_context():
+            dataset = Dataset.query.filter_by(id=dataset_id).first()
+            if not dataset:
+                raise ValueError("no dataset found")
+            keyword = Keyword(dataset)
+            keyword.create(documents)
+            if dataset.indexing_technique != "high_quality":
+                document_ids = [document.metadata["doc_id"] for document in documents]
+                db.session.query(DocumentSegment).filter(
+                    DocumentSegment.document_id == document_id,
+                    DocumentSegment.dataset_id == dataset_id,
+                    DocumentSegment.index_node_id.in_(document_ids),
+                    DocumentSegment.status == "indexing",
+                ).update(
+                    {
+                        DocumentSegment.status: "completed",
+                        DocumentSegment.enabled: True,
+                        DocumentSegment.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+                    }
+                )
+
+                db.session.commit()
+
+    def _process_chunk(
+        self, flask_app, index_processor, chunk_documents, dataset, dataset_document, embedding_model_instance
+    ):
+        with flask_app.app_context():
+            # check document is paused
+            self._check_document_paused_status(dataset_document.id)
+
+            tokens = 0
+            if embedding_model_instance:
+                tokens += sum(
+                    embedding_model_instance.get_text_embedding_num_tokens([document.page_content])
+                    for document in chunk_documents
+                )
+
+            # load index
+            index_processor.load(dataset, chunk_documents, with_keywords=False)
+
+            document_ids = [document.metadata["doc_id"] for document in chunk_documents]
+            db.session.query(DocumentSegment).filter(
+                DocumentSegment.document_id == dataset_document.id,
+                DocumentSegment.dataset_id == dataset.id,
+                DocumentSegment.index_node_id.in_(document_ids),
+                DocumentSegment.status == "indexing",
+            ).update(
+                {
+                    DocumentSegment.status: "completed",
+                    DocumentSegment.enabled: True,
+                    DocumentSegment.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+                }
+            )
+
+            db.session.commit()
+
+            return tokens
+
+    @staticmethod
+    def _check_document_paused_status(document_id: str):
+        indexing_cache_key = "document_{}_is_paused".format(document_id)
+        result = redis_client.get(indexing_cache_key)
+        if result:
+            raise DocumentIsPausedError()
+
+    @staticmethod
+    def _update_document_index_status(
+        document_id: str, after_indexing_status: str, extra_update_params: Optional[dict] = None
+    ) -> None:
+        """
+        Update the document indexing status.
+        """
+        count = DatasetDocument.query.filter_by(id=document_id, is_paused=True).count()
+        if count > 0:
+            raise DocumentIsPausedError()
+        document = DatasetDocument.query.filter_by(id=document_id).first()
+        if not document:
+            raise DocumentIsDeletedPausedError()
+
+        update_params = {DatasetDocument.indexing_status: after_indexing_status}
+
+        if extra_update_params:
+            update_params.update(extra_update_params)
+
+        DatasetDocument.query.filter_by(id=document_id).update(update_params)
+        db.session.commit()
+
+    @staticmethod
+    def _update_segments_by_document(dataset_document_id: str, update_params: dict) -> None:
+        """
+        Update the document segment by document id.
+        """
+        DocumentSegment.query.filter_by(document_id=dataset_document_id).update(update_params)
+        db.session.commit()
+
+    def _transform(
+        self,
+        index_processor: BaseIndexProcessor,
+        dataset: Dataset,
+        text_docs: list[Document],
+        doc_language: str,
+        process_rule: dict,
+    ) -> list[Document]:
+        # get embedding model instance
+        embedding_model_instance = None
+        if dataset.indexing_technique == "high_quality":
+            if dataset.embedding_model_provider:
+                embedding_model_instance = self.model_manager.get_model_instance(
+                    tenant_id=dataset.tenant_id,
+                    provider=dataset.embedding_model_provider,
+                    model_type=ModelType.TEXT_EMBEDDING,
+                    model=dataset.embedding_model,
+                )
+            else:
+                embedding_model_instance = self.model_manager.get_default_model_instance(
+                    tenant_id=dataset.tenant_id,
+                    model_type=ModelType.TEXT_EMBEDDING,
+                )
+
+        documents = index_processor.transform(
+            text_docs,
+            embedding_model_instance=embedding_model_instance,
+            process_rule=process_rule,
+            tenant_id=dataset.tenant_id,
+            doc_language=doc_language,
+        )
+
+        return documents
+
+    def _load_segments(self, dataset, dataset_document, documents):
+        # save node to document segment
+        doc_store = DatasetDocumentStore(
+            dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id
+        )
+
+        # add document segments
+        doc_store.add_documents(docs=documents, save_child=dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX)
+
+        # update document status to indexing
+        cur_time = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+        self._update_document_index_status(
+            document_id=dataset_document.id,
+            after_indexing_status="indexing",
+            extra_update_params={
+                DatasetDocument.cleaning_completed_at: cur_time,
+                DatasetDocument.splitting_completed_at: cur_time,
+            },
+        )
+
+        # update segment status to indexing
+        self._update_segments_by_document(
+            dataset_document_id=dataset_document.id,
+            update_params={
+                DocumentSegment.status: "indexing",
+                DocumentSegment.indexing_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+            },
+        )
+        pass
+
+
+class DocumentIsPausedError(Exception):
+    pass
+
+
+class DocumentIsDeletedPausedError(Exception):
+    pass

api/core/model_manager.py

@@ -0,0 +1,559 @@
+import logging
+from collections.abc import Callable, Generator, Iterable, Sequence
+from typing import IO, Any, Optional, Union, cast
+
+from configs import dify_config
+from core.entities.embedding_type import EmbeddingInputType
+from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
+from core.entities.provider_entities import ModelLoadBalancingConfiguration
+from core.errors.error import ProviderTokenNotInitError
+from core.model_runtime.callbacks.base_callback import Callback
+from core.model_runtime.entities.llm_entities import LLMResult
+from core.model_runtime.entities.message_entities import PromptMessage, PromptMessageTool
+from core.model_runtime.entities.model_entities import ModelType
+from core.model_runtime.entities.rerank_entities import RerankResult
+from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
+from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeConnectionError, InvokeRateLimitError
+from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
+from core.model_runtime.model_providers.__base.moderation_model import ModerationModel
+from core.model_runtime.model_providers.__base.rerank_model import RerankModel
+from core.model_runtime.model_providers.__base.speech2text_model import Speech2TextModel
+from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
+from core.model_runtime.model_providers.__base.tts_model import TTSModel
+from core.provider_manager import ProviderManager
+from extensions.ext_redis import redis_client
+from models.provider import ProviderType
+
+logger = logging.getLogger(__name__)
+
+
+class ModelInstance:
+    """
+    Model instance class
+    """
+
+    def __init__(self, provider_model_bundle: ProviderModelBundle, model: str) -> None:
+        self.provider_model_bundle = provider_model_bundle
+        self.model = model
+        self.provider = provider_model_bundle.configuration.provider.provider
+        self.credentials = self._fetch_credentials_from_bundle(provider_model_bundle, model)
+        self.model_type_instance = self.provider_model_bundle.model_type_instance
+        self.load_balancing_manager = self._get_load_balancing_manager(
+            configuration=provider_model_bundle.configuration,
+            model_type=provider_model_bundle.model_type_instance.model_type,
+            model=model,
+            credentials=self.credentials,
+        )
+
+    @staticmethod
+    def _fetch_credentials_from_bundle(provider_model_bundle: ProviderModelBundle, model: str) -> dict:
+        """
+        Fetch credentials from provider model bundle
+        :param provider_model_bundle: provider model bundle
+        :param model: model name
+        :return:
+        """
+        configuration = provider_model_bundle.configuration
+        model_type = provider_model_bundle.model_type_instance.model_type
+        credentials = configuration.get_current_credentials(model_type=model_type, model=model)
+
+        if credentials is None:
+            raise ProviderTokenNotInitError(f"Model {model} credentials is not initialized.")
+
+        return credentials
+
+    @staticmethod
+    def _get_load_balancing_manager(
+        configuration: ProviderConfiguration, model_type: ModelType, model: str, credentials: dict
+    ) -> Optional["LBModelManager"]:
+        """
+        Get load balancing model credentials
+        :param configuration: provider configuration
+        :param model_type: model type
+        :param model: model name
+        :param credentials: model credentials
+        :return:
+        """
+        if configuration.model_settings and configuration.using_provider_type == ProviderType.CUSTOM:
+            current_model_setting = None
+            # check if model is disabled by admin
+            for model_setting in configuration.model_settings:
+                if model_setting.model_type == model_type and model_setting.model == model:
+                    current_model_setting = model_setting
+                    break
+
+            # check if load balancing is enabled
+            if current_model_setting and current_model_setting.load_balancing_configs:
+                # use load balancing proxy to choose credentials
+                lb_model_manager = LBModelManager(
+                    tenant_id=configuration.tenant_id,
+                    provider=configuration.provider.provider,
+                    model_type=model_type,
+                    model=model,
+                    load_balancing_configs=current_model_setting.load_balancing_configs,
+                    managed_credentials=credentials if configuration.custom_configuration.provider else None,
+                )
+
+                return lb_model_manager
+
+        return None
+
+    def invoke_llm(
+        self,
+        prompt_messages: Sequence[PromptMessage],
+        model_parameters: Optional[dict] = None,
+        tools: Sequence[PromptMessageTool] | None = None,
+        stop: Optional[Sequence[str]] = None,
+        stream: bool = True,
+        user: Optional[str] = None,
+        callbacks: Optional[list[Callback]] = None,
+    ) -> Union[LLMResult, Generator]:
+        """
+        Invoke large language model
+
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        :param callbacks: callbacks
+        :return: full response or stream response chunk generator result
+        """
+        if not isinstance(self.model_type_instance, LargeLanguageModel):
+            raise Exception("Model type instance is not LargeLanguageModel")
+
+        self.model_type_instance = cast(LargeLanguageModel, self.model_type_instance)
+        return cast(
+            Union[LLMResult, Generator],
+            self._round_robin_invoke(
+                function=self.model_type_instance.invoke,
+                model=self.model,
+                credentials=self.credentials,
+                prompt_messages=prompt_messages,
+                model_parameters=model_parameters,
+                tools=tools,
+                stop=stop,
+                stream=stream,
+                user=user,
+                callbacks=callbacks,
+            ),
+        )
+
+    def get_llm_num_tokens(
+        self, prompt_messages: list[PromptMessage], tools: Optional[list[PromptMessageTool]] = None
+    ) -> int:
+        """
+        Get number of tokens for llm
+
+        :param prompt_messages: prompt messages
+        :param tools: tools for tool calling
+        :return:
+        """
+        if not isinstance(self.model_type_instance, LargeLanguageModel):
+            raise Exception("Model type instance is not LargeLanguageModel")
+
+        self.model_type_instance = cast(LargeLanguageModel, self.model_type_instance)
+        return cast(
+            int,
+            self._round_robin_invoke(
+                function=self.model_type_instance.get_num_tokens,
+                model=self.model,
+                credentials=self.credentials,
+                prompt_messages=prompt_messages,
+                tools=tools,
+            ),
+        )
+
+    def invoke_text_embedding(
+        self, texts: list[str], user: Optional[str] = None, input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT
+    ) -> TextEmbeddingResult:
+        """
+        Invoke large language model
+
+        :param texts: texts to embed
+        :param user: unique user id
+        :param input_type: input type
+        :return: embeddings result
+        """
+        if not isinstance(self.model_type_instance, TextEmbeddingModel):
+            raise Exception("Model type instance is not TextEmbeddingModel")
+
+        self.model_type_instance = cast(TextEmbeddingModel, self.model_type_instance)
+        return cast(
+            TextEmbeddingResult,
+            self._round_robin_invoke(
+                function=self.model_type_instance.invoke,
+                model=self.model,
+                credentials=self.credentials,
+                texts=texts,
+                user=user,
+                input_type=input_type,
+            ),
+        )
+
+    def get_text_embedding_num_tokens(self, texts: list[str]) -> int:
+        """
+        Get number of tokens for text embedding
+
+        :param texts: texts to embed
+        :return:
+        """
+        if not isinstance(self.model_type_instance, TextEmbeddingModel):
+            raise Exception("Model type instance is not TextEmbeddingModel")
+
+        self.model_type_instance = cast(TextEmbeddingModel, self.model_type_instance)
+        return cast(
+            int,
+            self._round_robin_invoke(
+                function=self.model_type_instance.get_num_tokens,
+                model=self.model,
+                credentials=self.credentials,
+                texts=texts,
+            ),
+        )
+
+    def invoke_rerank(
+        self,
+        query: str,
+        docs: list[str],
+        score_threshold: Optional[float] = None,
+        top_n: Optional[int] = None,
+        user: Optional[str] = None,
+    ) -> RerankResult:
+        """
+        Invoke rerank model
+
+        :param query: search query
+        :param docs: docs for reranking
+        :param score_threshold: score threshold
+        :param top_n: top n
+        :param user: unique user id
+        :return: rerank result
+        """
+        if not isinstance(self.model_type_instance, RerankModel):
+            raise Exception("Model type instance is not RerankModel")
+
+        self.model_type_instance = cast(RerankModel, self.model_type_instance)
+        return cast(
+            RerankResult,
+            self._round_robin_invoke(
+                function=self.model_type_instance.invoke,
+                model=self.model,
+                credentials=self.credentials,
+                query=query,
+                docs=docs,
+                score_threshold=score_threshold,
+                top_n=top_n,
+                user=user,
+            ),
+        )
+
+    def invoke_moderation(self, text: str, user: Optional[str] = None) -> bool:
+        """
+        Invoke moderation model
+
+        :param text: text to moderate
+        :param user: unique user id
+        :return: false if text is safe, true otherwise
+        """
+        if not isinstance(self.model_type_instance, ModerationModel):
+            raise Exception("Model type instance is not ModerationModel")
+
+        self.model_type_instance = cast(ModerationModel, self.model_type_instance)
+        return cast(
+            bool,
+            self._round_robin_invoke(
+                function=self.model_type_instance.invoke,
+                model=self.model,
+                credentials=self.credentials,
+                text=text,
+                user=user,
+            ),
+        )
+
+    def invoke_speech2text(self, file: IO[bytes], user: Optional[str] = None) -> str:
+        """
+        Invoke large language model
+
+        :param file: audio file
+        :param user: unique user id
+        :return: text for given audio file
+        """
+        if not isinstance(self.model_type_instance, Speech2TextModel):
+            raise Exception("Model type instance is not Speech2TextModel")
+
+        self.model_type_instance = cast(Speech2TextModel, self.model_type_instance)
+        return cast(
+            str,
+            self._round_robin_invoke(
+                function=self.model_type_instance.invoke,
+                model=self.model,
+                credentials=self.credentials,
+                file=file,
+                user=user,
+            ),
+        )
+
+    def invoke_tts(self, content_text: str, tenant_id: str, voice: str, user: Optional[str] = None) -> Iterable[bytes]:
+        """
+        Invoke large language tts model
+
+        :param content_text: text content to be translated
+        :param tenant_id: user tenant id
+        :param voice: model timbre
+        :param user: unique user id
+        :return: text for given audio file
+        """
+        if not isinstance(self.model_type_instance, TTSModel):
+            raise Exception("Model type instance is not TTSModel")
+
+        self.model_type_instance = cast(TTSModel, self.model_type_instance)
+        return cast(
+            Iterable[bytes],
+            self._round_robin_invoke(
+                function=self.model_type_instance.invoke,
+                model=self.model,
+                credentials=self.credentials,
+                content_text=content_text,
+                user=user,
+                tenant_id=tenant_id,
+                voice=voice,
+            ),
+        )
+
+    def _round_robin_invoke(self, function: Callable[..., Any], *args, **kwargs) -> Any:
+        """
+        Round-robin invoke
+        :param function: function to invoke
+        :param args: function args
+        :param kwargs: function kwargs
+        :return:
+        """
+        if not self.load_balancing_manager:
+            return function(*args, **kwargs)
+
+        last_exception: Union[InvokeRateLimitError, InvokeAuthorizationError, InvokeConnectionError, None] = None
+        while True:
+            lb_config = self.load_balancing_manager.fetch_next()
+            if not lb_config:
+                if not last_exception:
+                    raise ProviderTokenNotInitError("Model credentials is not initialized.")
+                else:
+                    raise last_exception
+
+            try:
+                if "credentials" in kwargs:
+                    del kwargs["credentials"]
+                return function(*args, **kwargs, credentials=lb_config.credentials)
+            except InvokeRateLimitError as e:
+                # expire in 60 seconds
+                self.load_balancing_manager.cooldown(lb_config, expire=60)
+                last_exception = e
+                continue
+            except (InvokeAuthorizationError, InvokeConnectionError) as e:
+                # expire in 10 seconds
+                self.load_balancing_manager.cooldown(lb_config, expire=10)
+                last_exception = e
+                continue
+            except Exception as e:
+                raise e
+
+    def get_tts_voices(self, language: Optional[str] = None) -> list:
+        """
+        Invoke large language tts model voices
+
+        :param language: tts language
+        :return: tts model voices
+        """
+        if not isinstance(self.model_type_instance, TTSModel):
+            raise Exception("Model type instance is not TTSModel")
+
+        self.model_type_instance = cast(TTSModel, self.model_type_instance)
+        return self.model_type_instance.get_tts_model_voices(
+            model=self.model, credentials=self.credentials, language=language
+        )
+
+
+class ModelManager:
+    def __init__(self) -> None:
+        self._provider_manager = ProviderManager()
+
+    def get_model_instance(self, tenant_id: str, provider: str, model_type: ModelType, model: str) -> ModelInstance:
+        """
+        Get model instance
+        :param tenant_id: tenant id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :return:
+        """
+        if not provider:
+            return self.get_default_model_instance(tenant_id, model_type)
+
+        provider_model_bundle = self._provider_manager.get_provider_model_bundle(
+            tenant_id=tenant_id, provider=provider, model_type=model_type
+        )
+
+        return ModelInstance(provider_model_bundle, model)
+
+    def get_default_provider_model_name(self, tenant_id: str, model_type: ModelType) -> tuple[str, str]:
+        """
+        Return first provider and the first model in the provider
+        :param tenant_id: tenant id
+        :param model_type: model type
+        :return: provider name, model name
+        """
+        return self._provider_manager.get_first_provider_first_model(tenant_id, model_type)
+
+    def get_default_model_instance(self, tenant_id: str, model_type: ModelType) -> ModelInstance:
+        """
+        Get default model instance
+        :param tenant_id: tenant id
+        :param model_type: model type
+        :return:
+        """
+        default_model_entity = self._provider_manager.get_default_model(tenant_id=tenant_id, model_type=model_type)
+
+        if not default_model_entity:
+            raise ProviderTokenNotInitError(f"Default model not found for {model_type}")
+
+        return self.get_model_instance(
+            tenant_id=tenant_id,
+            provider=default_model_entity.provider.provider,
+            model_type=model_type,
+            model=default_model_entity.model,
+        )
+
+
+class LBModelManager:
+    def __init__(
+        self,
+        tenant_id: str,
+        provider: str,
+        model_type: ModelType,
+        model: str,
+        load_balancing_configs: list[ModelLoadBalancingConfiguration],
+        managed_credentials: Optional[dict] = None,
+    ) -> None:
+        """
+        Load balancing model manager
+        :param tenant_id: tenant_id
+        :param provider: provider
+        :param model_type: model_type
+        :param model: model name
+        :param load_balancing_configs: all load balancing configurations
+        :param managed_credentials: credentials if load balancing configuration name is __inherit__
+        """
+        self._tenant_id = tenant_id
+        self._provider = provider
+        self._model_type = model_type
+        self._model = model
+        self._load_balancing_configs = load_balancing_configs
+
+        for load_balancing_config in self._load_balancing_configs[:]:  # Iterate over a shallow copy of the list
+            if load_balancing_config.name == "__inherit__":
+                if not managed_credentials:
+                    # remove __inherit__ if managed credentials is not provided
+                    self._load_balancing_configs.remove(load_balancing_config)
+                else:
+                    load_balancing_config.credentials = managed_credentials
+
+    def fetch_next(self) -> Optional[ModelLoadBalancingConfiguration]:
+        """
+        Get next model load balancing config
+        Strategy: Round Robin
+        :return:
+        """
+        cache_key = "model_lb_index:{}:{}:{}:{}".format(
+            self._tenant_id, self._provider, self._model_type.value, self._model
+        )
+
+        cooldown_load_balancing_configs = []
+        max_index = len(self._load_balancing_configs)
+
+        while True:
+            current_index = redis_client.incr(cache_key)
+            current_index = cast(int, current_index)
+            if current_index >= 10000000:
+                current_index = 1
+                redis_client.set(cache_key, current_index)
+
+            redis_client.expire(cache_key, 3600)
+            if current_index > max_index:
+                current_index = current_index % max_index
+
+            real_index = current_index - 1
+            if real_index > max_index:
+                real_index = 0
+
+            config: ModelLoadBalancingConfiguration = self._load_balancing_configs[real_index]
+
+            if self.in_cooldown(config):
+                cooldown_load_balancing_configs.append(config)
+                if len(cooldown_load_balancing_configs) >= len(self._load_balancing_configs):
+                    # all configs are in cooldown
+                    return None
+
+                continue
+
+            if dify_config.DEBUG:
+                logger.info(
+                    f"Model LB\nid: {config.id}\nname:{config.name}\n"
+                    f"tenant_id: {self._tenant_id}\nprovider: {self._provider}\n"
+                    f"model_type: {self._model_type.value}\nmodel: {self._model}"
+                )
+
+            return config
+
+        return None
+
+    def cooldown(self, config: ModelLoadBalancingConfiguration, expire: int = 60) -> None:
+        """
+        Cooldown model load balancing config
+        :param config: model load balancing config
+        :param expire: cooldown time
+        :return:
+        """
+        cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format(
+            self._tenant_id, self._provider, self._model_type.value, self._model, config.id
+        )
+
+        redis_client.setex(cooldown_cache_key, expire, "true")
+
+    def in_cooldown(self, config: ModelLoadBalancingConfiguration) -> bool:
+        """
+        Check if model load balancing config is in cooldown
+        :param config: model load balancing config
+        :return:
+        """
+        cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format(
+            self._tenant_id, self._provider, self._model_type.value, self._model, config.id
+        )
+
+        res: bool = redis_client.exists(cooldown_cache_key)
+        return res
+
+    @staticmethod
+    def get_config_in_cooldown_and_ttl(
+        tenant_id: str, provider: str, model_type: ModelType, model: str, config_id: str
+    ) -> tuple[bool, int]:
+        """
+        Get model load balancing config is in cooldown and ttl
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :param config_id: model load balancing config id
+        :return:
+        """
+        cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format(
+            tenant_id, provider, model_type.value, model, config_id
+        )
+
+        ttl = redis_client.ttl(cooldown_cache_key)
+        if ttl == -2:
+            return False, 0
+
+        ttl = cast(int, ttl)
+        return True, ttl

api/core/moderation/api/__builtin__

@@ -0,0 +1 @@
+3

api/core/moderation/api/api.py

@@ -0,0 +1,96 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from core.extension.api_based_extension_requestor import APIBasedExtensionPoint, APIBasedExtensionRequestor
+from core.helper.encrypter import decrypt_token
+from core.moderation.base import Moderation, ModerationAction, ModerationInputsResult, ModerationOutputsResult
+from extensions.ext_database import db
+from models.api_based_extension import APIBasedExtension
+
+
+class ModerationInputParams(BaseModel):
+    app_id: str = ""
+    inputs: dict = {}
+    query: str = ""
+
+
+class ModerationOutputParams(BaseModel):
+    app_id: str = ""
+    text: str
+
+
+class ApiModeration(Moderation):
+    name: str = "api"
+
+    @classmethod
+    def validate_config(cls, tenant_id: str, config: dict) -> None:
+        """
+        Validate the incoming form config data.
+
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        cls._validate_inputs_and_outputs_config(config, False)
+
+        api_based_extension_id = config.get("api_based_extension_id")
+        if not api_based_extension_id:
+            raise ValueError("api_based_extension_id is required")
+
+        extension = cls._get_api_based_extension(tenant_id, api_based_extension_id)
+        if not extension:
+            raise ValueError("API-based Extension not found. Please check it again.")
+
+    def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult:
+        flagged = False
+        preset_response = ""
+        if self.config is None:
+            raise ValueError("The config is not set.")
+
+        if self.config["inputs_config"]["enabled"]:
+            params = ModerationInputParams(app_id=self.app_id, inputs=inputs, query=query)
+
+            result = self._get_config_by_requestor(APIBasedExtensionPoint.APP_MODERATION_INPUT, params.model_dump())
+            return ModerationInputsResult(**result)
+
+        return ModerationInputsResult(
+            flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response
+        )
+
+    def moderation_for_outputs(self, text: str) -> ModerationOutputsResult:
+        flagged = False
+        preset_response = ""
+        if self.config is None:
+            raise ValueError("The config is not set.")
+
+        if self.config["outputs_config"]["enabled"]:
+            params = ModerationOutputParams(app_id=self.app_id, text=text)
+
+            result = self._get_config_by_requestor(APIBasedExtensionPoint.APP_MODERATION_OUTPUT, params.model_dump())
+            return ModerationOutputsResult(**result)
+
+        return ModerationOutputsResult(
+            flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response
+        )
+
+    def _get_config_by_requestor(self, extension_point: APIBasedExtensionPoint, params: dict) -> dict:
+        if self.config is None:
+            raise ValueError("The config is not set.")
+        extension = self._get_api_based_extension(self.tenant_id, self.config.get("api_based_extension_id", ""))
+        if not extension:
+            raise ValueError("API-based Extension not found. Please check it again.")
+        requestor = APIBasedExtensionRequestor(extension.api_endpoint, decrypt_token(self.tenant_id, extension.api_key))
+
+        result = requestor.request(extension_point, params)
+        return result
+
+    @staticmethod
+    def _get_api_based_extension(tenant_id: str, api_based_extension_id: str) -> Optional[APIBasedExtension]:
+        extension = (
+            db.session.query(APIBasedExtension)
+            .filter(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id)
+            .first()
+        )
+
+        return extension

api/core/moderation/base.py

@@ -0,0 +1,115 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel
+
+from core.extension.extensible import Extensible, ExtensionModule
+
+
+class ModerationAction(Enum):
+    DIRECT_OUTPUT = "direct_output"
+    OVERRIDDEN = "overridden"
+
+
+class ModerationInputsResult(BaseModel):
+    flagged: bool = False
+    action: ModerationAction
+    preset_response: str = ""
+    inputs: dict = {}
+    query: str = ""
+
+
+class ModerationOutputsResult(BaseModel):
+    flagged: bool = False
+    action: ModerationAction
+    preset_response: str = ""
+    text: str = ""
+
+
+class Moderation(Extensible, ABC):
+    """
+    The base class of moderation.
+    """
+
+    module: ExtensionModule = ExtensionModule.MODERATION
+
+    def __init__(self, app_id: str, tenant_id: str, config: Optional[dict] = None) -> None:
+        super().__init__(tenant_id, config)
+        self.app_id = app_id
+
+    @classmethod
+    @abstractmethod
+    def validate_config(cls, tenant_id: str, config: dict) -> None:
+        """
+        Validate the incoming form config data.
+
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult:
+        """
+        Moderation for inputs.
+        After the user inputs, this method will be called to perform sensitive content review
+        on the user inputs and return the processed results.
+
+        :param inputs: user inputs
+        :param query: query string (required in chat app)
+        :return:
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def moderation_for_outputs(self, text: str) -> ModerationOutputsResult:
+        """
+        Moderation for outputs.
+        When LLM outputs content, the front end will pass the output content (may be segmented)
+        to this method for sensitive content review, and the output content will be shielded if the review fails.
+
+        :param text: LLM output content
+        :return:
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def _validate_inputs_and_outputs_config(cls, config: dict, is_preset_response_required: bool) -> None:
+        # inputs_config
+        inputs_config = config.get("inputs_config")
+        if not isinstance(inputs_config, dict):
+            raise ValueError("inputs_config must be a dict")
+
+        # outputs_config
+        outputs_config = config.get("outputs_config")
+        if not isinstance(outputs_config, dict):
+            raise ValueError("outputs_config must be a dict")
+
+        inputs_config_enabled = inputs_config.get("enabled")
+        outputs_config_enabled = outputs_config.get("enabled")
+        if not inputs_config_enabled and not outputs_config_enabled:
+            raise ValueError("At least one of inputs_config or outputs_config must be enabled")
+
+        # preset_response
+        if not is_preset_response_required:
+            return
+
+        if inputs_config_enabled:
+            if not inputs_config.get("preset_response"):
+                raise ValueError("inputs_config.preset_response is required")
+
+            if len(inputs_config.get("preset_response", 0)) > 100:
+                raise ValueError("inputs_config.preset_response must be less than 100 characters")
+
+        if outputs_config_enabled:
+            if not outputs_config.get("preset_response"):
+                raise ValueError("outputs_config.preset_response is required")
+
+            if len(outputs_config.get("preset_response", 0)) > 100:
+                raise ValueError("outputs_config.preset_response must be less than 100 characters")
+
+
+class ModerationError(Exception):
+    pass

api/core/moderation/factory.py

@@ -0,0 +1,49 @@
+from core.extension.extensible import ExtensionModule
+from core.moderation.base import Moderation, ModerationInputsResult, ModerationOutputsResult
+from extensions.ext_code_based_extension import code_based_extension
+
+
+class ModerationFactory:
+    __extension_instance: Moderation
+
+    def __init__(self, name: str, app_id: str, tenant_id: str, config: dict) -> None:
+        extension_class = code_based_extension.extension_class(ExtensionModule.MODERATION, name)
+        self.__extension_instance = extension_class(app_id, tenant_id, config)
+
+    @classmethod
+    def validate_config(cls, name: str, tenant_id: str, config: dict) -> None:
+        """
+        Validate the incoming form config data.
+
+        :param name: the name of extension
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        code_based_extension.validate_form_schema(ExtensionModule.MODERATION, name, config)
+        extension_class = code_based_extension.extension_class(ExtensionModule.MODERATION, name)
+        # FIXME: mypy error, try to fix it instead of using type: ignore
+        extension_class.validate_config(tenant_id, config)  # type: ignore
+
+    def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult:
+        """
+        Moderation for inputs.
+        After the user inputs, this method will be called to perform sensitive content review
+        on the user inputs and return the processed results.
+
+        :param inputs: user inputs
+        :param query: query string (required in chat app)
+        :return:
+        """
+        return self.__extension_instance.moderation_for_inputs(inputs, query)
+
+    def moderation_for_outputs(self, text: str) -> ModerationOutputsResult:
+        """
+        Moderation for outputs.
+        When LLM outputs content, the front end will pass the output content (may be segmented)
+        to this method for sensitive content review, and the output content will be shielded if the review fails.
+
+        :param text: LLM output content
+        :return:
+        """
+        return self.__extension_instance.moderation_for_outputs(text)

api/core/moderation/input_moderation.py

@@ -0,0 +1,71 @@
+import logging
+from collections.abc import Mapping
+from typing import Any, Optional
+
+from core.app.app_config.entities import AppConfig
+from core.moderation.base import ModerationAction, ModerationError
+from core.moderation.factory import ModerationFactory
+from core.ops.entities.trace_entity import TraceTaskName
+from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
+from core.ops.utils import measure_time
+
+logger = logging.getLogger(__name__)
+
+
+class InputModeration:
+    def check(
+        self,
+        app_id: str,
+        tenant_id: str,
+        app_config: AppConfig,
+        inputs: Mapping[str, Any],
+        query: str,
+        message_id: str,
+        trace_manager: Optional[TraceQueueManager] = None,
+    ) -> tuple[bool, Mapping[str, Any], str]:
+        """
+        Process sensitive_word_avoidance.
+        :param app_id: app id
+        :param tenant_id: tenant id
+        :param app_config: app config
+        :param inputs: inputs
+        :param query: query
+        :param message_id: message id
+        :param trace_manager: trace manager
+        :return:
+        """
+        inputs = dict(inputs)
+        if not app_config.sensitive_word_avoidance:
+            return False, inputs, query
+
+        sensitive_word_avoidance_config = app_config.sensitive_word_avoidance
+        moderation_type = sensitive_word_avoidance_config.type
+
+        moderation_factory = ModerationFactory(
+            name=moderation_type, app_id=app_id, tenant_id=tenant_id, config=sensitive_word_avoidance_config.config
+        )
+
+        with measure_time() as timer:
+            moderation_result = moderation_factory.moderation_for_inputs(inputs, query)
+
+        if trace_manager:
+            trace_manager.add_trace_task(
+                TraceTask(
+                    TraceTaskName.MODERATION_TRACE,
+                    message_id=message_id,
+                    moderation_result=moderation_result,
+                    inputs=inputs,
+                    timer=timer,
+                )
+            )
+
+        if not moderation_result.flagged:
+            return False, inputs, query
+
+        if moderation_result.action == ModerationAction.DIRECT_OUTPUT:
+            raise ModerationError(moderation_result.preset_response)
+        elif moderation_result.action == ModerationAction.OVERRIDDEN:
+            inputs = moderation_result.inputs
+            query = moderation_result.query
+
+        return True, inputs, query

api/core/moderation/keywords/__builtin__

@@ -0,0 +1 @@
+2

api/core/moderation/keywords/keywords.py

@@ -0,0 +1,73 @@
+from collections.abc import Sequence
+from typing import Any
+
+from core.moderation.base import Moderation, ModerationAction, ModerationInputsResult, ModerationOutputsResult
+
+
+class KeywordsModeration(Moderation):
+    name: str = "keywords"
+
+    @classmethod
+    def validate_config(cls, tenant_id: str, config: dict) -> None:
+        """
+        Validate the incoming form config data.
+
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        cls._validate_inputs_and_outputs_config(config, True)
+
+        if not config.get("keywords"):
+            raise ValueError("keywords is required")
+
+        if len(config.get("keywords", [])) > 10000:
+            raise ValueError("keywords length must be less than 10000")
+
+        keywords_row_len = config["keywords"].split("\n")
+        if len(keywords_row_len) > 100:
+            raise ValueError("the number of rows for the keywords must be less than 100")
+
+    def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult:
+        flagged = False
+        preset_response = ""
+        if self.config is None:
+            raise ValueError("The config is not set.")
+
+        if self.config["inputs_config"]["enabled"]:
+            preset_response = self.config["inputs_config"]["preset_response"]
+
+            if query:
+                inputs["query__"] = query
+
+            # Filter out empty values
+            keywords_list = [keyword for keyword in self.config["keywords"].split("\n") if keyword]
+
+            flagged = self._is_violated(inputs, keywords_list)
+
+        return ModerationInputsResult(
+            flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response
+        )
+
+    def moderation_for_outputs(self, text: str) -> ModerationOutputsResult:
+        flagged = False
+        preset_response = ""
+        if self.config is None:
+            raise ValueError("The config is not set.")
+
+        if self.config["outputs_config"]["enabled"]:
+            # Filter out empty values
+            keywords_list = [keyword for keyword in self.config["keywords"].split("\n") if keyword]
+
+            flagged = self._is_violated({"text": text}, keywords_list)
+            preset_response = self.config["outputs_config"]["preset_response"]
+
+        return ModerationOutputsResult(
+            flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response
+        )
+
+    def _is_violated(self, inputs: dict, keywords_list: list) -> bool:
+        return any(self._check_keywords_in_value(keywords_list, value) for value in inputs.values())
+
+    def _check_keywords_in_value(self, keywords_list: Sequence[str], value: Any) -> bool:
+        return any(keyword.lower() in str(value).lower() for keyword in keywords_list)

api/core/moderation/openai_moderation/__builtin__

@@ -0,0 +1 @@
+1

api/core/moderation/openai_moderation/openai_moderation.py

@@ -0,0 +1,60 @@
+from core.model_manager import ModelManager
+from core.model_runtime.entities.model_entities import ModelType
+from core.moderation.base import Moderation, ModerationAction, ModerationInputsResult, ModerationOutputsResult
+
+
+class OpenAIModeration(Moderation):
+    name: str = "openai_moderation"
+
+    @classmethod
+    def validate_config(cls, tenant_id: str, config: dict) -> None:
+        """
+        Validate the incoming form config data.
+
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        cls._validate_inputs_and_outputs_config(config, True)
+
+    def moderation_for_inputs(self, inputs: dict, query: str = "") -> ModerationInputsResult:
+        flagged = False
+        preset_response = ""
+        if self.config is None:
+            raise ValueError("The config is not set.")
+
+        if self.config["inputs_config"]["enabled"]:
+            preset_response = self.config["inputs_config"]["preset_response"]
+
+            if query:
+                inputs["query__"] = query
+            flagged = self._is_violated(inputs)
+
+        return ModerationInputsResult(
+            flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response
+        )
+
+    def moderation_for_outputs(self, text: str) -> ModerationOutputsResult:
+        flagged = False
+        preset_response = ""
+        if self.config is None:
+            raise ValueError("The config is not set.")
+
+        if self.config["outputs_config"]["enabled"]:
+            flagged = self._is_violated({"text": text})
+            preset_response = self.config["outputs_config"]["preset_response"]
+
+        return ModerationOutputsResult(
+            flagged=flagged, action=ModerationAction.DIRECT_OUTPUT, preset_response=preset_response
+        )
+
+    def _is_violated(self, inputs: dict):
+        text = "\n".join(str(inputs.values()))
+        model_manager = ModelManager()
+        model_instance = model_manager.get_model_instance(
+            tenant_id=self.tenant_id, provider="openai", model_type=ModelType.MODERATION, model="text-moderation-stable"
+        )
+
+        openai_moderation = model_instance.invoke_moderation(text=text)
+
+        return openai_moderation

api/core/moderation/output_moderation.py

@@ -0,0 +1,131 @@
+import logging
+import threading
+import time
+from typing import Any, Optional
+
+from flask import Flask, current_app
+from pydantic import BaseModel, ConfigDict
+
+from configs import dify_config
+from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
+from core.app.entities.queue_entities import QueueMessageReplaceEvent
+from core.moderation.base import ModerationAction, ModerationOutputsResult
+from core.moderation.factory import ModerationFactory
+
+logger = logging.getLogger(__name__)
+
+
+class ModerationRule(BaseModel):
+    type: str
+    config: dict[str, Any]
+
+
+class OutputModeration(BaseModel):
+    tenant_id: str
+    app_id: str
+
+    rule: ModerationRule
+    queue_manager: AppQueueManager
+
+    thread: Optional[threading.Thread] = None
+    thread_running: bool = True
+    buffer: str = ""
+    is_final_chunk: bool = False
+    final_output: Optional[str] = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def should_direct_output(self) -> bool:
+        return self.final_output is not None
+
+    def get_final_output(self) -> str:
+        return self.final_output or ""
+
+    def append_new_token(self, token: str) -> None:
+        self.buffer += token
+
+        if not self.thread:
+            self.thread = self.start_thread()
+
+    def moderation_completion(self, completion: str, public_event: bool = False) -> str:
+        self.buffer = completion
+        self.is_final_chunk = True
+
+        result = self.moderation(tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=completion)
+
+        if not result or not result.flagged:
+            return completion
+
+        if result.action == ModerationAction.DIRECT_OUTPUT:
+            final_output = result.preset_response
+        else:
+            final_output = result.text
+
+        if public_event:
+            self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE)
+
+        return final_output
+
+    def start_thread(self) -> threading.Thread:
+        buffer_size = dify_config.MODERATION_BUFFER_SIZE
+        thread = threading.Thread(
+            target=self.worker,
+            kwargs={
+                "flask_app": current_app._get_current_object(),  # type: ignore
+                "buffer_size": buffer_size if buffer_size > 0 else dify_config.MODERATION_BUFFER_SIZE,
+            },
+        )
+
+        thread.start()
+
+        return thread
+
+    def stop_thread(self):
+        if self.thread and self.thread.is_alive():
+            self.thread_running = False
+
+    def worker(self, flask_app: Flask, buffer_size: int):
+        with flask_app.app_context():
+            current_length = 0
+            while self.thread_running:
+                moderation_buffer = self.buffer
+                buffer_length = len(moderation_buffer)
+                if not self.is_final_chunk:
+                    chunk_length = buffer_length - current_length
+                    if 0 <= chunk_length < buffer_size:
+                        time.sleep(1)
+                        continue
+
+                current_length = buffer_length
+
+                result = self.moderation(
+                    tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=moderation_buffer
+                )
+
+                if not result or not result.flagged:
+                    continue
+
+                if result.action == ModerationAction.DIRECT_OUTPUT:
+                    final_output = result.preset_response
+                    self.final_output = final_output
+                else:
+                    final_output = result.text + self.buffer[len(moderation_buffer) :]
+
+                # trigger replace event
+                if self.thread_running:
+                    self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE)
+
+                if result.action == ModerationAction.DIRECT_OUTPUT:
+                    break
+
+    def moderation(self, tenant_id: str, app_id: str, moderation_buffer: str) -> Optional[ModerationOutputsResult]:
+        try:
+            moderation_factory = ModerationFactory(
+                name=self.rule.type, app_id=app_id, tenant_id=tenant_id, config=self.rule.config
+            )
+
+            result: ModerationOutputsResult = moderation_factory.moderation_for_outputs(moderation_buffer)
+            return result
+        except Exception as e:
+            logger.exception(f"Moderation Output error, app_id: {app_id}")
+
+        return None

api/core/ops/base_trace_instance.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+
+from core.ops.entities.config_entity import BaseTracingConfig
+from core.ops.entities.trace_entity import BaseTraceInfo
+
+
+class BaseTraceInstance(ABC):
+    """
+    Base trace instance for ops trace services
+    """
+
+    @abstractmethod
+    def __init__(self, trace_config: BaseTracingConfig):
+        """
+        Abstract initializer for the trace instance.
+        Distribute trace tasks by matching entities
+        """
+        self.trace_config = trace_config
+
+    @abstractmethod
+    def trace(self, trace_info: BaseTraceInfo):
+        """
+        Abstract method to trace activities.
+        Subclasses must implement specific tracing logic for activities.
+        """
+        ...

api/core/ops/entities/config_entity.py

@@ -0,0 +1,92 @@
+from enum import Enum
+
+from pydantic import BaseModel, ValidationInfo, field_validator
+
+
+class TracingProviderEnum(Enum):
+    LANGFUSE = "langfuse"
+    LANGSMITH = "langsmith"
+    OPIK = "opik"
+
+
+class BaseTracingConfig(BaseModel):
+    """
+    Base model class for tracing
+    """
+
+    ...
+
+
+class LangfuseConfig(BaseTracingConfig):
+    """
+    Model class for Langfuse tracing config.
+    """
+
+    public_key: str
+    secret_key: str
+    host: str = "https://api.langfuse.com"
+
+    @field_validator("host")
+    @classmethod
+    def set_value(cls, v, info: ValidationInfo):
+        if v is None or v == "":
+            v = "https://api.langfuse.com"
+        if not v.startswith("https://") and not v.startswith("http://"):
+            raise ValueError("host must start with https:// or http://")
+
+        return v
+
+
+class LangSmithConfig(BaseTracingConfig):
+    """
+    Model class for Langsmith tracing config.
+    """
+
+    api_key: str
+    project: str
+    endpoint: str = "https://api.smith.langchain.com"
+
+    @field_validator("endpoint")
+    @classmethod
+    def set_value(cls, v, info: ValidationInfo):
+        if v is None or v == "":
+            v = "https://api.smith.langchain.com"
+        if not v.startswith("https://"):
+            raise ValueError("endpoint must start with https://")
+
+        return v
+
+
+class OpikConfig(BaseTracingConfig):
+    """
+    Model class for Opik tracing config.
+    """
+
+    api_key: str | None = None
+    project: str | None = None
+    workspace: str | None = None
+    url: str = "https://www.comet.com/opik/api/"
+
+    @field_validator("project")
+    @classmethod
+    def project_validator(cls, v, info: ValidationInfo):
+        if v is None or v == "":
+            v = "Default Project"
+
+        return v
+
+    @field_validator("url")
+    @classmethod
+    def url_validator(cls, v, info: ValidationInfo):
+        if v is None or v == "":
+            v = "https://www.comet.com/opik/api/"
+        if not v.startswith(("https://", "http://")):
+            raise ValueError("url must start with https:// or http://")
+        if not v.endswith("/api/"):
+            raise ValueError("url should ends with /api/")
+
+        return v
+
+
+OPS_FILE_PATH = "ops_trace/"
+OPS_TRACE_FAILED_KEY = "FAILED_OPS_TRACE"

api/core/ops/entities/trace_entity.py

@@ -0,0 +1,134 @@
+from collections.abc import Mapping
+from datetime import datetime
+from enum import StrEnum
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, field_validator
+
+
+class BaseTraceInfo(BaseModel):
+    message_id: Optional[str] = None
+    message_data: Optional[Any] = None
+    inputs: Optional[Union[str, dict[str, Any], list]] = None
+    outputs: Optional[Union[str, dict[str, Any], list]] = None
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    metadata: dict[str, Any]
+
+    @field_validator("inputs", "outputs")
+    @classmethod
+    def ensure_type(cls, v):
+        if v is None:
+            return None
+        if isinstance(v, str | dict | list):
+            return v
+        return ""
+
+    class Config:
+        json_encoders = {
+            datetime: lambda v: v.isoformat(),
+        }
+
+
+class WorkflowTraceInfo(BaseTraceInfo):
+    workflow_data: Any
+    conversation_id: Optional[str] = None
+    workflow_app_log_id: Optional[str] = None
+    workflow_id: str
+    tenant_id: str
+    workflow_run_id: str
+    workflow_run_elapsed_time: Union[int, float]
+    workflow_run_status: str
+    workflow_run_inputs: Mapping[str, Any]
+    workflow_run_outputs: Mapping[str, Any]
+    workflow_run_version: str
+    error: Optional[str] = None
+    total_tokens: int
+    file_list: list[str]
+    query: str
+    metadata: dict[str, Any]
+
+
+class MessageTraceInfo(BaseTraceInfo):
+    conversation_model: str
+    message_tokens: int
+    answer_tokens: int
+    total_tokens: int
+    error: Optional[str] = None
+    file_list: Optional[Union[str, dict[str, Any], list]] = None
+    message_file_data: Optional[Any] = None
+    conversation_mode: str
+
+
+class ModerationTraceInfo(BaseTraceInfo):
+    flagged: bool
+    action: str
+    preset_response: str
+    query: str
+
+
+class SuggestedQuestionTraceInfo(BaseTraceInfo):
+    total_tokens: int
+    status: Optional[str] = None
+    error: Optional[str] = None
+    from_account_id: Optional[str] = None
+    agent_based: Optional[bool] = None
+    from_source: Optional[str] = None
+    model_provider: Optional[str] = None
+    model_id: Optional[str] = None
+    suggested_question: list[str]
+    level: str
+    status_message: Optional[str] = None
+    workflow_run_id: Optional[str] = None
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+class DatasetRetrievalTraceInfo(BaseTraceInfo):
+    documents: Any
+
+
+class ToolTraceInfo(BaseTraceInfo):
+    tool_name: str
+    tool_inputs: dict[str, Any]
+    tool_outputs: str
+    metadata: dict[str, Any]
+    message_file_data: Any
+    error: Optional[str] = None
+    tool_config: dict[str, Any]
+    time_cost: Union[int, float]
+    tool_parameters: dict[str, Any]
+    file_url: Union[str, None, list]
+
+
+class GenerateNameTraceInfo(BaseTraceInfo):
+    conversation_id: Optional[str] = None
+    tenant_id: str
+
+
+class TaskData(BaseModel):
+    app_id: str
+    trace_info_type: str
+    trace_info: Any
+
+
+trace_info_info_map = {
+    "WorkflowTraceInfo": WorkflowTraceInfo,
+    "MessageTraceInfo": MessageTraceInfo,
+    "ModerationTraceInfo": ModerationTraceInfo,
+    "SuggestedQuestionTraceInfo": SuggestedQuestionTraceInfo,
+    "DatasetRetrievalTraceInfo": DatasetRetrievalTraceInfo,
+    "ToolTraceInfo": ToolTraceInfo,
+    "GenerateNameTraceInfo": GenerateNameTraceInfo,
+}
+
+
+class TraceTaskName(StrEnum):
+    CONVERSATION_TRACE = "conversation"
+    WORKFLOW_TRACE = "workflow"
+    MESSAGE_TRACE = "message"
+    MODERATION_TRACE = "moderation"
+    SUGGESTED_QUESTION_TRACE = "suggested_question"
+    DATASET_RETRIEVAL_TRACE = "dataset_retrieval"
+    TOOL_TRACE = "tool"
+    GENERATE_NAME_TRACE = "generate_conversation_name"

api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py

@@ -0,0 +1,282 @@
+from datetime import datetime
+from enum import StrEnum
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic_core.core_schema import ValidationInfo
+
+from core.ops.utils import replace_text_with_content
+
+
+def validate_input_output(v, field_name):
+    """
+    Validate input output
+    :param v:
+    :param field_name:
+    :return:
+    """
+    if v == {} or v is None:
+        return v
+    if isinstance(v, str):
+        return [
+            {
+                "role": "assistant" if field_name == "output" else "user",
+                "content": v,
+            }
+        ]
+    elif isinstance(v, list):
+        if len(v) > 0 and isinstance(v[0], dict):
+            v = replace_text_with_content(data=v)
+            return v
+        else:
+            return [
+                {
+                    "role": "assistant" if field_name == "output" else "user",
+                    "content": str(v),
+                }
+            ]
+
+    return v
+
+
+class LevelEnum(StrEnum):
+    DEBUG = "DEBUG"
+    WARNING = "WARNING"
+    ERROR = "ERROR"
+    DEFAULT = "DEFAULT"
+
+
+class LangfuseTrace(BaseModel):
+    """
+    Langfuse trace model
+    """
+
+    id: Optional[str] = Field(
+        default=None,
+        description="The id of the trace can be set, defaults to a random id. Used to link traces to external systems "
+        "or when creating a distributed trace. Traces are upserted on id.",
+    )
+    name: Optional[str] = Field(
+        default=None,
+        description="Identifier of the trace. Useful for sorting/filtering in the UI.",
+    )
+    input: Optional[Union[str, dict[str, Any], list, None]] = Field(
+        default=None, description="The input of the trace. Can be any JSON object."
+    )
+    output: Optional[Union[str, dict[str, Any], list, None]] = Field(
+        default=None, description="The output of the trace. Can be any JSON object."
+    )
+    metadata: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="Additional metadata of the trace. Can be any JSON object. Metadata is merged when being updated "
+        "via the API.",
+    )
+    user_id: Optional[str] = Field(
+        default=None,
+        description="The id of the user that triggered the execution. Used to provide user-level analytics.",
+    )
+    session_id: Optional[str] = Field(
+        default=None,
+        description="Used to group multiple traces into a session in Langfuse. Use your own session/thread identifier.",
+    )
+    version: Optional[str] = Field(
+        default=None,
+        description="The version of the trace type. Used to understand how changes to the trace type affect metrics. "
+        "Useful in debugging.",
+    )
+    release: Optional[str] = Field(
+        default=None,
+        description="The release identifier of the current deployment. Used to understand how changes of different "
+        "deployments affect metrics. Useful in debugging.",
+    )
+    tags: Optional[list[str]] = Field(
+        default=None,
+        description="Tags are used to categorize or label traces. Traces can be filtered by tags in the UI and GET "
+        "API. Tags can also be changed in the UI. Tags are merged and never deleted via the API.",
+    )
+    public: Optional[bool] = Field(
+        default=None,
+        description="You can make a trace public to share it via a public link. This allows others to view the trace "
+        "without needing to log in or be members of your Langfuse project.",
+    )
+
+    @field_validator("input", "output")
+    @classmethod
+    def ensure_dict(cls, v, info: ValidationInfo):
+        field_name = info.field_name
+        return validate_input_output(v, field_name)
+
+
+class LangfuseSpan(BaseModel):
+    """
+    Langfuse span model
+    """
+
+    id: Optional[str] = Field(
+        default=None,
+        description="The id of the span can be set, otherwise a random id is generated. Spans are upserted on id.",
+    )
+    session_id: Optional[str] = Field(
+        default=None,
+        description="Used to group multiple spans into a session in Langfuse. Use your own session/thread identifier.",
+    )
+    trace_id: Optional[str] = Field(
+        default=None,
+        description="The id of the trace the span belongs to. Used to link spans to traces.",
+    )
+    user_id: Optional[str] = Field(
+        default=None,
+        description="The id of the user that triggered the execution. Used to provide user-level analytics.",
+    )
+    start_time: Optional[datetime | str] = Field(
+        default_factory=datetime.now,
+        description="The time at which the span started, defaults to the current time.",
+    )
+    end_time: Optional[datetime | str] = Field(
+        default=None,
+        description="The time at which the span ended. Automatically set by span.end().",
+    )
+    name: Optional[str] = Field(
+        default=None,
+        description="Identifier of the span. Useful for sorting/filtering in the UI.",
+    )
+    metadata: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="Additional metadata of the span. Can be any JSON object. Metadata is merged when being updated "
+        "via the API.",
+    )
+    level: Optional[str] = Field(
+        default=None,
+        description="The level of the span. Can be DEBUG, DEFAULT, WARNING or ERROR. Used for sorting/filtering of "
+        "traces with elevated error levels and for highlighting in the UI.",
+    )
+    status_message: Optional[str] = Field(
+        default=None,
+        description="The status message of the span. Additional field for context of the event. E.g. the error "
+        "message of an error event.",
+    )
+    input: Optional[Union[str, dict[str, Any], list, None]] = Field(
+        default=None, description="The input of the span. Can be any JSON object."
+    )
+    output: Optional[Union[str, dict[str, Any], list, None]] = Field(
+        default=None, description="The output of the span. Can be any JSON object."
+    )
+    version: Optional[str] = Field(
+        default=None,
+        description="The version of the span type. Used to understand how changes to the span type affect metrics. "
+        "Useful in debugging.",
+    )
+    parent_observation_id: Optional[str] = Field(
+        default=None,
+        description="The id of the observation the span belongs to. Used to link spans to observations.",
+    )
+
+    @field_validator("input", "output")
+    @classmethod
+    def ensure_dict(cls, v, info: ValidationInfo):
+        field_name = info.field_name
+        return validate_input_output(v, field_name)
+
+
+class UnitEnum(StrEnum):
+    CHARACTERS = "CHARACTERS"
+    TOKENS = "TOKENS"
+    SECONDS = "SECONDS"
+    MILLISECONDS = "MILLISECONDS"
+    IMAGES = "IMAGES"
+
+
+class GenerationUsage(BaseModel):
+    promptTokens: Optional[int] = None
+    completionTokens: Optional[int] = None
+    total: Optional[int] = None
+    input: Optional[int] = None
+    output: Optional[int] = None
+    unit: Optional[UnitEnum] = None
+    inputCost: Optional[float] = None
+    outputCost: Optional[float] = None
+    totalCost: Optional[float] = None
+
+    @field_validator("input", "output")
+    @classmethod
+    def ensure_dict(cls, v, info: ValidationInfo):
+        field_name = info.field_name
+        return validate_input_output(v, field_name)
+
+
+class LangfuseGeneration(BaseModel):
+    id: Optional[str] = Field(
+        default=None,
+        description="The id of the generation can be set, defaults to random id.",
+    )
+    trace_id: Optional[str] = Field(
+        default=None,
+        description="The id of the trace the generation belongs to. Used to link generations to traces.",
+    )
+    parent_observation_id: Optional[str] = Field(
+        default=None,
+        description="The id of the observation the generation belongs to. Used to link generations to observations.",
+    )
+    name: Optional[str] = Field(
+        default=None,
+        description="Identifier of the generation. Useful for sorting/filtering in the UI.",
+    )
+    start_time: Optional[datetime | str] = Field(
+        default_factory=datetime.now,
+        description="The time at which the generation started, defaults to the current time.",
+    )
+    completion_start_time: Optional[datetime | str] = Field(
+        default=None,
+        description="The time at which the completion started (streaming). Set it to get latency analytics broken "
+        "down into time until completion started and completion duration.",
+    )
+    end_time: Optional[datetime | str] = Field(
+        default=None,
+        description="The time at which the generation ended. Automatically set by generation.end().",
+    )
+    model: Optional[str] = Field(default=None, description="The name of the model used for the generation.")
+    model_parameters: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="The parameters of the model used for the generation; can be any key-value pairs.",
+    )
+    input: Optional[Any] = Field(
+        default=None,
+        description="The prompt used for the generation. Can be any string or JSON object.",
+    )
+    output: Optional[Any] = Field(
+        default=None,
+        description="The completion generated by the model. Can be any string or JSON object.",
+    )
+    usage: Optional[GenerationUsage] = Field(
+        default=None,
+        description="The usage object supports the OpenAi structure with tokens and a more generic version with "
+        "detailed costs and units.",
+    )
+    metadata: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="Additional metadata of the generation. Can be any JSON object. Metadata is merged when being "
+        "updated via the API.",
+    )
+    level: Optional[LevelEnum] = Field(
+        default=None,
+        description="The level of the generation. Can be DEBUG, DEFAULT, WARNING or ERROR. Used for sorting/filtering "
+        "of traces with elevated error levels and for highlighting in the UI.",
+    )
+    status_message: Optional[str] = Field(
+        default=None,
+        description="The status message of the generation. Additional field for context of the event. E.g. the error "
+        "message of an error event.",
+    )
+    version: Optional[str] = Field(
+        default=None,
+        description="The version of the generation type. Used to understand how changes to the span type affect "
+        "metrics. Useful in debugging.",
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    @field_validator("input", "output")
+    @classmethod
+    def ensure_dict(cls, v, info: ValidationInfo):
+        field_name = info.field_name
+        return validate_input_output(v, field_name)

api/core/ops/langfuse_trace/langfuse_trace.py

@@ -0,0 +1,455 @@
+import json
+import logging
+import os
+from datetime import datetime, timedelta
+from typing import Optional
+
+from langfuse import Langfuse  # type: ignore
+
+from core.ops.base_trace_instance import BaseTraceInstance
+from core.ops.entities.config_entity import LangfuseConfig
+from core.ops.entities.trace_entity import (
+    BaseTraceInfo,
+    DatasetRetrievalTraceInfo,
+    GenerateNameTraceInfo,
+    MessageTraceInfo,
+    ModerationTraceInfo,
+    SuggestedQuestionTraceInfo,
+    ToolTraceInfo,
+    TraceTaskName,
+    WorkflowTraceInfo,
+)
+from core.ops.langfuse_trace.entities.langfuse_trace_entity import (
+    GenerationUsage,
+    LangfuseGeneration,
+    LangfuseSpan,
+    LangfuseTrace,
+    LevelEnum,
+    UnitEnum,
+)
+from core.ops.utils import filter_none_values
+from extensions.ext_database import db
+from models.model import EndUser
+from models.workflow import WorkflowNodeExecution
+
+logger = logging.getLogger(__name__)
+
+
+class LangFuseDataTrace(BaseTraceInstance):
+    def __init__(
+        self,
+        langfuse_config: LangfuseConfig,
+    ):
+        super().__init__(langfuse_config)
+        self.langfuse_client = Langfuse(
+            public_key=langfuse_config.public_key,
+            secret_key=langfuse_config.secret_key,
+            host=langfuse_config.host,
+        )
+        self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
+
+    def trace(self, trace_info: BaseTraceInfo):
+        if isinstance(trace_info, WorkflowTraceInfo):
+            self.workflow_trace(trace_info)
+        if isinstance(trace_info, MessageTraceInfo):
+            self.message_trace(trace_info)
+        if isinstance(trace_info, ModerationTraceInfo):
+            self.moderation_trace(trace_info)
+        if isinstance(trace_info, SuggestedQuestionTraceInfo):
+            self.suggested_question_trace(trace_info)
+        if isinstance(trace_info, DatasetRetrievalTraceInfo):
+            self.dataset_retrieval_trace(trace_info)
+        if isinstance(trace_info, ToolTraceInfo):
+            self.tool_trace(trace_info)
+        if isinstance(trace_info, GenerateNameTraceInfo):
+            self.generate_name_trace(trace_info)
+
+    def workflow_trace(self, trace_info: WorkflowTraceInfo):
+        trace_id = trace_info.workflow_run_id
+        user_id = trace_info.metadata.get("user_id")
+        metadata = trace_info.metadata
+        metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id
+
+        if trace_info.message_id:
+            trace_id = trace_info.message_id
+            name = TraceTaskName.MESSAGE_TRACE.value
+            trace_data = LangfuseTrace(
+                id=trace_id,
+                user_id=user_id,
+                name=name,
+                input=dict(trace_info.workflow_run_inputs),
+                output=dict(trace_info.workflow_run_outputs),
+                metadata=metadata,
+                session_id=trace_info.conversation_id,
+                tags=["message", "workflow"],
+            )
+            self.add_trace(langfuse_trace_data=trace_data)
+            workflow_span_data = LangfuseSpan(
+                id=trace_info.workflow_run_id,
+                name=TraceTaskName.WORKFLOW_TRACE.value,
+                input=dict(trace_info.workflow_run_inputs),
+                output=dict(trace_info.workflow_run_outputs),
+                trace_id=trace_id,
+                start_time=trace_info.start_time,
+                end_time=trace_info.end_time,
+                metadata=metadata,
+                level=LevelEnum.DEFAULT if trace_info.error == "" else LevelEnum.ERROR,
+                status_message=trace_info.error or "",
+            )
+            self.add_span(langfuse_span_data=workflow_span_data)
+        else:
+            trace_data = LangfuseTrace(
+                id=trace_id,
+                user_id=user_id,
+                name=TraceTaskName.WORKFLOW_TRACE.value,
+                input=dict(trace_info.workflow_run_inputs),
+                output=dict(trace_info.workflow_run_outputs),
+                metadata=metadata,
+                session_id=trace_info.conversation_id,
+                tags=["workflow"],
+            )
+            self.add_trace(langfuse_trace_data=trace_data)
+
+        # through workflow_run_id get all_nodes_execution
+        workflow_nodes_execution_id_records = (
+            db.session.query(WorkflowNodeExecution.id)
+            .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id)
+            .all()
+        )
+
+        for node_execution_id_record in workflow_nodes_execution_id_records:
+            node_execution = (
+                db.session.query(
+                    WorkflowNodeExecution.id,
+                    WorkflowNodeExecution.tenant_id,
+                    WorkflowNodeExecution.app_id,
+                    WorkflowNodeExecution.title,
+                    WorkflowNodeExecution.node_type,
+                    WorkflowNodeExecution.status,
+                    WorkflowNodeExecution.inputs,
+                    WorkflowNodeExecution.outputs,
+                    WorkflowNodeExecution.created_at,
+                    WorkflowNodeExecution.elapsed_time,
+                    WorkflowNodeExecution.process_data,
+                    WorkflowNodeExecution.execution_metadata,
+                )
+                .filter(WorkflowNodeExecution.id == node_execution_id_record.id)
+                .first()
+            )
+
+            if not node_execution:
+                continue
+
+            node_execution_id = node_execution.id
+            tenant_id = node_execution.tenant_id
+            app_id = node_execution.app_id
+            node_name = node_execution.title
+            node_type = node_execution.node_type
+            status = node_execution.status
+            if node_type == "llm":
+                inputs = (
+                    json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {}
+                )
+            else:
+                inputs = json.loads(node_execution.inputs) if node_execution.inputs else {}
+            outputs = json.loads(node_execution.outputs) if node_execution.outputs else {}
+            created_at = node_execution.created_at or datetime.now()
+            elapsed_time = node_execution.elapsed_time
+            finished_at = created_at + timedelta(seconds=elapsed_time)
+
+            metadata = json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {}
+            metadata.update(
+                {
+                    "workflow_run_id": trace_info.workflow_run_id,
+                    "node_execution_id": node_execution_id,
+                    "tenant_id": tenant_id,
+                    "app_id": app_id,
+                    "node_name": node_name,
+                    "node_type": node_type,
+                    "status": status,
+                }
+            )
+            process_data = json.loads(node_execution.process_data) if node_execution.process_data else {}
+            model_provider = process_data.get("model_provider", None)
+            model_name = process_data.get("model_name", None)
+            if model_provider is not None and model_name is not None:
+                metadata.update(
+                    {
+                        "model_provider": model_provider,
+                        "model_name": model_name,
+                    }
+                )
+
+            # add span
+            if trace_info.message_id:
+                span_data = LangfuseSpan(
+                    id=node_execution_id,
+                    name=node_type,
+                    input=inputs,
+                    output=outputs,
+                    trace_id=trace_id,
+                    start_time=created_at,
+                    end_time=finished_at,
+                    metadata=metadata,
+                    level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR),
+                    status_message=trace_info.error or "",
+                    parent_observation_id=trace_info.workflow_run_id,
+                )
+            else:
+                span_data = LangfuseSpan(
+                    id=node_execution_id,
+                    name=node_type,
+                    input=inputs,
+                    output=outputs,
+                    trace_id=trace_id,
+                    start_time=created_at,
+                    end_time=finished_at,
+                    metadata=metadata,
+                    level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR),
+                    status_message=trace_info.error or "",
+                )
+
+            self.add_span(langfuse_span_data=span_data)
+
+            if process_data and process_data.get("model_mode") == "chat":
+                total_token = metadata.get("total_tokens", 0)
+                # add generation
+                generation_usage = GenerationUsage(
+                    total=total_token,
+                )
+
+                node_generation_data = LangfuseGeneration(
+                    name="llm",
+                    trace_id=trace_id,
+                    model=process_data.get("model_name"),
+                    parent_observation_id=node_execution_id,
+                    start_time=created_at,
+                    end_time=finished_at,
+                    input=inputs,
+                    output=outputs,
+                    metadata=metadata,
+                    level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR),
+                    status_message=trace_info.error or "",
+                    usage=generation_usage,
+                )
+
+                self.add_generation(langfuse_generation_data=node_generation_data)
+
+    def message_trace(self, trace_info: MessageTraceInfo, **kwargs):
+        # get message file data
+        file_list = trace_info.file_list
+        metadata = trace_info.metadata
+        message_data = trace_info.message_data
+        if message_data is None:
+            return
+        message_id = message_data.id
+
+        user_id = message_data.from_account_id
+        if message_data.from_end_user_id:
+            end_user_data: Optional[EndUser] = (
+                db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first()
+            )
+            if end_user_data is not None:
+                user_id = end_user_data.session_id
+                metadata["user_id"] = user_id
+
+        trace_data = LangfuseTrace(
+            id=message_id,
+            user_id=user_id,
+            name=TraceTaskName.MESSAGE_TRACE.value,
+            input={
+                "message": trace_info.inputs,
+                "files": file_list,
+                "message_tokens": trace_info.message_tokens,
+                "answer_tokens": trace_info.answer_tokens,
+                "total_tokens": trace_info.total_tokens,
+                "error": trace_info.error,
+                "provider_response_latency": message_data.provider_response_latency,
+                "created_at": trace_info.start_time,
+            },
+            output=trace_info.outputs,
+            metadata=metadata,
+            session_id=message_data.conversation_id,
+            tags=["message", str(trace_info.conversation_mode)],
+            version=None,
+            release=None,
+            public=None,
+        )
+        self.add_trace(langfuse_trace_data=trace_data)
+
+        # start add span
+        generation_usage = GenerationUsage(
+            input=trace_info.message_tokens,
+            output=trace_info.answer_tokens,
+            total=trace_info.total_tokens,
+            unit=UnitEnum.TOKENS,
+            totalCost=message_data.total_price,
+        )
+
+        langfuse_generation_data = LangfuseGeneration(
+            name="llm",
+            trace_id=message_id,
+            start_time=trace_info.start_time,
+            end_time=trace_info.end_time,
+            model=message_data.model_id,
+            input=trace_info.inputs,
+            output=message_data.answer,
+            metadata=metadata,
+            level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR),
+            status_message=message_data.error or "",
+            usage=generation_usage,
+        )
+
+        self.add_generation(langfuse_generation_data)
+
+    def moderation_trace(self, trace_info: ModerationTraceInfo):
+        if trace_info.message_data is None:
+            return
+        span_data = LangfuseSpan(
+            name=TraceTaskName.MODERATION_TRACE.value,
+            input=trace_info.inputs,
+            output={
+                "action": trace_info.action,
+                "flagged": trace_info.flagged,
+                "preset_response": trace_info.preset_response,
+                "inputs": trace_info.inputs,
+            },
+            trace_id=trace_info.message_id,
+            start_time=trace_info.start_time or trace_info.message_data.created_at,
+            end_time=trace_info.end_time or trace_info.message_data.created_at,
+            metadata=trace_info.metadata,
+        )
+
+        self.add_span(langfuse_span_data=span_data)
+
+    def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo):
+        message_data = trace_info.message_data
+        if message_data is None:
+            return
+        generation_usage = GenerationUsage(
+            total=len(str(trace_info.suggested_question)),
+            input=len(trace_info.inputs) if trace_info.inputs else 0,
+            output=len(trace_info.suggested_question),
+            unit=UnitEnum.CHARACTERS,
+        )
+
+        generation_data = LangfuseGeneration(
+            name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value,
+            input=trace_info.inputs,
+            output=str(trace_info.suggested_question),
+            trace_id=trace_info.message_id,
+            start_time=trace_info.start_time,
+            end_time=trace_info.end_time,
+            metadata=trace_info.metadata,
+            level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR),
+            status_message=message_data.error or "",
+            usage=generation_usage,
+        )
+
+        self.add_generation(langfuse_generation_data=generation_data)
+
+    def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo):
+        if trace_info.message_data is None:
+            return
+        dataset_retrieval_span_data = LangfuseSpan(
+            name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value,
+            input=trace_info.inputs,
+            output={"documents": trace_info.documents},
+            trace_id=trace_info.message_id,
+            start_time=trace_info.start_time or trace_info.message_data.created_at,
+            end_time=trace_info.end_time or trace_info.message_data.updated_at,
+            metadata=trace_info.metadata,
+        )
+
+        self.add_span(langfuse_span_data=dataset_retrieval_span_data)
+
+    def tool_trace(self, trace_info: ToolTraceInfo):
+        tool_span_data = LangfuseSpan(
+            name=trace_info.tool_name,
+            input=trace_info.tool_inputs,
+            output=trace_info.tool_outputs,
+            trace_id=trace_info.message_id,
+            start_time=trace_info.start_time,
+            end_time=trace_info.end_time,
+            metadata=trace_info.metadata,
+            level=(LevelEnum.DEFAULT if trace_info.error == "" or trace_info.error is None else LevelEnum.ERROR),
+            status_message=trace_info.error,
+        )
+
+        self.add_span(langfuse_span_data=tool_span_data)
+
+    def generate_name_trace(self, trace_info: GenerateNameTraceInfo):
+        name_generation_trace_data = LangfuseTrace(
+            name=TraceTaskName.GENERATE_NAME_TRACE.value,
+            input=trace_info.inputs,
+            output=trace_info.outputs,
+            user_id=trace_info.tenant_id,
+            metadata=trace_info.metadata,
+            session_id=trace_info.conversation_id,
+        )
+
+        self.add_trace(langfuse_trace_data=name_generation_trace_data)
+
+        name_generation_span_data = LangfuseSpan(
+            name=TraceTaskName.GENERATE_NAME_TRACE.value,
+            input=trace_info.inputs,
+            output=trace_info.outputs,
+            trace_id=trace_info.conversation_id,
+            start_time=trace_info.start_time,
+            end_time=trace_info.end_time,
+            metadata=trace_info.metadata,
+        )
+        self.add_span(langfuse_span_data=name_generation_span_data)
+
+    def add_trace(self, langfuse_trace_data: Optional[LangfuseTrace] = None):
+        format_trace_data = filter_none_values(langfuse_trace_data.model_dump()) if langfuse_trace_data else {}
+        try:
+            self.langfuse_client.trace(**format_trace_data)
+            logger.debug("LangFuse Trace created successfully")
+        except Exception as e:
+            raise ValueError(f"LangFuse Failed to create trace: {str(e)}")
+
+    def add_span(self, langfuse_span_data: Optional[LangfuseSpan] = None):
+        format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}
+        try:
+            self.langfuse_client.span(**format_span_data)
+            logger.debug("LangFuse Span created successfully")
+        except Exception as e:
+            raise ValueError(f"LangFuse Failed to create span: {str(e)}")
+
+    def update_span(self, span, langfuse_span_data: Optional[LangfuseSpan] = None):
+        format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}
+
+        span.end(**format_span_data)
+
+    def add_generation(self, langfuse_generation_data: Optional[LangfuseGeneration] = None):
+        format_generation_data = (
+            filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {}
+        )
+        try:
+            self.langfuse_client.generation(**format_generation_data)
+            logger.debug("LangFuse Generation created successfully")
+        except Exception as e:
+            raise ValueError(f"LangFuse Failed to create generation: {str(e)}")
+
+    def update_generation(self, generation, langfuse_generation_data: Optional[LangfuseGeneration] = None):
+        format_generation_data = (
+            filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {}
+        )
+
+        generation.end(**format_generation_data)
+
+    def api_check(self):
+        try:
+            return self.langfuse_client.auth_check()
+        except Exception as e:
+            logger.debug(f"LangFuse API check failed: {str(e)}")
+            raise ValueError(f"LangFuse API check failed: {str(e)}")
+
+    def get_project_key(self):
+        try:
+            projects = self.langfuse_client.client.projects.get()
+            return projects.data[0].id
+        except Exception as e:
+            logger.debug(f"LangFuse get project key failed: {str(e)}")
+            raise ValueError(f"LangFuse get project key failed: {str(e)}")

api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py

@@ -0,0 +1,141 @@
+from datetime import datetime
+from enum import StrEnum
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, Field, field_validator
+from pydantic_core.core_schema import ValidationInfo
+
+from core.ops.utils import replace_text_with_content
+
+
+class LangSmithRunType(StrEnum):
+    tool = "tool"
+    chain = "chain"
+    llm = "llm"
+    retriever = "retriever"
+    embedding = "embedding"
+    prompt = "prompt"
+    parser = "parser"
+
+
+class LangSmithTokenUsage(BaseModel):
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+
+
+class LangSmithMultiModel(BaseModel):
+    file_list: Optional[list[str]] = Field(None, description="List of files")
+
+
+class LangSmithRunModel(LangSmithTokenUsage, LangSmithMultiModel):
+    name: Optional[str] = Field(..., description="Name of the run")
+    inputs: Optional[Union[str, dict[str, Any], list, None]] = Field(None, description="Inputs of the run")
+    outputs: Optional[Union[str, dict[str, Any], list, None]] = Field(None, description="Outputs of the run")
+    run_type: LangSmithRunType = Field(..., description="Type of the run")
+    start_time: Optional[datetime | str] = Field(None, description="Start time of the run")
+    end_time: Optional[datetime | str] = Field(None, description="End time of the run")
+    extra: Optional[dict[str, Any]] = Field(None, description="Extra information of the run")
+    error: Optional[str] = Field(None, description="Error message of the run")
+    serialized: Optional[dict[str, Any]] = Field(None, description="Serialized data of the run")
+    parent_run_id: Optional[str] = Field(None, description="Parent run ID")
+    events: Optional[list[dict[str, Any]]] = Field(None, description="Events associated with the run")
+    tags: Optional[list[str]] = Field(None, description="Tags associated with the run")
+    trace_id: Optional[str] = Field(None, description="Trace ID associated with the run")
+    dotted_order: Optional[str] = Field(None, description="Dotted order of the run")
+    id: Optional[str] = Field(None, description="ID of the run")
+    session_id: Optional[str] = Field(None, description="Session ID associated with the run")
+    session_name: Optional[str] = Field(None, description="Session name associated with the run")
+    reference_example_id: Optional[str] = Field(None, description="Reference example ID associated with the run")
+    input_attachments: Optional[dict[str, Any]] = Field(None, description="Input attachments of the run")
+    output_attachments: Optional[dict[str, Any]] = Field(None, description="Output attachments of the run")
+
+    @field_validator("inputs", "outputs")
+    @classmethod
+    def ensure_dict(cls, v, info: ValidationInfo):
+        field_name = info.field_name
+        values = info.data
+        if v == {} or v is None:
+            return v
+        usage_metadata = {
+            "input_tokens": values.get("input_tokens", 0),
+            "output_tokens": values.get("output_tokens", 0),
+            "total_tokens": values.get("total_tokens", 0),
+        }
+        file_list = values.get("file_list", [])
+        if isinstance(v, str):
+            if field_name == "inputs":
+                return {
+                    "messages": {
+                        "role": "user",
+                        "content": v,
+                        "usage_metadata": usage_metadata,
+                        "file_list": file_list,
+                    },
+                }
+            elif field_name == "outputs":
+                return {
+                    "choices": {
+                        "role": "ai",
+                        "content": v,
+                        "usage_metadata": usage_metadata,
+                        "file_list": file_list,
+                    },
+                }
+        elif isinstance(v, list):
+            data = {}
+            if len(v) > 0 and isinstance(v[0], dict):
+                # rename text to content
+                v = replace_text_with_content(data=v)
+                if field_name == "inputs":
+                    data = {
+                        "messages": v,
+                    }
+                elif field_name == "outputs":
+                    data = {
+                        "choices": {
+                            "role": "ai",
+                            "content": v,
+                            "usage_metadata": usage_metadata,
+                            "file_list": file_list,
+                        },
+                    }
+                return data
+            else:
+                return {
+                    "choices": {
+                        "role": "ai" if field_name == "outputs" else "user",
+                        "content": str(v),
+                        "usage_metadata": usage_metadata,
+                        "file_list": file_list,
+                    },
+                }
+        if isinstance(v, dict):
+            v["usage_metadata"] = usage_metadata
+            v["file_list"] = file_list
+            return v
+        return v
+
+    @classmethod
+    @field_validator("start_time", "end_time")
+    def format_time(cls, v, info: ValidationInfo):
+        if not isinstance(v, datetime):
+            raise ValueError(f"{info.field_name} must be a datetime object")
+        else:
+            return v.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+
+class LangSmithRunUpdateModel(BaseModel):
+    run_id: str = Field(..., description="ID of the run")
+    trace_id: Optional[str] = Field(None, description="Trace ID associated with the run")
+    dotted_order: Optional[str] = Field(None, description="Dotted order of the run")
+    parent_run_id: Optional[str] = Field(None, description="Parent run ID")
+    end_time: Optional[datetime | str] = Field(None, description="End time of the run")
+    error: Optional[str] = Field(None, description="Error message of the run")
+    inputs: Optional[dict[str, Any]] = Field(None, description="Inputs of the run")
+    outputs: Optional[dict[str, Any]] = Field(None, description="Outputs of the run")
+    events: Optional[list[dict[str, Any]]] = Field(None, description="Events associated with the run")
+    tags: Optional[list[str]] = Field(None, description="Tags associated with the run")
+    extra: Optional[dict[str, Any]] = Field(None, description="Extra information of the run")
+    input_attachments: Optional[dict[str, Any]] = Field(None, description="Input attachments of the run")
+    output_attachments: Optional[dict[str, Any]] = Field(None, description="Output attachments of the run")

api/core/ops/langsmith_trace/langsmith_trace.py

@@ -0,0 +1,524 @@
+import json
+import logging
+import os
+import uuid
+from datetime import datetime, timedelta
+from typing import Optional, cast
+
+from langsmith import Client
+from langsmith.schemas import RunBase
+
+from core.ops.base_trace_instance import BaseTraceInstance
+from core.ops.entities.config_entity import LangSmithConfig
+from core.ops.entities.trace_entity import (
+    BaseTraceInfo,
+    DatasetRetrievalTraceInfo,
+    GenerateNameTraceInfo,
+    MessageTraceInfo,
+    ModerationTraceInfo,
+    SuggestedQuestionTraceInfo,
+    ToolTraceInfo,
+    TraceTaskName,
+    WorkflowTraceInfo,
+)
+from core.ops.langsmith_trace.entities.langsmith_trace_entity import (
+    LangSmithRunModel,
+    LangSmithRunType,
+    LangSmithRunUpdateModel,
+)
+from core.ops.utils import filter_none_values, generate_dotted_order
+from extensions.ext_database import db
+from models.model import EndUser, MessageFile
+from models.workflow import WorkflowNodeExecution
+
+logger = logging.getLogger(__name__)
+
+
+class LangSmithDataTrace(BaseTraceInstance):
+    def __init__(
+        self,
+        langsmith_config: LangSmithConfig,
+    ):
+        super().__init__(langsmith_config)
+        self.langsmith_key = langsmith_config.api_key
+        self.project_name = langsmith_config.project
+        self.project_id = None
+        self.langsmith_client = Client(api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint)
+        self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
+
+    def trace(self, trace_info: BaseTraceInfo):
+        if isinstance(trace_info, WorkflowTraceInfo):
+            self.workflow_trace(trace_info)
+        if isinstance(trace_info, MessageTraceInfo):
+            self.message_trace(trace_info)
+        if isinstance(trace_info, ModerationTraceInfo):
+            self.moderation_trace(trace_info)
+        if isinstance(trace_info, SuggestedQuestionTraceInfo):
+            self.suggested_question_trace(trace_info)
+        if isinstance(trace_info, DatasetRetrievalTraceInfo):
+            self.dataset_retrieval_trace(trace_info)
+        if isinstance(trace_info, ToolTraceInfo):
+            self.tool_trace(trace_info)
+        if isinstance(trace_info, GenerateNameTraceInfo):
+            self.generate_name_trace(trace_info)
+
+    def workflow_trace(self, trace_info: WorkflowTraceInfo):
+        trace_id = trace_info.message_id or trace_info.workflow_run_id
+        if trace_info.start_time is None:
+            trace_info.start_time = datetime.now()
+        message_dotted_order = (
+            generate_dotted_order(trace_info.message_id, trace_info.start_time) if trace_info.message_id else None
+        )
+        workflow_dotted_order = generate_dotted_order(
+            trace_info.workflow_run_id,
+            trace_info.workflow_data.created_at,
+            message_dotted_order,
+        )
+        metadata = trace_info.metadata
+        metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id
+
+        if trace_info.message_id:
+            message_run = LangSmithRunModel(
+                id=trace_info.message_id,
+                name=TraceTaskName.MESSAGE_TRACE.value,
+                inputs=dict(trace_info.workflow_run_inputs),
+                outputs=dict(trace_info.workflow_run_outputs),
+                run_type=LangSmithRunType.chain,
+                start_time=trace_info.start_time,
+                end_time=trace_info.end_time,
+                extra={
+                    "metadata": metadata,
+                },
+                tags=["message", "workflow"],
+                error=trace_info.error,
+                trace_id=trace_id,
+                dotted_order=message_dotted_order,
+                file_list=[],
+                serialized=None,
+                parent_run_id=None,
+                events=[],
+                session_id=None,
+                session_name=None,
+                reference_example_id=None,
+                input_attachments={},
+                output_attachments={},
+            )
+            self.add_run(message_run)
+
+        langsmith_run = LangSmithRunModel(
+            file_list=trace_info.file_list,
+            total_tokens=trace_info.total_tokens,
+            id=trace_info.workflow_run_id,
+            name=TraceTaskName.WORKFLOW_TRACE.value,
+            inputs=dict(trace_info.workflow_run_inputs),
+            run_type=LangSmithRunType.tool,
+            start_time=trace_info.workflow_data.created_at,
+            end_time=trace_info.workflow_data.finished_at,
+            outputs=dict(trace_info.workflow_run_outputs),
+            extra={
+                "metadata": metadata,
+            },
+            error=trace_info.error,
+            tags=["workflow"],
+            parent_run_id=trace_info.message_id or None,
+            trace_id=trace_id,
+            dotted_order=workflow_dotted_order,
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+        )
+
+        self.add_run(langsmith_run)
+
+        # through workflow_run_id get all_nodes_execution
+        workflow_nodes_execution_id_records = (
+            db.session.query(WorkflowNodeExecution.id)
+            .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id)
+            .all()
+        )
+
+        for node_execution_id_record in workflow_nodes_execution_id_records:
+            node_execution = (
+                db.session.query(
+                    WorkflowNodeExecution.id,
+                    WorkflowNodeExecution.tenant_id,
+                    WorkflowNodeExecution.app_id,
+                    WorkflowNodeExecution.title,
+                    WorkflowNodeExecution.node_type,
+                    WorkflowNodeExecution.status,
+                    WorkflowNodeExecution.inputs,
+                    WorkflowNodeExecution.outputs,
+                    WorkflowNodeExecution.created_at,
+                    WorkflowNodeExecution.elapsed_time,
+                    WorkflowNodeExecution.process_data,
+                    WorkflowNodeExecution.execution_metadata,
+                )
+                .filter(WorkflowNodeExecution.id == node_execution_id_record.id)
+                .first()
+            )
+
+            if not node_execution:
+                continue
+
+            node_execution_id = node_execution.id
+            tenant_id = node_execution.tenant_id
+            app_id = node_execution.app_id
+            node_name = node_execution.title
+            node_type = node_execution.node_type
+            status = node_execution.status
+            if node_type == "llm":
+                inputs = (
+                    json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {}
+                )
+            else:
+                inputs = json.loads(node_execution.inputs) if node_execution.inputs else {}
+            outputs = json.loads(node_execution.outputs) if node_execution.outputs else {}
+            created_at = node_execution.created_at or datetime.now()
+            elapsed_time = node_execution.elapsed_time
+            finished_at = created_at + timedelta(seconds=elapsed_time)
+
+            execution_metadata = (
+                json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {}
+            )
+            node_total_tokens = execution_metadata.get("total_tokens", 0)
+            metadata = execution_metadata.copy()
+            metadata.update(
+                {
+                    "workflow_run_id": trace_info.workflow_run_id,
+                    "node_execution_id": node_execution_id,
+                    "tenant_id": tenant_id,
+                    "app_id": app_id,
+                    "app_name": node_name,
+                    "node_type": node_type,
+                    "status": status,
+                }
+            )
+
+            process_data = json.loads(node_execution.process_data) if node_execution.process_data else {}
+            if process_data and process_data.get("model_mode") == "chat":
+                run_type = LangSmithRunType.llm
+                metadata.update(
+                    {
+                        "ls_provider": process_data.get("model_provider", ""),
+                        "ls_model_name": process_data.get("model_name", ""),
+                    }
+                )
+            elif node_type == "knowledge-retrieval":
+                run_type = LangSmithRunType.retriever
+            else:
+                run_type = LangSmithRunType.tool
+
+            node_dotted_order = generate_dotted_order(node_execution_id, created_at, workflow_dotted_order)
+            langsmith_run = LangSmithRunModel(
+                total_tokens=node_total_tokens,
+                name=node_type,
+                inputs=inputs,
+                run_type=run_type,
+                start_time=created_at,
+                end_time=finished_at,
+                outputs=outputs,
+                file_list=trace_info.file_list,
+                extra={
+                    "metadata": metadata,
+                },
+                parent_run_id=trace_info.workflow_run_id,
+                tags=["node_execution"],
+                id=node_execution_id,
+                trace_id=trace_id,
+                dotted_order=node_dotted_order,
+                error="",
+                serialized=None,
+                events=[],
+                session_id=None,
+                session_name=None,
+                reference_example_id=None,
+                input_attachments={},
+                output_attachments={},
+            )
+
+            self.add_run(langsmith_run)
+
+    def message_trace(self, trace_info: MessageTraceInfo):
+        # get message file data
+        file_list = cast(list[str], trace_info.file_list) or []
+        message_file_data: Optional[MessageFile] = trace_info.message_file_data
+        file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else ""
+        file_list.append(file_url)
+        metadata = trace_info.metadata
+        message_data = trace_info.message_data
+        if message_data is None:
+            return
+        message_id = message_data.id
+
+        user_id = message_data.from_account_id
+        metadata["user_id"] = user_id
+
+        if message_data.from_end_user_id:
+            end_user_data: Optional[EndUser] = (
+                db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first()
+            )
+            if end_user_data is not None:
+                end_user_id = end_user_data.session_id
+                metadata["end_user_id"] = end_user_id
+
+        message_run = LangSmithRunModel(
+            input_tokens=trace_info.message_tokens,
+            output_tokens=trace_info.answer_tokens,
+            total_tokens=trace_info.total_tokens,
+            id=message_id,
+            name=TraceTaskName.MESSAGE_TRACE.value,
+            inputs=trace_info.inputs,
+            run_type=LangSmithRunType.chain,
+            start_time=trace_info.start_time,
+            end_time=trace_info.end_time,
+            outputs=message_data.answer,
+            extra={"metadata": metadata},
+            tags=["message", str(trace_info.conversation_mode)],
+            error=trace_info.error,
+            file_list=file_list,
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+            trace_id=None,
+            dotted_order=None,
+            parent_run_id=None,
+        )
+        self.add_run(message_run)
+
+        # create llm run parented to message run
+        llm_run = LangSmithRunModel(
+            input_tokens=trace_info.message_tokens,
+            output_tokens=trace_info.answer_tokens,
+            total_tokens=trace_info.total_tokens,
+            name="llm",
+            inputs=trace_info.inputs,
+            run_type=LangSmithRunType.llm,
+            start_time=trace_info.start_time,
+            end_time=trace_info.end_time,
+            outputs=message_data.answer,
+            extra={"metadata": metadata},
+            parent_run_id=message_id,
+            tags=["llm", str(trace_info.conversation_mode)],
+            error=trace_info.error,
+            file_list=file_list,
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+            trace_id=None,
+            dotted_order=None,
+            id=str(uuid.uuid4()),
+        )
+        self.add_run(llm_run)
+
+    def moderation_trace(self, trace_info: ModerationTraceInfo):
+        if trace_info.message_data is None:
+            return
+        langsmith_run = LangSmithRunModel(
+            name=TraceTaskName.MODERATION_TRACE.value,
+            inputs=trace_info.inputs,
+            outputs={
+                "action": trace_info.action,
+                "flagged": trace_info.flagged,
+                "preset_response": trace_info.preset_response,
+                "inputs": trace_info.inputs,
+            },
+            run_type=LangSmithRunType.tool,
+            extra={"metadata": trace_info.metadata},
+            tags=["moderation"],
+            parent_run_id=trace_info.message_id,
+            start_time=trace_info.start_time or trace_info.message_data.created_at,
+            end_time=trace_info.end_time or trace_info.message_data.updated_at,
+            id=str(uuid.uuid4()),
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+            trace_id=None,
+            dotted_order=None,
+            error="",
+            file_list=[],
+        )
+
+        self.add_run(langsmith_run)
+
+    def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo):
+        message_data = trace_info.message_data
+        if message_data is None:
+            return
+        suggested_question_run = LangSmithRunModel(
+            name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value,
+            inputs=trace_info.inputs,
+            outputs=trace_info.suggested_question,
+            run_type=LangSmithRunType.tool,
+            extra={"metadata": trace_info.metadata},
+            tags=["suggested_question"],
+            parent_run_id=trace_info.message_id,
+            start_time=trace_info.start_time or message_data.created_at,
+            end_time=trace_info.end_time or message_data.updated_at,
+            id=str(uuid.uuid4()),
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+            trace_id=None,
+            dotted_order=None,
+            error="",
+            file_list=[],
+        )
+
+        self.add_run(suggested_question_run)
+
+    def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo):
+        if trace_info.message_data is None:
+            return
+        dataset_retrieval_run = LangSmithRunModel(
+            name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value,
+            inputs=trace_info.inputs,
+            outputs={"documents": trace_info.documents},
+            run_type=LangSmithRunType.retriever,
+            extra={"metadata": trace_info.metadata},
+            tags=["dataset_retrieval"],
+            parent_run_id=trace_info.message_id,
+            start_time=trace_info.start_time or trace_info.message_data.created_at,
+            end_time=trace_info.end_time or trace_info.message_data.updated_at,
+            id=str(uuid.uuid4()),
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+            trace_id=None,
+            dotted_order=None,
+            error="",
+            file_list=[],
+        )
+
+        self.add_run(dataset_retrieval_run)
+
+    def tool_trace(self, trace_info: ToolTraceInfo):
+        tool_run = LangSmithRunModel(
+            name=trace_info.tool_name,
+            inputs=trace_info.tool_inputs,
+            outputs=trace_info.tool_outputs,
+            run_type=LangSmithRunType.tool,
+            extra={
+                "metadata": trace_info.metadata,
+            },
+            tags=["tool", trace_info.tool_name],
+            parent_run_id=trace_info.message_id,
+            start_time=trace_info.start_time,
+            end_time=trace_info.end_time,
+            file_list=[cast(str, trace_info.file_url)],
+            id=str(uuid.uuid4()),
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+            trace_id=None,
+            dotted_order=None,
+            error=trace_info.error or "",
+        )
+
+        self.add_run(tool_run)
+
+    def generate_name_trace(self, trace_info: GenerateNameTraceInfo):
+        name_run = LangSmithRunModel(
+            name=TraceTaskName.GENERATE_NAME_TRACE.value,
+            inputs=trace_info.inputs,
+            outputs=trace_info.outputs,
+            run_type=LangSmithRunType.tool,
+            extra={"metadata": trace_info.metadata},
+            tags=["generate_name"],
+            start_time=trace_info.start_time or datetime.now(),
+            end_time=trace_info.end_time or datetime.now(),
+            id=str(uuid.uuid4()),
+            serialized=None,
+            events=[],
+            session_id=None,
+            session_name=None,
+            reference_example_id=None,
+            input_attachments={},
+            output_attachments={},
+            trace_id=None,
+            dotted_order=None,
+            error="",
+            file_list=[],
+            parent_run_id=None,
+        )
+
+        self.add_run(name_run)
+
+    def add_run(self, run_data: LangSmithRunModel):
+        data = run_data.model_dump()
+        if self.project_id:
+            data["session_id"] = self.project_id
+        elif self.project_name:
+            data["session_name"] = self.project_name
+
+        data = filter_none_values(data)
+        try:
+            self.langsmith_client.create_run(**data)
+            logger.debug("LangSmith Run created successfully.")
+        except Exception as e:
+            raise ValueError(f"LangSmith Failed to create run: {str(e)}")
+
+    def update_run(self, update_run_data: LangSmithRunUpdateModel):
+        data = update_run_data.model_dump()
+        data = filter_none_values(data)
+        try:
+            self.langsmith_client.update_run(**data)
+            logger.debug("LangSmith Run updated successfully.")
+        except Exception as e:
+            raise ValueError(f"LangSmith Failed to update run: {str(e)}")
+
+    def api_check(self):
+        try:
+            random_project_name = f"test_project_{datetime.now().strftime('%Y%m%d%H%M%S')}"
+            self.langsmith_client.create_project(project_name=random_project_name)
+            self.langsmith_client.delete_project(project_name=random_project_name)
+            return True
+        except Exception as e:
+            logger.debug(f"LangSmith API check failed: {str(e)}")
+            raise ValueError(f"LangSmith API check failed: {str(e)}")
+
+    def get_project_url(self):
+        try:
+            run_data = RunBase(
+                id=uuid.uuid4(),
+                name="tool",
+                inputs={"input": "test"},
+                outputs={"output": "test"},
+                run_type=LangSmithRunType.tool,
+                start_time=datetime.now(),
+            )
+
+            project_url = self.langsmith_client.get_run_url(
+                run=run_data, project_id=self.project_id, project_name=self.project_name
+            )
+            return project_url.split("/r/")[0]
+        except Exception as e:
+            logger.debug(f"LangSmith get run url failed: {str(e)}")
+            raise ValueError(f"LangSmith get run url failed: {str(e)}")

api/core/ops/opik_trace/opik_trace.py

@@ -0,0 +1,469 @@
+import json
+import logging
+import os
+import uuid
+from datetime import datetime, timedelta
+from typing import Optional, cast
+
+from opik import Opik, Trace
+from opik.id_helpers import uuid4_to_uuid7
+
+from core.ops.base_trace_instance import BaseTraceInstance
+from core.ops.entities.config_entity import OpikConfig
+from core.ops.entities.trace_entity import (
+    BaseTraceInfo,
+    DatasetRetrievalTraceInfo,
+    GenerateNameTraceInfo,
+    MessageTraceInfo,
+    ModerationTraceInfo,
+    SuggestedQuestionTraceInfo,
+    ToolTraceInfo,
+    TraceTaskName,
+    WorkflowTraceInfo,
+)
+from extensions.ext_database import db
+from models.model import EndUser, MessageFile
+from models.workflow import WorkflowNodeExecution
+
+logger = logging.getLogger(__name__)
+
+
+def wrap_dict(key_name, data):
+    """Make sure that the input data is a dict"""
+    if not isinstance(data, dict):
+        return {key_name: data}
+
+    return data
+
+
+def wrap_metadata(metadata, **kwargs):
+    """Add common metatada to all Traces and Spans"""
+    metadata["created_from"] = "dify"
+
+    metadata.update(kwargs)
+
+    return metadata
+
+
+def prepare_opik_uuid(user_datetime: Optional[datetime], user_uuid: Optional[str]):
+    """Opik needs UUIDv7 while Dify uses UUIDv4 for identifier of most
+    messages and objects. The type-hints of BaseTraceInfo indicates that
+    objects start_time and message_id could be null which means we cannot map
+    it to a UUIDv7. Given that we have no way to identify that object
+    uniquely, generate a new random one UUIDv7 in that case.
+    """
+
+    if user_datetime is None:
+        user_datetime = datetime.now()
+
+    if user_uuid is None:
+        user_uuid = str(uuid.uuid4())
+
+    return uuid4_to_uuid7(user_datetime, user_uuid)
+
+
+class OpikDataTrace(BaseTraceInstance):
+    def __init__(
+        self,
+        opik_config: OpikConfig,
+    ):
+        super().__init__(opik_config)
+        self.opik_client = Opik(
+            project_name=opik_config.project,
+            workspace=opik_config.workspace,
+            host=opik_config.url,
+            api_key=opik_config.api_key,
+        )
+        self.project = opik_config.project
+        self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
+
+    def trace(self, trace_info: BaseTraceInfo):
+        if isinstance(trace_info, WorkflowTraceInfo):
+            self.workflow_trace(trace_info)
+        if isinstance(trace_info, MessageTraceInfo):
+            self.message_trace(trace_info)
+        if isinstance(trace_info, ModerationTraceInfo):
+            self.moderation_trace(trace_info)
+        if isinstance(trace_info, SuggestedQuestionTraceInfo):
+            self.suggested_question_trace(trace_info)
+        if isinstance(trace_info, DatasetRetrievalTraceInfo):
+            self.dataset_retrieval_trace(trace_info)
+        if isinstance(trace_info, ToolTraceInfo):
+            self.tool_trace(trace_info)
+        if isinstance(trace_info, GenerateNameTraceInfo):
+            self.generate_name_trace(trace_info)
+
+    def workflow_trace(self, trace_info: WorkflowTraceInfo):
+        dify_trace_id = trace_info.workflow_run_id
+        opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id)
+        workflow_metadata = wrap_metadata(
+            trace_info.metadata, message_id=trace_info.message_id, workflow_app_log_id=trace_info.workflow_app_log_id
+        )
+        root_span_id = None
+
+        if trace_info.message_id:
+            dify_trace_id = trace_info.message_id
+            opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id)
+
+            trace_data = {
+                "id": opik_trace_id,
+                "name": TraceTaskName.MESSAGE_TRACE.value,
+                "start_time": trace_info.start_time,
+                "end_time": trace_info.end_time,
+                "metadata": workflow_metadata,
+                "input": wrap_dict("input", trace_info.workflow_run_inputs),
+                "output": wrap_dict("output", trace_info.workflow_run_outputs),
+                "tags": ["message", "workflow"],
+                "project_name": self.project,
+            }
+            self.add_trace(trace_data)
+
+            root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id)
+            span_data = {
+                "id": root_span_id,
+                "parent_span_id": None,
+                "trace_id": opik_trace_id,
+                "name": TraceTaskName.WORKFLOW_TRACE.value,
+                "input": wrap_dict("input", trace_info.workflow_run_inputs),
+                "output": wrap_dict("output", trace_info.workflow_run_outputs),
+                "start_time": trace_info.start_time,
+                "end_time": trace_info.end_time,
+                "metadata": workflow_metadata,
+                "tags": ["workflow"],
+                "project_name": self.project,
+            }
+            self.add_span(span_data)
+        else:
+            trace_data = {
+                "id": opik_trace_id,
+                "name": TraceTaskName.MESSAGE_TRACE.value,
+                "start_time": trace_info.start_time,
+                "end_time": trace_info.end_time,
+                "metadata": workflow_metadata,
+                "input": wrap_dict("input", trace_info.workflow_run_inputs),
+                "output": wrap_dict("output", trace_info.workflow_run_outputs),
+                "tags": ["workflow"],
+                "project_name": self.project,
+            }
+            self.add_trace(trace_data)
+
+        # through workflow_run_id get all_nodes_execution
+        workflow_nodes_execution_id_records = (
+            db.session.query(WorkflowNodeExecution.id)
+            .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id)
+            .all()
+        )
+
+        for node_execution_id_record in workflow_nodes_execution_id_records:
+            node_execution = (
+                db.session.query(
+                    WorkflowNodeExecution.id,
+                    WorkflowNodeExecution.tenant_id,
+                    WorkflowNodeExecution.app_id,
+                    WorkflowNodeExecution.title,
+                    WorkflowNodeExecution.node_type,
+                    WorkflowNodeExecution.status,
+                    WorkflowNodeExecution.inputs,
+                    WorkflowNodeExecution.outputs,
+                    WorkflowNodeExecution.created_at,
+                    WorkflowNodeExecution.elapsed_time,
+                    WorkflowNodeExecution.process_data,
+                    WorkflowNodeExecution.execution_metadata,
+                )
+                .filter(WorkflowNodeExecution.id == node_execution_id_record.id)
+                .first()
+            )
+
+            if not node_execution:
+                continue
+
+            node_execution_id = node_execution.id
+            tenant_id = node_execution.tenant_id
+            app_id = node_execution.app_id
+            node_name = node_execution.title
+            node_type = node_execution.node_type
+            status = node_execution.status
+            if node_type == "llm":
+                inputs = (
+                    json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {}
+                )
+            else:
+                inputs = json.loads(node_execution.inputs) if node_execution.inputs else {}
+            outputs = json.loads(node_execution.outputs) if node_execution.outputs else {}
+            created_at = node_execution.created_at or datetime.now()
+            elapsed_time = node_execution.elapsed_time
+            finished_at = created_at + timedelta(seconds=elapsed_time)
+
+            execution_metadata = (
+                json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {}
+            )
+            metadata = execution_metadata.copy()
+            metadata.update(
+                {
+                    "workflow_run_id": trace_info.workflow_run_id,
+                    "node_execution_id": node_execution_id,
+                    "tenant_id": tenant_id,
+                    "app_id": app_id,
+                    "app_name": node_name,
+                    "node_type": node_type,
+                    "status": status,
+                }
+            )
+
+            process_data = json.loads(node_execution.process_data) if node_execution.process_data else {}
+
+            provider = None
+            model = None
+            total_tokens = 0
+            completion_tokens = 0
+            prompt_tokens = 0
+
+            if process_data and process_data.get("model_mode") == "chat":
+                run_type = "llm"
+                provider = process_data.get("model_provider", None)
+                model = process_data.get("model_name", "")
+                metadata.update(
+                    {
+                        "ls_provider": provider,
+                        "ls_model_name": model,
+                    }
+                )
+
+                try:
+                    if outputs.get("usage"):
+                        total_tokens = outputs["usage"].get("total_tokens", 0)
+                        prompt_tokens = outputs["usage"].get("prompt_tokens", 0)
+                        completion_tokens = outputs["usage"].get("completion_tokens", 0)
+                except Exception:
+                    logger.error("Failed to extract usage", exc_info=True)
+
+            else:
+                run_type = "tool"
+
+            parent_span_id = trace_info.workflow_app_log_id or trace_info.workflow_run_id
+
+            if not total_tokens:
+                total_tokens = execution_metadata.get("total_tokens", 0)
+
+            span_data = {
+                "trace_id": opik_trace_id,
+                "id": prepare_opik_uuid(created_at, node_execution_id),
+                "parent_span_id": prepare_opik_uuid(trace_info.start_time, parent_span_id),
+                "name": node_type,
+                "type": run_type,
+                "start_time": created_at,
+                "end_time": finished_at,
+                "metadata": wrap_metadata(metadata),
+                "input": wrap_dict("input", inputs),
+                "output": wrap_dict("output", outputs),
+                "tags": ["node_execution"],
+                "project_name": self.project,
+                "usage": {
+                    "total_tokens": total_tokens,
+                    "completion_tokens": completion_tokens,
+                    "prompt_tokens": prompt_tokens,
+                },
+                "model": model,
+                "provider": provider,
+            }
+
+            self.add_span(span_data)
+
+    def message_trace(self, trace_info: MessageTraceInfo):
+        # get message file data
+        file_list = cast(list[str], trace_info.file_list) or []
+        message_file_data: Optional[MessageFile] = trace_info.message_file_data
+
+        if message_file_data is not None:
+            file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else ""
+            file_list.append(file_url)
+
+        message_data = trace_info.message_data
+        if message_data is None:
+            return
+
+        metadata = trace_info.metadata
+        message_id = trace_info.message_id
+
+        user_id = message_data.from_account_id
+        metadata["user_id"] = user_id
+        metadata["file_list"] = file_list
+
+        if message_data.from_end_user_id:
+            end_user_data: Optional[EndUser] = (
+                db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first()
+            )
+            if end_user_data is not None:
+                end_user_id = end_user_data.session_id
+                metadata["end_user_id"] = end_user_id
+
+        trace_data = {
+            "id": prepare_opik_uuid(trace_info.start_time, message_id),
+            "name": TraceTaskName.MESSAGE_TRACE.value,
+            "start_time": trace_info.start_time,
+            "end_time": trace_info.end_time,
+            "metadata": wrap_metadata(metadata),
+            "input": trace_info.inputs,
+            "output": message_data.answer,
+            "tags": ["message", str(trace_info.conversation_mode)],
+            "project_name": self.project,
+        }
+        trace = self.add_trace(trace_data)
+
+        span_data = {
+            "trace_id": trace.id,
+            "name": "llm",
+            "type": "llm",
+            "start_time": trace_info.start_time,
+            "end_time": trace_info.end_time,
+            "metadata": wrap_metadata(metadata),
+            "input": {"input": trace_info.inputs},
+            "output": {"output": message_data.answer},
+            "tags": ["llm", str(trace_info.conversation_mode)],
+            "usage": {
+                "completion_tokens": trace_info.answer_tokens,
+                "prompt_tokens": trace_info.message_tokens,
+                "total_tokens": trace_info.total_tokens,
+            },
+            "project_name": self.project,
+        }
+        self.add_span(span_data)
+
+    def moderation_trace(self, trace_info: ModerationTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        start_time = trace_info.start_time or trace_info.message_data.created_at
+
+        span_data = {
+            "trace_id": prepare_opik_uuid(start_time, trace_info.message_id),
+            "name": TraceTaskName.MODERATION_TRACE.value,
+            "type": "tool",
+            "start_time": start_time,
+            "end_time": trace_info.end_time or trace_info.message_data.updated_at,
+            "metadata": wrap_metadata(trace_info.metadata),
+            "input": wrap_dict("input", trace_info.inputs),
+            "output": {
+                "action": trace_info.action,
+                "flagged": trace_info.flagged,
+                "preset_response": trace_info.preset_response,
+                "inputs": trace_info.inputs,
+            },
+            "tags": ["moderation"],
+        }
+
+        self.add_span(span_data)
+
+    def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo):
+        message_data = trace_info.message_data
+        if message_data is None:
+            return
+
+        start_time = trace_info.start_time or message_data.created_at
+
+        span_data = {
+            "trace_id": prepare_opik_uuid(start_time, trace_info.message_id),
+            "name": TraceTaskName.SUGGESTED_QUESTION_TRACE.value,
+            "type": "tool",
+            "start_time": start_time,
+            "end_time": trace_info.end_time or message_data.updated_at,
+            "metadata": wrap_metadata(trace_info.metadata),
+            "input": wrap_dict("input", trace_info.inputs),
+            "output": wrap_dict("output", trace_info.suggested_question),
+            "tags": ["suggested_question"],
+        }
+
+        self.add_span(span_data)
+
+    def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo):
+        if trace_info.message_data is None:
+            return
+
+        start_time = trace_info.start_time or trace_info.message_data.created_at
+
+        span_data = {
+            "trace_id": prepare_opik_uuid(start_time, trace_info.message_id),
+            "name": TraceTaskName.DATASET_RETRIEVAL_TRACE.value,
+            "type": "tool",
+            "start_time": start_time,
+            "end_time": trace_info.end_time or trace_info.message_data.updated_at,
+            "metadata": wrap_metadata(trace_info.metadata),
+            "input": wrap_dict("input", trace_info.inputs),
+            "output": {"documents": trace_info.documents},
+            "tags": ["dataset_retrieval"],
+        }
+
+        self.add_span(span_data)
+
+    def tool_trace(self, trace_info: ToolTraceInfo):
+        span_data = {
+            "trace_id": prepare_opik_uuid(trace_info.start_time, trace_info.message_id),
+            "name": trace_info.tool_name,
+            "type": "tool",
+            "start_time": trace_info.start_time,
+            "end_time": trace_info.end_time,
+            "metadata": wrap_metadata(trace_info.metadata),
+            "input": wrap_dict("input", trace_info.tool_inputs),
+            "output": wrap_dict("output", trace_info.tool_outputs),
+            "tags": ["tool", trace_info.tool_name],
+        }
+
+        self.add_span(span_data)
+
+    def generate_name_trace(self, trace_info: GenerateNameTraceInfo):
+        trace_data = {
+            "id": prepare_opik_uuid(trace_info.start_time, trace_info.message_id),
+            "name": TraceTaskName.GENERATE_NAME_TRACE.value,
+            "start_time": trace_info.start_time,
+            "end_time": trace_info.end_time,
+            "metadata": wrap_metadata(trace_info.metadata),
+            "input": trace_info.inputs,
+            "output": trace_info.outputs,
+            "tags": ["generate_name"],
+            "project_name": self.project,
+        }
+
+        trace = self.add_trace(trace_data)
+
+        span_data = {
+            "trace_id": trace.id,
+            "name": TraceTaskName.GENERATE_NAME_TRACE.value,
+            "start_time": trace_info.start_time,
+            "end_time": trace_info.end_time,
+            "metadata": wrap_metadata(trace_info.metadata),
+            "input": wrap_dict("input", trace_info.inputs),
+            "output": wrap_dict("output", trace_info.outputs),
+            "tags": ["generate_name"],
+        }
+
+        self.add_span(span_data)
+
+    def add_trace(self, opik_trace_data: dict) -> Trace:
+        try:
+            trace = self.opik_client.trace(**opik_trace_data)
+            logger.debug("Opik Trace created successfully")
+            return trace
+        except Exception as e:
+            raise ValueError(f"Opik Failed to create trace: {str(e)}")
+
+    def add_span(self, opik_span_data: dict):
+        try:
+            self.opik_client.span(**opik_span_data)
+            logger.debug("Opik Span created successfully")
+        except Exception as e:
+            raise ValueError(f"Opik Failed to create span: {str(e)}")
+
+    def api_check(self):
+        try:
+            self.opik_client.auth_check()
+            return True
+        except Exception as e:
+            logger.info(f"Opik API check failed: {str(e)}", exc_info=True)
+            raise ValueError(f"Opik API check failed: {str(e)}")
+
+    def get_project_url(self):
+        try:
+            return self.opik_client.get_project_url(project_name=self.project)
+        except Exception as e:
+            logger.info(f"Opik get run url failed: {str(e)}", exc_info=True)
+            raise ValueError(f"Opik get run url failed: {str(e)}")

api/core/ops/ops_trace_manager.py

@@ -0,0 +1,811 @@
+import json
+import logging
+import os
+import queue
+import threading
+import time
+from datetime import timedelta
+from typing import Any, Optional, Union
+from uuid import UUID, uuid4
+
+from flask import current_app
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
+from core.ops.entities.config_entity import (
+    OPS_FILE_PATH,
+    LangfuseConfig,
+    LangSmithConfig,
+    OpikConfig,
+    TracingProviderEnum,
+)
+from core.ops.entities.trace_entity import (
+    DatasetRetrievalTraceInfo,
+    GenerateNameTraceInfo,
+    MessageTraceInfo,
+    ModerationTraceInfo,
+    SuggestedQuestionTraceInfo,
+    TaskData,
+    ToolTraceInfo,
+    TraceTaskName,
+    WorkflowTraceInfo,
+)
+from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
+from core.ops.langsmith_trace.langsmith_trace import LangSmithDataTrace
+from core.ops.opik_trace.opik_trace import OpikDataTrace
+from core.ops.utils import get_message_data
+from extensions.ext_database import db
+from extensions.ext_storage import storage
+from models.model import App, AppModelConfig, Conversation, Message, MessageFile, TraceAppConfig
+from models.workflow import WorkflowAppLog, WorkflowRun
+from tasks.ops_trace_task import process_trace_tasks
+
+provider_config_map: dict[str, dict[str, Any]] = {
+    TracingProviderEnum.LANGFUSE.value: {
+        "config_class": LangfuseConfig,
+        "secret_keys": ["public_key", "secret_key"],
+        "other_keys": ["host", "project_key"],
+        "trace_instance": LangFuseDataTrace,
+    },
+    TracingProviderEnum.LANGSMITH.value: {
+        "config_class": LangSmithConfig,
+        "secret_keys": ["api_key"],
+        "other_keys": ["project", "endpoint"],
+        "trace_instance": LangSmithDataTrace,
+    },
+    TracingProviderEnum.OPIK.value: {
+        "config_class": OpikConfig,
+        "secret_keys": ["api_key"],
+        "other_keys": ["project", "url", "workspace"],
+        "trace_instance": OpikDataTrace,
+    },
+}
+
+
+class OpsTraceManager:
+    @classmethod
+    def encrypt_tracing_config(
+        cls, tenant_id: str, tracing_provider: str, tracing_config: dict, current_trace_config=None
+    ):
+        """
+        Encrypt tracing config.
+        :param tenant_id: tenant id
+        :param tracing_provider: tracing provider
+        :param tracing_config: tracing config dictionary to be encrypted
+        :param current_trace_config: current tracing configuration for keeping existing values
+        :return: encrypted tracing configuration
+        """
+        # Get the configuration class and the keys that require encryption
+        config_class, secret_keys, other_keys = (
+            provider_config_map[tracing_provider]["config_class"],
+            provider_config_map[tracing_provider]["secret_keys"],
+            provider_config_map[tracing_provider]["other_keys"],
+        )
+
+        new_config = {}
+        # Encrypt necessary keys
+        for key in secret_keys:
+            if key in tracing_config:
+                if "*" in tracing_config[key]:
+                    # If the key contains '*', retain the original value from the current config
+                    new_config[key] = current_trace_config.get(key, tracing_config[key])
+                else:
+                    # Otherwise, encrypt the key
+                    new_config[key] = encrypt_token(tenant_id, tracing_config[key])
+
+        for key in other_keys:
+            new_config[key] = tracing_config.get(key, "")
+
+        # Create a new instance of the config class with the new configuration
+        encrypted_config = config_class(**new_config)
+        return encrypted_config.model_dump()
+
+    @classmethod
+    def decrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_config: dict):
+        """
+        Decrypt tracing config
+        :param tenant_id: tenant id
+        :param tracing_provider: tracing provider
+        :param tracing_config: tracing config
+        :return:
+        """
+        config_class, secret_keys, other_keys = (
+            provider_config_map[tracing_provider]["config_class"],
+            provider_config_map[tracing_provider]["secret_keys"],
+            provider_config_map[tracing_provider]["other_keys"],
+        )
+        new_config = {}
+        for key in secret_keys:
+            if key in tracing_config:
+                new_config[key] = decrypt_token(tenant_id, tracing_config[key])
+
+        for key in other_keys:
+            new_config[key] = tracing_config.get(key, "")
+
+        return config_class(**new_config).model_dump()
+
+    @classmethod
+    def obfuscated_decrypt_token(cls, tracing_provider: str, decrypt_tracing_config: dict):
+        """
+        Decrypt tracing config
+        :param tracing_provider: tracing provider
+        :param decrypt_tracing_config: tracing config
+        :return:
+        """
+        config_class, secret_keys, other_keys = (
+            provider_config_map[tracing_provider]["config_class"],
+            provider_config_map[tracing_provider]["secret_keys"],
+            provider_config_map[tracing_provider]["other_keys"],
+        )
+        new_config = {}
+        for key in secret_keys:
+            if key in decrypt_tracing_config:
+                new_config[key] = obfuscated_token(decrypt_tracing_config[key])
+
+        for key in other_keys:
+            new_config[key] = decrypt_tracing_config.get(key, "")
+        return config_class(**new_config).model_dump()
+
+    @classmethod
+    def get_decrypted_tracing_config(cls, app_id: str, tracing_provider: str):
+        """
+        Get decrypted tracing config
+        :param app_id: app id
+        :param tracing_provider: tracing provider
+        :return:
+        """
+        trace_config_data: Optional[TraceAppConfig] = (
+            db.session.query(TraceAppConfig)
+            .filter(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider)
+            .first()
+        )
+
+        if not trace_config_data:
+            return None
+
+        # decrypt_token
+        app = db.session.query(App).filter(App.id == app_id).first()
+        if not app:
+            raise ValueError("App not found")
+
+        tenant_id = app.tenant_id
+        decrypt_tracing_config = cls.decrypt_tracing_config(
+            tenant_id, tracing_provider, trace_config_data.tracing_config
+        )
+
+        return decrypt_tracing_config
+
+    @classmethod
+    def get_ops_trace_instance(
+        cls,
+        app_id: Optional[Union[UUID, str]] = None,
+    ):
+        """
+        Get ops trace through model config
+        :param app_id: app_id
+        :return:
+        """
+        if isinstance(app_id, UUID):
+            app_id = str(app_id)
+
+        if app_id is None:
+            return None
+
+        app: Optional[App] = db.session.query(App).filter(App.id == app_id).first()
+
+        if app is None:
+            return None
+
+        app_ops_trace_config = json.loads(app.tracing) if app.tracing else None
+
+        if app_ops_trace_config is None:
+            return None
+
+        tracing_provider = app_ops_trace_config.get("tracing_provider")
+
+        if tracing_provider is None or tracing_provider not in provider_config_map:
+            return None
+
+        # decrypt_token
+        decrypt_trace_config = cls.get_decrypted_tracing_config(app_id, tracing_provider)
+        if app_ops_trace_config.get("enabled"):
+            trace_instance, config_class = (
+                provider_config_map[tracing_provider]["trace_instance"],
+                provider_config_map[tracing_provider]["config_class"],
+            )
+            tracing_instance = trace_instance(config_class(**decrypt_trace_config))
+            return tracing_instance
+
+        return None
+
+    @classmethod
+    def get_app_config_through_message_id(cls, message_id: str):
+        app_model_config = None
+        message_data = db.session.query(Message).filter(Message.id == message_id).first()
+        if not message_data:
+            return None
+        conversation_id = message_data.conversation_id
+        conversation_data = db.session.query(Conversation).filter(Conversation.id == conversation_id).first()
+        if not conversation_data:
+            return None
+
+        if conversation_data.app_model_config_id:
+            app_model_config = (
+                db.session.query(AppModelConfig)
+                .filter(AppModelConfig.id == conversation_data.app_model_config_id)
+                .first()
+            )
+        elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs:
+            app_model_config = conversation_data.override_model_configs
+
+        return app_model_config
+
+    @classmethod
+    def update_app_tracing_config(cls, app_id: str, enabled: bool, tracing_provider: str):
+        """
+        Update app tracing config
+        :param app_id: app id
+        :param enabled: enabled
+        :param tracing_provider: tracing provider
+        :return:
+        """
+        # auth check
+        if tracing_provider not in provider_config_map and tracing_provider is not None:
+            raise ValueError(f"Invalid tracing provider: {tracing_provider}")
+
+        app_config: Optional[App] = db.session.query(App).filter(App.id == app_id).first()
+        if not app_config:
+            raise ValueError("App not found")
+        app_config.tracing = json.dumps(
+            {
+                "enabled": enabled,
+                "tracing_provider": tracing_provider,
+            }
+        )
+        db.session.commit()
+
+    @classmethod
+    def get_app_tracing_config(cls, app_id: str):
+        """
+        Get app tracing config
+        :param app_id: app id
+        :return:
+        """
+        app: Optional[App] = db.session.query(App).filter(App.id == app_id).first()
+        if not app:
+            raise ValueError("App not found")
+        if not app.tracing:
+            return {"enabled": False, "tracing_provider": None}
+        app_trace_config = json.loads(app.tracing)
+        return app_trace_config
+
+    @staticmethod
+    def check_trace_config_is_effective(tracing_config: dict, tracing_provider: str):
+        """
+        Check trace config is effective
+        :param tracing_config: tracing config
+        :param tracing_provider: tracing provider
+        :return:
+        """
+        config_type, trace_instance = (
+            provider_config_map[tracing_provider]["config_class"],
+            provider_config_map[tracing_provider]["trace_instance"],
+        )
+        tracing_config = config_type(**tracing_config)
+        return trace_instance(tracing_config).api_check()
+
+    @staticmethod
+    def get_trace_config_project_key(tracing_config: dict, tracing_provider: str):
+        """
+        get trace config is project key
+        :param tracing_config: tracing config
+        :param tracing_provider: tracing provider
+        :return:
+        """
+        config_type, trace_instance = (
+            provider_config_map[tracing_provider]["config_class"],
+            provider_config_map[tracing_provider]["trace_instance"],
+        )
+        tracing_config = config_type(**tracing_config)
+        return trace_instance(tracing_config).get_project_key()
+
+    @staticmethod
+    def get_trace_config_project_url(tracing_config: dict, tracing_provider: str):
+        """
+        get trace config is project key
+        :param tracing_config: tracing config
+        :param tracing_provider: tracing provider
+        :return:
+        """
+        config_type, trace_instance = (
+            provider_config_map[tracing_provider]["config_class"],
+            provider_config_map[tracing_provider]["trace_instance"],
+        )
+        tracing_config = config_type(**tracing_config)
+        return trace_instance(tracing_config).get_project_url()
+
+
+class TraceTask:
+    def __init__(
+        self,
+        trace_type: Any,
+        message_id: Optional[str] = None,
+        workflow_run: Optional[WorkflowRun] = None,
+        conversation_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        timer: Optional[Any] = None,
+        **kwargs,
+    ):
+        self.trace_type = trace_type
+        self.message_id = message_id
+        self.workflow_run_id = workflow_run.id if workflow_run else None
+        self.conversation_id = conversation_id
+        self.user_id = user_id
+        self.timer = timer
+        self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
+        self.app_id = None
+
+        self.kwargs = kwargs
+
+    def execute(self):
+        return self.preprocess()
+
+    def preprocess(self):
+        preprocess_map = {
+            TraceTaskName.CONVERSATION_TRACE: lambda: self.conversation_trace(**self.kwargs),
+            TraceTaskName.WORKFLOW_TRACE: lambda: self.workflow_trace(
+                workflow_run_id=self.workflow_run_id, conversation_id=self.conversation_id, user_id=self.user_id
+            ),
+            TraceTaskName.MESSAGE_TRACE: lambda: self.message_trace(message_id=self.message_id),
+            TraceTaskName.MODERATION_TRACE: lambda: self.moderation_trace(
+                message_id=self.message_id, timer=self.timer, **self.kwargs
+            ),
+            TraceTaskName.SUGGESTED_QUESTION_TRACE: lambda: self.suggested_question_trace(
+                message_id=self.message_id, timer=self.timer, **self.kwargs
+            ),
+            TraceTaskName.DATASET_RETRIEVAL_TRACE: lambda: self.dataset_retrieval_trace(
+                message_id=self.message_id, timer=self.timer, **self.kwargs
+            ),
+            TraceTaskName.TOOL_TRACE: lambda: self.tool_trace(
+                message_id=self.message_id, timer=self.timer, **self.kwargs
+            ),
+            TraceTaskName.GENERATE_NAME_TRACE: lambda: self.generate_name_trace(
+                conversation_id=self.conversation_id, timer=self.timer, **self.kwargs
+            ),
+        }
+
+        return preprocess_map.get(self.trace_type, lambda: None)()
+
+    # process methods for different trace types
+    def conversation_trace(self, **kwargs):
+        return kwargs
+
+    def workflow_trace(
+        self,
+        *,
+        workflow_run_id: str | None,
+        conversation_id: str | None,
+        user_id: str | None,
+    ):
+        if not workflow_run_id:
+            return {}
+
+        with Session(db.engine) as session:
+            workflow_run_stmt = select(WorkflowRun).where(WorkflowRun.id == workflow_run_id)
+            workflow_run = session.scalars(workflow_run_stmt).first()
+            if not workflow_run:
+                raise ValueError("Workflow run not found")
+
+            workflow_id = workflow_run.workflow_id
+            tenant_id = workflow_run.tenant_id
+            workflow_run_id = workflow_run.id
+            workflow_run_elapsed_time = workflow_run.elapsed_time
+            workflow_run_status = workflow_run.status
+            workflow_run_inputs = workflow_run.inputs_dict
+            workflow_run_outputs = workflow_run.outputs_dict
+            workflow_run_version = workflow_run.version
+            error = workflow_run.error or ""
+
+            total_tokens = workflow_run.total_tokens
+
+            file_list = workflow_run_inputs.get("sys.file") or []
+            query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or ""
+
+            # get workflow_app_log_id
+            workflow_app_log_data_stmt = select(WorkflowAppLog.id).where(
+                WorkflowAppLog.tenant_id == tenant_id,
+                WorkflowAppLog.app_id == workflow_run.app_id,
+                WorkflowAppLog.workflow_run_id == workflow_run.id,
+            )
+            workflow_app_log_id = session.scalar(workflow_app_log_data_stmt)
+            # get message_id
+            message_id = None
+            if conversation_id:
+                message_data_stmt = select(Message.id).where(
+                    Message.conversation_id == conversation_id,
+                    Message.workflow_run_id == workflow_run_id,
+                )
+                message_id = session.scalar(message_data_stmt)
+
+            metadata = {
+                "workflow_id": workflow_id,
+                "conversation_id": conversation_id,
+                "workflow_run_id": workflow_run_id,
+                "tenant_id": tenant_id,
+                "elapsed_time": workflow_run_elapsed_time,
+                "status": workflow_run_status,
+                "version": workflow_run_version,
+                "total_tokens": total_tokens,
+                "file_list": file_list,
+                "triggered_form": workflow_run.triggered_from,
+                "user_id": user_id,
+            }
+
+            workflow_trace_info = WorkflowTraceInfo(
+                workflow_data=workflow_run.to_dict(),
+                conversation_id=conversation_id,
+                workflow_id=workflow_id,
+                tenant_id=tenant_id,
+                workflow_run_id=workflow_run_id,
+                workflow_run_elapsed_time=workflow_run_elapsed_time,
+                workflow_run_status=workflow_run_status,
+                workflow_run_inputs=workflow_run_inputs,
+                workflow_run_outputs=workflow_run_outputs,
+                workflow_run_version=workflow_run_version,
+                error=error,
+                total_tokens=total_tokens,
+                file_list=file_list,
+                query=query,
+                metadata=metadata,
+                workflow_app_log_id=workflow_app_log_id,
+                message_id=message_id,
+                start_time=workflow_run.created_at,
+                end_time=workflow_run.finished_at,
+            )
+        return workflow_trace_info
+
+    def message_trace(self, message_id: str | None):
+        if not message_id:
+            return {}
+        message_data = get_message_data(message_id)
+        if not message_data:
+            return {}
+        conversation_mode_stmt = select(Conversation.mode).where(Conversation.id == message_data.conversation_id)
+        conversation_mode = db.session.scalars(conversation_mode_stmt).all()
+        if not conversation_mode or len(conversation_mode) == 0:
+            return {}
+        conversation_mode = conversation_mode[0]
+        created_at = message_data.created_at
+        inputs = message_data.message
+
+        # get message file data
+        message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first()
+        file_list = []
+        if message_file_data and message_file_data.url is not None:
+            file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else ""
+            file_list.append(file_url)
+
+        metadata = {
+            "conversation_id": message_data.conversation_id,
+            "ls_provider": message_data.model_provider,
+            "ls_model_name": message_data.model_id,
+            "status": message_data.status,
+            "from_end_user_id": message_data.from_end_user_id,
+            "from_account_id": message_data.from_account_id,
+            "agent_based": message_data.agent_based,
+            "workflow_run_id": message_data.workflow_run_id,
+            "from_source": message_data.from_source,
+            "message_id": message_id,
+        }
+
+        message_tokens = message_data.message_tokens
+
+        message_trace_info = MessageTraceInfo(
+            message_id=message_id,
+            message_data=message_data.to_dict(),
+            conversation_model=conversation_mode,
+            message_tokens=message_tokens,
+            answer_tokens=message_data.answer_tokens,
+            total_tokens=message_tokens + message_data.answer_tokens,
+            error=message_data.error or "",
+            inputs=inputs,
+            outputs=message_data.answer,
+            file_list=file_list,
+            start_time=created_at,
+            end_time=created_at + timedelta(seconds=message_data.provider_response_latency),
+            metadata=metadata,
+            message_file_data=message_file_data,
+            conversation_mode=conversation_mode,
+        )
+
+        return message_trace_info
+
+    def moderation_trace(self, message_id, timer, **kwargs):
+        moderation_result = kwargs.get("moderation_result")
+        if not moderation_result:
+            return {}
+        inputs = kwargs.get("inputs")
+        message_data = get_message_data(message_id)
+        if not message_data:
+            return {}
+        metadata = {
+            "message_id": message_id,
+            "action": moderation_result.action,
+            "preset_response": moderation_result.preset_response,
+            "query": moderation_result.query,
+        }
+
+        # get workflow_app_log_id
+        workflow_app_log_id = None
+        if message_data.workflow_run_id:
+            workflow_app_log_data = (
+                db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first()
+            )
+            workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
+
+        moderation_trace_info = ModerationTraceInfo(
+            message_id=workflow_app_log_id or message_id,
+            inputs=inputs,
+            message_data=message_data.to_dict(),
+            flagged=moderation_result.flagged,
+            action=moderation_result.action,
+            preset_response=moderation_result.preset_response,
+            query=moderation_result.query,
+            start_time=timer.get("start"),
+            end_time=timer.get("end"),
+            metadata=metadata,
+        )
+
+        return moderation_trace_info
+
+    def suggested_question_trace(self, message_id, timer, **kwargs):
+        suggested_question = kwargs.get("suggested_question", [])
+        message_data = get_message_data(message_id)
+        if not message_data:
+            return {}
+        metadata = {
+            "message_id": message_id,
+            "ls_provider": message_data.model_provider,
+            "ls_model_name": message_data.model_id,
+            "status": message_data.status,
+            "from_end_user_id": message_data.from_end_user_id,
+            "from_account_id": message_data.from_account_id,
+            "agent_based": message_data.agent_based,
+            "workflow_run_id": message_data.workflow_run_id,
+            "from_source": message_data.from_source,
+        }
+
+        # get workflow_app_log_id
+        workflow_app_log_id = None
+        if message_data.workflow_run_id:
+            workflow_app_log_data = (
+                db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first()
+            )
+            workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
+
+        suggested_question_trace_info = SuggestedQuestionTraceInfo(
+            message_id=workflow_app_log_id or message_id,
+            message_data=message_data.to_dict(),
+            inputs=message_data.message,
+            outputs=message_data.answer,
+            start_time=timer.get("start"),
+            end_time=timer.get("end"),
+            metadata=metadata,
+            total_tokens=message_data.message_tokens + message_data.answer_tokens,
+            status=message_data.status,
+            error=message_data.error,
+            from_account_id=message_data.from_account_id,
+            agent_based=message_data.agent_based,
+            from_source=message_data.from_source,
+            model_provider=message_data.model_provider,
+            model_id=message_data.model_id,
+            suggested_question=suggested_question,
+            level=message_data.status,
+            status_message=message_data.error,
+        )
+
+        return suggested_question_trace_info
+
+    def dataset_retrieval_trace(self, message_id, timer, **kwargs):
+        documents = kwargs.get("documents")
+        message_data = get_message_data(message_id)
+        if not message_data:
+            return {}
+
+        metadata = {
+            "message_id": message_id,
+            "ls_provider": message_data.model_provider,
+            "ls_model_name": message_data.model_id,
+            "status": message_data.status,
+            "from_end_user_id": message_data.from_end_user_id,
+            "from_account_id": message_data.from_account_id,
+            "agent_based": message_data.agent_based,
+            "workflow_run_id": message_data.workflow_run_id,
+            "from_source": message_data.from_source,
+        }
+
+        dataset_retrieval_trace_info = DatasetRetrievalTraceInfo(
+            message_id=message_id,
+            inputs=message_data.query or message_data.inputs,
+            documents=[doc.model_dump() for doc in documents] if documents else [],
+            start_time=timer.get("start"),
+            end_time=timer.get("end"),
+            metadata=metadata,
+            message_data=message_data.to_dict(),
+        )
+
+        return dataset_retrieval_trace_info
+
+    def tool_trace(self, message_id, timer, **kwargs):
+        tool_name = kwargs.get("tool_name", "")
+        tool_inputs = kwargs.get("tool_inputs", {})
+        tool_outputs = kwargs.get("tool_outputs", {})
+        message_data = get_message_data(message_id)
+        if not message_data:
+            return {}
+        tool_config = {}
+        time_cost = 0
+        error = None
+        tool_parameters = {}
+        created_time = message_data.created_at
+        end_time = message_data.updated_at
+        agent_thoughts = message_data.agent_thoughts
+        for agent_thought in agent_thoughts:
+            if tool_name in agent_thought.tools:
+                created_time = agent_thought.created_at
+                tool_meta_data = agent_thought.tool_meta.get(tool_name, {})
+                tool_config = tool_meta_data.get("tool_config", {})
+                time_cost = tool_meta_data.get("time_cost", 0)
+                end_time = created_time + timedelta(seconds=time_cost)
+                error = tool_meta_data.get("error", "")
+                tool_parameters = tool_meta_data.get("tool_parameters", {})
+        metadata = {
+            "message_id": message_id,
+            "tool_name": tool_name,
+            "tool_inputs": tool_inputs,
+            "tool_outputs": tool_outputs,
+            "tool_config": tool_config,
+            "time_cost": time_cost,
+            "error": error,
+            "tool_parameters": tool_parameters,
+        }
+
+        file_url = ""
+        message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first()
+        if message_file_data:
+            message_file_id = message_file_data.id if message_file_data else None
+            type = message_file_data.type
+            created_by_role = message_file_data.created_by_role
+            created_user_id = message_file_data.created_by
+            file_url = f"{self.file_base_url}/{message_file_data.url}"
+
+            metadata.update(
+                {
+                    "message_file_id": message_file_id,
+                    "created_by_role": created_by_role,
+                    "created_user_id": created_user_id,
+                    "type": type,
+                }
+            )
+
+        tool_trace_info = ToolTraceInfo(
+            message_id=message_id,
+            message_data=message_data.to_dict(),
+            tool_name=tool_name,
+            start_time=timer.get("start") if timer else created_time,
+            end_time=timer.get("end") if timer else end_time,
+            tool_inputs=tool_inputs,
+            tool_outputs=tool_outputs,
+            metadata=metadata,
+            message_file_data=message_file_data,
+            error=error,
+            inputs=message_data.message,
+            outputs=message_data.answer,
+            tool_config=tool_config,
+            time_cost=time_cost,
+            tool_parameters=tool_parameters,
+            file_url=file_url,
+        )
+
+        return tool_trace_info
+
+    def generate_name_trace(self, conversation_id, timer, **kwargs):
+        generate_conversation_name = kwargs.get("generate_conversation_name")
+        inputs = kwargs.get("inputs")
+        tenant_id = kwargs.get("tenant_id")
+        if not tenant_id:
+            return {}
+        start_time = timer.get("start")
+        end_time = timer.get("end")
+
+        metadata = {
+            "conversation_id": conversation_id,
+            "tenant_id": tenant_id,
+        }
+
+        generate_name_trace_info = GenerateNameTraceInfo(
+            conversation_id=conversation_id,
+            inputs=inputs,
+            outputs=generate_conversation_name,
+            start_time=start_time,
+            end_time=end_time,
+            metadata=metadata,
+            tenant_id=tenant_id,
+        )
+
+        return generate_name_trace_info
+
+
+trace_manager_timer: Optional[threading.Timer] = None
+trace_manager_queue: queue.Queue = queue.Queue()
+trace_manager_interval = int(os.getenv("TRACE_QUEUE_MANAGER_INTERVAL", 5))
+trace_manager_batch_size = int(os.getenv("TRACE_QUEUE_MANAGER_BATCH_SIZE", 100))
+
+
+class TraceQueueManager:
+    def __init__(self, app_id=None, user_id=None):
+        global trace_manager_timer
+
+        self.app_id = app_id
+        self.user_id = user_id
+        self.trace_instance = OpsTraceManager.get_ops_trace_instance(app_id)
+        self.flask_app = current_app._get_current_object()  # type: ignore
+        if trace_manager_timer is None:
+            self.start_timer()
+
+    def add_trace_task(self, trace_task: TraceTask):
+        global trace_manager_timer, trace_manager_queue
+        try:
+            if self.trace_instance:
+                trace_task.app_id = self.app_id
+                trace_manager_queue.put(trace_task)
+        except Exception as e:
+            logging.exception(f"Error adding trace task, trace_type {trace_task.trace_type}")
+        finally:
+            self.start_timer()
+
+    def collect_tasks(self):
+        global trace_manager_queue
+        tasks: list[TraceTask] = []
+        while len(tasks) < trace_manager_batch_size and not trace_manager_queue.empty():
+            task = trace_manager_queue.get_nowait()
+            tasks.append(task)
+            trace_manager_queue.task_done()
+        return tasks
+
+    def run(self):
+        try:
+            tasks = self.collect_tasks()
+            if tasks:
+                self.send_to_celery(tasks)
+        except Exception as e:
+            logging.exception("Error processing trace tasks")
+
+    def start_timer(self):
+        global trace_manager_timer
+        if trace_manager_timer is None or not trace_manager_timer.is_alive():
+            trace_manager_timer = threading.Timer(trace_manager_interval, self.run)
+            trace_manager_timer.name = f"trace_manager_timer_{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}"
+            trace_manager_timer.daemon = False
+            trace_manager_timer.start()
+
+    def send_to_celery(self, tasks: list[TraceTask]):
+        with self.flask_app.app_context():
+            for task in tasks:
+                if task.app_id is None:
+                    continue
+                file_id = uuid4().hex
+                trace_info = task.execute()
+                task_data = TaskData(
+                    app_id=task.app_id,
+                    trace_info_type=type(trace_info).__name__,
+                    trace_info=trace_info.model_dump() if trace_info else None,
+                )
+                file_path = f"{OPS_FILE_PATH}{task.app_id}/{file_id}.json"
+                storage.save(file_path, task_data.model_dump_json().encode("utf-8"))
+                file_info = {
+                    "file_id": file_id,
+                    "app_id": task.app_id,
+                }
+                process_trace_tasks.delay(file_info)

api/core/ops/utils.py

@@ -0,0 +1,62 @@
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Optional, Union
+
+from extensions.ext_database import db
+from models.model import Message
+
+
+def filter_none_values(data: dict):
+    new_data = {}
+    for key, value in data.items():
+        if value is None:
+            continue
+        if isinstance(value, datetime):
+            new_data[key] = value.isoformat()
+        else:
+            new_data[key] = value
+    return new_data
+
+
+def get_message_data(message_id: str):
+    return db.session.query(Message).filter(Message.id == message_id).first()
+
+
+@contextmanager
+def measure_time():
+    timing_info = {"start": datetime.now(), "end": None}
+    try:
+        yield timing_info
+    finally:
+        timing_info["end"] = datetime.now()
+
+
+def replace_text_with_content(data):
+    if isinstance(data, dict):
+        new_data = {}
+        for key, value in data.items():
+            if key == "text":
+                new_data["content"] = value
+            else:
+                new_data[key] = replace_text_with_content(value)
+        return new_data
+    elif isinstance(data, list):
+        return [replace_text_with_content(item) for item in data]
+    else:
+        return data
+
+
+def generate_dotted_order(
+    run_id: str, start_time: Union[str, datetime], parent_dotted_order: Optional[str] = None
+) -> str:
+    """
+    generate dotted_order for langsmith
+    """
+    start_time = datetime.fromisoformat(start_time) if isinstance(start_time, str) else start_time
+    timestamp = start_time.strftime("%Y%m%dT%H%M%S%f")[:-3] + "Z"
+    current_segment = f"{timestamp}{run_id}"
+
+    if parent_dotted_order is None:
+        return current_segment
+
+    return f"{parent_dotted_order}.{current_segment}"

api/core/prompt/advanced_prompt_transform.py

@@ -0,0 +1,287 @@
+from collections.abc import Mapping, Sequence
+from typing import Optional, cast
+
+from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
+from core.file import file_manager
+from core.file.models import File
+from core.helper.code_executor.jinja2.jinja2_formatter import Jinja2Formatter
+from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_runtime.entities import (
+    AssistantPromptMessage,
+    PromptMessage,
+    PromptMessageContent,
+    PromptMessageRole,
+    SystemPromptMessage,
+    TextPromptMessageContent,
+    UserPromptMessage,
+)
+from core.model_runtime.entities.message_entities import ImagePromptMessageContent
+from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
+from core.prompt.prompt_transform import PromptTransform
+from core.prompt.utils.prompt_template_parser import PromptTemplateParser
+from core.workflow.entities.variable_pool import VariablePool
+
+
+class AdvancedPromptTransform(PromptTransform):
+    """
+    Advanced Prompt Transform for Workflow LLM Node.
+    """
+
+    def __init__(
+        self,
+        with_variable_tmpl: bool = False,
+        image_detail_config: ImagePromptMessageContent.DETAIL = ImagePromptMessageContent.DETAIL.LOW,
+    ) -> None:
+        self.with_variable_tmpl = with_variable_tmpl
+        self.image_detail_config = image_detail_config
+
+    def get_prompt(
+        self,
+        *,
+        prompt_template: Sequence[ChatModelMessage] | CompletionModelPromptTemplate,
+        inputs: Mapping[str, str],
+        query: str,
+        files: Sequence[File],
+        context: Optional[str],
+        memory_config: Optional[MemoryConfig],
+        memory: Optional[TokenBufferMemory],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> list[PromptMessage]:
+        prompt_messages = []
+
+        if isinstance(prompt_template, CompletionModelPromptTemplate):
+            prompt_messages = self._get_completion_model_prompt_messages(
+                prompt_template=prompt_template,
+                inputs=inputs,
+                query=query,
+                files=files,
+                context=context,
+                memory_config=memory_config,
+                memory=memory,
+                model_config=model_config,
+            )
+        elif isinstance(prompt_template, list) and all(isinstance(item, ChatModelMessage) for item in prompt_template):
+            prompt_messages = self._get_chat_model_prompt_messages(
+                prompt_template=prompt_template,
+                inputs=inputs,
+                query=query,
+                files=files,
+                context=context,
+                memory_config=memory_config,
+                memory=memory,
+                model_config=model_config,
+            )
+
+        return prompt_messages
+
+    def _get_completion_model_prompt_messages(
+        self,
+        prompt_template: CompletionModelPromptTemplate,
+        inputs: Mapping[str, str],
+        query: Optional[str],
+        files: Sequence[File],
+        context: Optional[str],
+        memory_config: Optional[MemoryConfig],
+        memory: Optional[TokenBufferMemory],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> list[PromptMessage]:
+        """
+        Get completion model prompt messages.
+        """
+        raw_prompt = prompt_template.text
+
+        prompt_messages: list[PromptMessage] = []
+
+        if prompt_template.edition_type == "basic" or not prompt_template.edition_type:
+            parser = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl)
+            prompt_inputs: Mapping[str, str] = {k: inputs[k] for k in parser.variable_keys if k in inputs}
+
+            prompt_inputs = self._set_context_variable(context, parser, prompt_inputs)
+
+            if memory and memory_config and memory_config.role_prefix:
+                role_prefix = memory_config.role_prefix
+                prompt_inputs = self._set_histories_variable(
+                    memory=memory,
+                    memory_config=memory_config,
+                    raw_prompt=raw_prompt,
+                    role_prefix=role_prefix,
+                    parser=parser,
+                    prompt_inputs=prompt_inputs,
+                    model_config=model_config,
+                )
+
+            if query:
+                prompt_inputs = self._set_query_variable(query, parser, prompt_inputs)
+
+            prompt = parser.format(prompt_inputs)
+        else:
+            prompt = raw_prompt
+            prompt_inputs = inputs
+
+            prompt = Jinja2Formatter.format(prompt, prompt_inputs)
+
+        if files:
+            prompt_message_contents: list[PromptMessageContent] = []
+            prompt_message_contents.append(TextPromptMessageContent(data=prompt))
+            for file in files:
+                prompt_message_contents.append(file_manager.to_prompt_message_content(file))
+
+            prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
+        else:
+            prompt_messages.append(UserPromptMessage(content=prompt))
+
+        return prompt_messages
+
+    def _get_chat_model_prompt_messages(
+        self,
+        prompt_template: list[ChatModelMessage],
+        inputs: Mapping[str, str],
+        query: Optional[str],
+        files: Sequence[File],
+        context: Optional[str],
+        memory_config: Optional[MemoryConfig],
+        memory: Optional[TokenBufferMemory],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> list[PromptMessage]:
+        """
+        Get chat model prompt messages.
+        """
+        prompt_messages: list[PromptMessage] = []
+        for prompt_item in prompt_template:
+            raw_prompt = prompt_item.text
+
+            if prompt_item.edition_type == "basic" or not prompt_item.edition_type:
+                if self.with_variable_tmpl:
+                    vp = VariablePool()
+                    for k, v in inputs.items():
+                        if k.startswith("#"):
+                            vp.add(k[1:-1].split("."), v)
+                    raw_prompt = raw_prompt.replace("{{#context#}}", context or "")
+                    prompt = vp.convert_template(raw_prompt).text
+                else:
+                    parser = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl)
+                    prompt_inputs: Mapping[str, str] = {k: inputs[k] for k in parser.variable_keys if k in inputs}
+                    prompt_inputs = self._set_context_variable(
+                        context=context, parser=parser, prompt_inputs=prompt_inputs
+                    )
+                    prompt = parser.format(prompt_inputs)
+            elif prompt_item.edition_type == "jinja2":
+                prompt = raw_prompt
+                prompt_inputs = inputs
+                prompt = Jinja2Formatter.format(template=prompt, inputs=prompt_inputs)
+            else:
+                raise ValueError(f"Invalid edition type: {prompt_item.edition_type}")
+
+            if prompt_item.role == PromptMessageRole.USER:
+                prompt_messages.append(UserPromptMessage(content=prompt))
+            elif prompt_item.role == PromptMessageRole.SYSTEM and prompt:
+                prompt_messages.append(SystemPromptMessage(content=prompt))
+            elif prompt_item.role == PromptMessageRole.ASSISTANT:
+                prompt_messages.append(AssistantPromptMessage(content=prompt))
+
+        if query and memory_config and memory_config.query_prompt_template:
+            parser = PromptTemplateParser(
+                template=memory_config.query_prompt_template, with_variable_tmpl=self.with_variable_tmpl
+            )
+            prompt_inputs = {k: inputs[k] for k in parser.variable_keys if k in inputs}
+            prompt_inputs["#sys.query#"] = query
+
+            prompt_inputs = self._set_context_variable(context, parser, prompt_inputs)
+
+            query = parser.format(prompt_inputs)
+
+        if memory and memory_config:
+            prompt_messages = self._append_chat_histories(memory, memory_config, prompt_messages, model_config)
+
+            if files and query is not None:
+                prompt_message_contents: list[PromptMessageContent] = []
+                prompt_message_contents.append(TextPromptMessageContent(data=query))
+                for file in files:
+                    prompt_message_contents.append(file_manager.to_prompt_message_content(file))
+                prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
+            else:
+                prompt_messages.append(UserPromptMessage(content=query))
+        elif files:
+            if not query:
+                # get last message
+                last_message = prompt_messages[-1] if prompt_messages else None
+                if last_message and last_message.role == PromptMessageRole.USER:
+                    # get last user message content and add files
+                    prompt_message_contents = [TextPromptMessageContent(data=cast(str, last_message.content))]
+                    for file in files:
+                        prompt_message_contents.append(file_manager.to_prompt_message_content(file))
+
+                    last_message.content = prompt_message_contents
+                else:
+                    prompt_message_contents = [TextPromptMessageContent(data="")]  # not for query
+                    for file in files:
+                        prompt_message_contents.append(file_manager.to_prompt_message_content(file))
+
+                    prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
+            else:
+                prompt_message_contents = [TextPromptMessageContent(data=query)]
+                for file in files:
+                    prompt_message_contents.append(file_manager.to_prompt_message_content(file))
+
+                prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
+        elif query:
+            prompt_messages.append(UserPromptMessage(content=query))
+
+        return prompt_messages
+
+    def _set_context_variable(
+        self, context: str | None, parser: PromptTemplateParser, prompt_inputs: Mapping[str, str]
+    ) -> Mapping[str, str]:
+        prompt_inputs = dict(prompt_inputs)
+        if "#context#" in parser.variable_keys:
+            if context:
+                prompt_inputs["#context#"] = context
+            else:
+                prompt_inputs["#context#"] = ""
+
+        return prompt_inputs
+
+    def _set_query_variable(
+        self, query: str, parser: PromptTemplateParser, prompt_inputs: Mapping[str, str]
+    ) -> Mapping[str, str]:
+        prompt_inputs = dict(prompt_inputs)
+        if "#query#" in parser.variable_keys:
+            if query:
+                prompt_inputs["#query#"] = query
+            else:
+                prompt_inputs["#query#"] = ""
+
+        return prompt_inputs
+
+    def _set_histories_variable(
+        self,
+        memory: TokenBufferMemory,
+        memory_config: MemoryConfig,
+        raw_prompt: str,
+        role_prefix: MemoryConfig.RolePrefix,
+        parser: PromptTemplateParser,
+        prompt_inputs: Mapping[str, str],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> Mapping[str, str]:
+        prompt_inputs = dict(prompt_inputs)
+        if "#histories#" in parser.variable_keys:
+            if memory:
+                inputs = {"#histories#": "", **prompt_inputs}
+                parser = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl)
+                prompt_inputs = {k: inputs[k] for k in parser.variable_keys if k in inputs}
+                tmp_human_message = UserPromptMessage(content=parser.format(prompt_inputs))
+
+                rest_tokens = self._calculate_rest_token([tmp_human_message], model_config)
+
+                histories = self._get_history_messages_from_memory(
+                    memory=memory,
+                    memory_config=memory_config,
+                    max_token_limit=rest_tokens,
+                    human_prefix=role_prefix.user,
+                    ai_prefix=role_prefix.assistant,
+                )
+                prompt_inputs["#histories#"] = histories
+            else:
+                prompt_inputs["#histories#"] = ""
+
+        return prompt_inputs

api/core/prompt/agent_history_prompt_transform.py

@@ -0,0 +1,80 @@
+from typing import Optional, cast
+
+from core.app.entities.app_invoke_entities import (
+    ModelConfigWithCredentialsEntity,
+)
+from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_runtime.entities.message_entities import (
+    PromptMessage,
+    SystemPromptMessage,
+    UserPromptMessage,
+)
+from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
+from core.prompt.prompt_transform import PromptTransform
+
+
+class AgentHistoryPromptTransform(PromptTransform):
+    """
+    History Prompt Transform for Agent App
+    """
+
+    def __init__(
+        self,
+        model_config: ModelConfigWithCredentialsEntity,
+        prompt_messages: list[PromptMessage],
+        history_messages: list[PromptMessage],
+        memory: Optional[TokenBufferMemory] = None,
+    ):
+        self.model_config = model_config
+        self.prompt_messages = prompt_messages
+        self.history_messages = history_messages
+        self.memory = memory
+
+    def get_prompt(self) -> list[PromptMessage]:
+        prompt_messages: list[PromptMessage] = []
+        num_system = 0
+        for prompt_message in self.history_messages:
+            if isinstance(prompt_message, SystemPromptMessage):
+                prompt_messages.append(prompt_message)
+                num_system += 1
+
+        if not self.memory:
+            return prompt_messages
+
+        max_token_limit = self._calculate_rest_token(self.prompt_messages, self.model_config)
+
+        model_type_instance = self.model_config.provider_model_bundle.model_type_instance
+        model_type_instance = cast(LargeLanguageModel, model_type_instance)
+
+        curr_message_tokens = model_type_instance.get_num_tokens(
+            self.memory.model_instance.model, self.memory.model_instance.credentials, self.history_messages
+        )
+        if curr_message_tokens <= max_token_limit:
+            return self.history_messages
+
+        # number of prompt has been appended in current message
+        num_prompt = 0
+        # append prompt messages in desc order
+        for prompt_message in self.history_messages[::-1]:
+            if isinstance(prompt_message, SystemPromptMessage):
+                continue
+            prompt_messages.append(prompt_message)
+            num_prompt += 1
+            # a message is start with UserPromptMessage
+            if isinstance(prompt_message, UserPromptMessage):
+                curr_message_tokens = model_type_instance.get_num_tokens(
+                    self.memory.model_instance.model, self.memory.model_instance.credentials, prompt_messages
+                )
+                # if current message token is overflow, drop all the prompts in current message and break
+                if curr_message_tokens > max_token_limit:
+                    prompt_messages = prompt_messages[:-num_prompt]
+                    break
+                num_prompt = 0
+        # return prompt messages in asc order
+        message_prompts = prompt_messages[num_system:]
+        message_prompts.reverse()
+
+        # merge system and message prompt
+        prompt_messages = prompt_messages[:num_system]
+        prompt_messages.extend(message_prompts)
+        return prompt_messages

api/core/prompt/entities/advanced_prompt_entities.py

@@ -0,0 +1,50 @@
+from typing import Literal, Optional
+
+from pydantic import BaseModel
+
+from core.model_runtime.entities.message_entities import PromptMessageRole
+
+
+class ChatModelMessage(BaseModel):
+    """
+    Chat Message.
+    """
+
+    text: str
+    role: PromptMessageRole
+    edition_type: Optional[Literal["basic", "jinja2"]] = None
+
+
+class CompletionModelPromptTemplate(BaseModel):
+    """
+    Completion Model Prompt Template.
+    """
+
+    text: str
+    edition_type: Optional[Literal["basic", "jinja2"]] = None
+
+
+class MemoryConfig(BaseModel):
+    """
+    Memory Config.
+    """
+
+    class RolePrefix(BaseModel):
+        """
+        Role Prefix.
+        """
+
+        user: str
+        assistant: str
+
+    class WindowConfig(BaseModel):
+        """
+        Window Config.
+        """
+
+        enabled: bool
+        size: Optional[int] = None
+
+    role_prefix: Optional[RolePrefix] = None
+    window: WindowConfig
+    query_prompt_template: Optional[str] = None

api/core/prompt/prompt_templates/advanced_prompt_templates.py

@@ -0,0 +1,45 @@
+CONTEXT = "Use the following context as your learned knowledge, inside <context></context> XML tags.\n\n<context>\n{{#context#}}\n</context>\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n"  # noqa: E501
+
+BAICHUAN_CONTEXT = "用户在与一个客观的助手对话。助手会尊重找到的材料，给出全面专业的解释，但不会过度演绎。同时回答中不会暴露引用的材料：\n\n```\n{{#context#}}\n```\n\n"  # noqa: E501
+
+CHAT_APP_COMPLETION_PROMPT_CONFIG = {
+    "completion_prompt_config": {
+        "prompt": {
+            "text": "{{#pre_prompt#}}\nHere are the chat histories between human and assistant, inside <histories></histories> XML tags.\n\n<histories>\n{{#histories#}}\n</histories>\n\n\nHuman: {{#query#}}\n\nAssistant: "  # noqa: E501
+        },
+        "conversation_histories_role": {"user_prefix": "Human", "assistant_prefix": "Assistant"},
+    },
+    "stop": ["Human:"],
+}
+
+CHAT_APP_CHAT_PROMPT_CONFIG = {"chat_prompt_config": {"prompt": [{"role": "system", "text": "{{#pre_prompt#}}"}]}}
+
+COMPLETION_APP_CHAT_PROMPT_CONFIG = {"chat_prompt_config": {"prompt": [{"role": "user", "text": "{{#pre_prompt#}}"}]}}
+
+COMPLETION_APP_COMPLETION_PROMPT_CONFIG = {
+    "completion_prompt_config": {"prompt": {"text": "{{#pre_prompt#}}"}},
+    "stop": ["Human:"],
+}
+
+BAICHUAN_CHAT_APP_COMPLETION_PROMPT_CONFIG = {
+    "completion_prompt_config": {
+        "prompt": {
+            "text": "{{#pre_prompt#}}\n\n用户和助手的历史对话内容如下：\n```\n{{#histories#}}\n```\n\n\n\n用户：{{#query#}}"  # noqa: E501
+        },
+        "conversation_histories_role": {"user_prefix": "用户", "assistant_prefix": "助手"},
+    },
+    "stop": ["用户:"],
+}
+
+BAICHUAN_CHAT_APP_CHAT_PROMPT_CONFIG = {
+    "chat_prompt_config": {"prompt": [{"role": "system", "text": "{{#pre_prompt#}}"}]}
+}
+
+BAICHUAN_COMPLETION_APP_CHAT_PROMPT_CONFIG = {
+    "chat_prompt_config": {"prompt": [{"role": "user", "text": "{{#pre_prompt#}}"}]}
+}
+
+BAICHUAN_COMPLETION_APP_COMPLETION_PROMPT_CONFIG = {
+    "completion_prompt_config": {"prompt": {"text": "{{#pre_prompt#}}"}},
+    "stop": ["用户:"],
+}

api/core/prompt/prompt_templates/baichuan_chat.json

@@ -0,0 +1,13 @@
+{
+  "human_prefix": "用户",
+  "assistant_prefix": "助手",
+  "context_prompt": "用户在与一个客观的助手对话。助手会尊重找到的材料，给出全面专业的解释，但不会过度演绎。同时回答中不会暴露引用的材料：\n\n```\n{{#context#}}\n```\n\n",
+  "histories_prompt": "用户和助手的历史对话内容如下：\n```\n{{#histories#}}\n```\n\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt",
+    "histories_prompt"
+  ],
+  "query_prompt": "\n\n用户：{{#query#}}",
+  "stops": ["用户："]
+}

api/core/prompt/prompt_templates/baichuan_completion.json

@@ -0,0 +1,9 @@
+{
+  "context_prompt": "用户在与一个客观的助手对话。助手会尊重找到的材料，给出全面专业的解释，但不会过度演绎。同时回答中不会暴露引用的材料：\n\n```\n{{#context#}}\n```\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt"
+  ],
+  "query_prompt": "{{#query#}}",
+  "stops": null
+}

api/core/prompt/prompt_templates/common_chat.json

@@ -0,0 +1,13 @@
+{
+  "human_prefix": "Human",
+  "assistant_prefix": "Assistant",
+  "context_prompt": "Use the following context as your learned knowledge, inside <context></context> XML tags.\n\n<context>\n{{#context#}}\n</context>\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n\n",
+  "histories_prompt": "Here is the chat histories between human and assistant, inside <histories></histories> XML tags.\n\n<histories>\n{{#histories#}}\n</histories>\n\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt",
+    "histories_prompt"
+  ],
+  "query_prompt": "\n\nHuman: {{#query#}}\n\nAssistant: ",
+  "stops": ["\nHuman:", "</histories>"]
+}

api/core/prompt/prompt_templates/common_completion.json

@@ -0,0 +1,9 @@
+{
+  "context_prompt": "Use the following context as your learned knowledge, inside <context></context> XML tags.\n\n<context>\n{{#context#}}\n</context>\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt"
+  ],
+  "query_prompt": "{{#query#}}",
+  "stops": null
+}

api/core/prompt/prompt_transform.py

@@ -0,0 +1,90 @@
+from typing import Any, Optional
+
+from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
+from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_manager import ModelInstance
+from core.model_runtime.entities.message_entities import PromptMessage
+from core.model_runtime.entities.model_entities import ModelPropertyKey
+from core.prompt.entities.advanced_prompt_entities import MemoryConfig
+
+
+class PromptTransform:
+    def _append_chat_histories(
+        self,
+        memory: TokenBufferMemory,
+        memory_config: MemoryConfig,
+        prompt_messages: list[PromptMessage],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> list[PromptMessage]:
+        rest_tokens = self._calculate_rest_token(prompt_messages, model_config)
+        histories = self._get_history_messages_list_from_memory(memory, memory_config, rest_tokens)
+        prompt_messages.extend(histories)
+
+        return prompt_messages
+
+    def _calculate_rest_token(
+        self, prompt_messages: list[PromptMessage], model_config: ModelConfigWithCredentialsEntity
+    ) -> int:
+        rest_tokens = 2000
+
+        model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
+        if model_context_tokens:
+            model_instance = ModelInstance(
+                provider_model_bundle=model_config.provider_model_bundle, model=model_config.model
+            )
+
+            curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
+
+            max_tokens = 0
+            for parameter_rule in model_config.model_schema.parameter_rules:
+                if parameter_rule.name == "max_tokens" or (
+                    parameter_rule.use_template and parameter_rule.use_template == "max_tokens"
+                ):
+                    max_tokens = (
+                        model_config.parameters.get(parameter_rule.name)
+                        or model_config.parameters.get(parameter_rule.use_template or "")
+                    ) or 0
+
+            rest_tokens = model_context_tokens - max_tokens - curr_message_tokens
+            rest_tokens = max(rest_tokens, 0)
+
+        return rest_tokens
+
+    def _get_history_messages_from_memory(
+        self,
+        memory: TokenBufferMemory,
+        memory_config: MemoryConfig,
+        max_token_limit: int,
+        human_prefix: Optional[str] = None,
+        ai_prefix: Optional[str] = None,
+    ) -> str:
+        """Get memory messages."""
+        kwargs: dict[str, Any] = {"max_token_limit": max_token_limit}
+
+        if human_prefix:
+            kwargs["human_prefix"] = human_prefix
+
+        if ai_prefix:
+            kwargs["ai_prefix"] = ai_prefix
+
+        if memory_config.window.enabled and memory_config.window.size is not None and memory_config.window.size > 0:
+            kwargs["message_limit"] = memory_config.window.size
+
+        return memory.get_history_prompt_text(**kwargs)
+
+    def _get_history_messages_list_from_memory(
+        self, memory: TokenBufferMemory, memory_config: MemoryConfig, max_token_limit: int
+    ) -> list[PromptMessage]:
+        """Get memory messages."""
+        return list(
+            memory.get_history_prompt_messages(
+                max_token_limit=max_token_limit,
+                message_limit=memory_config.window.size
+                if (
+                    memory_config.window.enabled
+                    and memory_config.window.size is not None
+                    and memory_config.window.size > 0
+                )
+                else None,
+            )
+        )

api/core/prompt/simple_prompt_transform.py

@@ -0,0 +1,327 @@
+import enum
+import json
+import os
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Optional, cast
+
+from core.app.app_config.entities import PromptTemplateEntity
+from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
+from core.file import file_manager
+from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_runtime.entities.message_entities import (
+    PromptMessage,
+    PromptMessageContent,
+    SystemPromptMessage,
+    TextPromptMessageContent,
+    UserPromptMessage,
+)
+from core.prompt.entities.advanced_prompt_entities import MemoryConfig
+from core.prompt.prompt_transform import PromptTransform
+from core.prompt.utils.prompt_template_parser import PromptTemplateParser
+from models.model import AppMode
+
+if TYPE_CHECKING:
+    from core.file.models import File
+
+
+class ModelMode(enum.StrEnum):
+    COMPLETION = "completion"
+    CHAT = "chat"
+
+    @classmethod
+    def value_of(cls, value: str) -> "ModelMode":
+        """
+        Get value of given mode.
+
+        :param value: mode value
+        :return: mode
+        """
+        for mode in cls:
+            if mode.value == value:
+                return mode
+        raise ValueError(f"invalid mode value {value}")
+
+
+prompt_file_contents: dict[str, Any] = {}
+
+
+class SimplePromptTransform(PromptTransform):
+    """
+    Simple Prompt Transform for Chatbot App Basic Mode.
+    """
+
+    def get_prompt(
+        self,
+        app_mode: AppMode,
+        prompt_template_entity: PromptTemplateEntity,
+        inputs: Mapping[str, str],
+        query: str,
+        files: Sequence["File"],
+        context: Optional[str],
+        memory: Optional[TokenBufferMemory],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> tuple[list[PromptMessage], Optional[list[str]]]:
+        inputs = {key: str(value) for key, value in inputs.items()}
+
+        model_mode = ModelMode.value_of(model_config.mode)
+        if model_mode == ModelMode.CHAT:
+            prompt_messages, stops = self._get_chat_model_prompt_messages(
+                app_mode=app_mode,
+                pre_prompt=prompt_template_entity.simple_prompt_template or "",
+                inputs=inputs,
+                query=query,
+                files=files,
+                context=context,
+                memory=memory,
+                model_config=model_config,
+            )
+        else:
+            prompt_messages, stops = self._get_completion_model_prompt_messages(
+                app_mode=app_mode,
+                pre_prompt=prompt_template_entity.simple_prompt_template or "",
+                inputs=inputs,
+                query=query,
+                files=files,
+                context=context,
+                memory=memory,
+                model_config=model_config,
+            )
+
+        return prompt_messages, stops
+
+    def get_prompt_str_and_rules(
+        self,
+        app_mode: AppMode,
+        model_config: ModelConfigWithCredentialsEntity,
+        pre_prompt: str,
+        inputs: dict,
+        query: Optional[str] = None,
+        context: Optional[str] = None,
+        histories: Optional[str] = None,
+    ) -> tuple[str, dict]:
+        # get prompt template
+        prompt_template_config = self.get_prompt_template(
+            app_mode=app_mode,
+            provider=model_config.provider,
+            model=model_config.model,
+            pre_prompt=pre_prompt,
+            has_context=context is not None,
+            query_in_prompt=query is not None,
+            with_memory_prompt=histories is not None,
+        )
+
+        variables = {k: inputs[k] for k in prompt_template_config["custom_variable_keys"] if k in inputs}
+
+        for v in prompt_template_config["special_variable_keys"]:
+            # support #context#, #query# and #histories#
+            if v == "#context#":
+                variables["#context#"] = context or ""
+            elif v == "#query#":
+                variables["#query#"] = query or ""
+            elif v == "#histories#":
+                variables["#histories#"] = histories or ""
+
+        prompt_template = prompt_template_config["prompt_template"]
+        prompt = prompt_template.format(variables)
+
+        return prompt, prompt_template_config["prompt_rules"]
+
+    def get_prompt_template(
+        self,
+        app_mode: AppMode,
+        provider: str,
+        model: str,
+        pre_prompt: str,
+        has_context: bool,
+        query_in_prompt: bool,
+        with_memory_prompt: bool = False,
+    ) -> dict:
+        prompt_rules = self._get_prompt_rule(app_mode=app_mode, provider=provider, model=model)
+
+        custom_variable_keys = []
+        special_variable_keys = []
+
+        prompt = ""
+        for order in prompt_rules["system_prompt_orders"]:
+            if order == "context_prompt" and has_context:
+                prompt += prompt_rules["context_prompt"]
+                special_variable_keys.append("#context#")
+            elif order == "pre_prompt" and pre_prompt:
+                prompt += pre_prompt + "\n"
+                pre_prompt_template = PromptTemplateParser(template=pre_prompt)
+                custom_variable_keys = pre_prompt_template.variable_keys
+            elif order == "histories_prompt" and with_memory_prompt:
+                prompt += prompt_rules["histories_prompt"]
+                special_variable_keys.append("#histories#")
+
+        if query_in_prompt:
+            prompt += prompt_rules.get("query_prompt", "{{#query#}}")
+            special_variable_keys.append("#query#")
+
+        return {
+            "prompt_template": PromptTemplateParser(template=prompt),
+            "custom_variable_keys": custom_variable_keys,
+            "special_variable_keys": special_variable_keys,
+            "prompt_rules": prompt_rules,
+        }
+
+    def _get_chat_model_prompt_messages(
+        self,
+        app_mode: AppMode,
+        pre_prompt: str,
+        inputs: dict,
+        query: str,
+        context: Optional[str],
+        files: Sequence["File"],
+        memory: Optional[TokenBufferMemory],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> tuple[list[PromptMessage], Optional[list[str]]]:
+        prompt_messages: list[PromptMessage] = []
+
+        # get prompt
+        prompt, _ = self.get_prompt_str_and_rules(
+            app_mode=app_mode,
+            model_config=model_config,
+            pre_prompt=pre_prompt,
+            inputs=inputs,
+            query=None,
+            context=context,
+        )
+
+        if prompt and query:
+            prompt_messages.append(SystemPromptMessage(content=prompt))
+
+        if memory:
+            prompt_messages = self._append_chat_histories(
+                memory=memory,
+                memory_config=MemoryConfig(
+                    window=MemoryConfig.WindowConfig(
+                        enabled=False,
+                    )
+                ),
+                prompt_messages=prompt_messages,
+                model_config=model_config,
+            )
+
+        if query:
+            prompt_messages.append(self.get_last_user_message(query, files))
+        else:
+            prompt_messages.append(self.get_last_user_message(prompt, files))
+
+        return prompt_messages, None
+
+    def _get_completion_model_prompt_messages(
+        self,
+        app_mode: AppMode,
+        pre_prompt: str,
+        inputs: dict,
+        query: str,
+        context: Optional[str],
+        files: Sequence["File"],
+        memory: Optional[TokenBufferMemory],
+        model_config: ModelConfigWithCredentialsEntity,
+    ) -> tuple[list[PromptMessage], Optional[list[str]]]:
+        # get prompt
+        prompt, prompt_rules = self.get_prompt_str_and_rules(
+            app_mode=app_mode,
+            model_config=model_config,
+            pre_prompt=pre_prompt,
+            inputs=inputs,
+            query=query,
+            context=context,
+        )
+
+        if memory:
+            tmp_human_message = UserPromptMessage(content=prompt)
+
+            rest_tokens = self._calculate_rest_token([tmp_human_message], model_config)
+            histories = self._get_history_messages_from_memory(
+                memory=memory,
+                memory_config=MemoryConfig(
+                    window=MemoryConfig.WindowConfig(
+                        enabled=False,
+                    )
+                ),
+                max_token_limit=rest_tokens,
+                human_prefix=prompt_rules.get("human_prefix", "Human"),
+                ai_prefix=prompt_rules.get("assistant_prefix", "Assistant"),
+            )
+
+            # get prompt
+            prompt, prompt_rules = self.get_prompt_str_and_rules(
+                app_mode=app_mode,
+                model_config=model_config,
+                pre_prompt=pre_prompt,
+                inputs=inputs,
+                query=query,
+                context=context,
+                histories=histories,
+            )
+
+        stops = prompt_rules.get("stops")
+        if stops is not None and len(stops) == 0:
+            stops = None
+
+        return [self.get_last_user_message(prompt, files)], stops
+
+    def get_last_user_message(self, prompt: str, files: Sequence["File"]) -> UserPromptMessage:
+        if files:
+            prompt_message_contents: list[PromptMessageContent] = []
+            prompt_message_contents.append(TextPromptMessageContent(data=prompt))
+            for file in files:
+                prompt_message_contents.append(file_manager.to_prompt_message_content(file))
+
+            prompt_message = UserPromptMessage(content=prompt_message_contents)
+        else:
+            prompt_message = UserPromptMessage(content=prompt)
+
+        return prompt_message
+
+    def _get_prompt_rule(self, app_mode: AppMode, provider: str, model: str) -> dict:
+        """
+        Get simple prompt rule.
+        :param app_mode: app mode
+        :param provider: model provider
+        :param model: model name
+        :return:
+        """
+        prompt_file_name = self._prompt_file_name(app_mode=app_mode, provider=provider, model=model)
+
+        # Check if the prompt file is already loaded
+        if prompt_file_name in prompt_file_contents:
+            return cast(dict, prompt_file_contents[prompt_file_name])
+
+        # Get the absolute path of the subdirectory
+        prompt_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "prompt_templates")
+        json_file_path = os.path.join(prompt_path, f"{prompt_file_name}.json")
+
+        # Open the JSON file and read its content
+        with open(json_file_path, encoding="utf-8") as json_file:
+            content = json.load(json_file)
+
+            # Store the content of the prompt file
+            prompt_file_contents[prompt_file_name] = content
+
+            return cast(dict, content)
+
+    def _prompt_file_name(self, app_mode: AppMode, provider: str, model: str) -> str:
+        # baichuan
+        is_baichuan = False
+        if provider == "baichuan":
+            is_baichuan = True
+        else:
+            baichuan_supported_providers = ["huggingface_hub", "openllm", "xinference"]
+            if provider in baichuan_supported_providers and "baichuan" in model.lower():
+                is_baichuan = True
+
+        if is_baichuan:
+            if app_mode == AppMode.COMPLETION:
+                return "baichuan_completion"
+            else:
+                return "baichuan_chat"
+
+        # common
+        if app_mode == AppMode.COMPLETION:
+            return "common_completion"
+        else:
+            return "common_chat"