From e49ca765721cfb8965e25e1d14c7d918e10e8a11 Mon Sep 17 00:00:00 2001 From: chiehw Date: Tue, 22 Apr 2025 17:07:22 +0800 Subject: [PATCH] feat: support var-selector params & surpport retriever_resources messsage --- api/core/agent/plugin_entities.py | 1 + api/core/entities/parameter_entities.py | 1 + api/core/plugin/entities/parameters.py | 6 ++++++ api/core/tools/entities/tool_entities.py | 9 ++++++++- api/core/workflow/nodes/tool/tool_node.py | 8 +++++++- 5 files changed, 23 insertions(+), 2 deletions(-) diff --git a/api/core/agent/plugin_entities.py b/api/core/agent/plugin_entities.py index 9c722baa23..308aea9c5d 100644 --- a/api/core/agent/plugin_entities.py +++ b/api/core/agent/plugin_entities.py @@ -41,6 +41,7 @@ class AgentStrategyParameter(PluginParameter): APP_SELECTOR = CommonParameterType.APP_SELECTOR.value MODEL_SELECTOR = CommonParameterType.MODEL_SELECTOR.value TOOLS_SELECTOR = CommonParameterType.TOOLS_SELECTOR.value + VAR_SELECTOR = CommonParameterType.VAR_SELECTOR.value # deprecated, should not use. SYSTEM_FILES = CommonParameterType.SYSTEM_FILES.value diff --git a/api/core/entities/parameter_entities.py b/api/core/entities/parameter_entities.py index 36800bc263..013e874306 100644 --- a/api/core/entities/parameter_entities.py +++ b/api/core/entities/parameter_entities.py @@ -14,6 +14,7 @@ class CommonParameterType(StrEnum): APP_SELECTOR = "app-selector" MODEL_SELECTOR = "model-selector" TOOLS_SELECTOR = "array[tools]" + VAR_SELECTOR = "var-selector" # TOOL_SELECTOR = "tool-selector" diff --git a/api/core/plugin/entities/parameters.py b/api/core/plugin/entities/parameters.py index 895dd0d0fc..9ee1b5cfb9 100644 --- a/api/core/plugin/entities/parameters.py +++ b/api/core/plugin/entities/parameters.py @@ -5,6 +5,7 @@ from pydantic import BaseModel, Field, field_validator from core.entities.parameter_entities import CommonParameterType from core.tools.entities.common_entities import I18nObject +from core.workflow.nodes.base.entities import NumberType class PluginParameterOption(BaseModel): @@ -35,6 +36,7 @@ class PluginParameterType(enum.StrEnum): APP_SELECTOR = CommonParameterType.APP_SELECTOR.value MODEL_SELECTOR = CommonParameterType.MODEL_SELECTOR.value TOOLS_SELECTOR = CommonParameterType.TOOLS_SELECTOR.value + VAR_SELECTOR = CommonParameterType.VAR_SELECTOR.value # deprecated, should not use. SYSTEM_FILES = CommonParameterType.SYSTEM_FILES.value @@ -134,6 +136,10 @@ def cast_parameter_value(typ: enum.StrEnum, value: Any, /): if value and not isinstance(value, list): raise ValueError("The tools selector must be a list.") return value + case PluginParameterType.VAR_SELECTOR: + if value and not isinstance(value, str | dict | list | NumberType): + raise ValueError("The var selector must be a string, dictionary, list or number.") + return value case _: return str(value) except ValueError: diff --git a/api/core/tools/entities/tool_entities.py b/api/core/tools/entities/tool_entities.py index 03047c0545..b72103490d 100644 --- a/api/core/tools/entities/tool_entities.py +++ b/api/core/tools/entities/tool_entities.py @@ -176,6 +176,10 @@ class ToolInvokeMessage(BaseModel): data: Mapping[str, Any] = Field(..., description="Detailed log data") metadata: Optional[Mapping[str, Any]] = Field(default=None, description="The metadata of the log") + class RetrieverResourceMessage(BaseModel): + retriever_resources: list[dict] = Field(..., description="retriever resources") + context: str = Field(..., description="context") + class MessageType(Enum): TEXT = "text" IMAGE = "image" @@ -188,13 +192,15 @@ class ToolInvokeMessage(BaseModel): FILE = "file" LOG = "log" BLOB_CHUNK = "blob_chunk" + RETRIEVER_RESOURCES = "retriever_resources" type: MessageType = MessageType.TEXT """ plain text, image url or link url """ message: ( - JsonMessage | TextMessage | BlobChunkMessage | BlobMessage | LogMessage | FileMessage | None | VariableMessage + JsonMessage | TextMessage | BlobChunkMessage | BlobMessage | LogMessage | FileMessage | None | VariableMessage \ + | RetrieverResourceMessage ) meta: dict[str, Any] | None = None @@ -240,6 +246,7 @@ class ToolParameter(PluginParameter): FILES = PluginParameterType.FILES.value APP_SELECTOR = PluginParameterType.APP_SELECTOR.value MODEL_SELECTOR = PluginParameterType.MODEL_SELECTOR.value + VAR_SELECTOR = PluginParameterType.VAR_SELECTOR.value # deprecated, should not use. SYSTEM_FILES = PluginParameterType.SYSTEM_FILES.value diff --git a/api/core/workflow/nodes/tool/tool_node.py b/api/core/workflow/nodes/tool/tool_node.py index aaecc7b989..a63585b2b0 100644 --- a/api/core/workflow/nodes/tool/tool_node.py +++ b/api/core/workflow/nodes/tool/tool_node.py @@ -21,7 +21,7 @@ from core.workflow.enums import SystemVariableKey from core.workflow.graph_engine.entities.event import AgentLogEvent from core.workflow.nodes.base import BaseNode from core.workflow.nodes.enums import NodeType -from core.workflow.nodes.event import RunCompletedEvent, RunStreamChunkEvent +from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent from core.workflow.utils.variable_template_parser import VariableTemplateParser from extensions.ext_database import db from factories import file_factory @@ -359,6 +359,12 @@ class ToolNode(BaseNode[ToolNodeData]): agent_logs.append(agent_log) yield agent_log + elif message.type == ToolInvokeMessage.MessageType.RETRIEVER_RESOURCES: + assert isinstance(message.message, ToolInvokeMessage.RetrieverResourceMessage) + yield RunRetrieverResourceEvent( + retriever_resources=message.message.retriever_resources, + context=message.message.context, + ) yield RunCompletedEvent( run_result=NodeRunResult(