node.py 42 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049
  1. import json
  2. import logging
  3. from collections.abc import Generator, Mapping, Sequence
  4. from typing import TYPE_CHECKING, Any, Optional, cast
  5. from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
  6. from core.entities.model_entities import ModelStatus
  7. from core.entities.provider_entities import QuotaUnit
  8. from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
  9. from core.file import FileType, file_manager
  10. from core.helper.code_executor import CodeExecutor, CodeLanguage
  11. from core.memory.token_buffer_memory import TokenBufferMemory
  12. from core.model_manager import ModelInstance, ModelManager
  13. from core.model_runtime.entities import (
  14. ImagePromptMessageContent,
  15. PromptMessage,
  16. PromptMessageContentType,
  17. TextPromptMessageContent,
  18. )
  19. from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage
  20. from core.model_runtime.entities.message_entities import (
  21. AssistantPromptMessage,
  22. PromptMessageContent,
  23. PromptMessageRole,
  24. SystemPromptMessage,
  25. UserPromptMessage,
  26. )
  27. from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey, ModelType
  28. from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
  29. from core.model_runtime.utils.encoders import jsonable_encoder
  30. from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptTemplate, MemoryConfig
  31. from core.prompt.utils.prompt_message_util import PromptMessageUtil
  32. from core.variables import (
  33. ArrayAnySegment,
  34. ArrayFileSegment,
  35. ArraySegment,
  36. FileSegment,
  37. NoneSegment,
  38. ObjectSegment,
  39. StringSegment,
  40. )
  41. from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID
  42. from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult
  43. from core.workflow.entities.variable_entities import VariableSelector
  44. from core.workflow.entities.variable_pool import VariablePool
  45. from core.workflow.enums import SystemVariableKey
  46. from core.workflow.graph_engine.entities.event import InNodeEvent
  47. from core.workflow.nodes.base import BaseNode
  48. from core.workflow.nodes.enums import NodeType
  49. from core.workflow.nodes.event import (
  50. ModelInvokeCompletedEvent,
  51. NodeEvent,
  52. RunCompletedEvent,
  53. RunRetrieverResourceEvent,
  54. RunStreamChunkEvent,
  55. )
  56. from core.workflow.utils.variable_template_parser import VariableTemplateParser
  57. from extensions.ext_database import db
  58. from models.model import Conversation
  59. from models.provider import Provider, ProviderType
  60. from models.workflow import WorkflowNodeExecutionStatus
  61. from .entities import (
  62. LLMNodeChatModelMessage,
  63. LLMNodeCompletionModelPromptTemplate,
  64. LLMNodeData,
  65. ModelConfig,
  66. )
  67. from .exc import (
  68. InvalidContextStructureError,
  69. InvalidVariableTypeError,
  70. LLMModeRequiredError,
  71. LLMNodeError,
  72. MemoryRolePrefixRequiredError,
  73. ModelNotExistError,
  74. NoPromptFoundError,
  75. TemplateTypeNotSupportError,
  76. VariableNotFoundError,
  77. )
  78. if TYPE_CHECKING:
  79. from core.file.models import File
  80. logger = logging.getLogger(__name__)
  81. class LLMNode(BaseNode[LLMNodeData]):
  82. _node_data_cls = LLMNodeData
  83. _node_type = NodeType.LLM
  84. def _run(self) -> NodeRunResult | Generator[NodeEvent | InNodeEvent, None, None]:
  85. node_inputs = None
  86. process_data = None
  87. try:
  88. # init messages template
  89. self.node_data.prompt_template = self._transform_chat_messages(self.node_data.prompt_template)
  90. # fetch variables and fetch values from variable pool
  91. inputs = self._fetch_inputs(node_data=self.node_data)
  92. # fetch jinja2 inputs
  93. jinja_inputs = self._fetch_jinja_inputs(node_data=self.node_data)
  94. # merge inputs
  95. inputs.update(jinja_inputs)
  96. node_inputs = {}
  97. # fetch files
  98. files = (
  99. self._fetch_files(selector=self.node_data.vision.configs.variable_selector)
  100. if self.node_data.vision.enabled
  101. else []
  102. )
  103. if files:
  104. node_inputs["#files#"] = [file.to_dict() for file in files]
  105. # fetch context value
  106. generator = self._fetch_context(node_data=self.node_data)
  107. context = None
  108. for event in generator:
  109. if isinstance(event, RunRetrieverResourceEvent):
  110. context = event.context
  111. yield event
  112. if context:
  113. node_inputs["#context#"] = context
  114. # fetch model config
  115. model_instance, model_config = self._fetch_model_config(self.node_data.model)
  116. # fetch memory
  117. memory = self._fetch_memory(node_data_memory=self.node_data.memory, model_instance=model_instance)
  118. query = None
  119. if self.node_data.memory:
  120. query = self.node_data.memory.query_prompt_template
  121. if not query and (
  122. query_variable := self.graph_runtime_state.variable_pool.get(
  123. (SYSTEM_VARIABLE_NODE_ID, SystemVariableKey.QUERY)
  124. )
  125. ):
  126. query = query_variable.text
  127. prompt_messages, stop = self._fetch_prompt_messages(
  128. user_query=query,
  129. user_files=files,
  130. context=context,
  131. memory=memory,
  132. model_config=model_config,
  133. prompt_template=self.node_data.prompt_template,
  134. memory_config=self.node_data.memory,
  135. vision_enabled=self.node_data.vision.enabled,
  136. vision_detail=self.node_data.vision.configs.detail,
  137. variable_pool=self.graph_runtime_state.variable_pool,
  138. jinja2_variables=self.node_data.prompt_config.jinja2_variables,
  139. )
  140. process_data = {
  141. "model_mode": model_config.mode,
  142. "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
  143. model_mode=model_config.mode, prompt_messages=prompt_messages
  144. ),
  145. "model_provider": model_config.provider,
  146. "model_name": model_config.model,
  147. }
  148. # handle invoke result
  149. generator = self._invoke_llm(
  150. node_data_model=self.node_data.model,
  151. model_instance=model_instance,
  152. prompt_messages=prompt_messages,
  153. stop=stop,
  154. )
  155. result_text = ""
  156. usage = LLMUsage.empty_usage()
  157. finish_reason = None
  158. for event in generator:
  159. if isinstance(event, RunStreamChunkEvent):
  160. yield event
  161. elif isinstance(event, ModelInvokeCompletedEvent):
  162. result_text = event.text
  163. usage = event.usage
  164. finish_reason = event.finish_reason
  165. break
  166. except LLMNodeError as e:
  167. yield RunCompletedEvent(
  168. run_result=NodeRunResult(
  169. status=WorkflowNodeExecutionStatus.FAILED,
  170. error=str(e),
  171. inputs=node_inputs,
  172. process_data=process_data,
  173. error_type=type(e).__name__,
  174. )
  175. )
  176. return
  177. except Exception as e:
  178. yield RunCompletedEvent(
  179. run_result=NodeRunResult(
  180. status=WorkflowNodeExecutionStatus.FAILED,
  181. error=str(e),
  182. inputs=node_inputs,
  183. process_data=process_data,
  184. )
  185. )
  186. return
  187. outputs = {"text": result_text, "usage": jsonable_encoder(usage), "finish_reason": finish_reason}
  188. yield RunCompletedEvent(
  189. run_result=NodeRunResult(
  190. status=WorkflowNodeExecutionStatus.SUCCEEDED,
  191. inputs=node_inputs,
  192. process_data=process_data,
  193. outputs=outputs,
  194. metadata={
  195. NodeRunMetadataKey.TOTAL_TOKENS: usage.total_tokens,
  196. NodeRunMetadataKey.TOTAL_PRICE: usage.total_price,
  197. NodeRunMetadataKey.CURRENCY: usage.currency,
  198. },
  199. llm_usage=usage,
  200. )
  201. )
  202. def _invoke_llm(
  203. self,
  204. node_data_model: ModelConfig,
  205. model_instance: ModelInstance,
  206. prompt_messages: Sequence[PromptMessage],
  207. stop: Optional[Sequence[str]] = None,
  208. ) -> Generator[NodeEvent, None, None]:
  209. db.session.close()
  210. invoke_result = model_instance.invoke_llm(
  211. prompt_messages=prompt_messages,
  212. model_parameters=node_data_model.completion_params,
  213. stop=stop,
  214. stream=True,
  215. user=self.user_id,
  216. )
  217. # handle invoke result
  218. generator = self._handle_invoke_result(invoke_result=invoke_result)
  219. usage = LLMUsage.empty_usage()
  220. for event in generator:
  221. yield event
  222. if isinstance(event, ModelInvokeCompletedEvent):
  223. usage = event.usage
  224. # deduct quota
  225. self.deduct_llm_quota(tenant_id=self.tenant_id, model_instance=model_instance, usage=usage)
  226. def _handle_invoke_result(self, invoke_result: LLMResult | Generator) -> Generator[NodeEvent, None, None]:
  227. if isinstance(invoke_result, LLMResult):
  228. return
  229. model = None
  230. prompt_messages: list[PromptMessage] = []
  231. full_text = ""
  232. usage = None
  233. finish_reason = None
  234. for result in invoke_result:
  235. text = result.delta.message.content
  236. full_text += text
  237. yield RunStreamChunkEvent(chunk_content=text, from_variable_selector=[self.node_id, "text"])
  238. if not model:
  239. model = result.model
  240. if not prompt_messages:
  241. prompt_messages = result.prompt_messages
  242. if not usage and result.delta.usage:
  243. usage = result.delta.usage
  244. if not finish_reason and result.delta.finish_reason:
  245. finish_reason = result.delta.finish_reason
  246. if not usage:
  247. usage = LLMUsage.empty_usage()
  248. yield ModelInvokeCompletedEvent(text=full_text, usage=usage, finish_reason=finish_reason)
  249. def _transform_chat_messages(
  250. self, messages: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate, /
  251. ) -> Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate:
  252. if isinstance(messages, LLMNodeCompletionModelPromptTemplate):
  253. if messages.edition_type == "jinja2" and messages.jinja2_text:
  254. messages.text = messages.jinja2_text
  255. return messages
  256. for message in messages:
  257. if message.edition_type == "jinja2" and message.jinja2_text:
  258. message.text = message.jinja2_text
  259. return messages
  260. def _fetch_jinja_inputs(self, node_data: LLMNodeData) -> dict[str, str]:
  261. variables = {}
  262. if not node_data.prompt_config:
  263. return variables
  264. for variable_selector in node_data.prompt_config.jinja2_variables or []:
  265. variable_name = variable_selector.variable
  266. variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
  267. if variable is None:
  268. raise VariableNotFoundError(f"Variable {variable_selector.variable} not found")
  269. def parse_dict(input_dict: Mapping[str, Any]) -> str:
  270. """
  271. Parse dict into string
  272. """
  273. # check if it's a context structure
  274. if "metadata" in input_dict and "_source" in input_dict["metadata"] and "content" in input_dict:
  275. return input_dict["content"]
  276. # else, parse the dict
  277. try:
  278. return json.dumps(input_dict, ensure_ascii=False)
  279. except Exception:
  280. return str(input_dict)
  281. if isinstance(variable, ArraySegment):
  282. result = ""
  283. for item in variable.value:
  284. if isinstance(item, dict):
  285. result += parse_dict(item)
  286. else:
  287. result += str(item)
  288. result += "\n"
  289. value = result.strip()
  290. elif isinstance(variable, ObjectSegment):
  291. value = parse_dict(variable.value)
  292. else:
  293. value = variable.text
  294. variables[variable_name] = value
  295. return variables
  296. def _fetch_inputs(self, node_data: LLMNodeData) -> dict[str, Any]:
  297. inputs = {}
  298. prompt_template = node_data.prompt_template
  299. variable_selectors = []
  300. if isinstance(prompt_template, list):
  301. for prompt in prompt_template:
  302. variable_template_parser = VariableTemplateParser(template=prompt.text)
  303. variable_selectors.extend(variable_template_parser.extract_variable_selectors())
  304. elif isinstance(prompt_template, CompletionModelPromptTemplate):
  305. variable_template_parser = VariableTemplateParser(template=prompt_template.text)
  306. variable_selectors = variable_template_parser.extract_variable_selectors()
  307. for variable_selector in variable_selectors:
  308. variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
  309. if variable is None:
  310. raise VariableNotFoundError(f"Variable {variable_selector.variable} not found")
  311. if isinstance(variable, NoneSegment):
  312. inputs[variable_selector.variable] = ""
  313. inputs[variable_selector.variable] = variable.to_object()
  314. memory = node_data.memory
  315. if memory and memory.query_prompt_template:
  316. query_variable_selectors = VariableTemplateParser(
  317. template=memory.query_prompt_template
  318. ).extract_variable_selectors()
  319. for variable_selector in query_variable_selectors:
  320. variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
  321. if variable is None:
  322. raise VariableNotFoundError(f"Variable {variable_selector.variable} not found")
  323. if isinstance(variable, NoneSegment):
  324. continue
  325. inputs[variable_selector.variable] = variable.to_object()
  326. return inputs
  327. def _fetch_files(self, *, selector: Sequence[str]) -> Sequence["File"]:
  328. variable = self.graph_runtime_state.variable_pool.get(selector)
  329. if variable is None:
  330. return []
  331. elif isinstance(variable, FileSegment):
  332. return [variable.value]
  333. elif isinstance(variable, ArrayFileSegment):
  334. return variable.value
  335. elif isinstance(variable, NoneSegment | ArrayAnySegment):
  336. return []
  337. raise InvalidVariableTypeError(f"Invalid variable type: {type(variable)}")
  338. def _fetch_context(self, node_data: LLMNodeData):
  339. if not node_data.context.enabled:
  340. return
  341. if not node_data.context.variable_selector:
  342. return
  343. context_value_variable = self.graph_runtime_state.variable_pool.get(node_data.context.variable_selector)
  344. if context_value_variable:
  345. if isinstance(context_value_variable, StringSegment):
  346. yield RunRetrieverResourceEvent(retriever_resources=[], context=context_value_variable.value)
  347. elif isinstance(context_value_variable, ArraySegment):
  348. context_str = ""
  349. original_retriever_resource = []
  350. for item in context_value_variable.value:
  351. if isinstance(item, str):
  352. context_str += item + "\n"
  353. else:
  354. if "content" not in item:
  355. raise InvalidContextStructureError(f"Invalid context structure: {item}")
  356. context_str += item["content"] + "\n"
  357. retriever_resource = self._convert_to_original_retriever_resource(item)
  358. if retriever_resource:
  359. original_retriever_resource.append(retriever_resource)
  360. yield RunRetrieverResourceEvent(
  361. retriever_resources=original_retriever_resource, context=context_str.strip()
  362. )
  363. def _convert_to_original_retriever_resource(self, context_dict: dict) -> Optional[dict]:
  364. if (
  365. "metadata" in context_dict
  366. and "_source" in context_dict["metadata"]
  367. and context_dict["metadata"]["_source"] == "knowledge"
  368. ):
  369. metadata = context_dict.get("metadata", {})
  370. source = {
  371. "position": metadata.get("position"),
  372. "dataset_id": metadata.get("dataset_id"),
  373. "dataset_name": metadata.get("dataset_name"),
  374. "document_id": metadata.get("document_id"),
  375. "document_name": metadata.get("document_name"),
  376. "data_source_type": metadata.get("document_data_source_type"),
  377. "segment_id": metadata.get("segment_id"),
  378. "retriever_from": metadata.get("retriever_from"),
  379. "score": metadata.get("score"),
  380. "hit_count": metadata.get("segment_hit_count"),
  381. "word_count": metadata.get("segment_word_count"),
  382. "segment_position": metadata.get("segment_position"),
  383. "index_node_hash": metadata.get("segment_index_node_hash"),
  384. "content": context_dict.get("content"),
  385. "page": metadata.get("page"),
  386. }
  387. return source
  388. return None
  389. def _fetch_model_config(
  390. self, node_data_model: ModelConfig
  391. ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]:
  392. model_name = node_data_model.name
  393. provider_name = node_data_model.provider
  394. model_manager = ModelManager()
  395. model_instance = model_manager.get_model_instance(
  396. tenant_id=self.tenant_id, model_type=ModelType.LLM, provider=provider_name, model=model_name
  397. )
  398. provider_model_bundle = model_instance.provider_model_bundle
  399. model_type_instance = model_instance.model_type_instance
  400. model_type_instance = cast(LargeLanguageModel, model_type_instance)
  401. model_credentials = model_instance.credentials
  402. # check model
  403. provider_model = provider_model_bundle.configuration.get_provider_model(
  404. model=model_name, model_type=ModelType.LLM
  405. )
  406. if provider_model is None:
  407. raise ModelNotExistError(f"Model {model_name} not exist.")
  408. if provider_model.status == ModelStatus.NO_CONFIGURE:
  409. raise ProviderTokenNotInitError(f"Model {model_name} credentials is not initialized.")
  410. elif provider_model.status == ModelStatus.NO_PERMISSION:
  411. raise ModelCurrentlyNotSupportError(f"Dify Hosted OpenAI {model_name} currently not support.")
  412. elif provider_model.status == ModelStatus.QUOTA_EXCEEDED:
  413. raise QuotaExceededError(f"Model provider {provider_name} quota exceeded.")
  414. # model config
  415. completion_params = node_data_model.completion_params
  416. stop = []
  417. if "stop" in completion_params:
  418. stop = completion_params["stop"]
  419. del completion_params["stop"]
  420. # get model mode
  421. model_mode = node_data_model.mode
  422. if not model_mode:
  423. raise LLMModeRequiredError("LLM mode is required.")
  424. model_schema = model_type_instance.get_model_schema(model_name, model_credentials)
  425. if not model_schema:
  426. raise ModelNotExistError(f"Model {model_name} not exist.")
  427. return model_instance, ModelConfigWithCredentialsEntity(
  428. provider=provider_name,
  429. model=model_name,
  430. model_schema=model_schema,
  431. mode=model_mode,
  432. provider_model_bundle=provider_model_bundle,
  433. credentials=model_credentials,
  434. parameters=completion_params,
  435. stop=stop,
  436. )
  437. def _fetch_memory(
  438. self, node_data_memory: Optional[MemoryConfig], model_instance: ModelInstance
  439. ) -> Optional[TokenBufferMemory]:
  440. if not node_data_memory:
  441. return None
  442. # get conversation id
  443. conversation_id_variable = self.graph_runtime_state.variable_pool.get(
  444. ["sys", SystemVariableKey.CONVERSATION_ID.value]
  445. )
  446. if not isinstance(conversation_id_variable, StringSegment):
  447. return None
  448. conversation_id = conversation_id_variable.value
  449. # get conversation
  450. conversation = (
  451. db.session.query(Conversation)
  452. .filter(Conversation.app_id == self.app_id, Conversation.id == conversation_id)
  453. .first()
  454. )
  455. if not conversation:
  456. return None
  457. memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance)
  458. return memory
  459. def _fetch_prompt_messages(
  460. self,
  461. *,
  462. user_query: str | None = None,
  463. user_files: Sequence["File"],
  464. context: str | None = None,
  465. memory: TokenBufferMemory | None = None,
  466. model_config: ModelConfigWithCredentialsEntity,
  467. prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate,
  468. memory_config: MemoryConfig | None = None,
  469. vision_enabled: bool = False,
  470. vision_detail: ImagePromptMessageContent.DETAIL,
  471. variable_pool: VariablePool,
  472. jinja2_variables: Sequence[VariableSelector],
  473. ) -> tuple[Sequence[PromptMessage], Optional[Sequence[str]]]:
  474. prompt_messages = []
  475. if isinstance(prompt_template, list):
  476. # For chat model
  477. prompt_messages.extend(
  478. _handle_list_messages(
  479. messages=prompt_template,
  480. context=context,
  481. jinja2_variables=jinja2_variables,
  482. variable_pool=variable_pool,
  483. vision_detail_config=vision_detail,
  484. )
  485. )
  486. # Get memory messages for chat mode
  487. memory_messages = _handle_memory_chat_mode(
  488. memory=memory,
  489. memory_config=memory_config,
  490. model_config=model_config,
  491. )
  492. # Extend prompt_messages with memory messages
  493. prompt_messages.extend(memory_messages)
  494. # Add current query to the prompt messages
  495. if user_query:
  496. message = LLMNodeChatModelMessage(
  497. text=user_query,
  498. role=PromptMessageRole.USER,
  499. edition_type="basic",
  500. )
  501. prompt_messages.extend(
  502. _handle_list_messages(
  503. messages=[message],
  504. context="",
  505. jinja2_variables=[],
  506. variable_pool=variable_pool,
  507. vision_detail_config=vision_detail,
  508. )
  509. )
  510. elif isinstance(prompt_template, LLMNodeCompletionModelPromptTemplate):
  511. # For completion model
  512. prompt_messages.extend(
  513. _handle_completion_template(
  514. template=prompt_template,
  515. context=context,
  516. jinja2_variables=jinja2_variables,
  517. variable_pool=variable_pool,
  518. )
  519. )
  520. # Get memory text for completion model
  521. memory_text = _handle_memory_completion_mode(
  522. memory=memory,
  523. memory_config=memory_config,
  524. model_config=model_config,
  525. )
  526. # Insert histories into the prompt
  527. prompt_content = prompt_messages[0].content
  528. # For issue #11247 - Check if prompt content is a string or a list
  529. prompt_content_type = type(prompt_content)
  530. if prompt_content_type == str:
  531. if "#histories#" in prompt_content:
  532. prompt_content = prompt_content.replace("#histories#", memory_text)
  533. else:
  534. prompt_content = memory_text + "\n" + prompt_content
  535. prompt_messages[0].content = prompt_content
  536. elif prompt_content_type == list:
  537. for content_item in prompt_content:
  538. if content_item.type == PromptMessageContentType.TEXT:
  539. if "#histories#" in content_item.data:
  540. content_item.data = content_item.data.replace("#histories#", memory_text)
  541. else:
  542. content_item.data = memory_text + "\n" + content_item.data
  543. else:
  544. raise ValueError("Invalid prompt content type")
  545. # Add current query to the prompt message
  546. if user_query:
  547. if prompt_content_type == str:
  548. prompt_content = prompt_messages[0].content.replace("#sys.query#", user_query)
  549. prompt_messages[0].content = prompt_content
  550. elif prompt_content_type == list:
  551. for content_item in prompt_content:
  552. if content_item.type == PromptMessageContentType.TEXT:
  553. content_item.data = user_query + "\n" + content_item.data
  554. else:
  555. raise ValueError("Invalid prompt content type")
  556. else:
  557. raise TemplateTypeNotSupportError(type_name=str(type(prompt_template)))
  558. if vision_enabled and user_files:
  559. file_prompts = []
  560. for file in user_files:
  561. file_prompt = file_manager.to_prompt_message_content(file, image_detail_config=vision_detail)
  562. file_prompts.append(file_prompt)
  563. if (
  564. len(prompt_messages) > 0
  565. and isinstance(prompt_messages[-1], UserPromptMessage)
  566. and isinstance(prompt_messages[-1].content, list)
  567. ):
  568. prompt_messages[-1] = UserPromptMessage(content=prompt_messages[-1].content + file_prompts)
  569. else:
  570. prompt_messages.append(UserPromptMessage(content=file_prompts))
  571. # Filter prompt messages
  572. filtered_prompt_messages = []
  573. for prompt_message in prompt_messages:
  574. if isinstance(prompt_message.content, list):
  575. prompt_message_content = []
  576. for content_item in prompt_message.content:
  577. # Skip content if features are not defined
  578. if not model_config.model_schema.features:
  579. if content_item.type != PromptMessageContentType.TEXT:
  580. continue
  581. prompt_message_content.append(content_item)
  582. continue
  583. # Skip content if corresponding feature is not supported
  584. if (
  585. (
  586. content_item.type == PromptMessageContentType.IMAGE
  587. and ModelFeature.VISION not in model_config.model_schema.features
  588. )
  589. or (
  590. content_item.type == PromptMessageContentType.DOCUMENT
  591. and ModelFeature.DOCUMENT not in model_config.model_schema.features
  592. )
  593. or (
  594. content_item.type == PromptMessageContentType.VIDEO
  595. and ModelFeature.VIDEO not in model_config.model_schema.features
  596. )
  597. or (
  598. content_item.type == PromptMessageContentType.AUDIO
  599. and ModelFeature.AUDIO not in model_config.model_schema.features
  600. )
  601. ):
  602. continue
  603. prompt_message_content.append(content_item)
  604. if len(prompt_message_content) == 1 and prompt_message_content[0].type == PromptMessageContentType.TEXT:
  605. prompt_message.content = prompt_message_content[0].data
  606. else:
  607. prompt_message.content = prompt_message_content
  608. if prompt_message.is_empty():
  609. continue
  610. filtered_prompt_messages.append(prompt_message)
  611. if len(filtered_prompt_messages) == 0:
  612. raise NoPromptFoundError(
  613. "No prompt found in the LLM configuration. "
  614. "Please ensure a prompt is properly configured before proceeding."
  615. )
  616. stop = model_config.stop
  617. return filtered_prompt_messages, stop
  618. @classmethod
  619. def deduct_llm_quota(cls, tenant_id: str, model_instance: ModelInstance, usage: LLMUsage) -> None:
  620. provider_model_bundle = model_instance.provider_model_bundle
  621. provider_configuration = provider_model_bundle.configuration
  622. if provider_configuration.using_provider_type != ProviderType.SYSTEM:
  623. return
  624. system_configuration = provider_configuration.system_configuration
  625. quota_unit = None
  626. for quota_configuration in system_configuration.quota_configurations:
  627. if quota_configuration.quota_type == system_configuration.current_quota_type:
  628. quota_unit = quota_configuration.quota_unit
  629. if quota_configuration.quota_limit == -1:
  630. return
  631. break
  632. used_quota = None
  633. if quota_unit:
  634. if quota_unit == QuotaUnit.TOKENS:
  635. used_quota = usage.total_tokens
  636. elif quota_unit == QuotaUnit.CREDITS:
  637. used_quota = 1
  638. if "gpt-4" in model_instance.model:
  639. used_quota = 20
  640. else:
  641. used_quota = 1
  642. if used_quota is not None and system_configuration.current_quota_type is not None:
  643. db.session.query(Provider).filter(
  644. Provider.tenant_id == tenant_id,
  645. Provider.provider_name == model_instance.provider,
  646. Provider.provider_type == ProviderType.SYSTEM.value,
  647. Provider.quota_type == system_configuration.current_quota_type.value,
  648. Provider.quota_limit > Provider.quota_used,
  649. ).update({"quota_used": Provider.quota_used + used_quota})
  650. db.session.commit()
  651. @classmethod
  652. def _extract_variable_selector_to_variable_mapping(
  653. cls,
  654. *,
  655. graph_config: Mapping[str, Any],
  656. node_id: str,
  657. node_data: LLMNodeData,
  658. ) -> Mapping[str, Sequence[str]]:
  659. prompt_template = node_data.prompt_template
  660. variable_selectors = []
  661. if isinstance(prompt_template, list) and all(
  662. isinstance(prompt, LLMNodeChatModelMessage) for prompt in prompt_template
  663. ):
  664. for prompt in prompt_template:
  665. if prompt.edition_type != "jinja2":
  666. variable_template_parser = VariableTemplateParser(template=prompt.text)
  667. variable_selectors.extend(variable_template_parser.extract_variable_selectors())
  668. elif isinstance(prompt_template, LLMNodeCompletionModelPromptTemplate):
  669. if prompt_template.edition_type != "jinja2":
  670. variable_template_parser = VariableTemplateParser(template=prompt_template.text)
  671. variable_selectors = variable_template_parser.extract_variable_selectors()
  672. else:
  673. raise InvalidVariableTypeError(f"Invalid prompt template type: {type(prompt_template)}")
  674. variable_mapping = {}
  675. for variable_selector in variable_selectors:
  676. variable_mapping[variable_selector.variable] = variable_selector.value_selector
  677. memory = node_data.memory
  678. if memory and memory.query_prompt_template:
  679. query_variable_selectors = VariableTemplateParser(
  680. template=memory.query_prompt_template
  681. ).extract_variable_selectors()
  682. for variable_selector in query_variable_selectors:
  683. variable_mapping[variable_selector.variable] = variable_selector.value_selector
  684. if node_data.context.enabled:
  685. variable_mapping["#context#"] = node_data.context.variable_selector
  686. if node_data.vision.enabled:
  687. variable_mapping["#files#"] = ["sys", SystemVariableKey.FILES.value]
  688. if node_data.memory:
  689. variable_mapping["#sys.query#"] = ["sys", SystemVariableKey.QUERY.value]
  690. if node_data.prompt_config:
  691. enable_jinja = False
  692. if isinstance(prompt_template, list):
  693. for prompt in prompt_template:
  694. if prompt.edition_type == "jinja2":
  695. enable_jinja = True
  696. break
  697. else:
  698. if prompt_template.edition_type == "jinja2":
  699. enable_jinja = True
  700. if enable_jinja:
  701. for variable_selector in node_data.prompt_config.jinja2_variables or []:
  702. variable_mapping[variable_selector.variable] = variable_selector.value_selector
  703. variable_mapping = {node_id + "." + key: value for key, value in variable_mapping.items()}
  704. return variable_mapping
  705. @classmethod
  706. def get_default_config(cls, filters: Optional[dict] = None) -> dict:
  707. return {
  708. "type": "llm",
  709. "config": {
  710. "prompt_templates": {
  711. "chat_model": {
  712. "prompts": [
  713. {"role": "system", "text": "You are a helpful AI assistant.", "edition_type": "basic"}
  714. ]
  715. },
  716. "completion_model": {
  717. "conversation_histories_role": {"user_prefix": "Human", "assistant_prefix": "Assistant"},
  718. "prompt": {
  719. "text": "Here are the chat histories between human and assistant, inside "
  720. "<histories></histories> XML tags.\n\n<histories>\n{{"
  721. "#histories#}}\n</histories>\n\n\nHuman: {{#sys.query#}}\n\nAssistant:",
  722. "edition_type": "basic",
  723. },
  724. "stop": ["Human:"],
  725. },
  726. }
  727. },
  728. }
  729. def _combine_message_content_with_role(*, contents: Sequence[PromptMessageContent], role: PromptMessageRole):
  730. match role:
  731. case PromptMessageRole.USER:
  732. return UserPromptMessage(content=contents)
  733. case PromptMessageRole.ASSISTANT:
  734. return AssistantPromptMessage(content=contents)
  735. case PromptMessageRole.SYSTEM:
  736. return SystemPromptMessage(content=contents)
  737. raise NotImplementedError(f"Role {role} is not supported")
  738. def _render_jinja2_message(
  739. *,
  740. template: str,
  741. jinjia2_variables: Sequence[VariableSelector],
  742. variable_pool: VariablePool,
  743. ):
  744. if not template:
  745. return ""
  746. jinjia2_inputs = {}
  747. for jinja2_variable in jinjia2_variables:
  748. variable = variable_pool.get(jinja2_variable.value_selector)
  749. jinjia2_inputs[jinja2_variable.variable] = variable.to_object() if variable else ""
  750. code_execute_resp = CodeExecutor.execute_workflow_code_template(
  751. language=CodeLanguage.JINJA2,
  752. code=template,
  753. inputs=jinjia2_inputs,
  754. )
  755. result_text = code_execute_resp["result"]
  756. return result_text
  757. def _handle_list_messages(
  758. *,
  759. messages: Sequence[LLMNodeChatModelMessage],
  760. context: Optional[str],
  761. jinja2_variables: Sequence[VariableSelector],
  762. variable_pool: VariablePool,
  763. vision_detail_config: ImagePromptMessageContent.DETAIL,
  764. ) -> Sequence[PromptMessage]:
  765. prompt_messages = []
  766. for message in messages:
  767. if message.edition_type == "jinja2":
  768. result_text = _render_jinja2_message(
  769. template=message.jinja2_text or "",
  770. jinjia2_variables=jinja2_variables,
  771. variable_pool=variable_pool,
  772. )
  773. prompt_message = _combine_message_content_with_role(
  774. contents=[TextPromptMessageContent(data=result_text)], role=message.role
  775. )
  776. prompt_messages.append(prompt_message)
  777. else:
  778. # Get segment group from basic message
  779. if context:
  780. template = message.text.replace("{#context#}", context)
  781. else:
  782. template = message.text
  783. segment_group = variable_pool.convert_template(template)
  784. # Process segments for images
  785. file_contents = []
  786. for segment in segment_group.value:
  787. if isinstance(segment, ArrayFileSegment):
  788. for file in segment.value:
  789. if file.type in {FileType.IMAGE, FileType.VIDEO, FileType.AUDIO, FileType.DOCUMENT}:
  790. file_content = file_manager.to_prompt_message_content(
  791. file, image_detail_config=vision_detail_config
  792. )
  793. file_contents.append(file_content)
  794. if isinstance(segment, FileSegment):
  795. file = segment.value
  796. if file.type in {FileType.IMAGE, FileType.VIDEO, FileType.AUDIO, FileType.DOCUMENT}:
  797. file_content = file_manager.to_prompt_message_content(
  798. file, image_detail_config=vision_detail_config
  799. )
  800. file_contents.append(file_content)
  801. # Create message with text from all segments
  802. plain_text = segment_group.text
  803. if plain_text:
  804. prompt_message = _combine_message_content_with_role(
  805. contents=[TextPromptMessageContent(data=plain_text)], role=message.role
  806. )
  807. prompt_messages.append(prompt_message)
  808. if file_contents:
  809. # Create message with image contents
  810. prompt_message = _combine_message_content_with_role(contents=file_contents, role=message.role)
  811. prompt_messages.append(prompt_message)
  812. return prompt_messages
  813. def _calculate_rest_token(
  814. *, prompt_messages: list[PromptMessage], model_config: ModelConfigWithCredentialsEntity
  815. ) -> int:
  816. rest_tokens = 2000
  817. model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
  818. if model_context_tokens:
  819. model_instance = ModelInstance(
  820. provider_model_bundle=model_config.provider_model_bundle, model=model_config.model
  821. )
  822. curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
  823. max_tokens = 0
  824. for parameter_rule in model_config.model_schema.parameter_rules:
  825. if parameter_rule.name == "max_tokens" or (
  826. parameter_rule.use_template and parameter_rule.use_template == "max_tokens"
  827. ):
  828. max_tokens = (
  829. model_config.parameters.get(parameter_rule.name)
  830. or model_config.parameters.get(str(parameter_rule.use_template))
  831. or 0
  832. )
  833. rest_tokens = model_context_tokens - max_tokens - curr_message_tokens
  834. rest_tokens = max(rest_tokens, 0)
  835. return rest_tokens
  836. def _handle_memory_chat_mode(
  837. *,
  838. memory: TokenBufferMemory | None,
  839. memory_config: MemoryConfig | None,
  840. model_config: ModelConfigWithCredentialsEntity,
  841. ) -> Sequence[PromptMessage]:
  842. memory_messages = []
  843. # Get messages from memory for chat model
  844. if memory and memory_config:
  845. rest_tokens = _calculate_rest_token(prompt_messages=[], model_config=model_config)
  846. memory_messages = memory.get_history_prompt_messages(
  847. max_token_limit=rest_tokens,
  848. message_limit=memory_config.window.size if memory_config.window.enabled else None,
  849. )
  850. return memory_messages
  851. def _handle_memory_completion_mode(
  852. *,
  853. memory: TokenBufferMemory | None,
  854. memory_config: MemoryConfig | None,
  855. model_config: ModelConfigWithCredentialsEntity,
  856. ) -> str:
  857. memory_text = ""
  858. # Get history text from memory for completion model
  859. if memory and memory_config:
  860. rest_tokens = _calculate_rest_token(prompt_messages=[], model_config=model_config)
  861. if not memory_config.role_prefix:
  862. raise MemoryRolePrefixRequiredError("Memory role prefix is required for completion model.")
  863. memory_text = memory.get_history_prompt_text(
  864. max_token_limit=rest_tokens,
  865. message_limit=memory_config.window.size if memory_config.window.enabled else None,
  866. human_prefix=memory_config.role_prefix.user,
  867. ai_prefix=memory_config.role_prefix.assistant,
  868. )
  869. return memory_text
  870. def _handle_completion_template(
  871. *,
  872. template: LLMNodeCompletionModelPromptTemplate,
  873. context: Optional[str],
  874. jinja2_variables: Sequence[VariableSelector],
  875. variable_pool: VariablePool,
  876. ) -> Sequence[PromptMessage]:
  877. """Handle completion template processing outside of LLMNode class.
  878. Args:
  879. template: The completion model prompt template
  880. context: Optional context string
  881. jinja2_variables: Variables for jinja2 template rendering
  882. variable_pool: Variable pool for template conversion
  883. Returns:
  884. Sequence of prompt messages
  885. """
  886. prompt_messages = []
  887. if template.edition_type == "jinja2":
  888. result_text = _render_jinja2_message(
  889. template=template.jinja2_text or "",
  890. jinjia2_variables=jinja2_variables,
  891. variable_pool=variable_pool,
  892. )
  893. else:
  894. if context:
  895. template_text = template.text.replace("{#context#}", context)
  896. else:
  897. template_text = template.text
  898. result_text = variable_pool.convert_template(template_text).text
  899. prompt_message = _combine_message_content_with_role(
  900. contents=[TextPromptMessageContent(data=result_text)], role=PromptMessageRole.USER
  901. )
  902. prompt_messages.append(prompt_message)
  903. return prompt_messages