advanced_prompt_transform.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237
  1. from typing import Optional, Union
  2. from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
  3. from core.file.file_obj import FileVar
  4. from core.memory.token_buffer_memory import TokenBufferMemory
  5. from core.model_runtime.entities.message_entities import (
  6. AssistantPromptMessage,
  7. PromptMessage,
  8. PromptMessageRole,
  9. SystemPromptMessage,
  10. TextPromptMessageContent,
  11. UserPromptMessage,
  12. )
  13. from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
  14. from core.prompt.prompt_transform import PromptTransform
  15. from core.prompt.simple_prompt_transform import ModelMode
  16. from core.prompt.utils.prompt_template_parser import PromptTemplateParser
  17. class AdvancedPromptTransform(PromptTransform):
  18. """
  19. Advanced Prompt Transform for Workflow LLM Node.
  20. """
  21. def __init__(self, with_variable_tmpl: bool = False) -> None:
  22. self.with_variable_tmpl = with_variable_tmpl
  23. def get_prompt(self, prompt_template: Union[list[ChatModelMessage], CompletionModelPromptTemplate],
  24. inputs: dict,
  25. query: str,
  26. files: list[FileVar],
  27. context: Optional[str],
  28. memory_config: Optional[MemoryConfig],
  29. memory: Optional[TokenBufferMemory],
  30. model_config: ModelConfigWithCredentialsEntity) -> list[PromptMessage]:
  31. inputs = {key: str(value) for key, value in inputs.items()}
  32. prompt_messages = []
  33. model_mode = ModelMode.value_of(model_config.mode)
  34. if model_mode == ModelMode.COMPLETION:
  35. prompt_messages = self._get_completion_model_prompt_messages(
  36. prompt_template=prompt_template,
  37. inputs=inputs,
  38. query=query,
  39. files=files,
  40. context=context,
  41. memory_config=memory_config,
  42. memory=memory,
  43. model_config=model_config
  44. )
  45. elif model_mode == ModelMode.CHAT:
  46. prompt_messages = self._get_chat_model_prompt_messages(
  47. prompt_template=prompt_template,
  48. inputs=inputs,
  49. query=query,
  50. files=files,
  51. context=context,
  52. memory_config=memory_config,
  53. memory=memory,
  54. model_config=model_config
  55. )
  56. return prompt_messages
  57. def _get_completion_model_prompt_messages(self,
  58. prompt_template: CompletionModelPromptTemplate,
  59. inputs: dict,
  60. query: Optional[str],
  61. files: list[FileVar],
  62. context: Optional[str],
  63. memory_config: Optional[MemoryConfig],
  64. memory: Optional[TokenBufferMemory],
  65. model_config: ModelConfigWithCredentialsEntity) -> list[PromptMessage]:
  66. """
  67. Get completion model prompt messages.
  68. """
  69. raw_prompt = prompt_template.text
  70. prompt_messages = []
  71. prompt_template = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl)
  72. prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
  73. prompt_inputs = self._set_context_variable(context, prompt_template, prompt_inputs)
  74. if memory and memory_config:
  75. role_prefix = memory_config.role_prefix
  76. prompt_inputs = self._set_histories_variable(
  77. memory=memory,
  78. memory_config=memory_config,
  79. raw_prompt=raw_prompt,
  80. role_prefix=role_prefix,
  81. prompt_template=prompt_template,
  82. prompt_inputs=prompt_inputs,
  83. model_config=model_config
  84. )
  85. if query:
  86. prompt_inputs = self._set_query_variable(query, prompt_template, prompt_inputs)
  87. prompt = prompt_template.format(
  88. prompt_inputs
  89. )
  90. if files:
  91. prompt_message_contents = [TextPromptMessageContent(data=prompt)]
  92. for file in files:
  93. prompt_message_contents.append(file.prompt_message_content)
  94. prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
  95. else:
  96. prompt_messages.append(UserPromptMessage(content=prompt))
  97. return prompt_messages
  98. def _get_chat_model_prompt_messages(self,
  99. prompt_template: list[ChatModelMessage],
  100. inputs: dict,
  101. query: Optional[str],
  102. files: list[FileVar],
  103. context: Optional[str],
  104. memory_config: Optional[MemoryConfig],
  105. memory: Optional[TokenBufferMemory],
  106. model_config: ModelConfigWithCredentialsEntity) -> list[PromptMessage]:
  107. """
  108. Get chat model prompt messages.
  109. """
  110. raw_prompt_list = prompt_template
  111. prompt_messages = []
  112. for prompt_item in raw_prompt_list:
  113. raw_prompt = prompt_item.text
  114. prompt_template = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl)
  115. prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
  116. prompt_inputs = self._set_context_variable(context, prompt_template, prompt_inputs)
  117. prompt = prompt_template.format(
  118. prompt_inputs
  119. )
  120. if prompt_item.role == PromptMessageRole.USER:
  121. prompt_messages.append(UserPromptMessage(content=prompt))
  122. elif prompt_item.role == PromptMessageRole.SYSTEM and prompt:
  123. prompt_messages.append(SystemPromptMessage(content=prompt))
  124. elif prompt_item.role == PromptMessageRole.ASSISTANT:
  125. prompt_messages.append(AssistantPromptMessage(content=prompt))
  126. if memory and memory_config:
  127. prompt_messages = self._append_chat_histories(memory, memory_config, prompt_messages, model_config)
  128. if files:
  129. prompt_message_contents = [TextPromptMessageContent(data=query)]
  130. for file in files:
  131. prompt_message_contents.append(file.prompt_message_content)
  132. prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
  133. else:
  134. prompt_messages.append(UserPromptMessage(content=query))
  135. elif files:
  136. if not query:
  137. # get last message
  138. last_message = prompt_messages[-1] if prompt_messages else None
  139. if last_message and last_message.role == PromptMessageRole.USER:
  140. # get last user message content and add files
  141. prompt_message_contents = [TextPromptMessageContent(data=last_message.content)]
  142. for file in files:
  143. prompt_message_contents.append(file.prompt_message_content)
  144. last_message.content = prompt_message_contents
  145. else:
  146. prompt_message_contents = [TextPromptMessageContent(data='')] # not for query
  147. for file in files:
  148. prompt_message_contents.append(file.prompt_message_content)
  149. prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
  150. else:
  151. prompt_message_contents = [TextPromptMessageContent(data=query)]
  152. for file in files:
  153. prompt_message_contents.append(file.prompt_message_content)
  154. prompt_messages.append(UserPromptMessage(content=prompt_message_contents))
  155. elif query:
  156. prompt_messages.append(UserPromptMessage(content=query))
  157. return prompt_messages
  158. def _set_context_variable(self, context: str, prompt_template: PromptTemplateParser, prompt_inputs: dict) -> dict:
  159. if '#context#' in prompt_template.variable_keys:
  160. if context:
  161. prompt_inputs['#context#'] = context
  162. else:
  163. prompt_inputs['#context#'] = ''
  164. return prompt_inputs
  165. def _set_query_variable(self, query: str, prompt_template: PromptTemplateParser, prompt_inputs: dict) -> dict:
  166. if '#query#' in prompt_template.variable_keys:
  167. if query:
  168. prompt_inputs['#query#'] = query
  169. else:
  170. prompt_inputs['#query#'] = ''
  171. return prompt_inputs
  172. def _set_histories_variable(self, memory: TokenBufferMemory,
  173. memory_config: MemoryConfig,
  174. raw_prompt: str,
  175. role_prefix: MemoryConfig.RolePrefix,
  176. prompt_template: PromptTemplateParser,
  177. prompt_inputs: dict,
  178. model_config: ModelConfigWithCredentialsEntity) -> dict:
  179. if '#histories#' in prompt_template.variable_keys:
  180. if memory:
  181. inputs = {'#histories#': '', **prompt_inputs}
  182. prompt_template = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl)
  183. prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs}
  184. tmp_human_message = UserPromptMessage(
  185. content=prompt_template.format(prompt_inputs)
  186. )
  187. rest_tokens = self._calculate_rest_token([tmp_human_message], model_config)
  188. histories = self._get_history_messages_from_memory(
  189. memory=memory,
  190. memory_config=memory_config,
  191. max_token_limit=rest_tokens,
  192. human_prefix=role_prefix.user,
  193. ai_prefix=role_prefix.assistant
  194. )
  195. prompt_inputs['#histories#'] = histories
  196. else:
  197. prompt_inputs['#histories#'] = ''
  198. return prompt_inputs