Initial commit

dify_0.15.3/api/events/event_handlers/__init__.py

@@ -0,0 +1,10 @@
+from .clean_when_dataset_deleted import handle
+from .clean_when_document_deleted import handle
+from .create_document_index import handle
+from .create_installed_app_when_app_created import handle
+from .create_site_record_when_app_created import handle
+from .deduct_quota_when_message_created import handle
+from .delete_tool_parameters_cache_when_sync_draft_workflow import handle
+from .update_app_dataset_join_when_app_model_config_updated import handle
+from .update_app_dataset_join_when_app_published_workflow_updated import handle
+from .update_provider_last_used_at_when_message_created import handle

dify_0.15.3/api/events/event_handlers/clean_when_dataset_deleted.py

@@ -0,0 +1,15 @@
+from events.dataset_event import dataset_was_deleted
+from tasks.clean_dataset_task import clean_dataset_task
+
+
+@dataset_was_deleted.connect
+def handle(sender, **kwargs):
+    dataset = sender
+    clean_dataset_task.delay(
+        dataset.id,
+        dataset.tenant_id,
+        dataset.indexing_technique,
+        dataset.index_struct,
+        dataset.collection_binding_id,
+        dataset.doc_form,
+    )

dify_0.15.3/api/events/event_handlers/clean_when_document_deleted.py

@@ -0,0 +1,11 @@
+from events.document_event import document_was_deleted
+from tasks.clean_document_task import clean_document_task
+
+
+@document_was_deleted.connect
+def handle(sender, **kwargs):
+    document_id = sender
+    dataset_id = kwargs.get("dataset_id")
+    doc_form = kwargs.get("doc_form")
+    file_id = kwargs.get("file_id")
+    clean_document_task.delay(document_id, dataset_id, doc_form, file_id)

dify_0.15.3/api/events/event_handlers/create_document_index.py

@@ -0,0 +1,49 @@
+import datetime
+import logging
+import time
+
+import click
+from werkzeug.exceptions import NotFound
+
+from core.indexing_runner import DocumentIsPausedError, IndexingRunner
+from events.event_handlers.document_index_event import document_index_created
+from extensions.ext_database import db
+from models.dataset import Document
+
+
+@document_index_created.connect
+def handle(sender, **kwargs):
+    dataset_id = sender
+    document_ids = kwargs.get("document_ids", [])
+    documents = []
+    start_at = time.perf_counter()
+    for document_id in document_ids:
+        logging.info(click.style("Start process document: {}".format(document_id), fg="green"))
+
+        document = (
+            db.session.query(Document)
+            .filter(
+                Document.id == document_id,
+                Document.dataset_id == dataset_id,
+            )
+            .first()
+        )
+
+        if not document:
+            raise NotFound("Document not found")
+
+        document.indexing_status = "parsing"
+        document.processing_started_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+        documents.append(document)
+        db.session.add(document)
+    db.session.commit()
+
+    try:
+        indexing_runner = IndexingRunner()
+        indexing_runner.run(documents)
+        end_at = time.perf_counter()
+        logging.info(click.style("Processed dataset: {} latency: {}".format(dataset_id, end_at - start_at), fg="green"))
+    except DocumentIsPausedError as ex:
+        logging.info(click.style(str(ex), fg="yellow"))
+    except Exception:
+        pass

dify_0.15.3/api/events/event_handlers/create_installed_app_when_app_created.py

@@ -0,0 +1,16 @@
+from events.app_event import app_was_created
+from extensions.ext_database import db
+from models.model import InstalledApp
+
+
+@app_was_created.connect
+def handle(sender, **kwargs):
+    """Create an installed app when an app is created."""
+    app = sender
+    installed_app = InstalledApp(
+        tenant_id=app.tenant_id,
+        app_id=app.id,
+        app_owner_tenant_id=app.tenant_id,
+    )
+    db.session.add(installed_app)
+    db.session.commit()

dify_0.15.3/api/events/event_handlers/create_site_record_when_app_created.py

@@ -0,0 +1,26 @@
+from events.app_event import app_was_created
+from extensions.ext_database import db
+from models.model import Site
+
+
+@app_was_created.connect
+def handle(sender, **kwargs):
+    """Create site record when an app is created."""
+    app = sender
+    account = kwargs.get("account")
+    if account is not None:
+        site = Site(
+            app_id=app.id,
+            title=app.name,
+            icon_type=app.icon_type,
+            icon=app.icon,
+            icon_background=app.icon_background,
+            default_language=account.interface_language,
+            customize_token_strategy="not_allow",
+            code=Site.generate_code(16),
+            created_by=app.created_by,
+            updated_by=app.updated_by,
+        )
+
+        db.session.add(site)
+        db.session.commit()

dify_0.15.3/api/events/event_handlers/deduct_quota_when_message_created.py

@@ -0,0 +1,53 @@
+from configs import dify_config
+from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity
+from core.entities.provider_entities import QuotaUnit
+from events.message_event import message_was_created
+from extensions.ext_database import db
+from models.provider import Provider, ProviderType
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    application_generate_entity = kwargs.get("application_generate_entity")
+
+    if not isinstance(application_generate_entity, ChatAppGenerateEntity | AgentChatAppGenerateEntity):
+        return
+
+    model_config = application_generate_entity.model_conf
+    provider_model_bundle = model_config.provider_model_bundle
+    provider_configuration = provider_model_bundle.configuration
+
+    if provider_configuration.using_provider_type != ProviderType.SYSTEM:
+        return
+
+    system_configuration = provider_configuration.system_configuration
+
+    quota_unit = None
+    for quota_configuration in system_configuration.quota_configurations:
+        if quota_configuration.quota_type == system_configuration.current_quota_type:
+            quota_unit = quota_configuration.quota_unit
+
+            if quota_configuration.quota_limit == -1:
+                return
+
+            break
+
+    used_quota = None
+    if quota_unit:
+        if quota_unit == QuotaUnit.TOKENS:
+            used_quota = message.message_tokens + message.answer_tokens
+        elif quota_unit == QuotaUnit.CREDITS:
+            used_quota = dify_config.get_model_credits(model_config.model)
+        else:
+            used_quota = 1
+
+    if used_quota is not None and system_configuration.current_quota_type is not None:
+        db.session.query(Provider).filter(
+            Provider.tenant_id == application_generate_entity.app_config.tenant_id,
+            Provider.provider_name == model_config.provider,
+            Provider.provider_type == ProviderType.SYSTEM.value,
+            Provider.quota_type == system_configuration.current_quota_type.value,
+            Provider.quota_limit > Provider.quota_used,
+        ).update({"quota_used": Provider.quota_used + used_quota})
+        db.session.commit()

dify_0.15.3/api/events/event_handlers/delete_tool_parameters_cache_when_sync_draft_workflow.py

@@ -0,0 +1,34 @@
+from core.tools.tool_manager import ToolManager
+from core.tools.utils.configuration import ToolParameterConfigurationManager
+from core.workflow.nodes import NodeType
+from core.workflow.nodes.tool.entities import ToolEntity
+from events.app_event import app_draft_workflow_was_synced
+
+
+@app_draft_workflow_was_synced.connect
+def handle(sender, **kwargs):
+    app = sender
+    synced_draft_workflow = kwargs.get("synced_draft_workflow")
+    if synced_draft_workflow is None:
+        return
+    for node_data in synced_draft_workflow.graph_dict.get("nodes", []):
+        if node_data.get("data", {}).get("type") == NodeType.TOOL.value:
+            try:
+                tool_entity = ToolEntity(**node_data["data"])
+                tool_runtime = ToolManager.get_tool_runtime(
+                    provider_type=tool_entity.provider_type,
+                    provider_id=tool_entity.provider_id,
+                    tool_name=tool_entity.tool_name,
+                    tenant_id=app.tenant_id,
+                )
+                manager = ToolParameterConfigurationManager(
+                    tenant_id=app.tenant_id,
+                    tool_runtime=tool_runtime,
+                    provider_name=tool_entity.provider_name,
+                    provider_type=tool_entity.provider_type,
+                    identity_id=f"WORKFLOW.{app.id}.{node_data.get('id')}",
+                )
+                manager.delete_tool_parameters_cache()
+            except:
+                # tool dose not exist
+                pass

dify_0.15.3/api/events/event_handlers/document_index_event.py

@@ -0,0 +1,4 @@
+from blinker import signal
+
+# sender: document
+document_index_created = signal("document-index-created")

dify_0.15.3/api/events/event_handlers/update_app_dataset_join_when_app_model_config_updated.py

@@ -0,0 +1,68 @@
+from events.app_event import app_model_config_was_updated
+from extensions.ext_database import db
+from models.dataset import AppDatasetJoin
+from models.model import AppModelConfig
+
+
+@app_model_config_was_updated.connect
+def handle(sender, **kwargs):
+    app = sender
+    app_model_config = kwargs.get("app_model_config")
+    if app_model_config is None:
+        return
+
+    dataset_ids = get_dataset_ids_from_model_config(app_model_config)
+
+    app_dataset_joins = db.session.query(AppDatasetJoin).filter(AppDatasetJoin.app_id == app.id).all()
+
+    removed_dataset_ids: set[str] = set()
+    if not app_dataset_joins:
+        added_dataset_ids = dataset_ids
+    else:
+        old_dataset_ids: set[str] = set()
+        old_dataset_ids.update(app_dataset_join.dataset_id for app_dataset_join in app_dataset_joins)
+
+        added_dataset_ids = dataset_ids - old_dataset_ids
+        removed_dataset_ids = old_dataset_ids - dataset_ids
+
+    if removed_dataset_ids:
+        for dataset_id in removed_dataset_ids:
+            db.session.query(AppDatasetJoin).filter(
+                AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id
+            ).delete()
+
+    if added_dataset_ids:
+        for dataset_id in added_dataset_ids:
+            app_dataset_join = AppDatasetJoin(app_id=app.id, dataset_id=dataset_id)
+            db.session.add(app_dataset_join)
+
+    db.session.commit()
+
+
+def get_dataset_ids_from_model_config(app_model_config: AppModelConfig) -> set[str]:
+    dataset_ids: set[str] = set()
+    if not app_model_config:
+        return dataset_ids
+
+    agent_mode = app_model_config.agent_mode_dict
+
+    tools = agent_mode.get("tools", []) or []
+    for tool in tools:
+        if len(list(tool.keys())) != 1:
+            continue
+
+        tool_type = list(tool.keys())[0]
+        tool_config = list(tool.values())[0]
+        if tool_type == "dataset":
+            dataset_ids.add(tool_config.get("id"))
+
+    # get dataset from dataset_configs
+    dataset_configs = app_model_config.dataset_configs_dict
+    datasets = dataset_configs.get("datasets", {}) or {}
+    for dataset in datasets.get("datasets", []) or []:
+        keys = list(dataset.keys())
+        if len(keys) == 1 and keys[0] == "dataset":
+            if dataset["dataset"].get("id"):
+                dataset_ids.add(dataset["dataset"].get("id"))
+
+    return dataset_ids

dify_0.15.3/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py

@@ -0,0 +1,67 @@
+from typing import cast
+
+from core.workflow.nodes import NodeType
+from core.workflow.nodes.knowledge_retrieval.entities import KnowledgeRetrievalNodeData
+from events.app_event import app_published_workflow_was_updated
+from extensions.ext_database import db
+from models.dataset import AppDatasetJoin
+from models.workflow import Workflow
+
+
+@app_published_workflow_was_updated.connect
+def handle(sender, **kwargs):
+    app = sender
+    published_workflow = kwargs.get("published_workflow")
+    published_workflow = cast(Workflow, published_workflow)
+
+    dataset_ids = get_dataset_ids_from_workflow(published_workflow)
+    app_dataset_joins = db.session.query(AppDatasetJoin).filter(AppDatasetJoin.app_id == app.id).all()
+
+    removed_dataset_ids: set[str] = set()
+    if not app_dataset_joins:
+        added_dataset_ids = dataset_ids
+    else:
+        old_dataset_ids: set[str] = set()
+        old_dataset_ids.update(app_dataset_join.dataset_id for app_dataset_join in app_dataset_joins)
+
+        added_dataset_ids = dataset_ids - old_dataset_ids
+        removed_dataset_ids = old_dataset_ids - dataset_ids
+
+    if removed_dataset_ids:
+        for dataset_id in removed_dataset_ids:
+            db.session.query(AppDatasetJoin).filter(
+                AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id
+            ).delete()
+
+    if added_dataset_ids:
+        for dataset_id in added_dataset_ids:
+            app_dataset_join = AppDatasetJoin(app_id=app.id, dataset_id=dataset_id)
+            db.session.add(app_dataset_join)
+
+    db.session.commit()
+
+
+def get_dataset_ids_from_workflow(published_workflow: Workflow) -> set[str]:
+    dataset_ids: set[str] = set()
+    graph = published_workflow.graph_dict
+    if not graph:
+        return dataset_ids
+
+    nodes = graph.get("nodes", [])
+
+    # fetch all knowledge retrieval nodes
+    knowledge_retrieval_nodes = [
+        node for node in nodes if node.get("data", {}).get("type") == NodeType.KNOWLEDGE_RETRIEVAL.value
+    ]
+
+    if not knowledge_retrieval_nodes:
+        return dataset_ids
+
+    for node in knowledge_retrieval_nodes:
+        try:
+            node_data = KnowledgeRetrievalNodeData(**node.get("data", {}))
+            dataset_ids.update(dataset_id for dataset_id in node_data.dataset_ids)
+        except Exception as e:
+            continue
+
+    return dataset_ids

dify_0.15.3/api/events/event_handlers/update_provider_last_used_at_when_message_created.py

@@ -0,0 +1,21 @@
+from datetime import UTC, datetime
+
+from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity
+from events.message_event import message_was_created
+from extensions.ext_database import db
+from models.provider import Provider
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    application_generate_entity = kwargs.get("application_generate_entity")
+
+    if not isinstance(application_generate_entity, ChatAppGenerateEntity | AgentChatAppGenerateEntity):
+        return
+
+    db.session.query(Provider).filter(
+        Provider.tenant_id == application_generate_entity.app_config.tenant_id,
+        Provider.provider_name == application_generate_entity.model_conf.provider,
+    ).update({"last_used": datetime.now(UTC).replace(tzinfo=None)})
+    db.session.commit()