dify

dify/api/core/external_data_tool/api/__builtin__

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

dify/api/core/external_data_tool/api/api.py

@@ -0,0 +1,92 @@
+from sqlalchemy import select
+
+from core.extension.api_based_extension_requestor import APIBasedExtensionRequestor
+from core.external_data_tool.base import ExternalDataTool
+from core.helper import encrypter
+from extensions.ext_database import db
+from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint
+
+
+class ApiExternalDataTool(ExternalDataTool):
+    """
+    The api external data tool.
+    """
+
+    name: str = "api"
+    """the unique name of external data tool"""
+
+    @classmethod
+    def validate_config(cls, tenant_id: str, config: dict):
+        """
+        Validate the incoming form config data.
+
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        # own validation logic
+        api_based_extension_id = config.get("api_based_extension_id")
+        if not api_based_extension_id:
+            raise ValueError("api_based_extension_id is required")
+        # get api_based_extension
+        stmt = select(APIBasedExtension).where(
+            APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id
+        )
+        api_based_extension = db.session.scalar(stmt)
+
+        if not api_based_extension:
+            raise ValueError("api_based_extension_id is invalid")
+
+    def query(self, inputs: dict, query: str | None = None) -> str:
+        """
+        Query the external data tool.
+
+        :param inputs: user inputs
+        :param query: the query of chat app
+        :return: the tool query result
+        """
+        # get params from config
+        if not self.config:
+            raise ValueError(f"config is required, config: {self.config}")
+        api_based_extension_id = self.config.get("api_based_extension_id")
+        assert api_based_extension_id is not None, "api_based_extension_id is required"
+        # get api_based_extension
+        stmt = select(APIBasedExtension).where(
+            APIBasedExtension.tenant_id == self.tenant_id, APIBasedExtension.id == api_based_extension_id
+        )
+        api_based_extension = db.session.scalar(stmt)
+
+        if not api_based_extension:
+            raise ValueError(
+                "[External data tool] API query failed, variable: {}, error: api_based_extension_id is invalid".format(
+                    self.variable
+                )
+            )
+
+        # decrypt api_key
+        api_key = encrypter.decrypt_token(tenant_id=self.tenant_id, token=api_based_extension.api_key)
+
+        try:
+            # request api
+            requestor = APIBasedExtensionRequestor(api_endpoint=api_based_extension.api_endpoint, api_key=api_key)
+        except Exception as e:
+            raise ValueError(f"[External data tool] API query failed, variable: {self.variable}, error: {e}")
+
+        response_json = requestor.request(
+            point=APIBasedExtensionPoint.APP_EXTERNAL_DATA_TOOL_QUERY,
+            params={"app_id": self.app_id, "tool_variable": self.variable, "inputs": inputs, "query": query},
+        )
+
+        if "result" not in response_json:
+            raise ValueError(
+                "[External data tool] API query failed, variable: {}, error: result not found in response".format(
+                    self.variable
+                )
+            )
+
+        if not isinstance(response_json["result"], str):
+            raise ValueError(
+                f"[External data tool] API query failed, variable: {self.variable}, error: result is not string"
+            )
+
+        return response_json["result"]

dify/api/core/external_data_tool/base.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+
+from core.extension.extensible import Extensible, ExtensionModule
+
+
+class ExternalDataTool(Extensible, ABC):
+    """
+    The base class of external data tool.
+    """
+
+    module: ExtensionModule = ExtensionModule.EXTERNAL_DATA_TOOL
+
+    app_id: str
+    """the id of app"""
+    variable: str
+    """the tool variable name of app tool"""
+
+    def __init__(self, tenant_id: str, app_id: str, variable: str, config: dict | None = None):
+        super().__init__(tenant_id, config)
+        self.app_id = app_id
+        self.variable = variable
+
+    @classmethod
+    @abstractmethod
+    def validate_config(cls, tenant_id: str, config: dict):
+        """
+        Validate the incoming form config data.
+
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def query(self, inputs: dict, query: str | None = None) -> str:
+        """
+        Query the external data tool.
+
+        :param inputs: user inputs
+        :param query: the query of chat app
+        :return: the tool query result
+        """
+        raise NotImplementedError

dify/api/core/external_data_tool/external_data_fetch.py

@@ -0,0 +1,89 @@
+import logging
+from collections.abc import Mapping
+from concurrent.futures import Future, ThreadPoolExecutor, as_completed
+from typing import Any
+
+from flask import Flask, current_app
+
+from core.app.app_config.entities import ExternalDataVariableEntity
+from core.external_data_tool.factory import ExternalDataToolFactory
+
+logger = logging.getLogger(__name__)
+
+
+class ExternalDataFetch:
+    def fetch(
+        self,
+        tenant_id: str,
+        app_id: str,
+        external_data_tools: list[ExternalDataVariableEntity],
+        inputs: Mapping[str, Any],
+        query: str,
+    ) -> Mapping[str, Any]:
+        """
+        Fill in variable inputs from external data tools if exists.
+
+        :param tenant_id: workspace id
+        :param app_id: app id
+        :param external_data_tools: external data tools configs
+        :param inputs: the inputs
+        :param query: the query
+        :return: the filled inputs
+        """
+        results: dict[str, Any] = {}
+        inputs = dict(inputs)
+        with ThreadPoolExecutor() as executor:
+            futures = {}
+            for tool in external_data_tools:
+                future: Future[tuple[str | None, str | None]] = executor.submit(
+                    self._query_external_data_tool,
+                    current_app._get_current_object(),  # type: ignore
+                    tenant_id,
+                    app_id,
+                    tool,
+                    inputs,
+                    query,
+                )
+
+                futures[future] = tool
+
+            for future in as_completed(futures):
+                tool_variable, result = future.result()
+                if tool_variable is not None:
+                    results[tool_variable] = result
+
+        inputs.update(results)
+        return inputs
+
+    def _query_external_data_tool(
+        self,
+        flask_app: Flask,
+        tenant_id: str,
+        app_id: str,
+        external_data_tool: ExternalDataVariableEntity,
+        inputs: Mapping[str, Any],
+        query: str,
+    ) -> tuple[str | None, str | None]:
+        """
+        Query external data tool.
+        :param flask_app: flask app
+        :param tenant_id: tenant id
+        :param app_id: app id
+        :param external_data_tool: external data tool
+        :param inputs: inputs
+        :param query: query
+        :return:
+        """
+        with flask_app.app_context():
+            tool_variable = external_data_tool.variable
+            tool_type = external_data_tool.type
+            tool_config = external_data_tool.config
+
+            external_data_tool_factory = ExternalDataToolFactory(
+                name=tool_type, tenant_id=tenant_id, app_id=app_id, variable=tool_variable, config=tool_config
+            )
+
+            # query external data tool
+            result = external_data_tool_factory.query(inputs=inputs, query=query)
+
+            return tool_variable, result

dify/api/core/external_data_tool/factory.py

@@ -0,0 +1,37 @@
+from collections.abc import Mapping
+from typing import Any, cast
+
+from core.extension.extensible import ExtensionModule
+from extensions.ext_code_based_extension import code_based_extension
+
+
+class ExternalDataToolFactory:
+    def __init__(self, name: str, tenant_id: str, app_id: str, variable: str, config: dict):
+        extension_class = code_based_extension.extension_class(ExtensionModule.EXTERNAL_DATA_TOOL, name)
+        self.__extension_instance = extension_class(
+            tenant_id=tenant_id, app_id=app_id, variable=variable, config=config
+        )
+
+    @classmethod
+    def validate_config(cls, name: str, tenant_id: str, config: dict):
+        """
+        Validate the incoming form config data.
+
+        :param name: the name of external data tool
+        :param tenant_id: the id of workspace
+        :param config: the form config data
+        :return:
+        """
+        extension_class = code_based_extension.extension_class(ExtensionModule.EXTERNAL_DATA_TOOL, name)
+        # FIXME mypy issue here, figure out how to fix it
+        extension_class.validate_config(tenant_id, config)  # type: ignore
+
+    def query(self, inputs: Mapping[str, Any], query: str | None = None) -> str:
+        """
+        Query the external data tool.
+
+        :param inputs: user inputs
+        :param query: the query of chat app
+        :return: the tool query result
+        """
+        return cast(str, self.__extension_instance.query(inputs, query))