refactor: tool

api/core/tools/utils/configuration.py

@@ -59,12 +59,11 @@ class ProviderConfigEncrypter(BaseModel):
             if field.type == BasicProviderConfig.Type.SECRET_INPUT:
                 if field_name in data:
                     if len(data[field_name]) > 6:
-                        data[field_name] = \
-                            data[field_name][:2] + \
-                            '*' * (len(data[field_name]) - 4) + \
-                            data[field_name][-2:]
+                        data[field_name] = (
+                            data[field_name][:2] + "*" * (len(data[field_name]) - 4) + data[field_name][-2:]
+                        )
                     else:
-                        data[field_name] = '*' * len(data[field_name])
+                        data[field_name] = "*" * len(data[field_name])
 
         return data
 
@@ -75,9 +74,9 @@ class ProviderConfigEncrypter(BaseModel):
         return a deep copy of credentials with decrypted values
         """
         cache = ToolProviderCredentialsCache(
-            tenant_id=self.tenant_id, 
-            identity_id=f'{self.provider_type}.{self.provider_identity}',
-            cache_type=ToolProviderCredentialsCacheType.PROVIDER
+            tenant_id=self.tenant_id,
+            identity_id=f"{self.provider_type}.{self.provider_identity}",
+            cache_type=ToolProviderCredentialsCacheType.PROVIDER,
         )
         cached_credentials = cache.get()
         if cached_credentials:
@@ -98,14 +97,14 @@ class ProviderConfigEncrypter(BaseModel):
 
     def delete_tool_credentials_cache(self):
         cache = ToolProviderCredentialsCache(
-            tenant_id=self.tenant_id, 
-            identity_id=f'{self.provider_type}.{self.provider_identity}',
-            cache_type=ToolProviderCredentialsCacheType.PROVIDER
+            tenant_id=self.tenant_id,
+            identity_id=f"{self.provider_type}.{self.provider_identity}",
+            cache_type=ToolProviderCredentialsCacheType.PROVIDER,
         )
         cache.delete()
 
 
-class ToolParameterConfigurationManager(BaseModel):
+class ToolParameterConfigurationManager:
     """
     Tool parameter configuration manager
     """
@@ -116,6 +115,15 @@ class ToolParameterConfigurationManager(BaseModel):
     provider_type: ToolProviderType
     identity_id: str
 
+    def __init__(
+        self, tenant_id: str, tool_runtime: Tool, provider_name: str, provider_type: ToolProviderType, identity_id: str
+    ) -> None:
+        self.tenant_id = tenant_id
+        self.tool_runtime = tool_runtime
+        self.provider_name = provider_name
+        self.provider_type = provider_type
+        self.identity_id = identity_id
+
     def _deep_copy(self, parameters: dict[str, Any]) -> dict[str, Any]:
         """
         deep copy parameters
@@ -127,7 +135,7 @@ class ToolParameterConfigurationManager(BaseModel):
         merge parameters
         """
         # get tool parameters
-        tool_parameters = self.tool_runtime.parameters or []
+        tool_parameters = self.tool_runtime.entity.parameters or []
         # get tool runtime parameters
         runtime_parameters = self.tool_runtime.get_runtime_parameters() or []
         # override parameters
@@ -203,8 +211,8 @@ class ToolParameterConfigurationManager(BaseModel):
         """
         cache = ToolParameterCache(
             tenant_id=self.tenant_id,
-            provider=f'{self.provider_type.value}.{self.provider_name}',
-            tool_name=self.tool_runtime.identity.name,
+            provider=f"{self.provider_type.value}.{self.provider_name}",
+            tool_name=self.tool_runtime.entity.identity.name,
             cache_type=ToolParameterCacheType.PARAMETER,
             identity_id=self.identity_id,
         )
@@ -236,8 +244,8 @@ class ToolParameterConfigurationManager(BaseModel):
     def delete_tool_parameters_cache(self):
         cache = ToolParameterCache(
             tenant_id=self.tenant_id,
-            provider=f'{self.provider_type.value}.{self.provider_name}',
-            tool_name=self.tool_runtime.identity.name,
+            provider=f"{self.provider_type.value}.{self.provider_name}",
+            tool_name=self.tool_runtime.entity.identity.name,
             cache_type=ToolParameterCacheType.PARAMETER,
             identity_id=self.identity_id,
         )

api/core/tools/utils/dataset_retriever_tool.py

@@ -6,9 +6,11 @@ from core.app.entities.app_invoke_entities import InvokeFrom
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
 from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
 from core.tools.__base.tool import Tool
+from core.tools.__base.tool_runtime import ToolRuntime
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_entities import (
     ToolDescription,
+    ToolEntity,
     ToolIdentity,
     ToolInvokeMessage,
     ToolParameter,
@@ -20,11 +22,15 @@ from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import Datas
 class DatasetRetrieverTool(Tool):
     retrieval_tool: DatasetRetrieverBaseTool
 
+    def __init__(self, entity: ToolEntity, runtime: ToolRuntime, retrieval_tool: DatasetRetrieverBaseTool) -> None:
+        super().__init__(entity, runtime)
+        self.retrieval_tool = retrieval_tool
+
     @staticmethod
     def get_dataset_tools(
         tenant_id: str,
         dataset_ids: list[str],
-        retrieve_config: DatasetRetrieveConfigEntity,
+        retrieve_config: DatasetRetrieveConfigEntity | None,
         return_resource: bool,
         invoke_from: InvokeFrom,
         hit_callback: DatasetIndexToolCallbackHandler,
@@ -54,7 +60,7 @@ class DatasetRetrieverTool(Tool):
         )
         if retrieval_tools is None or len(retrieval_tools) == 0:
             return []
-        
+
         # restore retrieve strategy
         retrieve_config.retrieve_strategy = original_retriever_mode
 
@@ -63,13 +69,14 @@ class DatasetRetrieverTool(Tool):
         for retrieval_tool in retrieval_tools:
             tool = DatasetRetrieverTool(
                 retrieval_tool=retrieval_tool,
-                identity=ToolIdentity(
-                    provider="", author="", name=retrieval_tool.name, label=I18nObject(en_US="", zh_Hans="")
+                entity=ToolEntity(
+                    identity=ToolIdentity(
+                        provider="", author="", name=retrieval_tool.name, label=I18nObject(en_US="", zh_Hans="")
+                    ),
+                    parameters=[],
+                    description=ToolDescription(human=I18nObject(en_US="", zh_Hans=""), llm=retrieval_tool.description),
                 ),
-                parameters=[],
-                is_team_authorization=True,
-                description=ToolDescription(human=I18nObject(en_US="", zh_Hans=""), llm=retrieval_tool.description),
-                runtime=DatasetRetrieverTool.Runtime(),
+                runtime=ToolRuntime(tenant_id=tenant_id),
             )
 
             tools.append(tool)
@@ -99,7 +106,7 @@ class DatasetRetrieverTool(Tool):
         """
         query = tool_parameters.get("query")
         if not query:
-            yield self.create_text_message(text='please input query')
+            yield self.create_text_message(text="please input query")
         else:
             # invoke dataset retriever tool
             result = self.retrieval_tool._run(query=query)