From 8495ed3348dfa5de0f34a0f6834b3ed0ef531cff Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Wed, 16 Oct 2024 15:10:50 +0800 Subject: [PATCH] add conversation id, app id and message id into plugin session --- api/core/agent/fc_agent_runner.py | 3 +++ api/core/plugin/manager/tool.py | 16 +++++++++-- api/core/tools/__base/tool.py | 33 ++++++++++++++++++++--- api/core/tools/custom_tool/tool.py | 11 ++++++-- api/core/tools/plugin_tool/tool.py | 9 ++++++- api/core/tools/tool_engine.py | 25 ++++++++++++++--- api/core/tools/workflow_as_tool/tool.py | 9 ++++++- api/core/workflow/nodes/tool/tool_node.py | 2 ++ 8 files changed, 94 insertions(+), 14 deletions(-) diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py index afdb1d70e2..cc4b1961ad 100644 --- a/api/core/agent/fc_agent_runner.py +++ b/api/core/agent/fc_agent_runner.py @@ -236,6 +236,9 @@ class FunctionCallAgentRunner(BaseAgentRunner): invoke_from=self.application_generate_entity.invoke_from, agent_tool_callback=self.agent_callback, trace_manager=trace_manager, + app_id=self.application_generate_entity.app_config.app_id, + message_id=self.message.id, + conversation_id=self.conversation.id, ) # publish files for message_file_id, save_as in message_files: diff --git a/api/core/plugin/manager/tool.py b/api/core/plugin/manager/tool.py index 4eba1d08a2..86df7e111b 100644 --- a/api/core/plugin/manager/tool.py +++ b/api/core/plugin/manager/tool.py @@ -1,5 +1,5 @@ from collections.abc import Generator -from typing import Any +from typing import Any, Optional from pydantic import BaseModel @@ -90,6 +90,9 @@ class PluginToolManager(BasePluginManager): tool_name: str, credentials: dict[str, Any], tool_parameters: dict[str, Any], + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, ) -> Generator[ToolInvokeMessage, None, None]: """ Invoke the tool with the given tenant, user, plugin, provider, name, credentials and parameters. @@ -103,6 +106,9 @@ class PluginToolManager(BasePluginManager): ToolInvokeMessage, data={ "user_id": user_id, + "conversation_id": conversation_id, + "app_id": app_id, + "message_id": message_id, "data": { "provider": provider_name, "tool": tool_name, @@ -154,6 +160,9 @@ class PluginToolManager(BasePluginManager): provider: str, credentials: dict[str, Any], tool: str, + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, ) -> list[ToolParameter]: """ get the runtime parameters of the tool @@ -168,7 +177,10 @@ class PluginToolManager(BasePluginManager): f"plugin/{tenant_id}/dispatch/tool/get_runtime_parameters", RuntimeParametersResponse, data={ - "user_id": "user_id", + "user_id": user_id, + "conversation_id": conversation_id, + "app_id": app_id, + "message_id": message_id, "data": { "provider": provider_name, "tool": tool, diff --git a/api/core/tools/__base/tool.py b/api/core/tools/__base/tool.py index d0bf2f0c31..8a96cd45a2 100644 --- a/api/core/tools/__base/tool.py +++ b/api/core/tools/__base/tool.py @@ -48,7 +48,14 @@ class Tool(ABC): :return: the tool provider type """ - def invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]: + def invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, + ) -> Generator[ToolInvokeMessage]: if self.runtime and self.runtime.runtime_parameters: tool_parameters.update(self.runtime.runtime_parameters) @@ -58,6 +65,9 @@ class Tool(ABC): result = self._invoke( user_id=user_id, tool_parameters=tool_parameters, + conversation_id=conversation_id, + app_id=app_id, + message_id=message_id, ) if isinstance(result, ToolInvokeMessage): @@ -91,11 +101,21 @@ class Tool(ABC): @abstractmethod def _invoke( - self, user_id: str, tool_parameters: dict[str, Any] + self, + user_id: str, + tool_parameters: dict[str, Any], + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, ) -> ToolInvokeMessage | list[ToolInvokeMessage] | Generator[ToolInvokeMessage, None, None]: pass - def get_runtime_parameters(self) -> list[ToolParameter]: + def get_runtime_parameters( + self, + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, + ) -> list[ToolParameter]: """ get the runtime parameters @@ -105,7 +125,12 @@ class Tool(ABC): """ return self.entity.parameters - def get_merged_runtime_parameters(self) -> list[ToolParameter]: + def get_merged_runtime_parameters( + self, + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, + ) -> list[ToolParameter]: """ get merged runtime parameters diff --git a/api/core/tools/custom_tool/tool.py b/api/core/tools/custom_tool/tool.py index e36c97a2de..a26b3763a2 100644 --- a/api/core/tools/custom_tool/tool.py +++ b/api/core/tools/custom_tool/tool.py @@ -1,7 +1,7 @@ import json from collections.abc import Generator from os import getenv -from typing import Any +from typing import Any, Optional from urllib.parse import urlencode import httpx @@ -287,7 +287,14 @@ class ApiTool(Tool): except ValueError as e: return value - def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage, None, None]: + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, + ) -> Generator[ToolInvokeMessage, None, None]: """ invoke http request """ diff --git a/api/core/tools/plugin_tool/tool.py b/api/core/tools/plugin_tool/tool.py index a4173a9630..7c3962a540 100644 --- a/api/core/tools/plugin_tool/tool.py +++ b/api/core/tools/plugin_tool/tool.py @@ -20,7 +20,14 @@ class PluginTool(Tool): def tool_provider_type(self) -> ToolProviderType: return ToolProviderType.PLUGIN - def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage, None, None]: + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, + ) -> Generator[ToolInvokeMessage, None, None]: manager = PluginToolManager() return manager.invoke( tenant_id=self.tenant_id, diff --git a/api/core/tools/tool_engine.py b/api/core/tools/tool_engine.py index 17be66035d..ab15970500 100644 --- a/api/core/tools/tool_engine.py +++ b/api/core/tools/tool_engine.py @@ -44,6 +44,9 @@ class ToolEngine: invoke_from: InvokeFrom, agent_tool_callback: DifyAgentCallbackHandler, trace_manager: Optional[TraceQueueManager] = None, + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, ) -> tuple[str, list[tuple[MessageFile, str]], ToolInvokeMeta]: """ Agent invokes the tool with the given arguments. @@ -66,7 +69,7 @@ class ToolEngine: # hit the callback handler agent_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) - messages = ToolEngine._invoke(tool, tool_parameters, user_id) + messages = ToolEngine._invoke(tool, tool_parameters, user_id, conversation_id, app_id, message_id) invocation_meta_dict: dict[str, ToolInvokeMeta] = {} def message_callback( @@ -138,6 +141,9 @@ class ToolEngine: workflow_tool_callback: DifyWorkflowCallbackHandler, workflow_call_depth: int, thread_pool_id: Optional[str] = None, + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, ) -> Generator[ToolInvokeMessage, None, None]: """ Workflow invokes the tool with the given arguments. @@ -153,7 +159,13 @@ class ToolEngine: if tool.runtime and tool.runtime.runtime_parameters: tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters} - response = tool.invoke(user_id=user_id, tool_parameters=tool_parameters) + response = tool.invoke( + user_id=user_id, + tool_parameters=tool_parameters, + conversation_id=conversation_id, + app_id=app_id, + message_id=message_id, + ) # hit the callback handler response = workflow_tool_callback.on_tool_execution( @@ -169,7 +181,12 @@ class ToolEngine: @staticmethod def _invoke( - tool: Tool, tool_parameters: dict, user_id: str + tool: Tool, + tool_parameters: dict, + user_id: str, + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, ) -> Generator[ToolInvokeMessage | ToolInvokeMeta, None, None]: """ Invoke the tool with the given arguments. @@ -190,7 +207,7 @@ class ToolEngine: }, ) try: - yield from tool.invoke(user_id, tool_parameters) + yield from tool.invoke(user_id, tool_parameters, conversation_id, app_id, message_id) except Exception as e: meta.error = str(e) raise ToolEngineInvokeError(meta) diff --git a/api/core/tools/workflow_as_tool/tool.py b/api/core/tools/workflow_as_tool/tool.py index 836855ba93..677e52b5ba 100644 --- a/api/core/tools/workflow_as_tool/tool.py +++ b/api/core/tools/workflow_as_tool/tool.py @@ -56,7 +56,14 @@ class WorkflowTool(Tool): """ return ToolProviderType.WORKFLOW - def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage, None, None]: + def _invoke( + self, + user_id: str, + tool_parameters: dict[str, Any], + conversation_id: Optional[str] = None, + app_id: Optional[str] = None, + message_id: Optional[str] = None, + ) -> Generator[ToolInvokeMessage, None, None]: """ invoke the tool """ diff --git a/api/core/workflow/nodes/tool/tool_node.py b/api/core/workflow/nodes/tool/tool_node.py index c65826dc3c..fea04f1fe9 100644 --- a/api/core/workflow/nodes/tool/tool_node.py +++ b/api/core/workflow/nodes/tool/tool_node.py @@ -73,6 +73,8 @@ class ToolNode(BaseNode): workflow_tool_callback=DifyWorkflowCallbackHandler(), workflow_call_depth=self.workflow_call_depth, thread_pool_id=self.thread_pool_id, + app_id=self.app_id, + # TODO: conversation id and message id ) except Exception as e: yield RunCompletedEvent(