Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 43 additions & 62 deletions apps/application/serializers/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,45 @@ def get_authorized_tool_ids(user_id: str, workspace_id: str, tool_ids: List[str]
return [tool_id for tool_id in tool_ids if tool_id in granted_tool_ids]


def resolve_import_tools(tool_list, user_id, workspace_id, force_new=False):
"""
处理导入包内的工具: 已存在且有权限的复用, 无权限或不存在则生成新id新建
@param tool_list: 导入包内的工具列表
@param user_id: 用户id
@param workspace_id: 工作空间id
@param force_new: 是否全部新建(工具模板导入)
@return: update_tool_map(包内旧id -> 最终id), tool_list(需要新建的工具, id已替换)
"""
tool_list = list({tool.get("id"): tool for tool in tool_list}.values())
update_tool_map = {}
create_tool_list = []
if len(tool_list) > 0:
tool_id_list = reduce(
lambda x, y: [*x, *y],
[[tool.get("id"), generate_uuid((tool.get("id") + workspace_id or ""))] for tool in tool_list],
[],
)
# 存在的工具列表
exits_tool_id_list = [
str(tool.id) for tool in QuerySet(Tool).filter(id__in=tool_id_list, workspace_id=workspace_id)
]
# 存在且有权限的工具列表
authorized_tool_ids = set(get_authorized_tool_ids(user_id, workspace_id, exits_tool_id_list))
for tool in tool_list:
old_id = tool.get("id")
generate_id = generate_uuid((old_id + workspace_id or ""))
reuse_id = next((i for i in [old_id, generate_id] if i in authorized_tool_ids), None)
if reuse_id is not None and not force_new:
# 有权限则复用
update_tool_map[old_id] = reuse_id
continue
# 无权限或不存在则生成新id新建
new_id = generate_id if not force_new and generate_id not in exits_tool_id_list else str(uuid.uuid7())
update_tool_map[old_id] = new_id
create_tool_list.append({**tool, "id": new_id})
return update_tool_map, create_tool_list


def get_authorized_application_ids(user_id: str, workspace_id: str, application_ids: List[str]) -> List[str]:
"""
返回 application_ids 中当前用户被授权绑定/使用的应用id。
Expand Down Expand Up @@ -924,44 +963,8 @@ def import_(self, instance: dict, is_import_tool, with_valid=True):
raise AppApiException(1001, _("Unsupported file format"))
application = mk_instance.application
tool_list = mk_instance.get_tool_list()
update_tool_map = {}
if len(tool_list) > 0:
tool_id_list = reduce(
lambda x, y: [*x, *y],
[[tool.get("id"), generate_uuid((tool.get("id") + workspace_id or ""))] for tool in tool_list],
[],
)
# 存在的工具列表
exits_tool_id_list = [
str(tool.id) for tool in QuerySet(Tool).filter(id__in=tool_id_list, workspace_id=workspace_id)
]
# 需要更新的工具集合
update_tool_map = {
tool.get("id"): generate_uuid((tool.get("id") + workspace_id or ""))
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
}

tool_list = [
{**tool, "id": update_tool_map.get(tool.get("id"))}
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
and not exits_tool_id_list.__contains__(generate_uuid((tool.get("id") + workspace_id or "")))
]
# 导入包内新建的工具由导入者本人持有,无需校验;仅需校验绑定到已存在工具的引用
existing_bound_tool_ids = [
tool_id for tool_id in get_bound_tool_ids(application) if tool_id not in update_tool_map
]
if existing_bound_tool_ids:
authorized_tool_ids = set(get_authorized_tool_ids(user_id, workspace_id, existing_bound_tool_ids))
unauthorized_tool_ids = [
tool_id for tool_id in existing_bound_tool_ids if tool_id not in authorized_tool_ids
]
if unauthorized_tool_ids:
message = lazy_format(
_("No permission to use tool(s): {tool_ids}"), tool_ids=", ".join(unauthorized_tool_ids)
)
raise AppApiException(403, str(message))
# 已存在且有权限的工具复用, 无权限或不存在则生成新id新建
update_tool_map, tool_list = resolve_import_tools(tool_list, user_id, workspace_id)
application_model = self.to_application(application, workspace_id, user_id, update_tool_map, folder_id)
tool_model_list = [self.to_tool(f, workspace_id, user_id) for f in tool_list]
application_model.save()
Expand Down Expand Up @@ -1584,30 +1587,8 @@ def update_template_workflow(self, instance: Dict, app: Application):
raise AppApiException(1001, _("Unsupported file format"))
application = mk_instance.application
tool_list = mk_instance.get_tool_list()
update_tool_map = {}
if len(tool_list) > 0:
tool_id_list = reduce(
lambda x, y: [*x, *y],
[[tool.get("id"), generate_uuid((tool.get("id") + app.workspace_id or ""))] for tool in tool_list],
[],
)
# 存在的工具列表
exits_tool_id_list = [
str(tool.id) for tool in QuerySet(Tool).filter(id__in=tool_id_list, workspace_id=app.workspace_id)
]
# 需要更新的工具集合
update_tool_map = {
tool.get("id"): generate_uuid((tool.get("id") + app.workspace_id or ""))
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
}

tool_list = [
{**tool, "id": update_tool_map.get(tool.get("id"))}
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
and not exits_tool_id_list.__contains__(generate_uuid((tool.get("id") + app.workspace_id or "")))
]
# 已存在且有权限的工具复用, 无权限或不存在则生成新id新建
update_tool_map, tool_list = resolve_import_tools(tool_list, self.data.get("user_id"), app.workspace_id)

tool_model_list = [self.to_tool(f, app.workspace_id, self.data.get("user_id")) for f in tool_list]
work_flow = application.get("work_flow")
Expand Down
35 changes: 8 additions & 27 deletions apps/knowledge/serializers/knowledge_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import base64
import json
import pickle
from functools import reduce
from typing import Dict, List

import requests
Expand All @@ -21,14 +20,18 @@
from application.flow.knowledge_workflow_manage import KnowledgeWorkflowManage
from application.flow.step_node import get_node
from application.flow.tools import save_workflow_mapping
from application.serializers.application import get_mcp_tools, validate_workflow_default_models
from application.serializers.application import (
get_mcp_tools,
resolve_import_tools,
validate_workflow_default_models,
)
from common.constants.cache_version import Cache_Version
from common.db.search import page_search
from common.exception.app_exception import AppApiException
from common.field.common import UploadedFileField
from common.result import result
from common.utils.common import bytes_to_uploaded_file
from common.utils.common import restricted_loads, generate_uuid
from common.utils.common import restricted_loads
from common.utils.logger import maxkb_logger
from common.utils.rsa_util import rsa_long_decrypt
from common.utils.tool_code import ToolExecutor
Expand Down Expand Up @@ -447,30 +450,8 @@ def import_(self, instance: dict, is_import_tool, with_valid=True):
raise AppApiException(1001, _("Unsupported file format"))
knowledge_workflow = kbwf_instance.knowledge_workflow
tool_list = kbwf_instance.get_tool_list()
update_tool_map = {}
if len(tool_list) > 0:
tool_id_list = reduce(
lambda x, y: [*x, *y],
[[tool.get("id"), generate_uuid((tool.get("id") + workspace_id or ""))] for tool in tool_list],
[],
)
# 存在的工具列表
exits_tool_id_list = [
str(tool.id) for tool in QuerySet(Tool).filter(id__in=tool_id_list, workspace_id=workspace_id)
]
# 需要更新的工具集合
update_tool_map = {
tool.get("id"): generate_uuid((tool.get("id") + workspace_id or ""))
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
}

tool_list = [
{**tool, "id": update_tool_map.get(tool.get("id"))}
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
and not exits_tool_id_list.__contains__(generate_uuid((tool.get("id") + workspace_id or "")))
]
# 已存在且有权限的工具复用, 无权限或不存在则生成新id新建
update_tool_map, tool_list = resolve_import_tools(tool_list, user_id, workspace_id)

work_flow = self.to_knowledge_workflow(
knowledge_workflow,
Expand Down
62 changes: 7 additions & 55 deletions apps/tools/serializers/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
import sys
import tempfile
import zipfile
from functools import reduce
from typing import Dict

import requests
Expand All @@ -22,7 +21,7 @@
from common.exception.app_exception import AppApiException
from common.field.common import UploadedImageField
from common.result import result
from common.utils.common import bytes_to_uploaded_file, common_convert_value, generate_uuid, get_file_content
from common.utils.common import bytes_to_uploaded_file, common_convert_value, get_file_content
from common.utils.logger import maxkb_logger
from common.utils.rsa_util import rsa_long_decrypt, rsa_long_encrypt
from common.utils.tool_code import ToolExecutor
Expand Down Expand Up @@ -86,19 +85,6 @@ def __init__(self, tool: dict, version: str):
ALLOWED_CLASSES = {("builtins", "dict"), ("uuid", "UUID"), ("tools.serializers.tool", "ToolInstance")}


class NewUUID:
def __init__(self):
self.uuid_dict = {}

def generate_uuid(self, _id):
_id = str(_id)
if _id in self.uuid_dict:
return self.uuid_dict.get(_id)
r = str(uuid.uuid7())
self.uuid_dict[_id] = r
return r


def to_dict(message, file_name):
return {
"line": message.line,
Expand Down Expand Up @@ -983,51 +969,17 @@ def import_workflow_tools(self, tool, workspace_id, user_id, folder_id, new_chil
2: 全部创建
@return:
"""
from application.serializers.application import resolve_import_tools

if new_child_policy == 0:
tool_list = []
else:
tool_list = tool.get("tool_list") or []

tool_list = {tool.get("id"): tool for tool in tool_list}.values()
update_tool_map = {}
if len(tool_list) > 0:
new_uuid = NewUUID()
tool_id_list = reduce(
lambda x, y: [*x, *y],
[
[
tool.get("id"),
new_uuid.generate_uuid(tool.get("id"))
if new_child_policy == 2
else generate_uuid((tool.get("id") + workspace_id or "")),
]
for tool in tool_list
],
[],
)
# 存在的工具列表
exits_tool_id_list = [
str(tool.id) for tool in QuerySet(Tool).filter(id__in=tool_id_list, workspace_id=workspace_id)
]
# 需要更新的工具集合
update_tool_map = {
tool.get("id"): new_uuid.generate_uuid(tool.get("id"))
if new_child_policy == 2
else generate_uuid((tool.get("id") + workspace_id or ""))
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
}

tool_list = [
{**tool, "id": update_tool_map.get(tool.get("id"))}
for tool in tool_list
if not exits_tool_id_list.__contains__(tool.get("id"))
and not exits_tool_id_list.__contains__(
new_uuid.generate_uuid(tool.get("id"))
if new_child_policy == 2
else generate_uuid((tool.get("id") + workspace_id or ""))
)
]
# 已存在且有权限的工具复用, 无权限或不存在则生成新id新建
update_tool_map, tool_list = resolve_import_tools(
tool_list, user_id, workspace_id, force_new=new_child_policy == 2
)

work_flow = self.to_tool_workflow(
tool.get("work_flow"),
Expand Down
Loading