base_agent_runner.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520
  1. import json
  2. import logging
  3. import uuid
  4. from collections.abc import Mapping, Sequence
  5. from typing import Optional, Union, cast
  6. from core.agent.entities import AgentEntity, AgentToolEntity
  7. from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
  8. from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig
  9. from core.app.apps.base_app_queue_manager import AppQueueManager
  10. from core.app.apps.base_app_runner import AppRunner
  11. from core.app.entities.app_invoke_entities import (
  12. AgentChatAppGenerateEntity,
  13. ModelConfigWithCredentialsEntity,
  14. )
  15. from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler
  16. from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
  17. from core.file.message_file_parser import MessageFileParser
  18. from core.memory.token_buffer_memory import TokenBufferMemory
  19. from core.model_manager import ModelInstance
  20. from core.model_runtime.entities.llm_entities import LLMUsage
  21. from core.model_runtime.entities.message_entities import (
  22. AssistantPromptMessage,
  23. PromptMessage,
  24. PromptMessageContent,
  25. PromptMessageTool,
  26. SystemPromptMessage,
  27. TextPromptMessageContent,
  28. ToolPromptMessage,
  29. UserPromptMessage,
  30. )
  31. from core.model_runtime.entities.model_entities import ModelFeature
  32. from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
  33. from core.prompt.utils.extract_thread_messages import extract_thread_messages
  34. from core.tools.__base.tool import Tool
  35. from core.tools.entities.tool_entities import (
  36. ToolParameter,
  37. )
  38. from core.tools.tool_manager import ToolManager
  39. from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool
  40. from core.tools.utils.tool_parameter_converter import ToolParameterConverter
  41. from extensions.ext_database import db
  42. from models.model import Conversation, Message, MessageAgentThought
  43. logger = logging.getLogger(__name__)
  44. class BaseAgentRunner(AppRunner):
  45. def __init__(
  46. self,
  47. tenant_id: str,
  48. application_generate_entity: AgentChatAppGenerateEntity,
  49. conversation: Conversation,
  50. app_config: AgentChatAppConfig,
  51. model_config: ModelConfigWithCredentialsEntity,
  52. config: AgentEntity,
  53. queue_manager: AppQueueManager,
  54. message: Message,
  55. user_id: str,
  56. model_instance: ModelInstance,
  57. memory: Optional[TokenBufferMemory] = None,
  58. prompt_messages: Optional[list[PromptMessage]] = None,
  59. ) -> None:
  60. """
  61. Agent runner
  62. :param tenant_id: tenant id
  63. :param application_generate_entity: application generate entity
  64. :param conversation: conversation
  65. :param app_config: app generate entity
  66. :param model_config: model config
  67. :param config: dataset config
  68. :param queue_manager: queue manager
  69. :param message: message
  70. :param user_id: user id
  71. :param memory: memory
  72. :param prompt_messages: prompt messages
  73. :param variables_pool: variables pool
  74. :param db_variables: db variables
  75. :param model_instance: model instance
  76. """
  77. self.tenant_id = tenant_id
  78. self.application_generate_entity = application_generate_entity
  79. self.conversation = conversation
  80. self.app_config = app_config
  81. self.model_config = model_config
  82. self.config = config
  83. self.queue_manager = queue_manager
  84. self.message = message
  85. self.user_id = user_id
  86. self.memory = memory
  87. self.history_prompt_messages = self.organize_agent_history(prompt_messages=prompt_messages or [])
  88. self.model_instance = model_instance
  89. # init callback
  90. self.agent_callback = DifyAgentCallbackHandler()
  91. # init dataset tools
  92. hit_callback = DatasetIndexToolCallbackHandler(
  93. queue_manager=queue_manager,
  94. app_id=self.app_config.app_id,
  95. message_id=message.id,
  96. user_id=user_id,
  97. invoke_from=self.application_generate_entity.invoke_from,
  98. )
  99. self.dataset_tools = DatasetRetrieverTool.get_dataset_tools(
  100. tenant_id=tenant_id,
  101. dataset_ids=app_config.dataset.dataset_ids if app_config.dataset else [],
  102. retrieve_config=app_config.dataset.retrieve_config if app_config.dataset else None,
  103. return_resource=app_config.additional_features.show_retrieve_source,
  104. invoke_from=application_generate_entity.invoke_from,
  105. hit_callback=hit_callback,
  106. )
  107. # get how many agent thoughts have been created
  108. self.agent_thought_count = (
  109. db.session.query(MessageAgentThought)
  110. .filter(
  111. MessageAgentThought.message_id == self.message.id,
  112. )
  113. .count()
  114. )
  115. db.session.close()
  116. # check if model supports stream tool call
  117. llm_model = cast(LargeLanguageModel, model_instance.model_type_instance)
  118. model_schema = llm_model.get_model_schema(model_instance.model, model_instance.credentials)
  119. if model_schema and ModelFeature.STREAM_TOOL_CALL in (model_schema.features or []):
  120. self.stream_tool_call = True
  121. else:
  122. self.stream_tool_call = False
  123. # check if model supports vision
  124. if model_schema and ModelFeature.VISION in (model_schema.features or []):
  125. self.files = application_generate_entity.files
  126. else:
  127. self.files = []
  128. self.query = None
  129. self._current_thoughts: list[PromptMessage] = []
  130. def _repack_app_generate_entity(
  131. self, app_generate_entity: AgentChatAppGenerateEntity
  132. ) -> AgentChatAppGenerateEntity:
  133. """
  134. Repack app generate entity
  135. """
  136. if app_generate_entity.app_config.prompt_template.simple_prompt_template is None:
  137. app_generate_entity.app_config.prompt_template.simple_prompt_template = ""
  138. return app_generate_entity
  139. def _convert_tool_to_prompt_message_tool(self, tool: AgentToolEntity) -> tuple[PromptMessageTool, Tool]:
  140. """
  141. convert tool to prompt message tool
  142. """
  143. tool_entity = ToolManager.get_agent_tool_runtime(
  144. tenant_id=self.tenant_id,
  145. app_id=self.app_config.app_id,
  146. agent_tool=tool,
  147. invoke_from=self.application_generate_entity.invoke_from,
  148. )
  149. assert tool_entity.entity.description
  150. message_tool = PromptMessageTool(
  151. name=tool.tool_name,
  152. description=tool_entity.entity.description.llm,
  153. parameters={
  154. "type": "object",
  155. "properties": {},
  156. "required": [],
  157. },
  158. )
  159. parameters = tool_entity.get_merged_runtime_parameters()
  160. for parameter in parameters:
  161. if parameter.form != ToolParameter.ToolParameterForm.LLM:
  162. continue
  163. parameter_type = ToolParameterConverter.get_parameter_type(parameter.type)
  164. enum = []
  165. if parameter.type == ToolParameter.ToolParameterType.SELECT:
  166. enum = [option.value for option in parameter.options]
  167. message_tool.parameters["properties"][parameter.name] = {
  168. "type": parameter_type,
  169. "description": parameter.llm_description or "",
  170. }
  171. if len(enum) > 0:
  172. message_tool.parameters["properties"][parameter.name]["enum"] = enum
  173. if parameter.required:
  174. message_tool.parameters["required"].append(parameter.name)
  175. return message_tool, tool_entity
  176. def _convert_dataset_retriever_tool_to_prompt_message_tool(self, tool: DatasetRetrieverTool) -> PromptMessageTool:
  177. """
  178. convert dataset retriever tool to prompt message tool
  179. """
  180. assert tool.entity.description
  181. prompt_tool = PromptMessageTool(
  182. name=tool.entity.identity.name,
  183. description=tool.entity.description.llm,
  184. parameters={
  185. "type": "object",
  186. "properties": {},
  187. "required": [],
  188. },
  189. )
  190. for parameter in tool.get_runtime_parameters():
  191. parameter_type = "string"
  192. prompt_tool.parameters["properties"][parameter.name] = {
  193. "type": parameter_type,
  194. "description": parameter.llm_description or "",
  195. }
  196. if parameter.required:
  197. if parameter.name not in prompt_tool.parameters["required"]:
  198. prompt_tool.parameters["required"].append(parameter.name)
  199. return prompt_tool
  200. def _init_prompt_tools(self) -> tuple[Mapping[str, Tool], Sequence[PromptMessageTool]]:
  201. """
  202. Init tools
  203. """
  204. tool_instances = {}
  205. prompt_messages_tools = []
  206. for tool in self.app_config.agent.tools or [] if self.app_config.agent else []:
  207. try:
  208. prompt_tool, tool_entity = self._convert_tool_to_prompt_message_tool(tool)
  209. except Exception:
  210. # api tool may be deleted
  211. continue
  212. # save tool entity
  213. tool_instances[tool.tool_name] = tool_entity
  214. # save prompt tool
  215. prompt_messages_tools.append(prompt_tool)
  216. # convert dataset tools into ModelRuntime Tool format
  217. for dataset_tool in self.dataset_tools:
  218. prompt_tool = self._convert_dataset_retriever_tool_to_prompt_message_tool(dataset_tool)
  219. # save prompt tool
  220. prompt_messages_tools.append(prompt_tool)
  221. # save tool entity
  222. tool_instances[dataset_tool.entity.identity.name] = dataset_tool
  223. return tool_instances, prompt_messages_tools
  224. def update_prompt_message_tool(self, tool: Tool, prompt_tool: PromptMessageTool) -> PromptMessageTool:
  225. """
  226. update prompt message tool
  227. """
  228. # try to get tool runtime parameters
  229. tool_runtime_parameters = tool.get_runtime_parameters() or []
  230. for parameter in tool_runtime_parameters:
  231. if parameter.form != ToolParameter.ToolParameterForm.LLM:
  232. continue
  233. parameter_type = ToolParameterConverter.get_parameter_type(parameter.type)
  234. enum = []
  235. if parameter.type == ToolParameter.ToolParameterType.SELECT:
  236. enum = [option.value for option in parameter.options]
  237. prompt_tool.parameters["properties"][parameter.name] = {
  238. "type": parameter_type,
  239. "description": parameter.llm_description or "",
  240. }
  241. if len(enum) > 0:
  242. prompt_tool.parameters["properties"][parameter.name]["enum"] = enum
  243. if parameter.required:
  244. if parameter.name not in prompt_tool.parameters["required"]:
  245. prompt_tool.parameters["required"].append(parameter.name)
  246. return prompt_tool
  247. def create_agent_thought(
  248. self, message_id: str, message: str, tool_name: str, tool_input: str, messages_ids: list[str]
  249. ) -> MessageAgentThought:
  250. """
  251. Create agent thought
  252. """
  253. thought = MessageAgentThought(
  254. message_id=message_id,
  255. message_chain_id=None,
  256. thought="",
  257. tool=tool_name,
  258. tool_labels_str="{}",
  259. tool_meta_str="{}",
  260. tool_input=tool_input,
  261. message=message,
  262. message_token=0,
  263. message_unit_price=0,
  264. message_price_unit=0,
  265. message_files=json.dumps(messages_ids) if messages_ids else "",
  266. answer="",
  267. observation="",
  268. answer_token=0,
  269. answer_unit_price=0,
  270. answer_price_unit=0,
  271. tokens=0,
  272. total_price=0,
  273. position=self.agent_thought_count + 1,
  274. currency="USD",
  275. latency=0,
  276. created_by_role="account",
  277. created_by=self.user_id,
  278. )
  279. db.session.add(thought)
  280. db.session.commit()
  281. db.session.refresh(thought)
  282. db.session.close()
  283. self.agent_thought_count += 1
  284. return thought
  285. def save_agent_thought(
  286. self,
  287. agent_thought: MessageAgentThought,
  288. tool_name: str | None,
  289. tool_input: Union[str, dict, None],
  290. thought: str | None,
  291. observation: Union[str, dict, None],
  292. tool_invoke_meta: Union[str, dict, None],
  293. answer: str | None,
  294. messages_ids: list[str],
  295. llm_usage: LLMUsage | None = None,
  296. ):
  297. """
  298. Save agent thought
  299. """
  300. updated_agent_thought = (
  301. db.session.query(MessageAgentThought).filter(MessageAgentThought.id == agent_thought.id).first()
  302. )
  303. if not updated_agent_thought:
  304. raise ValueError("agent thought not found")
  305. if thought is not None:
  306. updated_agent_thought.thought = thought
  307. if tool_name is not None:
  308. updated_agent_thought.tool = tool_name
  309. if tool_input is not None:
  310. if isinstance(tool_input, dict):
  311. try:
  312. tool_input = json.dumps(tool_input, ensure_ascii=False)
  313. except Exception as e:
  314. tool_input = json.dumps(tool_input)
  315. updated_agent_thought.tool_input = tool_input
  316. if observation is not None:
  317. if isinstance(observation, dict):
  318. try:
  319. observation = json.dumps(observation, ensure_ascii=False)
  320. except Exception as e:
  321. observation = json.dumps(observation)
  322. updated_agent_thought.observation = observation
  323. if answer is not None:
  324. updated_agent_thought.answer = answer
  325. if messages_ids is not None and len(messages_ids) > 0:
  326. updated_agent_thought.message_files = json.dumps(messages_ids)
  327. if llm_usage:
  328. updated_agent_thought.message_token = llm_usage.prompt_tokens
  329. updated_agent_thought.message_price_unit = llm_usage.prompt_price_unit
  330. updated_agent_thought.message_unit_price = llm_usage.prompt_unit_price
  331. updated_agent_thought.answer_token = llm_usage.completion_tokens
  332. updated_agent_thought.answer_price_unit = llm_usage.completion_price_unit
  333. updated_agent_thought.answer_unit_price = llm_usage.completion_unit_price
  334. updated_agent_thought.tokens = llm_usage.total_tokens
  335. updated_agent_thought.total_price = llm_usage.total_price
  336. # check if tool labels is not empty
  337. labels = updated_agent_thought.tool_labels or {}
  338. tools = updated_agent_thought.tool.split(";") if updated_agent_thought.tool else []
  339. for tool in tools:
  340. if not tool:
  341. continue
  342. if tool not in labels:
  343. tool_label = ToolManager.get_tool_label(tool)
  344. if tool_label:
  345. labels[tool] = tool_label.to_dict()
  346. else:
  347. labels[tool] = {"en_US": tool, "zh_Hans": tool}
  348. updated_agent_thought.tool_labels_str = json.dumps(labels)
  349. if tool_invoke_meta is not None:
  350. if isinstance(tool_invoke_meta, dict):
  351. try:
  352. tool_invoke_meta = json.dumps(tool_invoke_meta, ensure_ascii=False)
  353. except Exception as e:
  354. tool_invoke_meta = json.dumps(tool_invoke_meta)
  355. updated_agent_thought.tool_meta_str = tool_invoke_meta
  356. db.session.commit()
  357. db.session.close()
  358. def organize_agent_history(self, prompt_messages: list[PromptMessage]) -> list[PromptMessage]:
  359. """
  360. Organize agent history
  361. """
  362. result = []
  363. # check if there is a system message in the beginning of the conversation
  364. for prompt_message in prompt_messages:
  365. if isinstance(prompt_message, SystemPromptMessage):
  366. result.append(prompt_message)
  367. messages: list[Message] = (
  368. db.session.query(Message)
  369. .filter(
  370. Message.conversation_id == self.message.conversation_id,
  371. )
  372. .order_by(Message.created_at.desc())
  373. .all()
  374. )
  375. messages = list(reversed(extract_thread_messages(messages)))
  376. for message in messages:
  377. if message.id == self.message.id:
  378. continue
  379. result.append(self.organize_agent_user_prompt(message))
  380. agent_thoughts: list[MessageAgentThought] = message.agent_thoughts
  381. if agent_thoughts:
  382. for agent_thought in agent_thoughts:
  383. tools = agent_thought.tool
  384. if tools:
  385. tools = tools.split(";")
  386. tool_calls: list[AssistantPromptMessage.ToolCall] = []
  387. tool_call_response: list[ToolPromptMessage] = []
  388. try:
  389. tool_inputs = json.loads(agent_thought.tool_input)
  390. except Exception as e:
  391. tool_inputs = {tool: {} for tool in tools}
  392. try:
  393. tool_responses = json.loads(agent_thought.observation)
  394. except Exception as e:
  395. tool_responses = dict.fromkeys(tools, agent_thought.observation)
  396. for tool in tools:
  397. # generate a uuid for tool call
  398. tool_call_id = str(uuid.uuid4())
  399. tool_calls.append(
  400. AssistantPromptMessage.ToolCall(
  401. id=tool_call_id,
  402. type="function",
  403. function=AssistantPromptMessage.ToolCall.ToolCallFunction(
  404. name=tool,
  405. arguments=json.dumps(tool_inputs.get(tool, {})),
  406. ),
  407. )
  408. )
  409. tool_call_response.append(
  410. ToolPromptMessage(
  411. content=tool_responses.get(tool, agent_thought.observation),
  412. name=tool,
  413. tool_call_id=tool_call_id,
  414. )
  415. )
  416. result.extend(
  417. [
  418. AssistantPromptMessage(
  419. content=agent_thought.thought,
  420. tool_calls=tool_calls,
  421. ),
  422. *tool_call_response,
  423. ]
  424. )
  425. if not tools:
  426. result.append(AssistantPromptMessage(content=agent_thought.thought))
  427. else:
  428. if message.answer:
  429. result.append(AssistantPromptMessage(content=message.answer))
  430. db.session.close()
  431. return result
  432. def organize_agent_user_prompt(self, message: Message) -> UserPromptMessage:
  433. message_file_parser = MessageFileParser(
  434. tenant_id=self.tenant_id,
  435. app_id=self.app_config.app_id,
  436. )
  437. files = message.message_files
  438. if files:
  439. assert message.app_model_config
  440. file_extra_config = FileUploadConfigManager.convert(message.app_model_config.to_dict())
  441. if file_extra_config:
  442. file_objs = message_file_parser.transform_message_files(files, file_extra_config)
  443. else:
  444. file_objs = []
  445. if not file_objs:
  446. return UserPromptMessage(content=message.query)
  447. else:
  448. prompt_message_contents: list[PromptMessageContent] = [TextPromptMessageContent(data=message.query)]
  449. for file_obj in file_objs:
  450. prompt_message_contents.append(file_obj.prompt_message_content)
  451. return UserPromptMessage(content=prompt_message_contents)
  452. else:
  453. return UserPromptMessage(content=message.query)