tool.py 79 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653
  1. # -*- coding: utf-8 -*-
  2. import asyncio
  3. import base64
  4. import io
  5. import json
  6. import os
  7. import pickle
  8. import re
  9. import tempfile
  10. import zipfile
  11. from functools import reduce
  12. from typing import Dict
  13. import requests
  14. import uuid_utils.compat as uuid
  15. from django.core import validators
  16. from django.core.cache import cache
  17. from django.db import transaction
  18. from django.db.models import QuerySet, Q, Subquery, OuterRef, CharField, Value, When, Case
  19. from django.http import HttpResponse
  20. from django.utils import timezone
  21. from django.utils.translation import gettext_lazy as _
  22. from langchain_core.messages import HumanMessage, AIMessage
  23. from langchain_mcp_adapters.client import MultiServerMCPClient
  24. from pylint.lint import Run
  25. from pylint.reporters import JSON2Reporter
  26. from rest_framework import serializers, status
  27. from application.models import Application
  28. from common.constants.cache_version import Cache_Version
  29. from common.database_model_manage.database_model_manage import DatabaseModelManage
  30. from common.db.search import page_search, native_page_search, native_search
  31. from common.exception.app_exception import AppApiException
  32. from common.field.common import UploadedImageField
  33. from common.result import result
  34. from common.utils.common import get_file_content, generate_uuid, bytes_to_uploaded_file
  35. from common.utils.logger import maxkb_logger
  36. from common.utils.rsa_util import rsa_long_decrypt, rsa_long_encrypt
  37. from common.utils.tool_code import ToolExecutor
  38. from knowledge.models import File, FileSourceType, Knowledge
  39. from maxkb.const import PROJECT_DIR, CONFIG
  40. from models_provider.models import Model
  41. from system_manage.models import AuthTargetType, WorkspaceUserResourcePermission
  42. from system_manage.models.resource_mapping import ResourceMapping
  43. from system_manage.serializers.resource_mapping_serializers import ResourceMappingSerializer
  44. from system_manage.serializers.user_resource_permission import UserResourcePermissionSerializer
  45. from tools.models import Tool, ToolScope, ToolFolder, ToolType, ToolRecord
  46. from tools.models.tool_workflow import ToolWorkflow
  47. from trigger.models import TriggerTask, Trigger
  48. from users.serializers.user import is_workspace_manage
  49. tool_executor = ToolExecutor()
  50. def hand_node(node, update_tool_map):
  51. if node.get('type') == 'tool-lib-node':
  52. tool_lib_id = (node.get('properties', {}).get('node_data', {}).get('tool_lib_id') or '')
  53. node.get('properties', {}).get('node_data', {})['tool_lib_id'] = update_tool_map.get(tool_lib_id, tool_lib_id)
  54. if node.get('type') == 'tool-workflow-lib-node':
  55. tool_lib_id = (node.get('properties', {}).get('node_data', {}).get('tool_lib_id') or '')
  56. node.get('properties', {}).get('node_data', {})['tool_lib_id'] = update_tool_map.get(tool_lib_id, tool_lib_id)
  57. if node.get('type') == 'search-knowledge-node':
  58. node.get('properties', {}).get('node_data', {})['knowledge_id_list'] = []
  59. if node.get('type') == 'ai-chat-node':
  60. node_data = node.get('properties', {}).get('node_data', {})
  61. mcp_tool_ids = node_data.get('mcp_tool_ids') or []
  62. node_data['mcp_tool_ids'] = [update_tool_map.get(tool_id,
  63. tool_id) for tool_id in mcp_tool_ids]
  64. tool_ids = node_data.get('tool_ids') or []
  65. node_data['tool_ids'] = [update_tool_map.get(tool_id,
  66. tool_id) for tool_id in tool_ids]
  67. skill_tool_ids = node_data.get('skill_tool_ids') or []
  68. node_data['skill_tool_ids'] = [update_tool_map.get(tool_id,
  69. tool_id) for tool_id in skill_tool_ids]
  70. if node.get('type') == 'mcp-node':
  71. mcp_tool_id = (node.get('properties', {}).get('node_data', {}).get('mcp_tool_id') or '')
  72. node.get('properties', {}).get('node_data', {})['mcp_tool_id'] = update_tool_map.get(mcp_tool_id,
  73. mcp_tool_id)
  74. class ToolInstance:
  75. def __init__(self, tool: dict, version: str):
  76. self.tool = tool
  77. self.version = version
  78. ALLOWED_CLASSES = {
  79. ("builtins", "dict"),
  80. ('uuid', 'UUID'),
  81. ("tools.serializers.tool", "ToolInstance")
  82. }
  83. class NewUUID:
  84. def __init__(self):
  85. self.uuid_dict = {}
  86. def generate_uuid(self, _id):
  87. _id = str(_id)
  88. if _id in self.uuid_dict:
  89. return self.uuid_dict.get(_id)
  90. r = str(uuid.uuid7())
  91. self.uuid_dict[_id] = r
  92. return r
  93. def to_dict(message, file_name):
  94. return {
  95. 'line': message.line,
  96. 'column': message.column,
  97. 'endLine': message.end_line,
  98. 'endColumn': message.end_column,
  99. 'message': (message.msg or "").replace(file_name, 'code'),
  100. 'type': message.category
  101. }
  102. def get_file_name():
  103. file_name = f"{uuid.uuid7()}"
  104. pylint_dir = os.path.join(PROJECT_DIR, 'data', 'pylint')
  105. if not os.path.exists(pylint_dir):
  106. os.makedirs(pylint_dir, 0o700, exist_ok=True)
  107. os.chmod(os.path.dirname(pylint_dir), 0o700)
  108. return os.path.join(pylint_dir, file_name)
  109. class RestrictedUnpickler(pickle.Unpickler):
  110. def find_class(self, folder, name):
  111. if (folder, name) in ALLOWED_CLASSES:
  112. return super().find_class(folder, name)
  113. raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
  114. (folder, name))
  115. def encryption(message: str):
  116. """
  117. 加密敏感字段数据 加密方式是 如果密码是 1234567890 那么给前端则是 123******890
  118. :param message:
  119. :return:
  120. """
  121. if type(message) != str:
  122. return message
  123. if message == "":
  124. return ""
  125. max_pre_len = 8
  126. max_post_len = 4
  127. message_len = len(message)
  128. pre_len = int(message_len / 5 * 2)
  129. post_len = int(message_len / 5 * 1)
  130. pre_str = "".join([message[index] for index in
  131. range(0,
  132. max_pre_len if pre_len > max_pre_len else 1 if pre_len <= 0 else int(
  133. pre_len))])
  134. end_str = "".join(
  135. [message[index] for index in
  136. range(message_len - (int(post_len) if pre_len < max_post_len else max_post_len),
  137. message_len)])
  138. content = "***************"
  139. return pre_str + content + end_str
  140. def validate_mcp_config(servers: Dict):
  141. async def validate():
  142. client = MultiServerMCPClient(servers)
  143. await client.get_tools()
  144. try:
  145. asyncio.run(validate())
  146. except Exception as e:
  147. maxkb_logger.error(f"validate mcp config error: {e}, servers: {servers}")
  148. raise serializers.ValidationError(_('MCP configuration is invalid'))
  149. class ToolModelSerializer(serializers.ModelSerializer):
  150. class Meta:
  151. model = Tool
  152. fields = ['id', 'name', 'icon', 'desc', 'code', 'input_field_list', 'init_field_list', 'init_params',
  153. 'scope', 'is_active', 'user_id', 'template_id', 'workspace_id', 'folder_id', 'tool_type', 'label',
  154. 'version', 'create_time', 'update_time']
  155. class ToolRecordModelSerializer(serializers.ModelSerializer):
  156. class Meta:
  157. model = ToolRecord
  158. fields = ['id', 'workspace_id', 'tool_id', 'source_type', 'source_id', 'meta', 'state', 'run_time',
  159. 'create_time', 'update_time']
  160. class ToolExportModelSerializer(serializers.ModelSerializer):
  161. class Meta:
  162. model = Tool
  163. fields = ['id', 'name', 'icon', 'desc', 'code', 'input_field_list', 'init_field_list',
  164. 'scope', 'is_active', 'user_id', 'template_id', 'workspace_id', 'folder_id', 'tool_type', 'label',
  165. 'create_time', 'update_time']
  166. class UploadedFileField(serializers.FileField):
  167. def __init__(self, **kwargs):
  168. super().__init__(**kwargs)
  169. def to_representation(self, value):
  170. return value
  171. class ToolInputField(serializers.Serializer):
  172. name = serializers.CharField(required=True, label=_('variable name'))
  173. is_required = serializers.BooleanField(required=True, label=_('required'))
  174. type = serializers.CharField(required=True, label=_('type'), validators=[
  175. validators.RegexValidator(regex=re.compile("^string|int|dict|array|float$"),
  176. message=_('fields only support string|int|dict|array|float'), code=500)
  177. ])
  178. source = serializers.CharField(required=True, label=_('source'), validators=[
  179. validators.RegexValidator(regex=re.compile("^custom|reference$"),
  180. message=_('The field only supports custom|reference'), code=500)
  181. ])
  182. class InitField(serializers.Serializer):
  183. field = serializers.CharField(required=True, label=_('field name'))
  184. label = serializers.CharField(required=True, label=_('field label'))
  185. required = serializers.BooleanField(required=True, label=_('required'))
  186. input_type = serializers.CharField(required=True, label=_('input type'))
  187. default_value = serializers.CharField(required=False, allow_null=True, allow_blank=True)
  188. show_default_value = serializers.BooleanField(required=False, default=False)
  189. props_info = serializers.DictField(required=False, default=dict)
  190. attrs = serializers.DictField(required=False, default=dict)
  191. class ToolCreateRequest(serializers.Serializer):
  192. name = serializers.CharField(required=True, label=_('tool name'))
  193. desc = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool description'))
  194. code = serializers.CharField(required=True, label=_('tool content'))
  195. input_field_list = serializers.ListField(required=False, default=list, label=_('input field list'))
  196. init_field_list = serializers.ListField(required=False, default=list, label=_('init field list'))
  197. is_active = serializers.BooleanField(required=False, label=_('Is active'))
  198. folder_id = serializers.CharField(required=False, allow_null=True)
  199. class ToolEditRequest(serializers.Serializer):
  200. name = serializers.CharField(required=False, label=_('tool name'), allow_null=True)
  201. desc = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool description'))
  202. code = serializers.CharField(required=False, label=_('tool content'), allow_null=True, )
  203. input_field_list = serializers.ListField(required=False, default=list, allow_null=True, label=_('input field list'))
  204. init_field_list = serializers.ListField(required=False, default=list, allow_null=True, label=_('init field list'))
  205. init_params = serializers.DictField(required=False, default=dict, allow_null=True, label=_('init params'))
  206. is_active = serializers.BooleanField(required=False, label=_('Is active'), allow_null=True, )
  207. folder_id = serializers.CharField(required=False, allow_null=True)
  208. class AddInternalToolRequest(serializers.Serializer):
  209. name = serializers.CharField(required=False, label=_("tool name"), allow_null=True, allow_blank=True)
  210. folder_id = serializers.CharField(required=False, allow_null=True, label=_("folder id"))
  211. class DebugField(serializers.Serializer):
  212. name = serializers.CharField(required=True, label=_('variable name'))
  213. value = serializers.CharField(required=False, allow_blank=True, allow_null=True, label=_('variable value'))
  214. class ToolDebugRequest(serializers.Serializer):
  215. code = serializers.CharField(required=True, label=_('tool content'))
  216. input_field_list = serializers.ListField(required=False, default=list, label=_('input field list'))
  217. init_field_list = serializers.ListField(required=False, default=list, label=_('init field list'))
  218. init_params = serializers.DictField(required=False, default=dict, label=_('init params'))
  219. debug_field_list = DebugField(required=True, many=True)
  220. class PylintInstance(serializers.Serializer):
  221. code = serializers.CharField(required=True, allow_null=True, allow_blank=True, label=_('function content'))
  222. class ToolSerializer(serializers.Serializer):
  223. class Query(serializers.Serializer):
  224. workspace_id = serializers.CharField(required=True, label=_('workspace id'))
  225. folder_id = serializers.CharField(required=False, allow_blank=True, allow_null=True, label=_('folder id'))
  226. name = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool name'))
  227. user_id = serializers.UUIDField(required=False, allow_null=True, label=_('user id'))
  228. scope = serializers.CharField(required=True, label=_('scope'))
  229. tool_type = serializers.CharField(required=False, label=_('tool type'), allow_null=True, allow_blank=True)
  230. create_user = serializers.UUIDField(required=False, label=_('create user'), allow_null=True)
  231. def get_query_set(self, workspace_manage, is_x_pack_ee):
  232. tool_query_set = QuerySet(Tool).filter(workspace_id=self.data.get('workspace_id'))
  233. folder_query_set = QuerySet(ToolFolder)
  234. default_query_set = QuerySet(Tool)
  235. workspace_id = self.data.get('workspace_id')
  236. user_id = self.data.get('user_id')
  237. scope = self.data.get('scope')
  238. tool_type = self.data.get('tool_type')
  239. desc = self.data.get('desc')
  240. name = self.data.get('name')
  241. folder_id = self.data.get('folder_id')
  242. create_user = self.data.get('create_user')
  243. if workspace_id is not None:
  244. folder_query_set = folder_query_set.filter(workspace_id=workspace_id)
  245. default_query_set = default_query_set.filter(workspace_id=workspace_id)
  246. if folder_id is not None and folder_id != workspace_id:
  247. folder_query_set = folder_query_set.filter(parent=folder_id)
  248. default_query_set = default_query_set.filter(folder_id=folder_id)
  249. if name is not None:
  250. folder_query_set = folder_query_set.filter(name__icontains=name)
  251. default_query_set = default_query_set.filter(name__icontains=name)
  252. if desc is not None:
  253. folder_query_set = folder_query_set.filter(desc__icontains=desc)
  254. default_query_set = default_query_set.filter(desc__icontains=desc)
  255. if create_user is not None:
  256. tool_query_set = tool_query_set.filter(user_id=create_user)
  257. folder_query_set = folder_query_set.filter(user_id=create_user)
  258. default_query_set = default_query_set.order_by("-create_time")
  259. if scope is not None:
  260. tool_query_set = tool_query_set.filter(scope=scope)
  261. if tool_type:
  262. tool_query_set = tool_query_set.filter(tool_type=tool_type)
  263. query_set_dict = {
  264. 'tool_query_set': tool_query_set,
  265. 'default_query_set': default_query_set,
  266. }
  267. if not workspace_manage:
  268. query_set_dict['workspace_user_resource_permission_query_set'] = QuerySet(
  269. WorkspaceUserResourcePermission).filter(
  270. auth_target_type="TOOL",
  271. workspace_id=workspace_id,
  272. user_id=user_id
  273. )
  274. return query_set_dict
  275. def get_authorized_query_set(self):
  276. default_query_set = QuerySet(Tool)
  277. tool_type = self.data.get('tool_type')
  278. desc = self.data.get('desc')
  279. name = self.data.get('name')
  280. create_user = self.data.get('create_user')
  281. default_query_set = default_query_set.filter(workspace_id='None')
  282. default_query_set = default_query_set.filter(scope=ToolScope.SHARED)
  283. if name is not None:
  284. default_query_set = default_query_set.filter(name__icontains=name)
  285. if desc is not None:
  286. default_query_set = default_query_set.filter(desc__icontains=desc)
  287. if create_user is not None:
  288. default_query_set = default_query_set.filter(user_id=create_user)
  289. if tool_type:
  290. default_query_set = default_query_set.filter(tool_type=tool_type)
  291. default_query_set = default_query_set.order_by("-create_time")
  292. return default_query_set
  293. @staticmethod
  294. def is_x_pack_ee():
  295. workspace_user_role_mapping_model = DatabaseModelManage.get_model("workspace_user_role_mapping")
  296. role_permission_mapping_model = DatabaseModelManage.get_model("role_permission_mapping_model")
  297. return workspace_user_role_mapping_model is not None and role_permission_mapping_model is not None
  298. def get_tools(self):
  299. self.is_valid(raise_exception=True)
  300. workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
  301. is_x_pack_ee = self.is_x_pack_ee()
  302. results = native_search(
  303. self.get_query_set(workspace_manage, is_x_pack_ee),
  304. get_file_content(
  305. os.path.join(
  306. PROJECT_DIR,
  307. "apps", "tools", 'sql',
  308. 'list_tool.sql' if workspace_manage else (
  309. 'list_tool_user_ee.sql' if is_x_pack_ee else 'list_tool_user.sql'
  310. )
  311. )
  312. ),
  313. )
  314. get_authorized_tool = DatabaseModelManage.get_model("get_authorized_tool")
  315. shared_queryset = QuerySet(Tool).none()
  316. if get_authorized_tool is not None:
  317. shared_queryset = self.get_authorized_query_set()
  318. shared_queryset = get_authorized_tool(shared_queryset, self.data.get('workspace_id'))
  319. return {
  320. 'shared_tools': [
  321. ToolModelSerializer(data).data for data in shared_queryset
  322. ],
  323. 'tools': [
  324. {
  325. **tool,
  326. 'input_field_list': json.loads(tool.get('input_field_list', '[]')),
  327. 'init_field_list': json.loads(tool.get('init_field_list', '[]')),
  328. } for tool in results if tool['resource_type'] == 'tool'
  329. ],
  330. }
  331. class Create(serializers.Serializer):
  332. user_id = serializers.UUIDField(required=True, label=_('user id'))
  333. workspace_id = serializers.CharField(required=True, label=_('workspace id'))
  334. @transaction.atomic
  335. def insert(self, instance, with_valid=True):
  336. if with_valid:
  337. self.is_valid(raise_exception=True)
  338. ToolCreateRequest(data=instance).is_valid(raise_exception=True)
  339. # 校验代码是否包括禁止的关键字
  340. if instance.get('tool_type') == ToolType.MCP:
  341. ToolExecutor().validate_mcp_transport(instance.get('code', ''))
  342. # 处理 work_flow_template
  343. if instance.get('work_flow_template') is not None:
  344. template_instance = instance.get('work_flow_template')
  345. download_url = template_instance.get('downloadUrl')
  346. # 查找匹配的版本名称
  347. res = requests.get(download_url, timeout=5)
  348. tool = ToolSerializer.Import(data={
  349. 'file': bytes_to_uploaded_file(res.content, 'file.tool'),
  350. 'user_id': self.data.get('user_id'),
  351. 'workspace_id': self.data.get('workspace_id'),
  352. 'folder_id': str(instance.get('folder_id', self.data.get('workspace_id'))),
  353. }).import_(name=instance.get('name'), source='template')
  354. try:
  355. requests.get(template_instance.get('downloadCallbackUrl'), timeout=5)
  356. except Exception as e:
  357. maxkb_logger.error(f"callback appstore tool download error: {e}")
  358. return tool
  359. tool_id = uuid.uuid7()
  360. Tool(
  361. id=tool_id,
  362. name=instance.get('name'),
  363. desc=instance.get('desc'),
  364. code=instance.get('code'),
  365. user_id=self.data.get('user_id'),
  366. workspace_id=self.data.get('workspace_id'),
  367. input_field_list=instance.get('input_field_list', []),
  368. init_field_list=instance.get('init_field_list', []),
  369. scope=instance.get('scope', ToolScope.WORKSPACE),
  370. tool_type=instance.get('tool_type', ToolType.CUSTOM),
  371. folder_id=instance.get('folder_id', self.data.get('workspace_id')),
  372. is_active=False
  373. ).save()
  374. # 自动授权给创建者
  375. UserResourcePermissionSerializer(data={
  376. 'workspace_id': self.data.get('workspace_id'),
  377. 'user_id': self.data.get('user_id'),
  378. 'auth_target_type': AuthTargetType.TOOL.value
  379. }).auth_resource(str(tool_id))
  380. if instance.get('tool_type') == ToolType.WORKFLOW:
  381. ToolWorkflow(id=uuid.uuid7(), tool_id=tool_id, work_flow=instance.get('work_flow', {})).save()
  382. # 如果是SKILL类型的工具,修改file表中对应的记录
  383. if instance.get('tool_type') == ToolType.SKILL:
  384. file_id = instance.get('code')
  385. old_file = QuerySet(File).filter(id=file_id).first()
  386. if old_file:
  387. # 创建新的文件副本,不复制实际文件内容
  388. new_file_id = uuid.uuid7()
  389. new_file = File(
  390. id=new_file_id,
  391. file_name=old_file.file_name,
  392. file_size=old_file.file_size,
  393. sha256_hash=old_file.sha256_hash,
  394. source_type=FileSourceType.TOOL,
  395. source_id=tool_id,
  396. meta=old_file.meta,
  397. )
  398. new_file.save(old_file.get_bytes())
  399. # 更新工具的code为新的文件id
  400. QuerySet(Tool).filter(id=tool_id).update(code=str(new_file_id))
  401. return ToolSerializer.Operate(data={
  402. 'id': tool_id, 'workspace_id': self.data.get('workspace_id')
  403. }).one()
  404. class TestConnection(serializers.Serializer):
  405. workspace_id = serializers.CharField(required=True, label=_('workspace id'))
  406. code = serializers.CharField(required=True, label=_('tool content'))
  407. def test_connection(self):
  408. self.is_valid(raise_exception=True)
  409. # 校验代码是否包括禁止的关键字
  410. ToolExecutor().validate_mcp_transport(self.data.get('code', ''))
  411. # 校验mcp json
  412. validate_mcp_config(json.loads(self.data.get('code')))
  413. return True
  414. class Debug(serializers.Serializer):
  415. user_id = serializers.UUIDField(required=True, label=_('user id'))
  416. workspace_id = serializers.CharField(required=True, label=_('workspace id'))
  417. def debug(self, debug_instance):
  418. self.is_valid(raise_exception=True)
  419. ToolDebugRequest(data=debug_instance).is_valid(raise_exception=True)
  420. input_field_list = debug_instance.get('input_field_list')
  421. code = debug_instance.get('code')
  422. debug_field_list = debug_instance.get('debug_field_list')
  423. init_params = debug_instance.get('init_params')
  424. params = {field.get('name'): self.convert_value(field.get('name'), field.get('value'), field.get('type'),
  425. field.get('is_required'))
  426. for field in
  427. [{'value': self.get_field_value(debug_field_list, field.get('name'), field.get('is_required')),
  428. **field} for field in
  429. input_field_list]}
  430. # 合并初始化参数
  431. if init_params is not None:
  432. all_params = init_params | params
  433. else:
  434. all_params = params
  435. return tool_executor.exec_code(code, all_params)
  436. @staticmethod
  437. def get_field_value(debug_field_list, name, is_required):
  438. result = [field for field in debug_field_list if field.get('name') == name]
  439. if len(result) > 0:
  440. return result[-1].get('value')
  441. if is_required:
  442. raise AppApiException(500, f"{name}" + _('field has no value set'))
  443. return None
  444. @staticmethod
  445. def convert_value(name: str, value: str, _type: str, is_required: bool):
  446. if not is_required and (value is None or (isinstance(value, str) and len(value.strip()) == 0)):
  447. return None
  448. try:
  449. if _type == 'int':
  450. return int(value)
  451. if _type == 'boolean':
  452. value = 0 if ['0', '[]'].__contains__(value) else value
  453. return bool(value)
  454. if _type == 'float':
  455. return float(value)
  456. if _type == 'dict':
  457. v = json.loads(value)
  458. if isinstance(v, dict):
  459. return v
  460. raise Exception(_('type error'))
  461. if _type == 'array':
  462. v = json.loads(value)
  463. if isinstance(v, list):
  464. return v
  465. raise Exception(_('type error'))
  466. return value
  467. except Exception as e:
  468. raise AppApiException(500, _('Field: {name} Type: {type} Value: {value} Type conversion error').format(
  469. name=name, type=_type, value=value
  470. ))
  471. class Operate(serializers.Serializer):
  472. id = serializers.UUIDField(required=True, label=_('tool id'))
  473. workspace_id = serializers.CharField(required=True, label=_('workspace id'))
  474. def is_one_valid(self, *, raise_exception=False):
  475. super().is_valid(raise_exception=True)
  476. workspace_id = self.data.get('workspace_id')
  477. query_set = QuerySet(Tool).filter(id=self.data.get('id'))
  478. if workspace_id:
  479. query_set = query_set.filter(workspace_id=workspace_id)
  480. if not query_set.exists():
  481. get_authorized_tool = DatabaseModelManage.get_model('get_authorized_tool')
  482. if get_authorized_tool:
  483. if not get_authorized_tool(QuerySet(Tool).filter(id=self.data.get('id')), workspace_id).exists():
  484. raise AppApiException(500, _('Tool id does not exist'))
  485. def is_valid(self, *, raise_exception=False):
  486. super().is_valid(raise_exception=True)
  487. workspace_id = self.data.get('workspace_id')
  488. query_set = QuerySet(Tool).filter(id=self.data.get('id'))
  489. if workspace_id:
  490. query_set = query_set.filter(workspace_id=workspace_id)
  491. if not query_set.exists():
  492. raise AppApiException(500, _('Tool id does not exist'))
  493. @transaction.atomic
  494. def edit(self, instance, with_valid=True):
  495. if with_valid:
  496. self.is_valid(raise_exception=True)
  497. ToolEditRequest(data=instance).is_valid(raise_exception=True)
  498. # 校验代码是否包括禁止的关键字
  499. if instance.get('tool_type') == ToolType.MCP:
  500. ToolExecutor().validate_mcp_transport(instance.get('code', ''))
  501. if not QuerySet(Tool).filter(id=self.data.get('id')).exists():
  502. raise serializers.ValidationError(_('Tool not found'))
  503. edit_field_list = ['name', 'desc', 'code', 'icon', 'input_field_list', 'init_field_list', 'init_params',
  504. 'is_active', 'folder_id']
  505. edit_dict = {field: instance.get(field) for field in edit_field_list if (
  506. field in instance and instance.get(field) is not None)}
  507. tool = QuerySet(Tool).filter(id=self.data.get('id')).first()
  508. if 'init_params' in edit_dict:
  509. if edit_dict['init_field_list'] is not None:
  510. rm_key = []
  511. for key in edit_dict['init_params']:
  512. if key not in [field['field'] for field in edit_dict['init_field_list']]:
  513. rm_key.append(key)
  514. for key in rm_key:
  515. edit_dict['init_params'].pop(key)
  516. if tool.init_params:
  517. old_init_params = json.loads(rsa_long_decrypt(tool.init_params))
  518. for key in edit_dict['init_params']:
  519. if key in old_init_params and edit_dict['init_params'][key] == encryption(old_init_params[key]):
  520. edit_dict['init_params'][key] = old_init_params[key]
  521. edit_dict['init_params'] = rsa_long_encrypt(json.dumps(edit_dict['init_params']))
  522. edit_dict['update_time'] = timezone.now()
  523. QuerySet(Tool).filter(id=self.data.get('id')).update(**edit_dict)
  524. if 'is_active' in instance:
  525. QuerySet(TriggerTask).filter(source_type="TOOL", source_id=self.data.get('id')).update(
  526. is_active=instance.get('is_active'))
  527. # 如果是SKILL类型的工具,修改file表中对应的记录
  528. if instance.get('tool_type') == ToolType.SKILL:
  529. old_file_id = tool.code
  530. file_id = instance.get('code')
  531. if old_file_id != file_id:
  532. QuerySet(File).filter(id=old_file_id).delete()
  533. QuerySet(File).filter(id=file_id).update(source_id=tool.id, source_type=FileSourceType.TOOL)
  534. return self.one()
  535. @transaction.atomic
  536. def delete(self):
  537. from trigger.handler.simple_tools import deploy
  538. from trigger.serializers.trigger import TriggerModelSerializer
  539. self.is_valid(raise_exception=True)
  540. tool = QuerySet(Tool).filter(id=self.data.get('id')).first()
  541. if tool.template_id is None and tool.icon != '':
  542. QuerySet(File).filter(id=tool.icon.split('/')[-1]).delete()
  543. if tool.tool_type == ToolType.SKILL:
  544. QuerySet(File).filter(id=tool.code).delete()
  545. QuerySet(WorkspaceUserResourcePermission).filter(target=tool.id).delete()
  546. QuerySet(Tool).filter(id=self.data.get('id')).delete()
  547. ResourceMapping.objects.filter(
  548. Q(target_id=self.data.get('id')) | Q(source_id=self.data.get('id'))).delete()
  549. QuerySet(ToolRecord).filter(tool_id=self.data.get('id')).delete()
  550. trigger_ids = list(
  551. QuerySet(TriggerTask).filter(
  552. source_type="TOOL", source_id=self.data.get('id')
  553. ).values('trigger_id').distinct()
  554. )
  555. QuerySet(TriggerTask).filter(source_type="TOOL", source_id=self.data.get('id')).delete()
  556. for trigger_id in trigger_ids:
  557. trigger = Trigger.objects.filter(id=trigger_id['trigger_id']).first()
  558. if trigger and trigger.is_active:
  559. deploy(TriggerModelSerializer(trigger).data, **{})
  560. def one(self):
  561. self.is_one_valid(raise_exception=True)
  562. tool = QuerySet(Tool).filter(id=self.data.get('id')).select_related('user').first()
  563. nick_name = tool.user.nick_name if tool and tool.user else None
  564. if tool.init_params:
  565. tool.init_params = json.loads(rsa_long_decrypt(tool.init_params))
  566. if tool.init_field_list:
  567. password_fields = [i["field"] for i in tool.init_field_list if
  568. i.get("input_type") == "PasswordInput"]
  569. if tool.init_params:
  570. for k in tool.init_params:
  571. if k in password_fields and tool.init_params[k]:
  572. tool.init_params[k] = encryption(tool.init_params[k])
  573. if tool.tool_type == 'SKILL':
  574. skill_file = QuerySet(File).filter(id=tool.code).first()
  575. skill_file_dict = {
  576. 'id': str(skill_file.id),
  577. 'name': skill_file.file_name,
  578. 'size': skill_file.file_size,
  579. } if skill_file else None
  580. work_flow = {}
  581. is_publish = False
  582. if tool.tool_type == 'WORKFLOW':
  583. tool_workflow = QuerySet(ToolWorkflow).filter(tool_id=tool.id).first()
  584. if tool_workflow:
  585. work_flow = tool_workflow.work_flow
  586. is_publish = tool_workflow.is_publish
  587. return {
  588. **ToolModelSerializer(tool).data,
  589. 'init_params': tool.init_params if tool.init_params else {},
  590. 'nick_name': nick_name,
  591. 'fileList': [skill_file_dict] if tool.tool_type == 'SKILL' else [],
  592. 'work_flow': work_flow,
  593. 'is_publish': is_publish
  594. }
  595. def get_child_tool_list(self, work_flow, response):
  596. from application.flow.tools import get_tool_id_list
  597. tool_id_list = get_tool_id_list(work_flow, False)
  598. tool_id_list = [tool_id for tool_id in tool_id_list if
  599. len([r for r in response if r.get('id') == tool_id]) == 0]
  600. tool_list = []
  601. if len(tool_id_list) > 0:
  602. tool_list = QuerySet(Tool).filter(id__in=tool_id_list).exclude(scope=ToolScope.SHARED)
  603. work_flow_tools = [tool for tool in tool_list if tool.tool_type == ToolType.WORKFLOW]
  604. if len(work_flow_tools) > 0:
  605. work_flow_tool_dict = {tw.tool_id: tw for tw in
  606. QuerySet(ToolWorkflow).filter(tool_id__in=[t.id for t in work_flow_tools])}
  607. for tool in tool_list:
  608. if tool.tool_type == ToolType.WORKFLOW:
  609. response.append({**ToolExportModelSerializer(tool).data,
  610. 'work_flow': work_flow_tool_dict.get(tool.id).work_flow})
  611. self.get_child_tool_list(work_flow_tool_dict.get(tool.id).work_flow, response)
  612. else:
  613. response.append(ToolExportModelSerializer(tool).data)
  614. else:
  615. for tool in tool_list:
  616. response.append(ToolExportModelSerializer(tool).data)
  617. skill_tools = [tool for tool in tool_list if tool.tool_type == ToolType.SKILL]
  618. for tool in skill_tools:
  619. skill_file = QuerySet(File).filter(id=tool.code).first()
  620. if skill_file:
  621. tool.code = base64.b64encode(skill_file.get_bytes()).decode('utf-8')
  622. response.append(ToolExportModelSerializer(tool).data)
  623. return response
  624. def export(self):
  625. try:
  626. self.is_valid()
  627. id = self.data.get('id')
  628. tool = QuerySet(Tool).filter(id=id).first()
  629. tool_dict = ToolExportModelSerializer(tool).data
  630. # 如果是SKILL类型的工具,校验文件是否存在
  631. if tool.tool_type == ToolType.SKILL:
  632. skill_file = QuerySet(File).filter(id=tool.code).first()
  633. if skill_file:
  634. tool_dict['code'] = base64.b64encode(skill_file.get_bytes()).decode('utf-8')
  635. if tool.tool_type == ToolType.WORKFLOW:
  636. workflow = QuerySet(ToolWorkflow).filter(tool_id=tool.id).first()
  637. if workflow:
  638. tool_dict['work_flow'] = workflow.work_flow
  639. tool_dict['tool_list'] = self.get_child_tool_list(workflow.work_flow, [])
  640. mk_instance = ToolInstance(tool_dict, 'v2')
  641. tool_pickle = pickle.dumps(mk_instance)
  642. response = HttpResponse(content_type='text/plain', content=tool_pickle)
  643. response['Content-Disposition'] = f'attachment; filename="{tool.name}.tool"'
  644. return response
  645. except Exception as e:
  646. return result.error(str(e), response_status=status.HTTP_500_INTERNAL_SERVER_ERROR)
  647. class Pylint(serializers.Serializer):
  648. def run(self, instance, is_valid=True):
  649. if is_valid:
  650. self.is_valid(raise_exception=True)
  651. PylintInstance(data=instance).is_valid(raise_exception=True)
  652. code = instance.get('code')
  653. file_name = get_file_name()
  654. with open(file_name, 'w') as file:
  655. file.write(code)
  656. reporter = JSON2Reporter(output=io.StringIO())
  657. Run([file_name,
  658. "--disable=line-too-long",
  659. '--module-rgx=[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}'],
  660. reporter=reporter, exit=False)
  661. os.remove(file_name)
  662. return [to_dict(m, os.path.basename(file_name)) for m in reporter.messages]
  663. class Import(serializers.Serializer):
  664. file = UploadedFileField(required=True, label=_("file"))
  665. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  666. workspace_id = serializers.CharField(required=True, label=_("workspace id"))
  667. folder_id = serializers.CharField(required=False, allow_null=True, label=_("folder id"))
  668. @staticmethod
  669. def to_tool_workflow(work_flow, update_tool_map):
  670. for node in work_flow.get('nodes', []):
  671. hand_node(node, update_tool_map)
  672. if node.get('type') == 'loop-node':
  673. for n in node.get('properties', {}).get('node_data', {}).get('loop_body', {}).get('nodes', []):
  674. hand_node(n, update_tool_map)
  675. return work_flow
  676. @staticmethod
  677. def to_tool(tool, workspace_id, user_id, folder_id):
  678. # 如果是技能类型的工具,需要将code保存为文件
  679. code = tool.get('code')
  680. if tool.get('tool_type') == ToolType.SKILL:
  681. skill_file_id = uuid.uuid7()
  682. skill_file = File(
  683. id=skill_file_id,
  684. file_name=f"{tool.get('name')}.zip",
  685. source_type=FileSourceType.TOOL,
  686. source_id=tool.get('id'),
  687. meta={}
  688. )
  689. skill_file.save(base64.b64decode(code))
  690. tool['code'] = skill_file_id
  691. return Tool(id=tool.get('id'),
  692. user_id=user_id,
  693. name=tool.get('name'),
  694. code=tool.get('code'),
  695. template_id=tool.get('template_id'),
  696. input_field_list=tool.get('input_field_list'),
  697. init_field_list=tool.get('init_field_list'),
  698. is_active=False if (len((tool.get('init_field_list') or [])) > 0 or tool.get(
  699. 'tool_type') == ToolType.WORKFLOW) else tool.get('is_active'),
  700. tool_type=tool.get('tool_type', 'CUSTOM') or 'CUSTOM',
  701. scope=ToolScope.SHARED if workspace_id == 'None' else ToolScope.WORKSPACE,
  702. folder_id=folder_id if folder_id else 'default' if workspace_id == 'None' else workspace_id,
  703. workspace_id=workspace_id)
  704. def import_workflow_tools(self, tool, workspace_id, user_id, folder_id, new_child_policy):
  705. """
  706. @param tool: 工具对象
  707. @param workspace_id: 工作空间id
  708. @param user_id: 用户id
  709. @param folder_id: 文件夹id
  710. @param new_child_policy: 子工具创建策略
  711. 0: 不创建
  712. 1: 对比创建: 如果存在就不创建 不存在则创建
  713. 2: 全部创建
  714. @return:
  715. """
  716. if new_child_policy == 0:
  717. tool_list = []
  718. else:
  719. tool_list = tool.get('tool_list') or []
  720. tool_list = {tool.get('id'): tool for tool in tool_list}.values()
  721. update_tool_map = {}
  722. if len(tool_list) > 0:
  723. new_uuid = NewUUID()
  724. tool_id_list = reduce(lambda x, y: [*x, *y],
  725. [[tool.get('id'), new_uuid.generate_uuid(
  726. tool.get('id')) if new_child_policy == 2 else generate_uuid(
  727. (tool.get('id') + workspace_id or ''))]
  728. for tool
  729. in
  730. tool_list], [])
  731. # 存在的工具列表
  732. exits_tool_id_list = [str(tool.id) for tool in
  733. QuerySet(Tool).filter(id__in=tool_id_list, workspace_id=workspace_id)]
  734. # 需要更新的工具集合
  735. update_tool_map = {tool.get('id'): new_uuid.generate_uuid(
  736. tool.get('id')) if new_child_policy == 2 else generate_uuid(
  737. (tool.get('id') + workspace_id or '')) for tool
  738. in
  739. tool_list if
  740. not exits_tool_id_list.__contains__(
  741. tool.get('id'))}
  742. tool_list = [{**tool, 'id': update_tool_map.get(tool.get('id'))} for tool in tool_list if
  743. not exits_tool_id_list.__contains__(
  744. tool.get('id')) and not exits_tool_id_list.__contains__(
  745. new_uuid.generate_uuid(
  746. tool.get('id')) if new_child_policy == 2 else generate_uuid(
  747. (tool.get('id') + workspace_id or '')))]
  748. work_flow = self.to_tool_workflow(
  749. tool.get('work_flow'),
  750. update_tool_map,
  751. )
  752. QuerySet(ToolWorkflow).update_or_create(tool_id=tool.get('id'),
  753. create_defaults={'id': uuid.uuid7(),
  754. 'tool_id': tool.get('id'),
  755. "workspace_id": workspace_id,
  756. 'work_flow': work_flow, },
  757. defaults={
  758. 'tool_id': tool.get('id'),
  759. 'workspace_id': workspace_id,
  760. 'work_flow': work_flow
  761. })
  762. tool_model_list = [self.to_tool(tool, workspace_id, user_id, folder_id) for tool in tool_list]
  763. workflow_tool_model_list = [{'tool_id': t.get('id'), 'workflow': self.to_tool_workflow(
  764. t.get('work_flow'),
  765. update_tool_map,
  766. )} for t in tool_list if t.get('tool_type') == ToolType.WORKFLOW]
  767. existing_records = QuerySet(ToolWorkflow).filter(
  768. tool_id__in=[wt.get('tool_id') for wt in workflow_tool_model_list],
  769. workspace_id=workspace_id)
  770. existing_map = {
  771. record.tool_id: record
  772. for record in existing_records
  773. }
  774. QuerySet(ToolWorkflow).bulk_create(
  775. [ToolWorkflow(work_flow=wt.get('workflow'), workspace_id=workspace_id,
  776. tool_id=wt.get('tool_id')) for wt in
  777. workflow_tool_model_list if wt.get('tool_id') not in existing_map])
  778. if len(tool_model_list) > 0:
  779. QuerySet(Tool).bulk_create(tool_model_list)
  780. UserResourcePermissionSerializer(data={
  781. 'workspace_id': self.data.get('workspace_id'),
  782. 'user_id': self.data.get('user_id'),
  783. 'auth_target_type': AuthTargetType.TOOL.value
  784. }).auth_resource_batch([t.id for t in tool_model_list])
  785. def update_template_workflow(self, tool_id: str):
  786. self.is_valid(raise_exception=True)
  787. tool_instance_bytes = self.data.get('file').read()
  788. try:
  789. tool_instance = RestrictedUnpickler(io.BytesIO(tool_instance_bytes)).load()
  790. except Exception as e:
  791. raise AppApiException(1001, _("Unsupported file format"))
  792. tool = tool_instance.tool
  793. tool['id'] = tool_id
  794. folder_id = self.data.get('folder_id')
  795. self.import_workflow_tools(tool, workspace_id=self.data.get('workspace_id'),
  796. user_id=self.data.get('user_id'),
  797. folder_id=folder_id, new_child_policy=2)
  798. return True
  799. @transaction.atomic
  800. def import_(self, scope=ToolScope.WORKSPACE, name=None, source=None):
  801. self.is_valid()
  802. user_id = self.data.get('user_id')
  803. tool_instance_bytes = self.data.get('file').read()
  804. try:
  805. tool_instance = RestrictedUnpickler(io.BytesIO(tool_instance_bytes)).load()
  806. except Exception as e:
  807. raise AppApiException(1001, _("Unsupported file format"))
  808. if self.data.get('folder_id') is None:
  809. folder_id = self.data.get('workspace_id')
  810. else:
  811. folder_id = self.data.get('folder_id')
  812. tool = tool_instance.tool
  813. tool_id = uuid.uuid7()
  814. code = tool.get('code')
  815. if tool.get('tool_type') == ToolType.SKILL:
  816. skill_file_id = uuid.uuid7()
  817. skill_file = File(
  818. id=skill_file_id,
  819. file_name=f"{tool.get('name')}.zip",
  820. source_type=FileSourceType.TOOL,
  821. source_id=tool_id,
  822. meta={}
  823. )
  824. skill_file.save(base64.b64decode(code))
  825. code = skill_file_id
  826. tool_model = Tool(
  827. id=tool_id,
  828. name=name or tool.get('name'),
  829. desc=tool.get('desc'),
  830. code=code,
  831. user_id=user_id,
  832. workspace_id=self.data.get('workspace_id'),
  833. input_field_list=tool.get('input_field_list'),
  834. init_field_list=tool.get('init_field_list', []),
  835. tool_type=tool.get('tool_type'),
  836. folder_id=folder_id,
  837. scope=scope,
  838. is_active=False
  839. )
  840. tool_model.save()
  841. if tool.get('tool_type') == ToolType.WORKFLOW:
  842. tool['id'] = tool_id
  843. self.import_workflow_tools(tool, workspace_id=self.data.get('workspace_id'), user_id=user_id,
  844. folder_id=folder_id, new_child_policy=2 if source == 'template' else 1)
  845. # 自动授权给创建者
  846. UserResourcePermissionSerializer(data={
  847. 'workspace_id': self.data.get('workspace_id'),
  848. 'user_id': self.data.get('user_id'),
  849. 'auth_target_type': AuthTargetType.TOOL.value
  850. }).auth_resource(str(tool_id))
  851. return ToolSerializer.Operate(data={
  852. 'id': tool_id, 'workspace_id': self.data.get('workspace_id')
  853. }).one()
  854. class IconOperate(serializers.Serializer):
  855. id = serializers.UUIDField(required=True, label=_("function ID"))
  856. workspace_id = serializers.CharField(required=True, label=_("workspace id"))
  857. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  858. image = UploadedImageField(required=True, label=_("picture"))
  859. def is_valid(self, *, raise_exception=False):
  860. super().is_valid(raise_exception=True)
  861. workspace_id = self.data.get('workspace_id')
  862. query_set = QuerySet(Tool).filter(id=self.data.get('id'))
  863. if workspace_id:
  864. query_set = query_set.filter(workspace_id=workspace_id)
  865. if not query_set.exists():
  866. raise AppApiException(500, _('Tool id does not exist'))
  867. def edit(self, with_valid=True):
  868. if with_valid:
  869. self.is_valid(raise_exception=True)
  870. tool = QuerySet(Tool).filter(id=self.data.get('id')).first()
  871. if tool is None:
  872. raise AppApiException(500, _('Function does not exist'))
  873. # 删除旧的图片
  874. if tool.icon != '':
  875. QuerySet(File).filter(id=tool.icon.split('/')[-1]).delete()
  876. if self.data.get('image') is None:
  877. tool.icon = ''
  878. else:
  879. meta = {
  880. 'debug': False
  881. }
  882. file_id = uuid.uuid7()
  883. file = File(
  884. id=file_id,
  885. file_name=self.data.get('image').name,
  886. source_type=FileSourceType.TOOL,
  887. source_id=tool.id,
  888. meta=meta
  889. )
  890. file.save(self.data.get('image').read())
  891. tool.icon = f'./oss/file/{file_id}'
  892. tool.save()
  893. return tool.icon
  894. class InternalTool(serializers.Serializer):
  895. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  896. name = serializers.CharField(required=False, label=_("tool name"), allow_null=True, allow_blank=True)
  897. def get_internal_tools(self):
  898. self.is_valid(raise_exception=True)
  899. query_set = QuerySet(Tool)
  900. if self.data.get('name', '') != '':
  901. query_set = query_set.filter(
  902. Q(name__icontains=self.data.get('name')) |
  903. Q(desc__icontains=self.data.get('name'))
  904. )
  905. query_set = query_set.filter(
  906. Q(scope=ToolScope.INTERNAL) &
  907. Q(is_active=True)
  908. )
  909. return ToolModelSerializer(query_set, many=True).data
  910. class AddInternalTool(serializers.Serializer):
  911. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  912. workspace_id = serializers.CharField(required=True, label=_("workspace id"))
  913. tool_id = serializers.UUIDField(required=True, label=_("tool id"))
  914. def add(self, instance, with_valid=True):
  915. if with_valid:
  916. self.is_valid(raise_exception=True)
  917. AddInternalToolRequest(data=instance).is_valid(raise_exception=True)
  918. internal_tool = QuerySet(Tool).filter(id=self.data.get('tool_id')).first()
  919. if internal_tool is None:
  920. raise AppApiException(500, _('Tool does not exist'))
  921. tool_id = uuid.uuid7()
  922. tool = Tool(
  923. id=tool_id,
  924. name=instance.get('name', internal_tool.name),
  925. desc=internal_tool.desc,
  926. code=internal_tool.code,
  927. user_id=self.data.get('user_id'),
  928. icon=internal_tool.icon,
  929. workspace_id=self.data.get('workspace_id'),
  930. input_field_list=internal_tool.input_field_list,
  931. init_field_list=internal_tool.init_field_list,
  932. scope=ToolScope.WORKSPACE,
  933. tool_type=ToolType.CUSTOM,
  934. folder_id=instance.get('folder_id', self.data.get('workspace_id')),
  935. template_id=internal_tool.id,
  936. label=internal_tool.label,
  937. is_active=False
  938. )
  939. tool.save()
  940. # 自动授权给创建者
  941. UserResourcePermissionSerializer(data={
  942. 'workspace_id': self.data.get('workspace_id'),
  943. 'user_id': self.data.get('user_id'),
  944. 'auth_target_type': AuthTargetType.TOOL.value
  945. }).auth_resource(str(tool_id))
  946. return ToolModelSerializer(tool).data
  947. class StoreTool(serializers.Serializer):
  948. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  949. name = serializers.CharField(required=False, label=_("tool name"), allow_null=True, allow_blank=True)
  950. def get_appstore_tools(self):
  951. self.is_valid(raise_exception=True)
  952. # 下载zip文件
  953. try:
  954. appstore_url = CONFIG.get('APPSTORE_URL', 'https://apps-assets.fit2cloud.com/stable/maxkb.json.zip')
  955. res = requests.get(appstore_url, timeout=5)
  956. res.raise_for_status()
  957. # 创建临时文件保存zip
  958. with tempfile.NamedTemporaryFile(delete=False, suffix='.zip') as temp_zip:
  959. temp_zip.write(res.content)
  960. temp_zip_path = temp_zip.name
  961. try:
  962. # 解压zip文件
  963. with zipfile.ZipFile(temp_zip_path, 'r') as zip_ref:
  964. # 获取zip中的第一个文件(假设只有一个json文件)
  965. json_filename = zip_ref.namelist()[0]
  966. json_content = zip_ref.read(json_filename)
  967. # 将json转换为字典
  968. tool_store = json.loads(json_content.decode('utf-8'))
  969. tag_dict = {tag['name']: tag['key'] for tag in tool_store['additionalProperties']['tags']}
  970. filter_apps = []
  971. for tool in tool_store['apps']:
  972. if self.data.get('name', '') != '':
  973. if self.data.get('name').lower() not in tool.get('name', '').lower():
  974. continue
  975. if not tool['downloadUrl'].endswith('.tool'):
  976. continue
  977. versions = tool.get('versions', [])
  978. tool['label'] = tag_dict[tool.get('tags')[0]] if tool.get('tags') else ''
  979. tool['version'] = next(
  980. (version.get('name') for version in versions if
  981. version.get('downloadUrl') == tool['downloadUrl']),
  982. )
  983. filter_apps.append(tool)
  984. tool_store['apps'] = filter_apps
  985. return tool_store
  986. finally:
  987. # 清理临时文件
  988. os.unlink(temp_zip_path)
  989. except Exception as e:
  990. maxkb_logger.error(f"fetch appstore tools error: {e}")
  991. return {'apps': [], 'additionalProperties': {'tags': []}}
  992. class AddStoreTool(serializers.Serializer):
  993. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  994. workspace_id = serializers.CharField(required=True, label=_("workspace id"))
  995. tool_id = serializers.CharField(required=True, label=_("tool id"))
  996. def add(self, instance: Dict, with_valid=True):
  997. if with_valid:
  998. self.is_valid(raise_exception=True)
  999. AddInternalToolRequest(data=instance).is_valid(raise_exception=True)
  1000. versions = instance.get('versions', [])
  1001. download_url = instance.get('download_url')
  1002. # 查找匹配的版本名称
  1003. version_name = next(
  1004. (version.get('name') for version in versions if version.get('downloadUrl') == download_url),
  1005. )
  1006. res = requests.get(download_url, timeout=5)
  1007. tool_data = RestrictedUnpickler(io.BytesIO(res.content)).load().tool
  1008. tool_id = uuid.uuid7()
  1009. # 如果是SKILL类型的工具,保存文件内容到file表,并将code替换为file_id
  1010. if tool_data.get('tool_type') == ToolType.SKILL:
  1011. skill_file_id = uuid.uuid7()
  1012. skill_file = File(
  1013. id=skill_file_id,
  1014. file_name=f"{tool_data.get('name')}.zip",
  1015. source_type=FileSourceType.TOOL,
  1016. source_id=tool_id,
  1017. meta={}
  1018. )
  1019. skill_file.save(base64.b64decode(tool_data.get('code')))
  1020. tool_data['code'] = skill_file_id
  1021. tool = Tool(
  1022. id=tool_id,
  1023. name=instance.get('name'),
  1024. desc=tool_data.get('desc'),
  1025. code=tool_data.get('code'),
  1026. user_id=self.data.get('user_id'),
  1027. icon=instance.get('icon', ''),
  1028. workspace_id=self.data.get('workspace_id'),
  1029. input_field_list=tool_data.get('input_field_list', []),
  1030. init_field_list=tool_data.get('init_field_list', []),
  1031. scope=ToolScope.WORKSPACE,
  1032. tool_type=tool_data.get('tool_type', ToolType.CUSTOM),
  1033. folder_id=instance.get('folder_id', self.data.get('workspace_id')),
  1034. template_id=self.data.get('tool_id'),
  1035. label=instance.get('label'),
  1036. version=version_name,
  1037. is_active=False
  1038. )
  1039. tool.save()
  1040. # 自动授权给创建者
  1041. UserResourcePermissionSerializer(data={
  1042. 'workspace_id': self.data.get('workspace_id'),
  1043. 'user_id': self.data.get('user_id'),
  1044. 'auth_target_type': AuthTargetType.TOOL.value
  1045. }).auth_resource(str(tool_id))
  1046. try:
  1047. requests.get(instance.get('download_callback_url'), timeout=5)
  1048. except Exception as e:
  1049. maxkb_logger.error(f"callback appstore tool download error: {e}")
  1050. return ToolModelSerializer(tool).data
  1051. class UpdateStoreTool(serializers.Serializer):
  1052. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  1053. workspace_id = serializers.CharField(required=True, label=_("workspace id"))
  1054. tool_id = serializers.UUIDField(required=True, label=_("tool id"))
  1055. download_url = serializers.CharField(required=True, label=_("download url"))
  1056. download_callback_url = serializers.CharField(required=True, label=_("download callback url"))
  1057. icon = serializers.CharField(required=True, label=_("icon"), allow_null=True, allow_blank=True)
  1058. versions = serializers.ListField(required=True, label=_("versions"), child=serializers.DictField())
  1059. def update_tool(self, with_valid=True):
  1060. if with_valid:
  1061. self.is_valid(raise_exception=True)
  1062. tool = QuerySet(Tool).filter(id=self.data.get('tool_id')).first()
  1063. if tool is None:
  1064. raise AppApiException(500, _('Tool does not exist'))
  1065. # 查找匹配的版本名称
  1066. version_name = next(
  1067. (version.get('name') for version in self.data.get('versions') if
  1068. version.get('downloadUrl') == self.data.get('download_url')),
  1069. )
  1070. res = requests.get(self.data.get('download_url'), timeout=5)
  1071. tool_data = RestrictedUnpickler(io.BytesIO(res.content)).load().tool
  1072. # 如果是SKILL类型的工具,保存文件内容到file表,并将code替换为file_id
  1073. if tool_data.get('tool_type') == ToolType.SKILL:
  1074. skill_file_id = uuid.uuid7()
  1075. skill_file = File(
  1076. id=skill_file_id,
  1077. file_name=f"{tool_data.get('name')}.zip",
  1078. source_type=FileSourceType.TOOL,
  1079. source_id=tool.id,
  1080. meta={}
  1081. )
  1082. skill_file.save(base64.b64decode(tool_data.get('code')))
  1083. tool_data['code'] = skill_file_id
  1084. tool.desc = tool_data.get('desc')
  1085. tool.code = tool_data.get('code')
  1086. tool.input_field_list = tool_data.get('input_field_list', [])
  1087. tool.init_field_list = tool_data.get('init_field_list', [])
  1088. tool.icon = self.data.get('icon', tool.icon)
  1089. tool.version = version_name
  1090. # tool.is_active = False
  1091. tool.save()
  1092. try:
  1093. requests.get(self.data.get('download_callback_url'), timeout=5)
  1094. except Exception as e:
  1095. maxkb_logger.error(f"callback appstore tool download error: {e}")
  1096. return ToolModelSerializer(tool).data
  1097. class ToolRecord(serializers.Serializer):
  1098. workspace_id = serializers.CharField(required=False, allow_null=True, label=_('workspace id'))
  1099. tool_id = serializers.UUIDField(required=True, label=_('tool id'))
  1100. record_id = serializers.UUIDField(required=False, allow_null=True, label=_('record id'))
  1101. source_name = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('source name'))
  1102. source_type = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('source type'))
  1103. state = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('state'))
  1104. class Operate(serializers.Serializer):
  1105. id = serializers.UUIDField(required=False, allow_null=True, label=_('record id'))
  1106. tool_id = serializers.UUIDField(required=True, label=_('tool id'))
  1107. workspace_id = serializers.CharField(required=False, allow_null=True, label=_('workspace id'))
  1108. def one(self):
  1109. self.is_valid(raise_exception=True)
  1110. tool_record = cache.get(Cache_Version.TOOL_WORKFLOW_EXECUTE.get_key(key=self.data.get('id')),
  1111. version=Cache_Version.TOOL_WORKFLOW_EXECUTE.get_version())
  1112. if tool_record:
  1113. return tool_record
  1114. tool_record = QuerySet(ToolRecord).filter(id=self.data.get('id'), tool_id=self.data.get('tool_id'),
  1115. workspace_id=self.data.get('workspace_id')).first()
  1116. if tool_record:
  1117. return {'id': tool_record.id,
  1118. 'tool_id': tool_record.tool_id,
  1119. 'workspace_id': tool_record.workspace_id,
  1120. 'source_type': tool_record.source_type,
  1121. 'source_id': tool_record.source_id,
  1122. 'meta': tool_record.meta,
  1123. 'state': tool_record.state,
  1124. 'run_time': tool_record.run_time}
  1125. raise AppApiException(500, _('Tool record does not exist'))
  1126. def one(self):
  1127. self.is_valid(raise_exception=True)
  1128. if self.data.get('record_id'):
  1129. page = self.get_tool_records(1, 1)
  1130. return page.get('records')[0]
  1131. return None
  1132. def get_tool_records(self, current_page: int, page_size: int):
  1133. self.is_valid(raise_exception=True)
  1134. application_subquery = Application.objects.filter(id=OuterRef('source_id')).values('name')[:1]
  1135. knowledge_subquery = Knowledge.objects.filter(id=OuterRef('source_id')).values('name')[:1]
  1136. trigger_subquery = Trigger.objects.filter(id=OuterRef('source_id')).values('name')[:1]
  1137. trigger_type_subquery = Trigger.objects.filter(id=OuterRef('source_id')).values('trigger_type')[:1]
  1138. query_set = QuerySet(ToolRecord)
  1139. query_set = query_set.filter(
  1140. tool_id=self.data.get('tool_id')
  1141. ).annotate(
  1142. source_name=Case(
  1143. When(source_type='APPLICATION', then=Subquery(application_subquery)),
  1144. When(source_type='KNOWLEDGE', then=Subquery(knowledge_subquery)),
  1145. When(source_type='TRIGGER', then=Subquery(trigger_subquery)),
  1146. default=Value(''),
  1147. output_field=CharField()
  1148. )
  1149. ).annotate(
  1150. trigger_type=Case(
  1151. When(source_type='TRIGGER', then=Subquery(trigger_type_subquery)),
  1152. default=Value(''),
  1153. output_field=CharField()
  1154. )
  1155. ).annotate(
  1156. tool_name=Subquery(
  1157. Tool.objects.filter(id=OuterRef('tool_id')).values('name')[:1]
  1158. )
  1159. ).annotate(
  1160. tool_icon=Subquery(
  1161. Tool.objects.filter(id=OuterRef('tool_id')).values('icon')[:1]
  1162. )
  1163. )
  1164. if self.data.get('source_type'):
  1165. query_set = query_set.filter(Q(source_type=self.data.get('source_type', '')))
  1166. if self.data.get('state'):
  1167. query_set = query_set.filter(Q(state=self.data.get('state', '')))
  1168. if self.data.get('source_name'):
  1169. query_set = query_set.filter(Q(tool_name__icontains=self.data.get('source_name', '')))
  1170. if self.data.get('record_id'):
  1171. query_set = query_set.filter(Q(id=self.data.get('record_id')))
  1172. if self.data.get('workspace_id'):
  1173. query_set = query_set.filter(Q(workspace_id=self.data.get('workspace_id')))
  1174. query_set = query_set.order_by('-create_time')
  1175. return page_search(
  1176. current_page, page_size, query_set,
  1177. lambda record: {
  1178. **ToolRecordModelSerializer(record).data,
  1179. 'source_name': record.tool_name,
  1180. 'tool_icon': record.tool_icon,
  1181. 'trigger_type': record.trigger_type,
  1182. }
  1183. )
  1184. class UploadSkillFile(serializers.Serializer):
  1185. file = UploadedFileField(required=True, label=_("file"))
  1186. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  1187. workspace_id = serializers.CharField(required=True, label=_("workspace id"))
  1188. def upload(self):
  1189. self.is_valid()
  1190. file = self.data.get('file')
  1191. if not file.name.endswith('.zip'):
  1192. raise AppApiException(1001, _("Unsupported file format"))
  1193. file_id = uuid.uuid7()
  1194. file = File(
  1195. id=file_id,
  1196. file_name=self.data.get('file').name,
  1197. meta={}
  1198. )
  1199. file.save(self.data.get('file').read())
  1200. return file_id
  1201. class DownloadSkillFile(serializers.Serializer):
  1202. user_id = serializers.UUIDField(required=True, label=_("User ID"))
  1203. workspace_id = serializers.CharField(required=True, label=_("workspace id"))
  1204. tool_id = serializers.CharField(required=True, label=_("tool id"))
  1205. def download(self):
  1206. self.is_valid(raise_exception=True)
  1207. tool = QuerySet(Tool).filter(
  1208. id=self.data.get('tool_id'), workspace_id=self.data.get('workspace_id'), tool_type=ToolType.SKILL
  1209. ).first()
  1210. if tool is None:
  1211. raise AppApiException(500, _('Tool does not exist'))
  1212. skill_file = QuerySet(File).filter(id=tool.code).first()
  1213. if skill_file is None:
  1214. raise AppApiException(500, _('Skill file does not exist'))
  1215. response = HttpResponse(content_type='application/zip', content=skill_file.get_bytes())
  1216. response['Content-Disposition'] = f'attachment; filename="{skill_file.file_name}"'
  1217. return response
  1218. class GenerateCodeSerializer(serializers.Serializer):
  1219. workspace_id = serializers.CharField(required=True, label=_('Workspace ID'))
  1220. model_id = serializers.UUIDField(required=True, label=_('Model ID'))
  1221. prompt = serializers.CharField(required=True, label=_('Prompt'))
  1222. messages = serializers.ListField(required=True, label=_('Messages'))
  1223. model_params_setting = serializers.DictField(required=False, default=dict, label=_('Model Params Setting'))
  1224. init_field_list = serializers.ListField(required=False, default=list, label=_('Init Field List'))
  1225. input_field_list = serializers.ListField(required=False, default=list, label=_('Input Field List'))
  1226. def generate_code(self):
  1227. from models_provider.tools import get_model_instance_by_model_workspace_id
  1228. from application.flow.tools import to_stream_response_simple
  1229. self.is_valid(raise_exception=True)
  1230. workspace_id = self.data.get('workspace_id')
  1231. model_id = self.data.get('model_id')
  1232. prompt = self.data.get('prompt')
  1233. messages = self.data.get('messages')
  1234. model_params_setting = self.data.get('model_params_setting')
  1235. init_field_list = self.data.get('init_field_list')
  1236. input_field_list = self.data.get('input_field_list')
  1237. message = messages[-1]['content']
  1238. q = prompt.replace(
  1239. "{userInput}", message
  1240. ).replace(
  1241. "{initFieldList}", json.dumps(init_field_list)
  1242. ).replace(
  1243. "{inputFieldList}", json.dumps(input_field_list)
  1244. )
  1245. messages[-1]['content'] = q
  1246. SUPPORTED_MODEL_TYPES = ["LLM"]
  1247. model_exist = QuerySet(Model).filter(
  1248. id=model_id,
  1249. model_type__in=SUPPORTED_MODEL_TYPES
  1250. ).exists()
  1251. if not model_exist:
  1252. raise Exception(_("Model does not exists or is not an LLM model"))
  1253. def process():
  1254. model = get_model_instance_by_model_workspace_id(
  1255. model_id=model_id, workspace_id=workspace_id, **model_params_setting
  1256. )
  1257. try:
  1258. for r in model.stream([
  1259. # SystemMessage(content=SYSTEM_ROLE),
  1260. *[
  1261. HumanMessage(
  1262. content=m.get('content')
  1263. ) if m.get('role') == 'user' else AIMessage(
  1264. content=m.get('content')
  1265. ) for m in messages
  1266. ]
  1267. ]):
  1268. yield 'data: ' + json.dumps({'content': r.content}) + '\n\n'
  1269. except Exception as e:
  1270. yield 'data: ' + json.dumps({'error': str(e)}) + '\n\n'
  1271. return to_stream_response_simple(process())
  1272. class ToolBatchOperateSerializer(serializers.Serializer):
  1273. workspace_id = serializers.CharField(required=True, label=_('workspace id'))
  1274. def is_valid(self, *, raise_exception=False):
  1275. super().is_valid(raise_exception=True)
  1276. @transaction.atomic
  1277. def batch_delete(self, instance: Dict, with_valid=True):
  1278. from knowledge.serializers.common import BatchSerializer
  1279. from trigger.handler.simple_tools import deploy
  1280. from trigger.serializers.trigger import TriggerModelSerializer
  1281. if with_valid:
  1282. BatchSerializer(data=instance).is_valid(model=Tool, raise_exception=True)
  1283. self.is_valid(raise_exception=True)
  1284. id_list = instance.get('id_list')
  1285. workspace_id = self.data.get('workspace_id')
  1286. tool_query_set = QuerySet(Tool).filter(id__in=id_list, workspace_id=workspace_id)
  1287. for tool in tool_query_set:
  1288. if tool.template_id is None and tool.icon != '':
  1289. QuerySet(File).filter(id=tool.icon.split('/')[-1]).delete()
  1290. if tool.tool_type == ToolType.SKILL:
  1291. QuerySet(File).filter(id=tool.code).delete()
  1292. QuerySet(WorkspaceUserResourcePermission).filter(target__in=id_list).delete()
  1293. QuerySet(ResourceMapping).filter(Q(target_id__in=id_list) | Q(source_id__in=id_list)).delete()
  1294. QuerySet(ToolRecord).filter(tool_id__in=id_list).delete()
  1295. trigger_ids = list(
  1296. QuerySet(TriggerTask).filter(
  1297. source_type="TOOL", source_id__in=id_list
  1298. ).values('trigger_id').distinct()
  1299. )
  1300. QuerySet(TriggerTask).filter(source_type="TOOL", source_id__in=id_list).delete()
  1301. for trigger_id in trigger_ids:
  1302. trigger = Trigger.objects.filter(id=trigger_id['trigger_id']).first()
  1303. if trigger and trigger.is_active:
  1304. deploy(TriggerModelSerializer(trigger).data, **{})
  1305. tool_query_set.delete()
  1306. return True
  1307. def batch_move(self, instance: Dict, with_valid=True):
  1308. from knowledge.serializers.common import BatchMoveSerializer
  1309. if with_valid:
  1310. BatchMoveSerializer(data=instance).is_valid(model=Tool, raise_exception=True)
  1311. self.is_valid(raise_exception=True)
  1312. id_list = instance.get('id_list')
  1313. folder_id = instance.get('folder_id')
  1314. workspace_id = self.data.get('workspace_id')
  1315. QuerySet(Tool).filter(id__in=id_list, workspace_id=workspace_id).update(folder_id=folder_id)
  1316. return True
  1317. class ToolTreeSerializer(serializers.Serializer):
  1318. class Query(serializers.Serializer):
  1319. workspace_id = serializers.CharField(required=True, label=_('workspace id'))
  1320. folder_id = serializers.CharField(required=True, label=_('folder id'))
  1321. name = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool name'))
  1322. user_id = serializers.UUIDField(required=False, allow_null=True, label=_('user id'))
  1323. scope = serializers.CharField(required=True, label=_('scope'))
  1324. tool_type = serializers.CharField(required=False, label=_('tool type'), allow_null=True, allow_blank=True)
  1325. tool_type_list = serializers.ListField(child=serializers.CharField(), required=False, label=_('tool type list'),
  1326. allow_null=True, allow_empty=True)
  1327. create_user = serializers.UUIDField(required=False, label=_('create user'), allow_null=True)
  1328. def page_tool(self, current_page: int, page_size: int):
  1329. self.is_valid(raise_exception=True)
  1330. folder_id = self.data.get('folder_id', self.data.get('workspace_id'))
  1331. root = ToolFolder.objects.filter(id=folder_id).first()
  1332. if not root:
  1333. raise serializers.ValidationError(_('Folder not found'))
  1334. # 使用MPTT的get_descendants()方法获取所有相关节点
  1335. all_folders = root.get_descendants(include_self=True)
  1336. if self.data.get('name'):
  1337. tools = QuerySet(Tool).filter(
  1338. Q(workspace_id=self.data.get('workspace_id')) &
  1339. Q(folder_id__in=all_folders) &
  1340. Q(user_id=self.data.get('user_id')) &
  1341. Q(name__contains=self.data.get('name'))
  1342. )
  1343. else:
  1344. tools = QuerySet(Tool).filter(
  1345. Q(workspace_id=self.data.get('workspace_id')) &
  1346. Q(folder_id__in=all_folders) &
  1347. Q(user_id=self.data.get('user_id'))
  1348. )
  1349. return page_search(current_page, page_size, tools, lambda record: ToolModelSerializer(record).data)
  1350. def get_query_set(self, workspace_manage, is_x_pack_ee):
  1351. tool_query_set = QuerySet(Tool).filter(workspace_id=self.data.get('workspace_id'))
  1352. folder_query_set = QuerySet(ToolFolder)
  1353. default_query_set = QuerySet(Tool)
  1354. workspace_id = self.data.get('workspace_id')
  1355. user_id = self.data.get('user_id')
  1356. scope = self.data.get('scope')
  1357. tool_type = self.data.get('tool_type')
  1358. desc = self.data.get('desc')
  1359. name = self.data.get('name')
  1360. folder_id = self.data.get('folder_id')
  1361. create_user = self.data.get('create_user')
  1362. if workspace_id is not None:
  1363. folder_query_set = folder_query_set.filter(workspace_id=workspace_id)
  1364. default_query_set = default_query_set.filter(workspace_id=workspace_id)
  1365. if folder_id is not None and folder_id != workspace_id:
  1366. folder_query_set = folder_query_set.filter(parent=folder_id)
  1367. default_query_set = default_query_set.filter(folder_id=folder_id)
  1368. if name is not None:
  1369. folder_query_set = folder_query_set.filter(name__icontains=name)
  1370. default_query_set = default_query_set.filter(name__icontains=name)
  1371. if desc is not None:
  1372. folder_query_set = folder_query_set.filter(desc__icontains=desc)
  1373. default_query_set = default_query_set.filter(desc__icontains=desc)
  1374. if create_user is not None:
  1375. tool_query_set = tool_query_set.filter(user_id=create_user)
  1376. folder_query_set = folder_query_set.filter(user_id=create_user)
  1377. default_query_set = default_query_set.order_by("-create_time")
  1378. if scope is not None:
  1379. tool_query_set = tool_query_set.filter(scope=scope)
  1380. tool_type_list = self.data.get('tool_type_list')
  1381. if tool_type_list:
  1382. tool_query_set = tool_query_set.filter(tool_type__in=tool_type_list)
  1383. elif tool_type:
  1384. tool_query_set = tool_query_set.filter(tool_type=tool_type)
  1385. query_set_dict = {
  1386. 'tool_query_set': tool_query_set,
  1387. 'default_query_set': default_query_set,
  1388. }
  1389. if not workspace_manage:
  1390. query_set_dict['workspace_user_resource_permission_query_set'] = QuerySet(
  1391. WorkspaceUserResourcePermission).filter(
  1392. auth_target_type="TOOL",
  1393. workspace_id=workspace_id,
  1394. user_id=user_id
  1395. )
  1396. return query_set_dict
  1397. @staticmethod
  1398. def is_x_pack_ee():
  1399. workspace_user_role_mapping_model = DatabaseModelManage.get_model("workspace_user_role_mapping")
  1400. role_permission_mapping_model = DatabaseModelManage.get_model("role_permission_mapping_model")
  1401. return workspace_user_role_mapping_model is not None and role_permission_mapping_model is not None
  1402. def page_tool_with_folders(self, current_page: int, page_size: int):
  1403. self.is_valid(raise_exception=True)
  1404. workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
  1405. is_x_pack_ee = self.is_x_pack_ee()
  1406. result = native_page_search(
  1407. current_page, page_size, self.get_query_set(workspace_manage, is_x_pack_ee),
  1408. get_file_content(
  1409. os.path.join(
  1410. PROJECT_DIR,
  1411. "apps", "tools", 'sql',
  1412. 'list_tool.sql' if workspace_manage else (
  1413. 'list_tool_user_ee.sql' if is_x_pack_ee else 'list_tool_user.sql'
  1414. )
  1415. )
  1416. ),
  1417. post_records_handler=lambda record: {
  1418. **record,
  1419. 'input_field_list': json.loads(record.get('input_field_list', '[]')),
  1420. 'init_field_list': json.loads(record.get('init_field_list', '[]')),
  1421. },
  1422. )
  1423. return ResourceMappingSerializer().get_resource_count(result)
  1424. def get_tools(self):
  1425. self.is_valid(raise_exception=True)
  1426. workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
  1427. is_x_pack_ee = self.is_x_pack_ee()
  1428. results = native_search(
  1429. self.get_query_set(workspace_manage, is_x_pack_ee),
  1430. get_file_content(
  1431. os.path.join(
  1432. PROJECT_DIR,
  1433. "apps", "tools", 'sql',
  1434. 'list_tool.sql' if workspace_manage else (
  1435. 'list_tool_user_ee.sql' if is_x_pack_ee else 'list_tool_user.sql'
  1436. )
  1437. )
  1438. ),
  1439. )
  1440. # 返回包含文件夹和工具的结构
  1441. return {
  1442. 'folders': [
  1443. folder for folder in results if folder['resource_type'] == 'folder'
  1444. ],
  1445. 'tools': [
  1446. {
  1447. **tool,
  1448. 'input_field_list': json.loads(tool.get('input_field_list', '[]')),
  1449. 'init_field_list': json.loads(tool.get('init_field_list', '[]')),
  1450. } for tool in results if tool['resource_type'] == 'tool'
  1451. ],
  1452. }