dify/api/core/datasource/datasource_engine.py

import json
from collections.abc import Generator, Iterable
from mimetypes import guess_type
from typing import Any, Optional, cast

from yarl import URL

from core.app.entities.app_invoke_entities import InvokeFrom
from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
from core.datasource.__base.datasource_plugin import DatasourcePlugin
from core.datasource.entities.datasource_entities import (
    DatasourceInvokeMessage,
    DatasourceInvokeMessageBinary,
)
from core.file import FileType
from core.file.models import FileTransferMethod
from extensions.ext_database import db
from models.enums import CreatedByRole
from models.model import Message, MessageFile


class DatasourceEngine:
    """
    Datasource runtime engine take care of the datasource executions.
    """

    @staticmethod
    def invoke_first_step(
        datasource: DatasourcePlugin,
        datasource_parameters: dict[str, Any],
        user_id: str,
        workflow_tool_callback: DifyWorkflowCallbackHandler,
        conversation_id: Optional[str] = None,
        app_id: Optional[str] = None,
        message_id: Optional[str] = None,
    ) -> Generator[DatasourceInvokeMessage, None, None]:
        """
        Workflow invokes the datasource with the given arguments.
        """
        try:
            # hit the callback handler
            workflow_tool_callback.on_datasource_start(
                datasource_name=datasource.entity.identity.name, datasource_inputs=datasource_parameters
            )

            if datasource.runtime and datasource.runtime.runtime_parameters:
                datasource_parameters = {**datasource.runtime.runtime_parameters, **datasource_parameters}

            response = datasource._invoke_first_step(
                user_id=user_id,
                datasource_parameters=datasource_parameters,
                conversation_id=conversation_id,
                app_id=app_id,
                message_id=message_id,
            )

            # hit the callback handler
            response = workflow_tool_callback.on_datasource_end(
                datasource_name=datasource.entity.identity.name,
                datasource_inputs=datasource_parameters,
                datasource_outputs=response,
            )

            return response
        except Exception as e:
            workflow_tool_callback.on_tool_error(e)
            raise e

    @staticmethod
    def invoke_second_step(
        datasource: DatasourcePlugin,
        datasource_parameters: dict[str, Any],
        user_id: str,
        workflow_tool_callback: DifyWorkflowCallbackHandler,
    ) -> Generator[DatasourceInvokeMessage, None, None]:
        """
        Workflow invokes the datasource with the given arguments.
        """
        try:
            response = datasource._invoke_second_step(
                user_id=user_id,
                datasource_parameters=datasource_parameters,
            )

            return response
        except Exception as e:
            workflow_tool_callback.on_tool_error(e)
            raise e

    @staticmethod
    def _convert_datasource_response_to_str(datasource_response: list[DatasourceInvokeMessage]) -> str:
        """
        Handle datasource response
        """
        result = ""
        for response in datasource_response:
            if response.type == DatasourceInvokeMessage.MessageType.TEXT:
                result += cast(DatasourceInvokeMessage.TextMessage, response.message).text
            elif response.type == DatasourceInvokeMessage.MessageType.LINK:
                result += (
                    f"result link: {cast(DatasourceInvokeMessage.TextMessage, response.message).text}."
                    + " please tell user to check it."
                )
            elif response.type in {
                DatasourceInvokeMessage.MessageType.IMAGE_LINK,
                DatasourceInvokeMessage.MessageType.IMAGE,
            }:
                result += (
                    "image has been created and sent to user already, "
                    + "you do not need to create it, just tell the user to check it now."
                )
            elif response.type == DatasourceInvokeMessage.MessageType.JSON:
                result = json.dumps(
                    cast(DatasourceInvokeMessage.JsonMessage, response.message).json_object, ensure_ascii=False
                )
            else:
                result += str(response.message)

        return result

    @staticmethod
    def _extract_datasource_response_binary_and_text(
        datasource_response: list[DatasourceInvokeMessage],
    ) -> Generator[DatasourceInvokeMessageBinary, None, None]:
        """
        Extract datasource response binary
        """
        for response in datasource_response:
            if response.type in {
                DatasourceInvokeMessage.MessageType.IMAGE_LINK,
                DatasourceInvokeMessage.MessageType.IMAGE,
            }:
                mimetype = None
                if not response.meta:
                    raise ValueError("missing meta data")
                if response.meta.get("mime_type"):
                    mimetype = response.meta.get("mime_type")
                else:
                    try:
                        url = URL(cast(DatasourceInvokeMessage.TextMessage, response.message).text)
                        extension = url.suffix
                        guess_type_result, _ = guess_type(f"a{extension}")
                        if guess_type_result:
                            mimetype = guess_type_result
                    except Exception:
                        pass

                if not mimetype:
                    mimetype = "image/jpeg"

                yield DatasourceInvokeMessageBinary(
                    mimetype=response.meta.get("mime_type", "image/jpeg"),
                    url=cast(DatasourceInvokeMessage.TextMessage, response.message).text,
                )
            elif response.type == DatasourceInvokeMessage.MessageType.BLOB:
                if not response.meta:
                    raise ValueError("missing meta data")

                yield DatasourceInvokeMessageBinary(
                    mimetype=response.meta.get("mime_type", "application/octet-stream"),
                    url=cast(DatasourceInvokeMessage.TextMessage, response.message).text,
                )
            elif response.type == DatasourceInvokeMessage.MessageType.LINK:
                # check if there is a mime type in meta
                if response.meta and "mime_type" in response.meta:
                    yield DatasourceInvokeMessageBinary(
                        mimetype=response.meta.get("mime_type", "application/octet-stream")
                        if response.meta
                        else "application/octet-stream",
                        url=cast(DatasourceInvokeMessage.TextMessage, response.message).text,
                    )

    @staticmethod
    def _create_message_files(
        datasource_messages: Iterable[DatasourceInvokeMessageBinary],
        agent_message: Message,
        invoke_from: InvokeFrom,
        user_id: str,
    ) -> list[str]:
        """
        Create message file

        :return: message file ids
        """
        result = []

        for message in datasource_messages:
            if "image" in message.mimetype:
                file_type = FileType.IMAGE
            elif "video" in message.mimetype:
                file_type = FileType.VIDEO
            elif "audio" in message.mimetype:
                file_type = FileType.AUDIO
            elif "text" in message.mimetype or "pdf" in message.mimetype:
                file_type = FileType.DOCUMENT
            else:
                file_type = FileType.CUSTOM

            # extract tool file id from url
            tool_file_id = message.url.split("/")[-1].split(".")[0]
            message_file = MessageFile(
                message_id=agent_message.id,
                type=file_type,
                transfer_method=FileTransferMethod.TOOL_FILE,
                belongs_to="assistant",
                url=message.url,
                upload_file_id=tool_file_id,
                created_by_role=(
                    CreatedByRole.ACCOUNT
                    if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}
                    else CreatedByRole.END_USER
                ),
                created_by=user_id,
            )

            db.session.add(message_file)
            db.session.commit()
            db.session.refresh(message_file)

            result.append(message_file.id)

        db.session.close()

        return result