Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.

app_dsl_service.py 32KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806
  1. import base64
  2. import hashlib
  3. import logging
  4. import uuid
  5. from collections.abc import Mapping
  6. from enum import StrEnum
  7. from typing import Optional
  8. from urllib.parse import urlparse
  9. from uuid import uuid4
  10. import yaml # type: ignore
  11. from Crypto.Cipher import AES
  12. from Crypto.Util.Padding import pad, unpad
  13. from packaging import version
  14. from packaging.version import parse as parse_version
  15. from pydantic import BaseModel, Field
  16. from sqlalchemy import select
  17. from sqlalchemy.orm import Session
  18. from core.helper import ssrf_proxy
  19. from core.model_runtime.utils.encoders import jsonable_encoder
  20. from core.plugin.entities.plugin import PluginDependency
  21. from core.workflow.nodes.enums import NodeType
  22. from core.workflow.nodes.knowledge_retrieval.entities import KnowledgeRetrievalNodeData
  23. from core.workflow.nodes.llm.entities import LLMNodeData
  24. from core.workflow.nodes.parameter_extractor.entities import ParameterExtractorNodeData
  25. from core.workflow.nodes.question_classifier.entities import QuestionClassifierNodeData
  26. from core.workflow.nodes.tool.entities import ToolNodeData
  27. from events.app_event import app_model_config_was_updated, app_was_created
  28. from extensions.ext_redis import redis_client
  29. from factories import variable_factory
  30. from models import Account, App, AppMode
  31. from models.model import AppModelConfig
  32. from models.workflow import Workflow
  33. from services.plugin.dependencies_analysis import DependenciesAnalysisService
  34. from services.workflow_draft_variable_service import WorkflowDraftVariableService
  35. from services.workflow_service import WorkflowService
  36. logger = logging.getLogger(__name__)
  37. IMPORT_INFO_REDIS_KEY_PREFIX = "app_import_info:"
  38. CHECK_DEPENDENCIES_REDIS_KEY_PREFIX = "app_check_dependencies:"
  39. IMPORT_INFO_REDIS_EXPIRY = 10 * 60 # 10 minutes
  40. DSL_MAX_SIZE = 10 * 1024 * 1024 # 10MB
  41. CURRENT_DSL_VERSION = "0.4.0"
  42. class ImportMode(StrEnum):
  43. YAML_CONTENT = "yaml-content"
  44. YAML_URL = "yaml-url"
  45. class ImportStatus(StrEnum):
  46. COMPLETED = "completed"
  47. COMPLETED_WITH_WARNINGS = "completed-with-warnings"
  48. PENDING = "pending"
  49. FAILED = "failed"
  50. class Import(BaseModel):
  51. id: str
  52. status: ImportStatus
  53. app_id: Optional[str] = None
  54. app_mode: Optional[str] = None
  55. current_dsl_version: str = CURRENT_DSL_VERSION
  56. imported_dsl_version: str = ""
  57. error: str = ""
  58. class CheckDependenciesResult(BaseModel):
  59. leaked_dependencies: list[PluginDependency] = Field(default_factory=list)
  60. def _check_version_compatibility(imported_version: str) -> ImportStatus:
  61. """Determine import status based on version comparison"""
  62. try:
  63. current_ver = version.parse(CURRENT_DSL_VERSION)
  64. imported_ver = version.parse(imported_version)
  65. except version.InvalidVersion:
  66. return ImportStatus.FAILED
  67. # If imported version is newer than current, always return PENDING
  68. if imported_ver > current_ver:
  69. return ImportStatus.PENDING
  70. # If imported version is older than current's major, return PENDING
  71. if imported_ver.major < current_ver.major:
  72. return ImportStatus.PENDING
  73. # If imported version is older than current's minor, return COMPLETED_WITH_WARNINGS
  74. if imported_ver.minor < current_ver.minor:
  75. return ImportStatus.COMPLETED_WITH_WARNINGS
  76. # If imported version equals or is older than current's micro, return COMPLETED
  77. return ImportStatus.COMPLETED
  78. class PendingData(BaseModel):
  79. import_mode: str
  80. yaml_content: str
  81. name: str | None
  82. description: str | None
  83. icon_type: str | None
  84. icon: str | None
  85. icon_background: str | None
  86. app_id: str | None
  87. class CheckDependenciesPendingData(BaseModel):
  88. dependencies: list[PluginDependency]
  89. app_id: str | None
  90. class AppDslService:
  91. def __init__(self, session: Session):
  92. self._session = session
  93. def import_app(
  94. self,
  95. *,
  96. account: Account,
  97. import_mode: str,
  98. yaml_content: Optional[str] = None,
  99. yaml_url: Optional[str] = None,
  100. name: Optional[str] = None,
  101. description: Optional[str] = None,
  102. icon_type: Optional[str] = None,
  103. icon: Optional[str] = None,
  104. icon_background: Optional[str] = None,
  105. app_id: Optional[str] = None,
  106. ) -> Import:
  107. """Import an app from YAML content or URL."""
  108. import_id = str(uuid.uuid4())
  109. # Validate import mode
  110. try:
  111. mode = ImportMode(import_mode)
  112. except ValueError:
  113. raise ValueError(f"Invalid import_mode: {import_mode}")
  114. # Get YAML content
  115. content: str = ""
  116. if mode == ImportMode.YAML_URL:
  117. if not yaml_url:
  118. return Import(
  119. id=import_id,
  120. status=ImportStatus.FAILED,
  121. error="yaml_url is required when import_mode is yaml-url",
  122. )
  123. try:
  124. parsed_url = urlparse(yaml_url)
  125. if (
  126. parsed_url.scheme == "https"
  127. and parsed_url.netloc == "github.com"
  128. and parsed_url.path.endswith((".yml", ".yaml"))
  129. ):
  130. yaml_url = yaml_url.replace("https://github.com", "https://raw.githubusercontent.com")
  131. yaml_url = yaml_url.replace("/blob/", "/")
  132. response = ssrf_proxy.get(yaml_url.strip(), follow_redirects=True, timeout=(10, 10))
  133. response.raise_for_status()
  134. content = response.content.decode()
  135. if len(content) > DSL_MAX_SIZE:
  136. return Import(
  137. id=import_id,
  138. status=ImportStatus.FAILED,
  139. error="File size exceeds the limit of 10MB",
  140. )
  141. if not content:
  142. return Import(
  143. id=import_id,
  144. status=ImportStatus.FAILED,
  145. error="Empty content from url",
  146. )
  147. except Exception as e:
  148. return Import(
  149. id=import_id,
  150. status=ImportStatus.FAILED,
  151. error=f"Error fetching YAML from URL: {str(e)}",
  152. )
  153. elif mode == ImportMode.YAML_CONTENT:
  154. if not yaml_content:
  155. return Import(
  156. id=import_id,
  157. status=ImportStatus.FAILED,
  158. error="yaml_content is required when import_mode is yaml-content",
  159. )
  160. content = yaml_content
  161. # Process YAML content
  162. try:
  163. # Parse YAML to validate format
  164. data = yaml.safe_load(content)
  165. if not isinstance(data, dict):
  166. return Import(
  167. id=import_id,
  168. status=ImportStatus.FAILED,
  169. error="Invalid YAML format: content must be a mapping",
  170. )
  171. # Validate and fix DSL version
  172. if not data.get("version"):
  173. data["version"] = "0.1.0"
  174. if not data.get("kind") or data.get("kind") != "app":
  175. data["kind"] = "app"
  176. imported_version = data.get("version", "0.1.0")
  177. # check if imported_version is a float-like string
  178. if not isinstance(imported_version, str):
  179. raise ValueError(f"Invalid version type, expected str, got {type(imported_version)}")
  180. status = _check_version_compatibility(imported_version)
  181. # Extract app data
  182. app_data = data.get("app")
  183. if not app_data:
  184. return Import(
  185. id=import_id,
  186. status=ImportStatus.FAILED,
  187. error="Missing app data in YAML content",
  188. )
  189. # If app_id is provided, check if it exists
  190. app = None
  191. if app_id:
  192. stmt = select(App).where(App.id == app_id, App.tenant_id == account.current_tenant_id)
  193. app = self._session.scalar(stmt)
  194. if not app:
  195. return Import(
  196. id=import_id,
  197. status=ImportStatus.FAILED,
  198. error="App not found",
  199. )
  200. if app.mode not in [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]:
  201. return Import(
  202. id=import_id,
  203. status=ImportStatus.FAILED,
  204. error="Only workflow or advanced chat apps can be overwritten",
  205. )
  206. # If major version mismatch, store import info in Redis
  207. if status == ImportStatus.PENDING:
  208. pending_data = PendingData(
  209. import_mode=import_mode,
  210. yaml_content=content,
  211. name=name,
  212. description=description,
  213. icon_type=icon_type,
  214. icon=icon,
  215. icon_background=icon_background,
  216. app_id=app_id,
  217. )
  218. redis_client.setex(
  219. f"{IMPORT_INFO_REDIS_KEY_PREFIX}{import_id}",
  220. IMPORT_INFO_REDIS_EXPIRY,
  221. pending_data.model_dump_json(),
  222. )
  223. return Import(
  224. id=import_id,
  225. status=status,
  226. app_id=app_id,
  227. imported_dsl_version=imported_version,
  228. )
  229. # Extract dependencies
  230. dependencies = data.get("dependencies", [])
  231. check_dependencies_pending_data = None
  232. if dependencies:
  233. check_dependencies_pending_data = [PluginDependency.model_validate(d) for d in dependencies]
  234. elif parse_version(imported_version) <= parse_version("0.1.5"):
  235. if "workflow" in data:
  236. graph = data.get("workflow", {}).get("graph", {})
  237. dependencies_list = self._extract_dependencies_from_workflow_graph(graph)
  238. else:
  239. dependencies_list = self._extract_dependencies_from_model_config(data.get("model_config", {}))
  240. check_dependencies_pending_data = DependenciesAnalysisService.generate_latest_dependencies(
  241. dependencies_list
  242. )
  243. # Create or update app
  244. app = self._create_or_update_app(
  245. app=app,
  246. data=data,
  247. account=account,
  248. name=name,
  249. description=description,
  250. icon_type=icon_type,
  251. icon=icon,
  252. icon_background=icon_background,
  253. dependencies=check_dependencies_pending_data,
  254. )
  255. draft_var_srv = WorkflowDraftVariableService(session=self._session)
  256. draft_var_srv.delete_workflow_variables(app_id=app.id)
  257. return Import(
  258. id=import_id,
  259. status=status,
  260. app_id=app.id,
  261. app_mode=app.mode,
  262. imported_dsl_version=imported_version,
  263. )
  264. except yaml.YAMLError as e:
  265. return Import(
  266. id=import_id,
  267. status=ImportStatus.FAILED,
  268. error=f"Invalid YAML format: {str(e)}",
  269. )
  270. except Exception as e:
  271. logger.exception("Failed to import app")
  272. return Import(
  273. id=import_id,
  274. status=ImportStatus.FAILED,
  275. error=str(e),
  276. )
  277. def confirm_import(self, *, import_id: str, account: Account) -> Import:
  278. """
  279. Confirm an import that requires confirmation
  280. """
  281. redis_key = f"{IMPORT_INFO_REDIS_KEY_PREFIX}{import_id}"
  282. pending_data = redis_client.get(redis_key)
  283. if not pending_data:
  284. return Import(
  285. id=import_id,
  286. status=ImportStatus.FAILED,
  287. error="Import information expired or does not exist",
  288. )
  289. try:
  290. if not isinstance(pending_data, str | bytes):
  291. return Import(
  292. id=import_id,
  293. status=ImportStatus.FAILED,
  294. error="Invalid import information",
  295. )
  296. pending_data = PendingData.model_validate_json(pending_data)
  297. data = yaml.safe_load(pending_data.yaml_content)
  298. app = None
  299. if pending_data.app_id:
  300. stmt = select(App).where(App.id == pending_data.app_id, App.tenant_id == account.current_tenant_id)
  301. app = self._session.scalar(stmt)
  302. # Create or update app
  303. app = self._create_or_update_app(
  304. app=app,
  305. data=data,
  306. account=account,
  307. name=pending_data.name,
  308. description=pending_data.description,
  309. icon_type=pending_data.icon_type,
  310. icon=pending_data.icon,
  311. icon_background=pending_data.icon_background,
  312. )
  313. # Delete import info from Redis
  314. redis_client.delete(redis_key)
  315. return Import(
  316. id=import_id,
  317. status=ImportStatus.COMPLETED,
  318. app_id=app.id,
  319. app_mode=app.mode,
  320. current_dsl_version=CURRENT_DSL_VERSION,
  321. imported_dsl_version=data.get("version", "0.1.0"),
  322. )
  323. except Exception as e:
  324. logger.exception("Error confirming import")
  325. return Import(
  326. id=import_id,
  327. status=ImportStatus.FAILED,
  328. error=str(e),
  329. )
  330. def check_dependencies(
  331. self,
  332. *,
  333. app_model: App,
  334. ) -> CheckDependenciesResult:
  335. """Check dependencies"""
  336. # Get dependencies from Redis
  337. redis_key = f"{CHECK_DEPENDENCIES_REDIS_KEY_PREFIX}{app_model.id}"
  338. dependencies = redis_client.get(redis_key)
  339. if not dependencies:
  340. return CheckDependenciesResult()
  341. # Extract dependencies
  342. dependencies = CheckDependenciesPendingData.model_validate_json(dependencies)
  343. # Get leaked dependencies
  344. leaked_dependencies = DependenciesAnalysisService.get_leaked_dependencies(
  345. tenant_id=app_model.tenant_id, dependencies=dependencies.dependencies
  346. )
  347. return CheckDependenciesResult(
  348. leaked_dependencies=leaked_dependencies,
  349. )
  350. def _create_or_update_app(
  351. self,
  352. *,
  353. app: Optional[App],
  354. data: dict,
  355. account: Account,
  356. name: Optional[str] = None,
  357. description: Optional[str] = None,
  358. icon_type: Optional[str] = None,
  359. icon: Optional[str] = None,
  360. icon_background: Optional[str] = None,
  361. dependencies: Optional[list[PluginDependency]] = None,
  362. ) -> App:
  363. """Create a new app or update an existing one."""
  364. app_data = data.get("app", {})
  365. app_mode = app_data.get("mode")
  366. if not app_mode:
  367. raise ValueError("loss app mode")
  368. app_mode = AppMode(app_mode)
  369. # Set icon type
  370. icon_type_value = icon_type or app_data.get("icon_type")
  371. if icon_type_value in ["emoji", "link", "image"]:
  372. icon_type = icon_type_value
  373. else:
  374. icon_type = "emoji"
  375. icon = icon or str(app_data.get("icon", ""))
  376. if app:
  377. # Update existing app
  378. app.name = name or app_data.get("name", app.name)
  379. app.description = description or app_data.get("description", app.description)
  380. app.icon_type = icon_type
  381. app.icon = icon
  382. app.icon_background = icon_background or app_data.get("icon_background", app.icon_background)
  383. app.updated_by = account.id
  384. else:
  385. if account.current_tenant_id is None:
  386. raise ValueError("Current tenant is not set")
  387. # Create new app
  388. app = App()
  389. app.id = str(uuid4())
  390. app.tenant_id = account.current_tenant_id
  391. app.mode = app_mode.value
  392. app.name = name or app_data.get("name", "")
  393. app.description = description or app_data.get("description", "")
  394. app.icon_type = icon_type
  395. app.icon = icon
  396. app.icon_background = icon_background or app_data.get("icon_background", "#FFFFFF")
  397. app.enable_site = True
  398. app.enable_api = True
  399. app.use_icon_as_answer_icon = app_data.get("use_icon_as_answer_icon", False)
  400. app.created_by = account.id
  401. app.updated_by = account.id
  402. self._session.add(app)
  403. self._session.commit()
  404. app_was_created.send(app, account=account)
  405. # save dependencies
  406. if dependencies:
  407. redis_client.setex(
  408. f"{CHECK_DEPENDENCIES_REDIS_KEY_PREFIX}{app.id}",
  409. IMPORT_INFO_REDIS_EXPIRY,
  410. CheckDependenciesPendingData(app_id=app.id, dependencies=dependencies).model_dump_json(),
  411. )
  412. # Initialize app based on mode
  413. if app_mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
  414. workflow_data = data.get("workflow")
  415. if not workflow_data or not isinstance(workflow_data, dict):
  416. raise ValueError("Missing workflow data for workflow/advanced chat app")
  417. environment_variables_list = workflow_data.get("environment_variables", [])
  418. environment_variables = [
  419. variable_factory.build_environment_variable_from_mapping(obj) for obj in environment_variables_list
  420. ]
  421. conversation_variables_list = workflow_data.get("conversation_variables", [])
  422. conversation_variables = [
  423. variable_factory.build_conversation_variable_from_mapping(obj) for obj in conversation_variables_list
  424. ]
  425. workflow_service = WorkflowService()
  426. current_draft_workflow = workflow_service.get_draft_workflow(app_model=app)
  427. if current_draft_workflow:
  428. unique_hash = current_draft_workflow.unique_hash
  429. else:
  430. unique_hash = None
  431. graph = workflow_data.get("graph", {})
  432. for node in graph.get("nodes", []):
  433. if node.get("data", {}).get("type", "") == NodeType.KNOWLEDGE_RETRIEVAL.value:
  434. dataset_ids = node["data"].get("dataset_ids", [])
  435. node["data"]["dataset_ids"] = [
  436. decrypted_id
  437. for dataset_id in dataset_ids
  438. if (decrypted_id := self.decrypt_dataset_id(encrypted_data=dataset_id, tenant_id=app.tenant_id))
  439. ]
  440. workflow_service.sync_draft_workflow(
  441. app_model=app,
  442. graph=workflow_data.get("graph", {}),
  443. features=workflow_data.get("features", {}),
  444. unique_hash=unique_hash,
  445. account=account,
  446. environment_variables=environment_variables,
  447. conversation_variables=conversation_variables,
  448. )
  449. elif app_mode in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.COMPLETION}:
  450. # Initialize model config
  451. model_config = data.get("model_config")
  452. if not model_config or not isinstance(model_config, dict):
  453. raise ValueError("Missing model_config for chat/agent-chat/completion app")
  454. # Initialize or update model config
  455. if not app.app_model_config:
  456. app_model_config = AppModelConfig().from_model_config_dict(model_config)
  457. app_model_config.id = str(uuid4())
  458. app_model_config.app_id = app.id
  459. app_model_config.created_by = account.id
  460. app_model_config.updated_by = account.id
  461. app.app_model_config_id = app_model_config.id
  462. self._session.add(app_model_config)
  463. app_model_config_was_updated.send(app, app_model_config=app_model_config)
  464. else:
  465. raise ValueError("Invalid app mode")
  466. return app
  467. @classmethod
  468. def export_dsl(cls, app_model: App, include_secret: bool = False, workflow_id: Optional[str] = None) -> str:
  469. """
  470. Export app
  471. :param app_model: App instance
  472. :param include_secret: Whether include secret variable
  473. :return:
  474. """
  475. app_mode = AppMode.value_of(app_model.mode)
  476. export_data = {
  477. "version": CURRENT_DSL_VERSION,
  478. "kind": "app",
  479. "app": {
  480. "name": app_model.name,
  481. "mode": app_model.mode,
  482. "icon": "🤖" if app_model.icon_type == "image" else app_model.icon,
  483. "icon_background": "#FFEAD5" if app_model.icon_type == "image" else app_model.icon_background,
  484. "description": app_model.description,
  485. "use_icon_as_answer_icon": app_model.use_icon_as_answer_icon,
  486. },
  487. }
  488. if app_mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
  489. cls._append_workflow_export_data(
  490. export_data=export_data, app_model=app_model, include_secret=include_secret, workflow_id=workflow_id
  491. )
  492. else:
  493. cls._append_model_config_export_data(export_data, app_model)
  494. return yaml.dump(export_data, allow_unicode=True) # type: ignore
  495. @classmethod
  496. def _append_workflow_export_data(
  497. cls, *, export_data: dict, app_model: App, include_secret: bool, workflow_id: Optional[str] = None
  498. ) -> None:
  499. """
  500. Append workflow export data
  501. :param export_data: export data
  502. :param app_model: App instance
  503. """
  504. workflow_service = WorkflowService()
  505. workflow = workflow_service.get_draft_workflow(app_model, workflow_id)
  506. if not workflow:
  507. raise ValueError("Missing draft workflow configuration, please check.")
  508. workflow_dict = workflow.to_dict(include_secret=include_secret)
  509. # TODO: refactor: we need a better way to filter workspace related data from nodes
  510. for node in workflow_dict.get("graph", {}).get("nodes", []):
  511. node_data = node.get("data", {})
  512. if not node_data:
  513. continue
  514. data_type = node_data.get("type", "")
  515. if data_type == NodeType.KNOWLEDGE_RETRIEVAL.value:
  516. dataset_ids = node_data.get("dataset_ids", [])
  517. node_data["dataset_ids"] = [
  518. cls.encrypt_dataset_id(dataset_id=dataset_id, tenant_id=app_model.tenant_id)
  519. for dataset_id in dataset_ids
  520. ]
  521. # filter credential id from tool node
  522. if not include_secret and data_type == NodeType.TOOL.value:
  523. node_data.pop("credential_id", None)
  524. # filter credential id from agent node
  525. if not include_secret and data_type == NodeType.AGENT.value:
  526. for tool in node_data.get("agent_parameters", {}).get("tools", {}).get("value", []):
  527. tool.pop("credential_id", None)
  528. export_data["workflow"] = workflow_dict
  529. dependencies = cls._extract_dependencies_from_workflow(workflow)
  530. export_data["dependencies"] = [
  531. jsonable_encoder(d.model_dump())
  532. for d in DependenciesAnalysisService.generate_dependencies(
  533. tenant_id=app_model.tenant_id, dependencies=dependencies
  534. )
  535. ]
  536. @classmethod
  537. def _append_model_config_export_data(cls, export_data: dict, app_model: App) -> None:
  538. """
  539. Append model config export data
  540. :param export_data: export data
  541. :param app_model: App instance
  542. """
  543. app_model_config = app_model.app_model_config
  544. if not app_model_config:
  545. raise ValueError("Missing app configuration, please check.")
  546. model_config = app_model_config.to_dict()
  547. # TODO: refactor: we need a better way to filter workspace related data from model config
  548. # filter credential id from model config
  549. for tool in model_config.get("agent_mode", {}).get("tools", []):
  550. tool.pop("credential_id", None)
  551. export_data["model_config"] = model_config
  552. dependencies = cls._extract_dependencies_from_model_config(app_model_config.to_dict())
  553. export_data["dependencies"] = [
  554. jsonable_encoder(d.model_dump())
  555. for d in DependenciesAnalysisService.generate_dependencies(
  556. tenant_id=app_model.tenant_id, dependencies=dependencies
  557. )
  558. ]
  559. @classmethod
  560. def _extract_dependencies_from_workflow(cls, workflow: Workflow) -> list[str]:
  561. """
  562. Extract dependencies from workflow
  563. :param workflow: Workflow instance
  564. :return: dependencies list format like ["langgenius/google"]
  565. """
  566. graph = workflow.graph_dict
  567. dependencies = cls._extract_dependencies_from_workflow_graph(graph)
  568. return dependencies
  569. @classmethod
  570. def _extract_dependencies_from_workflow_graph(cls, graph: Mapping) -> list[str]:
  571. """
  572. Extract dependencies from workflow graph
  573. :param graph: Workflow graph
  574. :return: dependencies list format like ["langgenius/google"]
  575. """
  576. dependencies = []
  577. for node in graph.get("nodes", []):
  578. try:
  579. typ = node.get("data", {}).get("type")
  580. match typ:
  581. case NodeType.TOOL.value:
  582. tool_entity = ToolNodeData(**node["data"])
  583. dependencies.append(
  584. DependenciesAnalysisService.analyze_tool_dependency(tool_entity.provider_id),
  585. )
  586. case NodeType.LLM.value:
  587. llm_entity = LLMNodeData(**node["data"])
  588. dependencies.append(
  589. DependenciesAnalysisService.analyze_model_provider_dependency(llm_entity.model.provider),
  590. )
  591. case NodeType.QUESTION_CLASSIFIER.value:
  592. question_classifier_entity = QuestionClassifierNodeData(**node["data"])
  593. dependencies.append(
  594. DependenciesAnalysisService.analyze_model_provider_dependency(
  595. question_classifier_entity.model.provider
  596. ),
  597. )
  598. case NodeType.PARAMETER_EXTRACTOR.value:
  599. parameter_extractor_entity = ParameterExtractorNodeData(**node["data"])
  600. dependencies.append(
  601. DependenciesAnalysisService.analyze_model_provider_dependency(
  602. parameter_extractor_entity.model.provider
  603. ),
  604. )
  605. case NodeType.KNOWLEDGE_RETRIEVAL.value:
  606. knowledge_retrieval_entity = KnowledgeRetrievalNodeData(**node["data"])
  607. if knowledge_retrieval_entity.retrieval_mode == "multiple":
  608. if knowledge_retrieval_entity.multiple_retrieval_config:
  609. if (
  610. knowledge_retrieval_entity.multiple_retrieval_config.reranking_mode
  611. == "reranking_model"
  612. ):
  613. if knowledge_retrieval_entity.multiple_retrieval_config.reranking_model:
  614. dependencies.append(
  615. DependenciesAnalysisService.analyze_model_provider_dependency(
  616. knowledge_retrieval_entity.multiple_retrieval_config.reranking_model.provider
  617. ),
  618. )
  619. elif (
  620. knowledge_retrieval_entity.multiple_retrieval_config.reranking_mode
  621. == "weighted_score"
  622. ):
  623. if knowledge_retrieval_entity.multiple_retrieval_config.weights:
  624. vector_setting = (
  625. knowledge_retrieval_entity.multiple_retrieval_config.weights.vector_setting
  626. )
  627. dependencies.append(
  628. DependenciesAnalysisService.analyze_model_provider_dependency(
  629. vector_setting.embedding_provider_name
  630. ),
  631. )
  632. elif knowledge_retrieval_entity.retrieval_mode == "single":
  633. model_config = knowledge_retrieval_entity.single_retrieval_config
  634. if model_config:
  635. dependencies.append(
  636. DependenciesAnalysisService.analyze_model_provider_dependency(
  637. model_config.model.provider
  638. ),
  639. )
  640. case _:
  641. # TODO: Handle default case or unknown node types
  642. pass
  643. except Exception as e:
  644. logger.exception("Error extracting node dependency", exc_info=e)
  645. return dependencies
  646. @classmethod
  647. def _extract_dependencies_from_model_config(cls, model_config: Mapping) -> list[str]:
  648. """
  649. Extract dependencies from model config
  650. :param model_config: model config dict
  651. :return: dependencies list format like ["langgenius/google"]
  652. """
  653. dependencies = []
  654. try:
  655. # completion model
  656. model_dict = model_config.get("model", {})
  657. if model_dict:
  658. dependencies.append(
  659. DependenciesAnalysisService.analyze_model_provider_dependency(model_dict.get("provider", ""))
  660. )
  661. # reranking model
  662. dataset_configs = model_config.get("dataset_configs", {})
  663. if dataset_configs:
  664. for dataset_config in dataset_configs.get("datasets", {}).get("datasets", []):
  665. if dataset_config.get("reranking_model"):
  666. dependencies.append(
  667. DependenciesAnalysisService.analyze_model_provider_dependency(
  668. dataset_config.get("reranking_model", {})
  669. .get("reranking_provider_name", {})
  670. .get("provider")
  671. )
  672. )
  673. # tools
  674. agent_configs = model_config.get("agent_mode", {})
  675. if agent_configs:
  676. for agent_config in agent_configs.get("tools", []):
  677. dependencies.append(
  678. DependenciesAnalysisService.analyze_tool_dependency(agent_config.get("provider_id"))
  679. )
  680. except Exception as e:
  681. logger.exception("Error extracting model config dependency", exc_info=e)
  682. return dependencies
  683. @classmethod
  684. def get_leaked_dependencies(cls, tenant_id: str, dsl_dependencies: list[dict]) -> list[PluginDependency]:
  685. """
  686. Returns the leaked dependencies in current workspace
  687. """
  688. dependencies = [PluginDependency(**dep) for dep in dsl_dependencies]
  689. if not dependencies:
  690. return []
  691. return DependenciesAnalysisService.get_leaked_dependencies(tenant_id=tenant_id, dependencies=dependencies)
  692. @staticmethod
  693. def _generate_aes_key(tenant_id: str) -> bytes:
  694. """Generate AES key based on tenant_id"""
  695. return hashlib.sha256(tenant_id.encode()).digest()
  696. @classmethod
  697. def encrypt_dataset_id(cls, dataset_id: str, tenant_id: str) -> str:
  698. """Encrypt dataset_id using AES-CBC mode"""
  699. key = cls._generate_aes_key(tenant_id)
  700. iv = key[:16]
  701. cipher = AES.new(key, AES.MODE_CBC, iv)
  702. ct_bytes = cipher.encrypt(pad(dataset_id.encode(), AES.block_size))
  703. return base64.b64encode(ct_bytes).decode()
  704. @classmethod
  705. def decrypt_dataset_id(cls, encrypted_data: str, tenant_id: str) -> str | None:
  706. """AES decryption"""
  707. try:
  708. key = cls._generate_aes_key(tenant_id)
  709. iv = key[:16]
  710. cipher = AES.new(key, AES.MODE_CBC, iv)
  711. pt = unpad(cipher.decrypt(base64.b64decode(encrypted_data)), AES.block_size)
  712. return pt.decode()
  713. except Exception:
  714. return None