From a175d6b2d74c2334fa8ae7341cd209380ec0b22d Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Tue, 10 Dec 2024 00:22:41 +0800 Subject: [PATCH 1/2] feat: agent management --- api/controllers/console/__init__.py | 1 + .../console/workspace/agent_providers.py | 35 +++++++++++++++++++ api/core/plugin/manager/agent.py | 7 ++-- api/core/workflow/nodes/agent/agent_node.py | 4 +-- api/services/agent_service.py | 32 ++++++++++++++--- 5 files changed, 71 insertions(+), 8 deletions(-) create mode 100644 api/controllers/console/workspace/agent_providers.py diff --git a/api/controllers/console/__init__.py b/api/controllers/console/__init__.py index 2592f3342a..ad4f562953 100644 --- a/api/controllers/console/__init__.py +++ b/api/controllers/console/__init__.py @@ -91,5 +91,6 @@ from .workspace import ( models, plugin, tool_providers, + agent_providers, workspace, ) diff --git a/api/controllers/console/workspace/agent_providers.py b/api/controllers/console/workspace/agent_providers.py new file mode 100644 index 0000000000..13343f2e35 --- /dev/null +++ b/api/controllers/console/workspace/agent_providers.py @@ -0,0 +1,35 @@ +from flask_login import current_user +from flask_restful import Resource +from controllers.console import api +from controllers.console.wraps import account_initialization_required, setup_required +from core.model_runtime.utils.encoders import jsonable_encoder +from libs.login import login_required +from services.agent_service import AgentService + + +class AgentProviderListApi(Resource): + @setup_required + @login_required + @account_initialization_required + def get(self): + user = current_user + + user_id = user.id + tenant_id = user.current_tenant_id + + return jsonable_encoder(AgentService.list_agent_providers(user_id, tenant_id)) + + +class AgentProviderApi(Resource): + @setup_required + @login_required + @account_initialization_required + def get(self, provider_name: str): + user = current_user + user_id = user.id + tenant_id = user.current_tenant_id + return jsonable_encoder(AgentService.get_agent_provider(user_id, tenant_id, provider_name)) + + +api.add_resource(AgentProviderListApi, "/workspaces/current/agent-providers") +api.add_resource(AgentProviderApi, "/workspaces/current/agent-provider/") diff --git a/api/core/plugin/manager/agent.py b/api/core/plugin/manager/agent.py index fadc6cbf6e..c99e71d7c7 100644 --- a/api/core/plugin/manager/agent.py +++ b/api/core/plugin/manager/agent.py @@ -18,9 +18,12 @@ class PluginAgentManager(BasePluginManager): def transformer(json_response: dict[str, Any]) -> dict: for provider in json_response.get("data", []): declaration = provider.get("declaration", {}) or {} + declaration["identity"]["name"] = ( + f"{provider.get('plugin_id')}/{declaration.get('identity', {}).get('name')}" + ) provider_name = declaration.get("identity", {}).get("name") - for tool in declaration.get("tools", []): - tool["identity"]["provider"] = provider_name + for strategy in declaration.get("strategies", []): + strategy["identity"]["provider"] = provider_name return json_response diff --git a/api/core/workflow/nodes/agent/agent_node.py b/api/core/workflow/nodes/agent/agent_node.py index 4da3355b89..ea9d8c2906 100644 --- a/api/core/workflow/nodes/agent/agent_node.py +++ b/api/core/workflow/nodes/agent/agent_node.py @@ -99,9 +99,9 @@ class AgentNode(ToolNode): Generate parameters based on the given tool parameters, variable pool, and node data. Args: - tool_parameters (Sequence[ToolParameter]): The list of tool parameters. + agent_parameters (Sequence[AgentParameter]): The list of agent parameters. variable_pool (VariablePool): The variable pool containing the variables. - node_data (ToolNodeData): The data associated with the tool node. + node_data (AgentNodeData): The data associated with the agent node. Returns: Mapping[str, Any]: A dictionary containing the generated parameters. diff --git a/api/services/agent_service.py b/api/services/agent_service.py index bc64316faa..78573bfaae 100644 --- a/api/services/agent_service.py +++ b/api/services/agent_service.py @@ -5,6 +5,7 @@ from flask_login import current_user import contexts from core.app.app_config.easy_ui_based_app.agent.manager import AgentConfigManager +from core.plugin.manager.agent import PluginAgentManager from core.tools.tool_manager import ToolManager from extensions.ext_database import db from models.account import Account @@ -63,6 +64,10 @@ class AgentService: timezone = pytz.timezone(current_user.timezone) + app_model_config = app_model.app_model_config + if not app_model_config: + raise ValueError("App model config not found") + result = { "meta": { "status": "success", @@ -70,15 +75,18 @@ class AgentService: "start_time": message.created_at.astimezone(timezone).isoformat(), "elapsed_time": message.provider_response_latency, "total_tokens": message.answer_tokens + message.message_tokens, - "agent_mode": app_model.app_model_config.agent_mode_dict.get("strategy", "react"), + "agent_mode": app_model_config.agent_mode_dict.get("strategy", "react"), "iterations": len(agent_thoughts), }, "iterations": [], "files": message.message_files, } - agent_config = AgentConfigManager.convert(app_model.app_model_config.to_dict()) - agent_tools = agent_config.tools + agent_config = AgentConfigManager.convert(app_model_config.to_dict()) + if not agent_config: + raise ValueError("Agent config not found") + + agent_tools = agent_config.tools or [] def find_agent_tool(tool_name: str): for agent_tool in agent_tools: @@ -90,7 +98,7 @@ class AgentService: tool_labels = agent_thought.tool_labels tool_meta = agent_thought.tool_meta tool_inputs = agent_thought.tool_inputs_dict - tool_outputs = agent_thought.tool_outputs_dict + tool_outputs = agent_thought.tool_outputs_dict or {} tool_calls = [] for tool in tools: tool_name = tool @@ -145,3 +153,19 @@ class AgentService: ) return result + + @classmethod + def list_agent_providers(cls, user_id: str, tenant_id: str): + """ + List agent providers + """ + manager = PluginAgentManager() + return manager.fetch_agent_providers(tenant_id) + + @classmethod + def get_agent_provider(cls, user_id: str, tenant_id: str, provider_name: str): + """ + Get agent provider + """ + manager = PluginAgentManager() + return manager.fetch_agent_provider(tenant_id, provider_name) From 93786f516caa1a794026966057812f17851db6f5 Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Tue, 10 Dec 2024 00:22:54 +0800 Subject: [PATCH 2/2] apply ruff --- api/controllers/console/__init__.py | 2 +- api/controllers/console/workspace/agent_providers.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/api/controllers/console/__init__.py b/api/controllers/console/__init__.py index ad4f562953..9f52b98ae7 100644 --- a/api/controllers/console/__init__.py +++ b/api/controllers/console/__init__.py @@ -84,6 +84,7 @@ from .tag import tags # Import workspace controllers from .workspace import ( account, + agent_providers, endpoint, load_balancing_config, members, @@ -91,6 +92,5 @@ from .workspace import ( models, plugin, tool_providers, - agent_providers, workspace, ) diff --git a/api/controllers/console/workspace/agent_providers.py b/api/controllers/console/workspace/agent_providers.py index 13343f2e35..88c37767e3 100644 --- a/api/controllers/console/workspace/agent_providers.py +++ b/api/controllers/console/workspace/agent_providers.py @@ -1,5 +1,6 @@ from flask_login import current_user from flask_restful import Resource + from controllers.console import api from controllers.console.wraps import account_initialization_required, setup_required from core.model_runtime.utils.encoders import jsonable_encoder