mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 19:59:50 +08:00
424 lines
17 KiB
Python
424 lines
17 KiB
Python
import json
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Generator
|
|
from typing import Union
|
|
|
|
from core.agent.base_agent_runner import BaseAgentRunner
|
|
from core.agent.entities import AgentScratchpadUnit
|
|
from core.agent.output_parser.cot_output_parser import CotAgentOutputParser
|
|
from core.app.apps.base_app_queue_manager import PublishFrom
|
|
from core.app.entities.queue_entities import QueueAgentThoughtEvent, QueueMessageEndEvent, QueueMessageFileEvent
|
|
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
|
|
from core.model_runtime.entities.message_entities import (
|
|
AssistantPromptMessage,
|
|
PromptMessage,
|
|
ToolPromptMessage,
|
|
UserPromptMessage,
|
|
)
|
|
from core.tools.entities.tool_entities import ToolInvokeMeta
|
|
from core.tools.tool.tool import Tool
|
|
from core.tools.tool_engine import ToolEngine
|
|
from models.model import Message
|
|
|
|
|
|
class CotAgentRunner(BaseAgentRunner, ABC):
|
|
_is_first_iteration = True
|
|
_ignore_observation_providers = ['wenxin']
|
|
_historic_prompt_messages: list[PromptMessage] = None
|
|
_agent_scratchpad: list[AgentScratchpadUnit] = None
|
|
_instruction: str = None
|
|
_query: str = None
|
|
_prompt_messages_tools: list[PromptMessage] = None
|
|
|
|
def run(self, message: Message,
|
|
query: str,
|
|
inputs: dict[str, str],
|
|
) -> Union[Generator, LLMResult]:
|
|
"""
|
|
Run Cot agent application
|
|
"""
|
|
app_generate_entity = self.application_generate_entity
|
|
self._repack_app_generate_entity(app_generate_entity)
|
|
self._init_react_state(query)
|
|
|
|
# check model mode
|
|
if 'Observation' not in app_generate_entity.model_config.stop:
|
|
if app_generate_entity.model_config.provider not in self._ignore_observation_providers:
|
|
app_generate_entity.model_config.stop.append('Observation')
|
|
|
|
app_config = self.app_config
|
|
|
|
# init instruction
|
|
inputs = inputs or {}
|
|
instruction = app_config.prompt_template.simple_prompt_template
|
|
self._instruction = self._fill_in_inputs_from_external_data_tools(instruction, inputs)
|
|
|
|
iteration_step = 1
|
|
max_iteration_steps = min(app_config.agent.max_iteration, 5) + 1
|
|
|
|
# convert tools into ModelRuntime Tool format
|
|
tool_instances, self._prompt_messages_tools = self._init_prompt_tools()
|
|
|
|
prompt_messages = self._organize_prompt_messages()
|
|
|
|
function_call_state = True
|
|
llm_usage = {
|
|
'usage': None
|
|
}
|
|
final_answer = ''
|
|
|
|
def increase_usage(final_llm_usage_dict: dict[str, LLMUsage], usage: LLMUsage):
|
|
if not final_llm_usage_dict['usage']:
|
|
final_llm_usage_dict['usage'] = usage
|
|
else:
|
|
llm_usage = final_llm_usage_dict['usage']
|
|
llm_usage.prompt_tokens += usage.prompt_tokens
|
|
llm_usage.completion_tokens += usage.completion_tokens
|
|
llm_usage.prompt_price += usage.prompt_price
|
|
llm_usage.completion_price += usage.completion_price
|
|
|
|
model_instance = self.model_instance
|
|
|
|
while function_call_state and iteration_step <= max_iteration_steps:
|
|
# continue to run until there is not any tool call
|
|
function_call_state = False
|
|
|
|
if iteration_step == max_iteration_steps:
|
|
# the last iteration, remove all tools
|
|
self._prompt_messages_tools = []
|
|
|
|
message_file_ids = []
|
|
|
|
agent_thought = self.create_agent_thought(
|
|
message_id=message.id,
|
|
message='',
|
|
tool_name='',
|
|
tool_input='',
|
|
messages_ids=message_file_ids
|
|
)
|
|
|
|
if iteration_step > 1:
|
|
self.queue_manager.publish(QueueAgentThoughtEvent(
|
|
agent_thought_id=agent_thought.id
|
|
), PublishFrom.APPLICATION_MANAGER)
|
|
|
|
# recalc llm max tokens
|
|
prompt_messages = self._organize_prompt_messages()
|
|
self.recalc_llm_max_tokens(self.model_config, prompt_messages)
|
|
# invoke model
|
|
chunks: Generator[LLMResultChunk, None, None] = model_instance.invoke_llm(
|
|
prompt_messages=prompt_messages,
|
|
model_parameters=app_generate_entity.model_config.parameters,
|
|
tools=[],
|
|
stop=app_generate_entity.model_config.stop,
|
|
stream=True,
|
|
user=self.user_id,
|
|
callbacks=[],
|
|
)
|
|
|
|
# check llm result
|
|
if not chunks:
|
|
raise ValueError("failed to invoke llm")
|
|
|
|
usage_dict = {}
|
|
react_chunks = CotAgentOutputParser.handle_react_stream_output(chunks)
|
|
scratchpad = AgentScratchpadUnit(
|
|
agent_response='',
|
|
thought='',
|
|
action_str='',
|
|
observation='',
|
|
action=None,
|
|
)
|
|
|
|
# publish agent thought if it's first iteration
|
|
if iteration_step == 1:
|
|
self.queue_manager.publish(QueueAgentThoughtEvent(
|
|
agent_thought_id=agent_thought.id
|
|
), PublishFrom.APPLICATION_MANAGER)
|
|
|
|
for chunk in react_chunks:
|
|
if isinstance(chunk, AgentScratchpadUnit.Action):
|
|
action = chunk
|
|
# detect action
|
|
scratchpad.agent_response += json.dumps(chunk.dict())
|
|
scratchpad.action_str = json.dumps(chunk.dict())
|
|
scratchpad.action = action
|
|
else:
|
|
scratchpad.agent_response += chunk
|
|
scratchpad.thought += chunk
|
|
yield LLMResultChunk(
|
|
model=self.model_config.model,
|
|
prompt_messages=prompt_messages,
|
|
system_fingerprint='',
|
|
delta=LLMResultChunkDelta(
|
|
index=0,
|
|
message=AssistantPromptMessage(
|
|
content=chunk
|
|
),
|
|
usage=None
|
|
)
|
|
)
|
|
|
|
scratchpad.thought = scratchpad.thought.strip() or 'I am thinking about how to help you'
|
|
self._agent_scratchpad.append(scratchpad)
|
|
|
|
# get llm usage
|
|
if 'usage' in usage_dict:
|
|
increase_usage(llm_usage, usage_dict['usage'])
|
|
else:
|
|
usage_dict['usage'] = LLMUsage.empty_usage()
|
|
|
|
self.save_agent_thought(
|
|
agent_thought=agent_thought,
|
|
tool_name=scratchpad.action.action_name if scratchpad.action else '',
|
|
tool_input={
|
|
scratchpad.action.action_name: scratchpad.action.action_input
|
|
} if scratchpad.action else {},
|
|
tool_invoke_meta={},
|
|
thought=scratchpad.thought,
|
|
observation='',
|
|
answer=scratchpad.agent_response,
|
|
messages_ids=[],
|
|
llm_usage=usage_dict['usage']
|
|
)
|
|
|
|
if not scratchpad.is_final():
|
|
self.queue_manager.publish(QueueAgentThoughtEvent(
|
|
agent_thought_id=agent_thought.id
|
|
), PublishFrom.APPLICATION_MANAGER)
|
|
|
|
if not scratchpad.action:
|
|
# failed to extract action, return final answer directly
|
|
final_answer = scratchpad.agent_response or ''
|
|
else:
|
|
if scratchpad.action.action_name.lower() == "final answer":
|
|
# action is final answer, return final answer directly
|
|
try:
|
|
if isinstance(scratchpad.action.action_input, dict):
|
|
final_answer = json.dumps(scratchpad.action.action_input)
|
|
elif isinstance(scratchpad.action.action_input, str):
|
|
final_answer = scratchpad.action.action_input
|
|
else:
|
|
final_answer = f'{scratchpad.action.action_input}'
|
|
except json.JSONDecodeError:
|
|
final_answer = f'{scratchpad.action.action_input}'
|
|
else:
|
|
function_call_state = True
|
|
# action is tool call, invoke tool
|
|
tool_invoke_response, tool_invoke_meta = self._handle_invoke_action(
|
|
action=scratchpad.action,
|
|
tool_instances=tool_instances,
|
|
message_file_ids=message_file_ids
|
|
)
|
|
scratchpad.observation = tool_invoke_response
|
|
scratchpad.agent_response = tool_invoke_response
|
|
|
|
self.save_agent_thought(
|
|
agent_thought=agent_thought,
|
|
tool_name=scratchpad.action.action_name,
|
|
tool_input={scratchpad.action.action_name: scratchpad.action.action_input},
|
|
thought=scratchpad.thought,
|
|
observation={scratchpad.action.action_name: tool_invoke_response},
|
|
tool_invoke_meta=tool_invoke_meta.to_dict(),
|
|
answer=scratchpad.agent_response,
|
|
messages_ids=message_file_ids,
|
|
llm_usage=usage_dict['usage']
|
|
)
|
|
|
|
self.queue_manager.publish(QueueAgentThoughtEvent(
|
|
agent_thought_id=agent_thought.id
|
|
), PublishFrom.APPLICATION_MANAGER)
|
|
|
|
# update prompt tool message
|
|
for prompt_tool in self._prompt_messages_tools:
|
|
self.update_prompt_message_tool(tool_instances[prompt_tool.name], prompt_tool)
|
|
|
|
iteration_step += 1
|
|
|
|
yield LLMResultChunk(
|
|
model=model_instance.model,
|
|
prompt_messages=prompt_messages,
|
|
delta=LLMResultChunkDelta(
|
|
index=0,
|
|
message=AssistantPromptMessage(
|
|
content=final_answer
|
|
),
|
|
usage=llm_usage['usage']
|
|
),
|
|
system_fingerprint=''
|
|
)
|
|
|
|
# save agent thought
|
|
self.save_agent_thought(
|
|
agent_thought=agent_thought,
|
|
tool_name='',
|
|
tool_input={},
|
|
tool_invoke_meta={},
|
|
thought=final_answer,
|
|
observation={},
|
|
answer=final_answer,
|
|
messages_ids=[]
|
|
)
|
|
|
|
self.update_db_variables(self.variables_pool, self.db_variables_pool)
|
|
# publish end event
|
|
self.queue_manager.publish(QueueMessageEndEvent(llm_result=LLMResult(
|
|
model=model_instance.model,
|
|
prompt_messages=prompt_messages,
|
|
message=AssistantPromptMessage(
|
|
content=final_answer
|
|
),
|
|
usage=llm_usage['usage'] if llm_usage['usage'] else LLMUsage.empty_usage(),
|
|
system_fingerprint=''
|
|
)), PublishFrom.APPLICATION_MANAGER)
|
|
|
|
def _handle_invoke_action(self, action: AgentScratchpadUnit.Action,
|
|
tool_instances: dict[str, Tool],
|
|
message_file_ids: list[str]) -> tuple[str, ToolInvokeMeta]:
|
|
"""
|
|
handle invoke action
|
|
:param action: action
|
|
:param tool_instances: tool instances
|
|
:return: observation, meta
|
|
"""
|
|
# action is tool call, invoke tool
|
|
tool_call_name = action.action_name
|
|
tool_call_args = action.action_input
|
|
tool_instance = tool_instances.get(tool_call_name)
|
|
|
|
if not tool_instance:
|
|
answer = f"there is not a tool named {tool_call_name}"
|
|
return answer, ToolInvokeMeta.error_instance(answer)
|
|
|
|
if isinstance(tool_call_args, str):
|
|
try:
|
|
tool_call_args = json.loads(tool_call_args)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# invoke tool
|
|
tool_invoke_response, message_files, tool_invoke_meta = ToolEngine.agent_invoke(
|
|
tool=tool_instance,
|
|
tool_parameters=tool_call_args,
|
|
user_id=self.user_id,
|
|
tenant_id=self.tenant_id,
|
|
message=self.message,
|
|
invoke_from=self.application_generate_entity.invoke_from,
|
|
agent_tool_callback=self.agent_callback
|
|
)
|
|
|
|
# publish files
|
|
for message_file, save_as in message_files:
|
|
if save_as:
|
|
self.variables_pool.set_file(tool_name=tool_call_name, value=message_file.id, name=save_as)
|
|
|
|
# publish message file
|
|
self.queue_manager.publish(QueueMessageFileEvent(
|
|
message_file_id=message_file.id
|
|
), PublishFrom.APPLICATION_MANAGER)
|
|
# add message file ids
|
|
message_file_ids.append(message_file.id)
|
|
|
|
return tool_invoke_response, tool_invoke_meta
|
|
|
|
def _convert_dict_to_action(self, action: dict) -> AgentScratchpadUnit.Action:
|
|
"""
|
|
convert dict to action
|
|
"""
|
|
return AgentScratchpadUnit.Action(
|
|
action_name=action['action'],
|
|
action_input=action['action_input']
|
|
)
|
|
|
|
def _fill_in_inputs_from_external_data_tools(self, instruction: str, inputs: dict) -> str:
|
|
"""
|
|
fill in inputs from external data tools
|
|
"""
|
|
for key, value in inputs.items():
|
|
try:
|
|
instruction = instruction.replace(f'{{{{{key}}}}}', str(value))
|
|
except Exception as e:
|
|
continue
|
|
|
|
return instruction
|
|
|
|
def _init_react_state(self, query) -> None:
|
|
"""
|
|
init agent scratchpad
|
|
"""
|
|
self._query = query
|
|
self._agent_scratchpad = []
|
|
self._historic_prompt_messages = self._organize_historic_prompt_messages()
|
|
|
|
@abstractmethod
|
|
def _organize_prompt_messages(self) -> list[PromptMessage]:
|
|
"""
|
|
organize prompt messages
|
|
"""
|
|
|
|
def _format_assistant_message(self, agent_scratchpad: list[AgentScratchpadUnit]) -> str:
|
|
"""
|
|
format assistant message
|
|
"""
|
|
message = ''
|
|
for scratchpad in agent_scratchpad:
|
|
if scratchpad.is_final():
|
|
message += f"Final Answer: {scratchpad.agent_response}"
|
|
else:
|
|
message += f"Thought: {scratchpad.thought}\n\n"
|
|
if scratchpad.action_str:
|
|
message += f"Action: {scratchpad.action_str}\n\n"
|
|
if scratchpad.observation:
|
|
message += f"Observation: {scratchpad.observation}\n\n"
|
|
|
|
return message
|
|
|
|
def _organize_historic_prompt_messages(self) -> list[PromptMessage]:
|
|
"""
|
|
organize historic prompt messages
|
|
"""
|
|
result: list[PromptMessage] = []
|
|
scratchpad: list[AgentScratchpadUnit] = []
|
|
current_scratchpad: AgentScratchpadUnit = None
|
|
|
|
for message in self.history_prompt_messages:
|
|
if isinstance(message, AssistantPromptMessage):
|
|
current_scratchpad = AgentScratchpadUnit(
|
|
agent_response=message.content,
|
|
thought=message.content or 'I am thinking about how to help you',
|
|
action_str='',
|
|
action=None,
|
|
observation=None,
|
|
)
|
|
if message.tool_calls:
|
|
try:
|
|
current_scratchpad.action = AgentScratchpadUnit.Action(
|
|
action_name=message.tool_calls[0].function.name,
|
|
action_input=json.loads(message.tool_calls[0].function.arguments)
|
|
)
|
|
current_scratchpad.action_str = json.dumps(
|
|
current_scratchpad.action.to_dict()
|
|
)
|
|
except:
|
|
pass
|
|
|
|
scratchpad.append(current_scratchpad)
|
|
elif isinstance(message, ToolPromptMessage):
|
|
if current_scratchpad:
|
|
current_scratchpad.observation = message.content
|
|
elif isinstance(message, UserPromptMessage):
|
|
result.append(message)
|
|
|
|
if scratchpad:
|
|
result.append(AssistantPromptMessage(
|
|
content=self._format_assistant_message(scratchpad)
|
|
))
|
|
|
|
scratchpad = []
|
|
|
|
if scratchpad:
|
|
result.append(AssistantPromptMessage(
|
|
content=self._format_assistant_message(scratchpad)
|
|
))
|
|
|
|
return result |