Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

rag_pipeline_dsl_service.py 42KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944
  1. import base64
  2. import hashlib
  3. import json
  4. import logging
  5. import uuid
  6. from collections.abc import Mapping
  7. from datetime import UTC, datetime
  8. from enum import StrEnum
  9. from typing import cast
  10. from urllib.parse import urlparse
  11. from uuid import uuid4
  12. import yaml # type: ignore
  13. from Crypto.Cipher import AES
  14. from Crypto.Util.Padding import pad, unpad
  15. from flask_login import current_user
  16. from packaging import version
  17. from pydantic import BaseModel, Field
  18. from sqlalchemy import select
  19. from sqlalchemy.orm import Session
  20. from core.helper import ssrf_proxy
  21. from core.helper.name_generator import generate_incremental_name
  22. from core.model_runtime.utils.encoders import jsonable_encoder
  23. from core.plugin.entities.plugin import PluginDependency
  24. from core.workflow.enums import NodeType
  25. from core.workflow.nodes.datasource.entities import DatasourceNodeData
  26. from core.workflow.nodes.knowledge_retrieval.entities import KnowledgeRetrievalNodeData
  27. from core.workflow.nodes.llm.entities import LLMNodeData
  28. from core.workflow.nodes.parameter_extractor.entities import ParameterExtractorNodeData
  29. from core.workflow.nodes.question_classifier.entities import QuestionClassifierNodeData
  30. from core.workflow.nodes.tool.entities import ToolNodeData
  31. from extensions.ext_redis import redis_client
  32. from factories import variable_factory
  33. from models import Account
  34. from models.dataset import Dataset, DatasetCollectionBinding, Pipeline
  35. from models.workflow import Workflow, WorkflowType
  36. from services.entities.knowledge_entities.rag_pipeline_entities import (
  37. IconInfo,
  38. KnowledgeConfiguration,
  39. RagPipelineDatasetCreateEntity,
  40. )
  41. from services.plugin.dependencies_analysis import DependenciesAnalysisService
  42. logger = logging.getLogger(__name__)
  43. IMPORT_INFO_REDIS_KEY_PREFIX = "app_import_info:"
  44. CHECK_DEPENDENCIES_REDIS_KEY_PREFIX = "app_check_dependencies:"
  45. IMPORT_INFO_REDIS_EXPIRY = 10 * 60 # 10 minutes
  46. DSL_MAX_SIZE = 10 * 1024 * 1024 # 10MB
  47. CURRENT_DSL_VERSION = "0.1.0"
  48. class ImportMode(StrEnum):
  49. YAML_CONTENT = "yaml-content"
  50. YAML_URL = "yaml-url"
  51. class ImportStatus(StrEnum):
  52. COMPLETED = "completed"
  53. COMPLETED_WITH_WARNINGS = "completed-with-warnings"
  54. PENDING = "pending"
  55. FAILED = "failed"
  56. class RagPipelineImportInfo(BaseModel):
  57. id: str
  58. status: ImportStatus
  59. pipeline_id: str | None = None
  60. current_dsl_version: str = CURRENT_DSL_VERSION
  61. imported_dsl_version: str = ""
  62. error: str = ""
  63. dataset_id: str | None = None
  64. class CheckDependenciesResult(BaseModel):
  65. leaked_dependencies: list[PluginDependency] = Field(default_factory=list)
  66. def _check_version_compatibility(imported_version: str) -> ImportStatus:
  67. """Determine import status based on version comparison"""
  68. try:
  69. current_ver = version.parse(CURRENT_DSL_VERSION)
  70. imported_ver = version.parse(imported_version)
  71. except version.InvalidVersion:
  72. return ImportStatus.FAILED
  73. # If imported version is newer than current, always return PENDING
  74. if imported_ver > current_ver:
  75. return ImportStatus.PENDING
  76. # If imported version is older than current's major, return PENDING
  77. if imported_ver.major < current_ver.major:
  78. return ImportStatus.PENDING
  79. # If imported version is older than current's minor, return COMPLETED_WITH_WARNINGS
  80. if imported_ver.minor < current_ver.minor:
  81. return ImportStatus.COMPLETED_WITH_WARNINGS
  82. # If imported version equals or is older than current's micro, return COMPLETED
  83. return ImportStatus.COMPLETED
  84. class RagPipelinePendingData(BaseModel):
  85. import_mode: str
  86. yaml_content: str
  87. pipeline_id: str | None
  88. class CheckDependenciesPendingData(BaseModel):
  89. dependencies: list[PluginDependency]
  90. pipeline_id: str | None
  91. class RagPipelineDslService:
  92. def __init__(self, session: Session):
  93. self._session = session
  94. def import_rag_pipeline(
  95. self,
  96. *,
  97. account: Account,
  98. import_mode: str,
  99. yaml_content: str | None = None,
  100. yaml_url: str | None = None,
  101. pipeline_id: str | None = None,
  102. dataset: Dataset | None = None,
  103. dataset_name: str | None = None,
  104. icon_info: IconInfo | None = None,
  105. ) -> RagPipelineImportInfo:
  106. """Import an app from YAML content or URL."""
  107. import_id = str(uuid.uuid4())
  108. # Validate import mode
  109. try:
  110. mode = ImportMode(import_mode)
  111. except ValueError:
  112. raise ValueError(f"Invalid import_mode: {import_mode}")
  113. # Get YAML content
  114. content: str = ""
  115. if mode == ImportMode.YAML_URL:
  116. if not yaml_url:
  117. return RagPipelineImportInfo(
  118. id=import_id,
  119. status=ImportStatus.FAILED,
  120. error="yaml_url is required when import_mode is yaml-url",
  121. )
  122. try:
  123. parsed_url = urlparse(yaml_url)
  124. if (
  125. parsed_url.scheme == "https"
  126. and parsed_url.netloc == "github.com"
  127. and parsed_url.path.endswith((".yml", ".yaml"))
  128. ):
  129. yaml_url = yaml_url.replace("https://github.com", "https://raw.githubusercontent.com")
  130. yaml_url = yaml_url.replace("/blob/", "/")
  131. response = ssrf_proxy.get(yaml_url.strip(), follow_redirects=True, timeout=(10, 10))
  132. response.raise_for_status()
  133. content = response.content.decode()
  134. if len(content) > DSL_MAX_SIZE:
  135. return RagPipelineImportInfo(
  136. id=import_id,
  137. status=ImportStatus.FAILED,
  138. error="File size exceeds the limit of 10MB",
  139. )
  140. if not content:
  141. return RagPipelineImportInfo(
  142. id=import_id,
  143. status=ImportStatus.FAILED,
  144. error="Empty content from url",
  145. )
  146. except Exception as e:
  147. return RagPipelineImportInfo(
  148. id=import_id,
  149. status=ImportStatus.FAILED,
  150. error=f"Error fetching YAML from URL: {str(e)}",
  151. )
  152. elif mode == ImportMode.YAML_CONTENT:
  153. if not yaml_content:
  154. return RagPipelineImportInfo(
  155. id=import_id,
  156. status=ImportStatus.FAILED,
  157. error="yaml_content is required when import_mode is yaml-content",
  158. )
  159. content = yaml_content
  160. # Process YAML content
  161. try:
  162. # Parse YAML to validate format
  163. data = yaml.safe_load(content)
  164. if not isinstance(data, dict):
  165. return RagPipelineImportInfo(
  166. id=import_id,
  167. status=ImportStatus.FAILED,
  168. error="Invalid YAML format: content must be a mapping",
  169. )
  170. # Validate and fix DSL version
  171. if not data.get("version"):
  172. data["version"] = "0.1.0"
  173. if not data.get("kind") or data.get("kind") != "rag_pipeline":
  174. data["kind"] = "rag_pipeline"
  175. imported_version = data.get("version", "0.1.0")
  176. # check if imported_version is a float-like string
  177. if not isinstance(imported_version, str):
  178. raise ValueError(f"Invalid version type, expected str, got {type(imported_version)}")
  179. status = _check_version_compatibility(imported_version)
  180. # Extract app data
  181. pipeline_data = data.get("rag_pipeline")
  182. if not pipeline_data:
  183. return RagPipelineImportInfo(
  184. id=import_id,
  185. status=ImportStatus.FAILED,
  186. error="Missing rag_pipeline data in YAML content",
  187. )
  188. # If app_id is provided, check if it exists
  189. pipeline = None
  190. if pipeline_id:
  191. stmt = select(Pipeline).where(
  192. Pipeline.id == pipeline_id,
  193. Pipeline.tenant_id == account.current_tenant_id,
  194. )
  195. pipeline = self._session.scalar(stmt)
  196. if not pipeline:
  197. return RagPipelineImportInfo(
  198. id=import_id,
  199. status=ImportStatus.FAILED,
  200. error="Pipeline not found",
  201. )
  202. dataset = pipeline.retrieve_dataset(session=self._session)
  203. # If major version mismatch, store import info in Redis
  204. if status == ImportStatus.PENDING:
  205. pending_data = RagPipelinePendingData(
  206. import_mode=import_mode,
  207. yaml_content=content,
  208. pipeline_id=pipeline_id,
  209. )
  210. redis_client.setex(
  211. f"{IMPORT_INFO_REDIS_KEY_PREFIX}{import_id}",
  212. IMPORT_INFO_REDIS_EXPIRY,
  213. pending_data.model_dump_json(),
  214. )
  215. return RagPipelineImportInfo(
  216. id=import_id,
  217. status=status,
  218. pipeline_id=pipeline_id,
  219. imported_dsl_version=imported_version,
  220. )
  221. # Extract dependencies
  222. dependencies = data.get("dependencies", [])
  223. check_dependencies_pending_data = None
  224. if dependencies:
  225. check_dependencies_pending_data = [PluginDependency.model_validate(d) for d in dependencies]
  226. # Create or update pipeline
  227. pipeline = self._create_or_update_pipeline(
  228. pipeline=pipeline,
  229. data=data,
  230. account=account,
  231. dependencies=check_dependencies_pending_data,
  232. )
  233. # create dataset
  234. name = pipeline.name or "Untitled"
  235. description = pipeline.description
  236. if icon_info:
  237. icon_type = icon_info.icon_type
  238. icon = icon_info.icon
  239. icon_background = icon_info.icon_background
  240. icon_url = icon_info.icon_url
  241. else:
  242. icon_type = data.get("rag_pipeline", {}).get("icon_type")
  243. icon = data.get("rag_pipeline", {}).get("icon")
  244. icon_background = data.get("rag_pipeline", {}).get("icon_background")
  245. icon_url = data.get("rag_pipeline", {}).get("icon_url")
  246. workflow = data.get("workflow", {})
  247. graph = workflow.get("graph", {})
  248. nodes = graph.get("nodes", [])
  249. dataset_id = None
  250. for node in nodes:
  251. if node.get("data", {}).get("type") == "knowledge-index":
  252. knowledge_configuration = KnowledgeConfiguration(**node.get("data", {}))
  253. if (
  254. dataset
  255. and pipeline.is_published
  256. and dataset.chunk_structure != knowledge_configuration.chunk_structure
  257. ):
  258. raise ValueError("Chunk structure is not compatible with the published pipeline")
  259. if not dataset:
  260. datasets = self._session.query(Dataset).filter_by(tenant_id=account.current_tenant_id).all()
  261. names = [dataset.name for dataset in datasets]
  262. generate_name = generate_incremental_name(names, name)
  263. dataset = Dataset(
  264. tenant_id=account.current_tenant_id,
  265. name=generate_name,
  266. description=description,
  267. icon_info={
  268. "icon_type": icon_type,
  269. "icon": icon,
  270. "icon_background": icon_background,
  271. "icon_url": icon_url,
  272. },
  273. indexing_technique=knowledge_configuration.indexing_technique,
  274. created_by=account.id,
  275. retrieval_model=knowledge_configuration.retrieval_model.model_dump(),
  276. runtime_mode="rag_pipeline",
  277. chunk_structure=knowledge_configuration.chunk_structure,
  278. )
  279. if knowledge_configuration.indexing_technique == "high_quality":
  280. dataset_collection_binding = (
  281. self._session.query(DatasetCollectionBinding)
  282. .where(
  283. DatasetCollectionBinding.provider_name
  284. == knowledge_configuration.embedding_model_provider,
  285. DatasetCollectionBinding.model_name == knowledge_configuration.embedding_model,
  286. DatasetCollectionBinding.type == "dataset",
  287. )
  288. .order_by(DatasetCollectionBinding.created_at)
  289. .first()
  290. )
  291. if not dataset_collection_binding:
  292. dataset_collection_binding = DatasetCollectionBinding(
  293. provider_name=knowledge_configuration.embedding_model_provider,
  294. model_name=knowledge_configuration.embedding_model,
  295. collection_name=Dataset.gen_collection_name_by_id(str(uuid.uuid4())),
  296. type="dataset",
  297. )
  298. self._session.add(dataset_collection_binding)
  299. self._session.commit()
  300. dataset_collection_binding_id = dataset_collection_binding.id
  301. dataset.collection_binding_id = dataset_collection_binding_id
  302. dataset.embedding_model = knowledge_configuration.embedding_model
  303. dataset.embedding_model_provider = knowledge_configuration.embedding_model_provider
  304. elif knowledge_configuration.indexing_technique == "economy":
  305. dataset.keyword_number = knowledge_configuration.keyword_number
  306. dataset.pipeline_id = pipeline.id
  307. self._session.add(dataset)
  308. self._session.commit()
  309. dataset_id = dataset.id
  310. if not dataset_id:
  311. raise ValueError("DSL is not valid, please check the Knowledge Index node.")
  312. return RagPipelineImportInfo(
  313. id=import_id,
  314. status=status,
  315. pipeline_id=pipeline.id,
  316. dataset_id=dataset_id,
  317. imported_dsl_version=imported_version,
  318. )
  319. except yaml.YAMLError as e:
  320. return RagPipelineImportInfo(
  321. id=import_id,
  322. status=ImportStatus.FAILED,
  323. error=f"Invalid YAML format: {str(e)}",
  324. )
  325. except Exception as e:
  326. logger.exception("Failed to import app")
  327. return RagPipelineImportInfo(
  328. id=import_id,
  329. status=ImportStatus.FAILED,
  330. error=str(e),
  331. )
  332. def confirm_import(self, *, import_id: str, account: Account) -> RagPipelineImportInfo:
  333. """
  334. Confirm an import that requires confirmation
  335. """
  336. redis_key = f"{IMPORT_INFO_REDIS_KEY_PREFIX}{import_id}"
  337. pending_data = redis_client.get(redis_key)
  338. if not pending_data:
  339. return RagPipelineImportInfo(
  340. id=import_id,
  341. status=ImportStatus.FAILED,
  342. error="Import information expired or does not exist",
  343. )
  344. try:
  345. if not isinstance(pending_data, str | bytes):
  346. return RagPipelineImportInfo(
  347. id=import_id,
  348. status=ImportStatus.FAILED,
  349. error="Invalid import information",
  350. )
  351. pending_data = RagPipelinePendingData.model_validate_json(pending_data)
  352. data = yaml.safe_load(pending_data.yaml_content)
  353. pipeline = None
  354. if pending_data.pipeline_id:
  355. stmt = select(Pipeline).where(
  356. Pipeline.id == pending_data.pipeline_id,
  357. Pipeline.tenant_id == account.current_tenant_id,
  358. )
  359. pipeline = self._session.scalar(stmt)
  360. # Create or update app
  361. pipeline = self._create_or_update_pipeline(
  362. pipeline=pipeline,
  363. data=data,
  364. account=account,
  365. )
  366. dataset = pipeline.retrieve_dataset(session=self._session)
  367. # create dataset
  368. name = pipeline.name
  369. description = pipeline.description
  370. icon_type = data.get("rag_pipeline", {}).get("icon_type")
  371. icon = data.get("rag_pipeline", {}).get("icon")
  372. icon_background = data.get("rag_pipeline", {}).get("icon_background")
  373. icon_url = data.get("rag_pipeline", {}).get("icon_url")
  374. workflow = data.get("workflow", {})
  375. graph = workflow.get("graph", {})
  376. nodes = graph.get("nodes", [])
  377. dataset_id = None
  378. for node in nodes:
  379. if node.get("data", {}).get("type") == "knowledge-index":
  380. knowledge_configuration = KnowledgeConfiguration(**node.get("data", {}))
  381. if not dataset:
  382. dataset = Dataset(
  383. tenant_id=account.current_tenant_id,
  384. name=name,
  385. description=description,
  386. icon_info={
  387. "icon_type": icon_type,
  388. "icon": icon,
  389. "icon_background": icon_background,
  390. "icon_url": icon_url,
  391. },
  392. indexing_technique=knowledge_configuration.indexing_technique,
  393. created_by=account.id,
  394. retrieval_model=knowledge_configuration.retrieval_model.model_dump(),
  395. runtime_mode="rag_pipeline",
  396. chunk_structure=knowledge_configuration.chunk_structure,
  397. )
  398. else:
  399. dataset.indexing_technique = knowledge_configuration.indexing_technique
  400. dataset.retrieval_model = knowledge_configuration.retrieval_model.model_dump()
  401. dataset.runtime_mode = "rag_pipeline"
  402. dataset.chunk_structure = knowledge_configuration.chunk_structure
  403. if knowledge_configuration.indexing_technique == "high_quality":
  404. dataset_collection_binding = (
  405. self._session.query(DatasetCollectionBinding)
  406. .where(
  407. DatasetCollectionBinding.provider_name
  408. == knowledge_configuration.embedding_model_provider,
  409. DatasetCollectionBinding.model_name == knowledge_configuration.embedding_model,
  410. DatasetCollectionBinding.type == "dataset",
  411. )
  412. .order_by(DatasetCollectionBinding.created_at)
  413. .first()
  414. )
  415. if not dataset_collection_binding:
  416. dataset_collection_binding = DatasetCollectionBinding(
  417. provider_name=knowledge_configuration.embedding_model_provider,
  418. model_name=knowledge_configuration.embedding_model,
  419. collection_name=Dataset.gen_collection_name_by_id(str(uuid.uuid4())),
  420. type="dataset",
  421. )
  422. self._session.add(dataset_collection_binding)
  423. self._session.commit()
  424. dataset_collection_binding_id = dataset_collection_binding.id
  425. dataset.collection_binding_id = dataset_collection_binding_id
  426. dataset.embedding_model = knowledge_configuration.embedding_model
  427. dataset.embedding_model_provider = knowledge_configuration.embedding_model_provider
  428. elif knowledge_configuration.indexing_technique == "economy":
  429. dataset.keyword_number = knowledge_configuration.keyword_number
  430. dataset.pipeline_id = pipeline.id
  431. self._session.add(dataset)
  432. self._session.commit()
  433. dataset_id = dataset.id
  434. if not dataset_id:
  435. raise ValueError("DSL is not valid, please check the Knowledge Index node.")
  436. # Delete import info from Redis
  437. redis_client.delete(redis_key)
  438. return RagPipelineImportInfo(
  439. id=import_id,
  440. status=ImportStatus.COMPLETED,
  441. pipeline_id=pipeline.id,
  442. dataset_id=dataset_id,
  443. current_dsl_version=CURRENT_DSL_VERSION,
  444. imported_dsl_version=data.get("version", "0.1.0"),
  445. )
  446. except Exception as e:
  447. logger.exception("Error confirming import")
  448. return RagPipelineImportInfo(
  449. id=import_id,
  450. status=ImportStatus.FAILED,
  451. error=str(e),
  452. )
  453. def check_dependencies(
  454. self,
  455. *,
  456. pipeline: Pipeline,
  457. ) -> CheckDependenciesResult:
  458. """Check dependencies"""
  459. # Get dependencies from Redis
  460. redis_key = f"{CHECK_DEPENDENCIES_REDIS_KEY_PREFIX}{pipeline.id}"
  461. dependencies = redis_client.get(redis_key)
  462. if not dependencies:
  463. return CheckDependenciesResult()
  464. # Extract dependencies
  465. dependencies = CheckDependenciesPendingData.model_validate_json(dependencies)
  466. # Get leaked dependencies
  467. leaked_dependencies = DependenciesAnalysisService.get_leaked_dependencies(
  468. tenant_id=pipeline.tenant_id, dependencies=dependencies.dependencies
  469. )
  470. return CheckDependenciesResult(
  471. leaked_dependencies=leaked_dependencies,
  472. )
  473. def _create_or_update_pipeline(
  474. self,
  475. *,
  476. pipeline: Pipeline | None,
  477. data: dict,
  478. account: Account,
  479. dependencies: list[PluginDependency] | None = None,
  480. ) -> Pipeline:
  481. """Create a new app or update an existing one."""
  482. if not account.current_tenant_id:
  483. raise ValueError("Tenant id is required")
  484. pipeline_data = data.get("rag_pipeline", {})
  485. # Initialize pipeline based on mode
  486. workflow_data = data.get("workflow")
  487. if not workflow_data or not isinstance(workflow_data, dict):
  488. raise ValueError("Missing workflow data for rag pipeline")
  489. environment_variables_list = workflow_data.get("environment_variables", [])
  490. environment_variables = [
  491. variable_factory.build_environment_variable_from_mapping(obj) for obj in environment_variables_list
  492. ]
  493. conversation_variables_list = workflow_data.get("conversation_variables", [])
  494. conversation_variables = [
  495. variable_factory.build_conversation_variable_from_mapping(obj) for obj in conversation_variables_list
  496. ]
  497. rag_pipeline_variables_list = workflow_data.get("rag_pipeline_variables", [])
  498. graph = workflow_data.get("graph", {})
  499. for node in graph.get("nodes", []):
  500. if node.get("data", {}).get("type", "") == NodeType.KNOWLEDGE_RETRIEVAL.value:
  501. dataset_ids = node["data"].get("dataset_ids", [])
  502. node["data"]["dataset_ids"] = [
  503. decrypted_id
  504. for dataset_id in dataset_ids
  505. if (
  506. decrypted_id := self.decrypt_dataset_id(
  507. encrypted_data=dataset_id,
  508. tenant_id=account.current_tenant_id,
  509. )
  510. )
  511. ]
  512. if pipeline:
  513. # Update existing pipeline
  514. pipeline.name = pipeline_data.get("name", pipeline.name)
  515. pipeline.description = pipeline_data.get("description", pipeline.description)
  516. pipeline.updated_by = account.id
  517. else:
  518. if account.current_tenant_id is None:
  519. raise ValueError("Current tenant is not set")
  520. # Create new app
  521. pipeline = Pipeline()
  522. pipeline.id = str(uuid4())
  523. pipeline.tenant_id = account.current_tenant_id
  524. pipeline.name = pipeline_data.get("name", "")
  525. pipeline.description = pipeline_data.get("description", "")
  526. pipeline.created_by = account.id
  527. pipeline.updated_by = account.id
  528. self._session.add(pipeline)
  529. self._session.commit()
  530. # save dependencies
  531. if dependencies:
  532. redis_client.setex(
  533. f"{CHECK_DEPENDENCIES_REDIS_KEY_PREFIX}{pipeline.id}",
  534. IMPORT_INFO_REDIS_EXPIRY,
  535. CheckDependenciesPendingData(pipeline_id=pipeline.id, dependencies=dependencies).model_dump_json(),
  536. )
  537. workflow = (
  538. self._session.query(Workflow)
  539. .where(
  540. Workflow.tenant_id == pipeline.tenant_id,
  541. Workflow.app_id == pipeline.id,
  542. Workflow.version == "draft",
  543. )
  544. .first()
  545. )
  546. # create draft workflow if not found
  547. if not workflow:
  548. workflow = Workflow(
  549. tenant_id=pipeline.tenant_id,
  550. app_id=pipeline.id,
  551. features="{}",
  552. type=WorkflowType.RAG_PIPELINE.value,
  553. version="draft",
  554. graph=json.dumps(graph),
  555. created_by=account.id,
  556. environment_variables=environment_variables,
  557. conversation_variables=conversation_variables,
  558. rag_pipeline_variables=rag_pipeline_variables_list,
  559. )
  560. self._session.add(workflow)
  561. self._session.flush()
  562. pipeline.workflow_id = workflow.id
  563. else:
  564. workflow.graph = json.dumps(graph)
  565. workflow.updated_by = account.id
  566. workflow.updated_at = datetime.now(UTC).replace(tzinfo=None)
  567. workflow.environment_variables = environment_variables
  568. workflow.conversation_variables = conversation_variables
  569. workflow.rag_pipeline_variables = rag_pipeline_variables_list
  570. # commit db session changes
  571. self._session.commit()
  572. return pipeline
  573. def export_rag_pipeline_dsl(self, pipeline: Pipeline, include_secret: bool = False) -> str:
  574. """
  575. Export pipeline
  576. :param pipeline: Pipeline instance
  577. :param include_secret: Whether include secret variable
  578. :return:
  579. """
  580. dataset = pipeline.retrieve_dataset(session=self._session)
  581. if not dataset:
  582. raise ValueError("Missing dataset for rag pipeline")
  583. icon_info = dataset.icon_info
  584. export_data = {
  585. "version": CURRENT_DSL_VERSION,
  586. "kind": "rag_pipeline",
  587. "rag_pipeline": {
  588. "name": dataset.name,
  589. "icon": icon_info.get("icon", "📙") if icon_info else "📙",
  590. "icon_type": icon_info.get("icon_type", "emoji") if icon_info else "emoji",
  591. "icon_background": icon_info.get("icon_background", "#FFEAD5") if icon_info else "#FFEAD5",
  592. "icon_url": icon_info.get("icon_url") if icon_info else None,
  593. "description": pipeline.description,
  594. },
  595. }
  596. self._append_workflow_export_data(export_data=export_data, pipeline=pipeline, include_secret=include_secret)
  597. return yaml.dump(export_data, allow_unicode=True) # type: ignore
  598. def _append_workflow_export_data(self, *, export_data: dict, pipeline: Pipeline, include_secret: bool) -> None:
  599. """
  600. Append workflow export data
  601. :param export_data: export data
  602. :param pipeline: Pipeline instance
  603. """
  604. workflow = (
  605. self._session.query(Workflow)
  606. .where(
  607. Workflow.tenant_id == pipeline.tenant_id,
  608. Workflow.app_id == pipeline.id,
  609. Workflow.version == "draft",
  610. )
  611. .first()
  612. )
  613. if not workflow:
  614. raise ValueError("Missing draft workflow configuration, please check.")
  615. workflow_dict = workflow.to_dict(include_secret=include_secret)
  616. for node in workflow_dict.get("graph", {}).get("nodes", []):
  617. node_data = node.get("data", {})
  618. if not node_data:
  619. continue
  620. data_type = node_data.get("type", "")
  621. if data_type == NodeType.KNOWLEDGE_RETRIEVAL.value:
  622. dataset_ids = node_data.get("dataset_ids", [])
  623. node["data"]["dataset_ids"] = [
  624. self.encrypt_dataset_id(dataset_id=dataset_id, tenant_id=pipeline.tenant_id)
  625. for dataset_id in dataset_ids
  626. ]
  627. # filter credential id from tool node
  628. if not include_secret and data_type == NodeType.TOOL.value:
  629. node_data.pop("credential_id", None)
  630. # filter credential id from agent node
  631. if not include_secret and data_type == NodeType.AGENT.value:
  632. for tool in node_data.get("agent_parameters", {}).get("tools", {}).get("value", []):
  633. tool.pop("credential_id", None)
  634. export_data["workflow"] = workflow_dict
  635. dependencies = self._extract_dependencies_from_workflow(workflow)
  636. export_data["dependencies"] = [
  637. jsonable_encoder(d.model_dump())
  638. for d in DependenciesAnalysisService.generate_dependencies(
  639. tenant_id=pipeline.tenant_id, dependencies=dependencies
  640. )
  641. ]
  642. def _extract_dependencies_from_workflow(self, workflow: Workflow) -> list[str]:
  643. """
  644. Extract dependencies from workflow
  645. :param workflow: Workflow instance
  646. :return: dependencies list format like ["langgenius/google"]
  647. """
  648. graph = workflow.graph_dict
  649. dependencies = self._extract_dependencies_from_workflow_graph(graph)
  650. return dependencies
  651. def _extract_dependencies_from_workflow_graph(self, graph: Mapping) -> list[str]:
  652. """
  653. Extract dependencies from workflow graph
  654. :param graph: Workflow graph
  655. :return: dependencies list format like ["langgenius/google"]
  656. """
  657. dependencies = []
  658. for node in graph.get("nodes", []):
  659. try:
  660. typ = node.get("data", {}).get("type")
  661. match typ:
  662. case NodeType.TOOL.value:
  663. tool_entity = ToolNodeData(**node["data"])
  664. dependencies.append(
  665. DependenciesAnalysisService.analyze_tool_dependency(tool_entity.provider_id),
  666. )
  667. case NodeType.DATASOURCE.value:
  668. datasource_entity = DatasourceNodeData(**node["data"])
  669. if datasource_entity.provider_type != "local_file":
  670. dependencies.append(datasource_entity.plugin_id)
  671. case NodeType.LLM.value:
  672. llm_entity = LLMNodeData(**node["data"])
  673. dependencies.append(
  674. DependenciesAnalysisService.analyze_model_provider_dependency(llm_entity.model.provider),
  675. )
  676. case NodeType.QUESTION_CLASSIFIER.value:
  677. question_classifier_entity = QuestionClassifierNodeData(**node["data"])
  678. dependencies.append(
  679. DependenciesAnalysisService.analyze_model_provider_dependency(
  680. question_classifier_entity.model.provider
  681. ),
  682. )
  683. case NodeType.PARAMETER_EXTRACTOR.value:
  684. parameter_extractor_entity = ParameterExtractorNodeData(**node["data"])
  685. dependencies.append(
  686. DependenciesAnalysisService.analyze_model_provider_dependency(
  687. parameter_extractor_entity.model.provider
  688. ),
  689. )
  690. case NodeType.KNOWLEDGE_INDEX.value:
  691. knowledge_index_entity = KnowledgeConfiguration(**node["data"])
  692. if knowledge_index_entity.indexing_technique == "high_quality":
  693. if knowledge_index_entity.embedding_model_provider:
  694. dependencies.append(
  695. DependenciesAnalysisService.analyze_model_provider_dependency(
  696. knowledge_index_entity.embedding_model_provider
  697. ),
  698. )
  699. if knowledge_index_entity.retrieval_model.reranking_mode == "reranking_model":
  700. if knowledge_index_entity.retrieval_model.reranking_enable:
  701. if (
  702. knowledge_index_entity.retrieval_model.reranking_model
  703. and knowledge_index_entity.retrieval_model.reranking_mode == "reranking_model"
  704. ):
  705. if knowledge_index_entity.retrieval_model.reranking_model.reranking_provider_name:
  706. dependencies.append(
  707. DependenciesAnalysisService.analyze_model_provider_dependency(
  708. knowledge_index_entity.retrieval_model.reranking_model.reranking_provider_name
  709. ),
  710. )
  711. case NodeType.KNOWLEDGE_RETRIEVAL.value:
  712. knowledge_retrieval_entity = KnowledgeRetrievalNodeData(**node["data"])
  713. if knowledge_retrieval_entity.retrieval_mode == "multiple":
  714. if knowledge_retrieval_entity.multiple_retrieval_config:
  715. if (
  716. knowledge_retrieval_entity.multiple_retrieval_config.reranking_mode
  717. == "reranking_model"
  718. ):
  719. if knowledge_retrieval_entity.multiple_retrieval_config.reranking_model:
  720. dependencies.append(
  721. DependenciesAnalysisService.analyze_model_provider_dependency(
  722. knowledge_retrieval_entity.multiple_retrieval_config.reranking_model.provider
  723. ),
  724. )
  725. elif (
  726. knowledge_retrieval_entity.multiple_retrieval_config.reranking_mode
  727. == "weighted_score"
  728. ):
  729. if knowledge_retrieval_entity.multiple_retrieval_config.weights:
  730. vector_setting = (
  731. knowledge_retrieval_entity.multiple_retrieval_config.weights.vector_setting
  732. )
  733. dependencies.append(
  734. DependenciesAnalysisService.analyze_model_provider_dependency(
  735. vector_setting.embedding_provider_name
  736. ),
  737. )
  738. elif knowledge_retrieval_entity.retrieval_mode == "single":
  739. model_config = knowledge_retrieval_entity.single_retrieval_config
  740. if model_config:
  741. dependencies.append(
  742. DependenciesAnalysisService.analyze_model_provider_dependency(
  743. model_config.model.provider
  744. ),
  745. )
  746. case _:
  747. # TODO: Handle default case or unknown node types
  748. pass
  749. except Exception as e:
  750. logger.exception("Error extracting node dependency", exc_info=e)
  751. return dependencies
  752. @classmethod
  753. def _extract_dependencies_from_model_config(cls, model_config: Mapping) -> list[str]:
  754. """
  755. Extract dependencies from model config
  756. :param model_config: model config dict
  757. :return: dependencies list format like ["langgenius/google"]
  758. """
  759. dependencies = []
  760. try:
  761. # completion model
  762. model_dict = model_config.get("model", {})
  763. if model_dict:
  764. dependencies.append(
  765. DependenciesAnalysisService.analyze_model_provider_dependency(model_dict.get("provider", ""))
  766. )
  767. # reranking model
  768. dataset_configs = model_config.get("dataset_configs", {})
  769. if dataset_configs:
  770. for dataset_config in dataset_configs.get("datasets", {}).get("datasets", []):
  771. if dataset_config.get("reranking_model"):
  772. dependencies.append(
  773. DependenciesAnalysisService.analyze_model_provider_dependency(
  774. dataset_config.get("reranking_model", {})
  775. .get("reranking_provider_name", {})
  776. .get("provider")
  777. )
  778. )
  779. # tools
  780. agent_configs = model_config.get("agent_mode", {})
  781. if agent_configs:
  782. for agent_config in agent_configs.get("tools", []):
  783. dependencies.append(
  784. DependenciesAnalysisService.analyze_tool_dependency(agent_config.get("provider_id"))
  785. )
  786. except Exception as e:
  787. logger.exception("Error extracting model config dependency", exc_info=e)
  788. return dependencies
  789. @classmethod
  790. def get_leaked_dependencies(cls, tenant_id: str, dsl_dependencies: list[dict]) -> list[PluginDependency]:
  791. """
  792. Returns the leaked dependencies in current workspace
  793. """
  794. dependencies = [PluginDependency(**dep) for dep in dsl_dependencies]
  795. if not dependencies:
  796. return []
  797. return DependenciesAnalysisService.get_leaked_dependencies(tenant_id=tenant_id, dependencies=dependencies)
  798. def _generate_aes_key(self, tenant_id: str) -> bytes:
  799. """Generate AES key based on tenant_id"""
  800. return hashlib.sha256(tenant_id.encode()).digest()
  801. def encrypt_dataset_id(self, dataset_id: str, tenant_id: str) -> str:
  802. """Encrypt dataset_id using AES-CBC mode"""
  803. key = self._generate_aes_key(tenant_id)
  804. iv = key[:16]
  805. cipher = AES.new(key, AES.MODE_CBC, iv)
  806. ct_bytes = cipher.encrypt(pad(dataset_id.encode(), AES.block_size))
  807. return base64.b64encode(ct_bytes).decode()
  808. def decrypt_dataset_id(self, encrypted_data: str, tenant_id: str) -> str | None:
  809. """AES decryption"""
  810. try:
  811. key = self._generate_aes_key(tenant_id)
  812. iv = key[:16]
  813. cipher = AES.new(key, AES.MODE_CBC, iv)
  814. pt = unpad(cipher.decrypt(base64.b64decode(encrypted_data)), AES.block_size)
  815. return pt.decode()
  816. except Exception:
  817. return None
  818. def create_rag_pipeline_dataset(
  819. self,
  820. tenant_id: str,
  821. rag_pipeline_dataset_create_entity: RagPipelineDatasetCreateEntity,
  822. ):
  823. if rag_pipeline_dataset_create_entity.name:
  824. # check if dataset name already exists
  825. if (
  826. self._session.query(Dataset)
  827. .filter_by(name=rag_pipeline_dataset_create_entity.name, tenant_id=tenant_id)
  828. .first()
  829. ):
  830. raise ValueError(f"Dataset with name {rag_pipeline_dataset_create_entity.name} already exists.")
  831. else:
  832. # generate a random name as Untitled 1 2 3 ...
  833. datasets = self._session.query(Dataset).filter_by(tenant_id=tenant_id).all()
  834. names = [dataset.name for dataset in datasets]
  835. rag_pipeline_dataset_create_entity.name = generate_incremental_name(
  836. names,
  837. "Untitled",
  838. )
  839. account = cast(Account, current_user)
  840. rag_pipeline_import_info: RagPipelineImportInfo = self.import_rag_pipeline(
  841. account=account,
  842. import_mode=ImportMode.YAML_CONTENT.value,
  843. yaml_content=rag_pipeline_dataset_create_entity.yaml_content,
  844. dataset=None,
  845. dataset_name=rag_pipeline_dataset_create_entity.name,
  846. icon_info=rag_pipeline_dataset_create_entity.icon_info,
  847. )
  848. return {
  849. "id": rag_pipeline_import_info.id,
  850. "dataset_id": rag_pipeline_import_info.dataset_id,
  851. "pipeline_id": rag_pipeline_import_info.pipeline_id,
  852. "status": rag_pipeline_import_info.status,
  853. "imported_dsl_version": rag_pipeline_import_info.imported_dsl_version,
  854. "current_dsl_version": rag_pipeline_import_info.current_dsl_version,
  855. "error": rag_pipeline_import_info.error,
  856. }