opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-05 00:31:19 +08:00
Code Issues Packages Projects Releases Wiki Activity
4ffcf3bf4c
dify/api/core/agent/base_agent_runner.py
dparkmit24 2c162109c5 test(agent): pin the colliding case and state why the rule is shaped this way 
The occurrence selector tells a per-call list from a single call's list
value by length alone. One test pinned the side where that is
unambiguous -- a 3-element list against 2 calls, replayed whole. The side
that decides whether the rule is safe is the other one: a tool called n
times whose single stored value is a list of length n. That case was
decided by the rule and described by no test, so it read as an oversight
rather than a decision.

Pin it at both readers. The replay reader in base_agent_runner and the
display reader on MessageAgentThought each get a two-call record whose
one stored value is a two-element list, asserting that call 1 reads
element 0 and call 2 reads element 1 -- what the rule does today. The
name says what the case concedes rather than what it asserts.

The docstrings say why the asymmetry is tolerable. observation values
are always str: ToolEngine.agent_invoke is typed
-> tuple[str, list[str], ToolInvokeMeta] and both runners store element
0, so a list under a tool name is not a shape any writer produces and
the length check is defensive there. tool_input values are json.loads of
the model's arguments with no shape check, so a legacy list-valued input
is possible in principle, and that is the side the collision can reach.

The selector is defined twice, identically, because models/ importing
from core/agent/ is the worse layering trade and the reverse is odd.
Neither copy is in the wrong place, so each now names the other and says
the duplication is deliberate -- enough for a future editor to find both.

No behaviour change: the condition, the ordering and the fallback are
untouched.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-05 18:57:51 -07:00

530 lines
21 KiB
Python
Raw Blame History

import json
import logging
import uuid
from collections import Counter
from decimal import Decimal
from typing import Any, Union, cast
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from core.agent.entities import AgentEntity, AgentToolEntity
from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig
from core.app.apps.base_app_queue_manager import AppQueueManager
from core.app.apps.base_app_runner import AppRunner
from core.app.entities.app_invoke_entities import (
AgentChatAppGenerateEntity,
ModelConfigWithCredentialsEntity,
)
from core.app.file_access import DatabaseFileAccessController
from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler
from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_manager import ModelInstance
from core.prompt.utils.extract_thread_messages import extract_thread_messages
from core.tools.__base.tool import Tool
from core.tools.tool_manager import ToolManager
from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool
from extensions.ext_database import db
from factories import file_factory
from graphon.file import file_manager
from graphon.model_runtime.entities import (
AssistantPromptMessage,
LLMUsage,
PromptMessage,
PromptMessageTool,
SystemPromptMessage,
TextPromptMessageContent,
ToolPromptMessage,
UserPromptMessage,
)
from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes
from graphon.model_runtime.entities.model_entities import ModelFeature
from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
from models.enums import CreatorUserRole
from models.model import Conversation, Message, MessageAgentThought, MessageFile, load_annotation_reply_config
logger = logging.getLogger(__name__)
_file_access_controller = DatabaseFileAccessController()
def _select_tool_occurrence(value: Any, occurrence: int, occurrences: int) -> Any:
"""
Pick one call's value out of a persisted agent thought payload.
A tool called several times in one turn stores one value per call, in call
order. Records written before those calls were kept apart store a single
value for the tool name, and every occurrence replays it — the behaviour
those records were written with.
`observation` values are always `str`: `ToolEngine.agent_invoke` is typed
`-> tuple[str, list[str], ToolInvokeMeta]` and the runners store element 0,
so a list there is not a shape any writer produces and the length check is
defensive. `tool_input` values are `json.loads` of the model's `arguments`
with no shape check, so a legacy list-valued input is possible in principle;
on a length collision the rule reads per call, not whole —
`test_a_legacy_list_of_matching_length_is_read_per_call_not_whole` pins it.
The same function is defined, identically, in `models/model.py`.
The duplication is deliberate: `models/` importing from `core/agent/` is the
worse layering trade and the reverse is odd, so neither copy is in the wrong
place and neither should move. A change to this rule must be applied in both.
"""
if occurrences > 1 and isinstance(value, list) and len(value) == occurrences:
return value[occurrence]
return value
class BaseAgentRunner(AppRunner):
def __init__(
self,
*,
session: Session,
tenant_id: str,
application_generate_entity: AgentChatAppGenerateEntity,
conversation: Conversation,
app_config: AgentChatAppConfig,
model_config: ModelConfigWithCredentialsEntity,
config: AgentEntity,
queue_manager: AppQueueManager,
message: Message,
user_id: str,
model_instance: ModelInstance,
memory: TokenBufferMemory | None = None,
prompt_messages: list[PromptMessage] | None = None,
):
self.tenant_id = tenant_id
self.application_generate_entity = application_generate_entity
self.conversation = conversation
self.app_config = app_config
self.model_config = model_config
self.config = config
self.queue_manager = queue_manager
self.message = message
self.user_id = user_id
self.memory = memory
self.history_prompt_messages = self.organize_agent_history(
session=session, prompt_messages=prompt_messages or []
)
self.model_instance = model_instance
# init callback
self.agent_callback = DifyAgentCallbackHandler()
# init dataset tools
hit_callback = DatasetIndexToolCallbackHandler(
queue_manager=queue_manager,
app_id=self.app_config.app_id,
message_id=message.id,
user_id=user_id,
invoke_from=self.application_generate_entity.invoke_from,
)
self.dataset_tools = DatasetRetrieverTool.get_dataset_tools(
session=session,
tenant_id=tenant_id,
dataset_ids=app_config.dataset.dataset_ids if app_config.dataset else [],
retrieve_config=app_config.dataset.retrieve_config if app_config.dataset else None,
return_resource=(
app_config.additional_features.show_retrieve_source if app_config.additional_features else False
),
invoke_from=application_generate_entity.invoke_from,
hit_callback=hit_callback,
user_id=user_id,
inputs=cast(dict, application_generate_entity.inputs),
)
# get how many agent thoughts have been created
self.agent_thought_count = (
session.scalar(
select(func.count())
.select_from(MessageAgentThought)
.where(
MessageAgentThought.message_id == self.message.id,
)
)
or 0
)
session.close()
# check if model supports stream tool call
llm_model = cast(LargeLanguageModel, model_instance.model_type_instance)
model_schema = llm_model.get_model_schema(model_instance.model_name, model_instance.credentials)
features = model_schema.features if model_schema and model_schema.features else []
self.stream_tool_call = ModelFeature.STREAM_TOOL_CALL in features
self.vision_enabled = ModelFeature.VISION in features
self.files = application_generate_entity.files if self.vision_enabled else []
self.query: str = ""
self._current_thoughts: list[PromptMessage] = []
def _repack_app_generate_entity(
self, app_generate_entity: AgentChatAppGenerateEntity
) -> AgentChatAppGenerateEntity:
"""
Repack app generate entity
"""
if app_generate_entity.app_config.prompt_template.simple_prompt_template is None:
app_generate_entity.app_config.prompt_template.simple_prompt_template = ""
return app_generate_entity
def _convert_tool_to_prompt_message_tool(self, tool: AgentToolEntity) -> tuple[PromptMessageTool, Tool]:
"""
convert tool to prompt message tool
"""
tool_entity = ToolManager.get_agent_tool_runtime(
tenant_id=self.tenant_id,
app_id=self.app_config.app_id,
agent_tool=tool,
user_id=self.user_id,
invoke_from=self.application_generate_entity.invoke_from,
)
assert tool_entity.entity.description
message_tool = PromptMessageTool(
name=tool.tool_name,
description=tool_entity.entity.description.llm,
parameters=tool_entity.get_llm_parameters_json_schema(),
)
return message_tool, tool_entity
def _convert_dataset_retriever_tool_to_prompt_message_tool(self, tool: DatasetRetrieverTool) -> PromptMessageTool:
"""
convert dataset retriever tool to prompt message tool
"""
assert tool.entity.description
prompt_tool = PromptMessageTool(
name=tool.entity.identity.name,
description=tool.entity.description.llm,
parameters={
"type": "object",
"properties": {},
"required": [],
},
)
for parameter in tool.get_runtime_parameters():
parameter_type = "string"
prompt_tool.parameters["properties"][parameter.name] = {
"type": parameter_type,
"description": parameter.llm_description or "",
}
if parameter.required:
if parameter.name not in prompt_tool.parameters["required"]:
prompt_tool.parameters["required"].append(parameter.name)
return prompt_tool
def _init_prompt_tools(self) -> tuple[dict[str, Tool], list[PromptMessageTool]]:
"""
Init tools
"""
tool_instances = {}
prompt_messages_tools = []
for tool in self.app_config.agent.tools or [] if self.app_config.agent else []:
try:
prompt_tool, tool_entity = self._convert_tool_to_prompt_message_tool(tool)
except Exception:
# api tool may be deleted
continue
# save tool entity
tool_instances[tool.tool_name] = tool_entity
# save prompt tool
prompt_messages_tools.append(prompt_tool)
# convert dataset tools into ModelRuntime Tool format
for dataset_tool in self.dataset_tools:
prompt_tool = self._convert_dataset_retriever_tool_to_prompt_message_tool(dataset_tool)
# save prompt tool
prompt_messages_tools.append(prompt_tool)
# save tool entity
tool_instances[dataset_tool.entity.identity.name] = dataset_tool
return tool_instances, prompt_messages_tools
def update_prompt_message_tool(self, tool: Tool, prompt_tool: PromptMessageTool) -> PromptMessageTool:
"""
update prompt message tool
"""
prompt_tool.parameters = tool.get_llm_parameters_json_schema()
return prompt_tool
def create_agent_thought(
self, message_id: str, message: str, tool_name: str, tool_input: str, messages_ids: list[str]
) -> str:
"""
Create agent thought
"""
thought = MessageAgentThought(
message_id=message_id,
message_chain_id=None,
tool_process_data=None,
thought="",
tool=tool_name,
tool_labels_str="{}",
tool_meta_str="{}",
tool_input=tool_input,
message=message,
message_token=0,
message_unit_price=Decimal(0),
message_price_unit=Decimal("0.001"),
message_files=json.dumps(messages_ids) if messages_ids else "",
answer="",
observation="",
answer_token=0,
answer_unit_price=Decimal(0),
answer_price_unit=Decimal("0.001"),
tokens=0,
total_price=Decimal(0),
position=self.agent_thought_count + 1,
currency="USD",
latency=0,
created_by_role=CreatorUserRole.ACCOUNT,
created_by=self.user_id,
)
db.session.add(thought)
db.session.commit()
agent_thought_id = str(thought.id)
self.agent_thought_count += 1
db.session.close()
return agent_thought_id
def save_agent_thought(
self,
agent_thought_id: str,
tool_name: str | None,
tool_input: Union[str, dict, None],
thought: str | None,
observation: Union[str, dict, None],
tool_invoke_meta: Union[str, dict, None],
answer: str | None,
messages_ids: list[str],
llm_usage: LLMUsage | None = None,
):
"""
Save agent thought
"""
stmt = select(MessageAgentThought).where(MessageAgentThought.id == agent_thought_id)
agent_thought = db.session.scalar(stmt)
if not agent_thought:
raise ValueError("agent thought not found")
if thought:
existing_thought = agent_thought.thought or ""
agent_thought.thought = f"{existing_thought}{thought}"
if tool_name:
agent_thought.tool = tool_name
if tool_input:
if isinstance(tool_input, dict):
try:
tool_input = json.dumps(tool_input, ensure_ascii=False)
except Exception:
tool_input = json.dumps(tool_input)
agent_thought.tool_input = tool_input
if observation:
if isinstance(observation, dict):
try:
observation = json.dumps(observation, ensure_ascii=False)
except Exception:
observation = json.dumps(observation)
agent_thought.observation = observation
if answer:
agent_thought.answer = answer
if messages_ids is not None and len(messages_ids) > 0:
agent_thought.message_files = json.dumps(messages_ids)
if llm_usage:
agent_thought.message_token = llm_usage.prompt_tokens
agent_thought.message_price_unit = llm_usage.prompt_price_unit
agent_thought.message_unit_price = llm_usage.prompt_unit_price
agent_thought.answer_token = llm_usage.completion_tokens
agent_thought.answer_price_unit = llm_usage.completion_price_unit
agent_thought.answer_unit_price = llm_usage.completion_unit_price
agent_thought.tokens = llm_usage.total_tokens
agent_thought.total_price = llm_usage.total_price
# check if tool labels is not empty
labels = agent_thought.tool_labels or {}
tools = agent_thought.tool.split(";") if agent_thought.tool else []
for tool in tools:
if not tool:
continue
if tool not in labels:
tool_label = ToolManager.get_tool_label(tool)
if tool_label:
labels[tool] = tool_label.to_dict()
else:
labels[tool] = {"en_US": tool, "zh_Hans": tool}
agent_thought.tool_labels_str = json.dumps(labels)
if tool_invoke_meta is not None:
if isinstance(tool_invoke_meta, dict):
try:
tool_invoke_meta = json.dumps(tool_invoke_meta, ensure_ascii=False)
except Exception:
tool_invoke_meta = json.dumps(tool_invoke_meta)
agent_thought.tool_meta_str = tool_invoke_meta
db.session.commit()
db.session.close()
def organize_agent_history(self, prompt_messages: list[PromptMessage], *, session: Session) -> list[PromptMessage]:
"""
Organize agent history
"""
result: list[PromptMessage] = []
# check if there is a system message in the beginning of the conversation
for prompt_message in prompt_messages:
if isinstance(prompt_message, SystemPromptMessage):
result.append(prompt_message)
messages = (
(
session.execute(
select(Message)
.where(Message.conversation_id == self.message.conversation_id)
.order_by(Message.created_at.desc())
)
)
.scalars()
.all()
)
messages = list(reversed(extract_thread_messages(messages)))
for message in messages:
if message.id == self.message.id:
continue
result.append(self.organize_agent_user_prompt(message, session=session))
agent_thoughts = message.agent_thoughts_with_session(session=session)
if agent_thoughts:
for agent_thought in agent_thoughts:
tool_names_raw = agent_thought.tool
if tool_names_raw:
tool_names = tool_names_raw.split(";")
tool_calls: list[AssistantPromptMessage.ToolCall] = []
tool_call_response: list[ToolPromptMessage] = []
tool_input_payload = agent_thought.tool_input
if tool_input_payload:
try:
tool_inputs = json.loads(tool_input_payload)
except Exception:
tool_inputs = {tool: {} for tool in tool_names}
else:
tool_inputs = {tool: {} for tool in tool_names}
observation_payload = agent_thought.observation
if observation_payload:
try:
tool_responses = json.loads(observation_payload)
except Exception:
tool_responses = dict.fromkeys(tool_names, observation_payload)
else:
tool_responses = dict.fromkeys(tool_names, observation_payload)
tool_occurrences = Counter(tool_names)
seen_tools: Counter[str] = Counter()
for tool in tool_names:
occurrence = seen_tools[tool]
seen_tools[tool] += 1
# generate a uuid for tool call
tool_call_id = str(uuid.uuid4())
tool_calls.append(
AssistantPromptMessage.ToolCall(
id=tool_call_id,
type="function",
function=AssistantPromptMessage.ToolCall.ToolCallFunction(
name=tool,
arguments=json.dumps(
_select_tool_occurrence(
tool_inputs.get(tool, {}), occurrence, tool_occurrences[tool]
)
),
),
)
)
tool_call_response.append(
ToolPromptMessage(
content=_select_tool_occurrence(
tool_responses.get(tool, agent_thought.observation),
occurrence,
tool_occurrences[tool],
),
name=tool,
tool_call_id=tool_call_id,
)
)
result.extend(
[
AssistantPromptMessage(
content=agent_thought.thought,
tool_calls=tool_calls,
),
*tool_call_response,
]
)
if not tool_names_raw:
result.append(AssistantPromptMessage(content=agent_thought.thought))
else:
if message.answer:
result.append(AssistantPromptMessage(content=message.answer))
session.close()
return result
def organize_agent_user_prompt(self, message: Message, *, session: Session) -> UserPromptMessage:
stmt = select(MessageFile).where(MessageFile.message_id == message.id)
files = session.scalars(stmt).all()
if not files:
return UserPromptMessage(content=message.query)
app_model_config = message.app_model_config_with_session(session=session)
if app_model_config:
annotation_reply = load_annotation_reply_config(session, app_model_config.app_id)
file_extra_config = FileUploadConfigManager.convert(
app_model_config.to_dict(annotation_reply=annotation_reply)
)
else:
file_extra_config = None
if not file_extra_config:
return UserPromptMessage(content=message.query)
image_detail_config = file_extra_config.image_config.detail if file_extra_config.image_config else None
image_detail_config = image_detail_config or ImagePromptMessageContent.DETAIL.LOW
file_objs = file_factory.build_from_message_files(
message_files=files,
tenant_id=self.tenant_id,
access_controller=_file_access_controller,
)
if not file_objs:
return UserPromptMessage(content=message.query)
prompt_message_contents: list[PromptMessageContentUnionTypes] = []
for file in file_objs:
prompt_message_contents.append(
file_manager.to_prompt_message_content(
file,
image_detail_config=image_detail_config,
)
)
prompt_message_contents.append(TextPromptMessageContent(data=message.query))
return UserPromptMessage(content=prompt_message_contents)
Reference in New Issue View Git Blame Copy Permalink