jprdonnelly/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-06-02 16:00:54 -04:00
Code Issues Packages Projects Releases Wiki Activity

feat: extract model runtime

Browse Source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN-
2026-03-15 15:34:47 +08:00
parent 3d5a29462e
commit fbb74a4af9
178 changed files with 4343 additions and 2134 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
api/dify_graph/entities/graph_init_params.py
View File

@@ -3,8 +3,6 @@ from typing import Any
from pydantic import BaseModel, Field
DIFY_RUN_CONTEXT_KEY = "_dify"
class GraphInitParams(BaseModel):
"""GraphInitParams encapsulates the configurations and contextual information

2
api/dify_graph/entities/workflow_execution.py
View File

@@ -14,7 +14,7 @@ from typing import Any
from pydantic import BaseModel, Field
from dify_graph.enums import WorkflowExecutionStatus, WorkflowType
from libs.datetime_utils import naive_utc_now
from dify_graph.utils.datetime_utils import naive_utc_now
class WorkflowExecution(BaseModel):

5
api/dify_graph/graph/graph.py
View File

@@ -10,7 +10,6 @@ from pydantic import TypeAdapter
from dify_graph.entities.graph_config import NodeConfigDict
from dify_graph.enums import ErrorStrategy, NodeExecutionType, NodeState
from dify_graph.nodes.base.node import Node
from libs.typing import is_str
from .edge import Edge
from .validation import get_graph_validator
@@ -102,7 +101,7 @@ class Graph:
source = edge_config.get("source")
target = edge_config.get("target")
if not is_str(source) or not is_str(target):
if not isinstance(source, str) or not isinstance(target, str):
continue
# Create edge
@@ -110,7 +109,7 @@ class Graph:
edge_counter += 1
source_handle = edge_config.get("sourceHandle", "source")
if not is_str(source_handle):
if not isinstance(source_handle, str):
continue
edge = Edge(

6
api/dify_graph/graph_events/node.py
View File

@@ -1,9 +1,9 @@
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Any
from pydantic import Field
from core.rag.entities.citation_metadata import RetrievalSourceMetadata
from dify_graph.entities.pause_reason import PauseReason
from .base import GraphNodeEventBase
@@ -30,7 +30,7 @@ class NodeRunStreamChunkEvent(GraphNodeEventBase):
class NodeRunRetrieverResourceEvent(GraphNodeEventBase):
retriever_resources: Sequence[RetrievalSourceMetadata] = Field(..., description="retriever resources")
retriever_resources: Sequence[Mapping[str, Any]] = Field(..., description="retriever resources")
context: str = Field(..., description="context")

14
api/dify_graph/model_runtime/entities/provider_entities.py
View File

@@ -93,10 +93,14 @@ class ModelCredentialSchema(BaseModel):
class SimpleProviderEntity(BaseModel):
"""
Simple model class for provider.
Simplified provider schema exposed to callers.
`provider` is the canonical runtime identifier. `provider_name` is an optional
compatibility alias for short-name lookups and is empty when no alias exists.
"""
provider: str
provider_name: str = ""
label: I18nObject
icon_small: I18nObject | None = None
icon_small_dark: I18nObject | None = None
@@ -115,10 +119,15 @@ class ProviderHelpEntity(BaseModel):
class ProviderEntity(BaseModel):
"""
Model class for provider.
Runtime-native provider schema.
`provider` is the canonical runtime identifier. `provider_name` is a
compatibility alias for callers that still resolve providers by short name and
is empty when no alias exists.
"""
provider: str
provider_name: str = ""
label: I18nObject
description: I18nObject | None = None
icon_small: I18nObject | None = None
@@ -153,6 +162,7 @@ class ProviderEntity(BaseModel):
"""
return SimpleProviderEntity(
provider=self.provider,
provider_name=self.provider_name,
label=self.label,
icon_small=self.icon_small,
supported_model_types=self.supported_model_types,

7
api/dify_graph/model_runtime/entities/rerank_entities.py
View File

@@ -1,6 +1,13 @@
from typing import TypedDict
from pydantic import BaseModel
class MultimodalRerankInput(TypedDict):
content: str
content_type: str
class RerankDocument(BaseModel):
"""
Model class for rerank document.

8
api/dify_graph/model_runtime/entities/text_embedding_entities.py
View File

@@ -1,10 +1,18 @@
from decimal import Decimal
from enum import StrEnum, auto
from pydantic import BaseModel
from dify_graph.model_runtime.entities.model_entities import ModelUsage
class EmbeddingInputType(StrEnum):
"""Embedding request input variants understood by the model runtime."""
DOCUMENT = auto()
QUERY = auto()
class EmbeddingUsage(ModelUsage):
"""
Model class for embedding usage.

133
api/dify_graph/model_runtime/model_providers/__base/ai_model.py
View File

@@ -1,12 +1,5 @@
import decimal
import hashlib
import logging
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from redis import RedisError
from configs import dify_config
from core.plugin.entities.plugin_daemon import PluginModelProviderEntity
from dify_graph.model_runtime.entities.common_entities import I18nObject
from dify_graph.model_runtime.entities.defaults import PARAMETER_RULE_TEMPLATE
from dify_graph.model_runtime.entities.model_entities import (
@@ -17,6 +10,7 @@ from dify_graph.model_runtime.entities.model_entities import (
PriceInfo,
PriceType,
)
from dify_graph.model_runtime.entities.provider_entities import ProviderEntity
from dify_graph.model_runtime.errors.invoke import (
InvokeAuthorizationError,
InvokeBadRequestError,
@@ -25,45 +19,61 @@ from dify_graph.model_runtime.errors.invoke import (
InvokeRateLimitError,
InvokeServerUnavailableError,
)
from extensions.ext_redis import redis_client
logger = logging.getLogger(__name__)
from dify_graph.model_runtime.runtime import ModelRuntime
class AIModel(BaseModel):
class AIModel:
"""
Base class for all models.
Runtime-facing base class for all model providers.
This stays a regular Python class because instances hold live collaborators
such as the provider schema and runtime adapter rather than user input that
benefits from Pydantic validation. Subclasses must pin ``model_type`` via a
class attribute; the base class is not meant to be instantiated directly.
"""
tenant_id: str = Field(description="Tenant ID")
model_type: ModelType = Field(description="Model type")
plugin_id: str = Field(description="Plugin ID")
provider_name: str = Field(description="Provider")
plugin_model_provider: PluginModelProviderEntity = Field(description="Plugin model provider")
started_at: float = Field(description="Invoke start time", default=0)
model_type: ModelType
provider_schema: ProviderEntity
model_runtime: ModelRuntime
started_at: float
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
def __init__(
self,
provider_schema: ProviderEntity,
model_runtime: ModelRuntime,
*,
started_at: float = 0,
) -> None:
if getattr(type(self), "model_type", None) is None:
raise TypeError("AIModel subclasses must define model_type as a class attribute")
self.model_type = type(self).model_type
self.provider_schema = provider_schema
self.model_runtime = model_runtime
self.started_at = started_at
@property
def provider(self) -> str:
return self.provider_schema.provider
@property
def provider_display_name(self) -> str:
return self.provider_schema.label.en_US
@property
def _invoke_error_mapping(self) -> dict[type[Exception], list[type[Exception]]]:
"""
Map model invoke error to unified error
The key is the error type thrown to the caller
The value is the error type thrown by the model,
which needs to be converted into a unified error type for the caller.
Map model invoke error to unified error.
:return: Invoke error mapping
The key is the error type thrown to the caller, and the value contains
runtime-facing exception types that should be normalized to it.
"""
from core.plugin.entities.plugin_daemon import PluginDaemonInnerError
return {
InvokeConnectionError: [InvokeConnectionError],
InvokeServerUnavailableError: [InvokeServerUnavailableError],
InvokeRateLimitError: [InvokeRateLimitError],
InvokeAuthorizationError: [InvokeAuthorizationError],
InvokeBadRequestError: [InvokeBadRequestError],
PluginDaemonInnerError: [PluginDaemonInnerError],
ValueError: [ValueError],
}
@@ -79,15 +89,18 @@ class AIModel(BaseModel):
if invoke_error == InvokeAuthorizationError:
return InvokeAuthorizationError(
description=(
f"[{self.provider_name}] Incorrect model credentials provided, please check and try again."
f"[{self.provider_display_name}] Incorrect model credentials provided, "
"please check and try again."
)
)
elif isinstance(invoke_error, InvokeError):
return InvokeError(description=f"[{self.provider_name}] {invoke_error.description}, {str(error)}")
return InvokeError(
description=f"[{self.provider_display_name}] {invoke_error.description}, {str(error)}"
)
else:
return error
return InvokeError(description=f"[{self.provider_name}] Error: {str(error)}")
return InvokeError(description=f"[{self.provider_display_name}] Error: {str(error)}")
def get_price(self, model: str, credentials: dict, price_type: PriceType, tokens: int) -> PriceInfo:
"""
@@ -144,65 +157,13 @@ class AIModel(BaseModel):
:param credentials: model credentials
:return: model schema
"""
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
cache_key = f"{self.tenant_id}:{self.plugin_id}:{self.provider_name}:{self.model_type.value}:{model}"
sorted_credentials = sorted(credentials.items()) if credentials else []
cache_key += ":".join([hashlib.md5(f"{k}:{v}".encode()).hexdigest() for k, v in sorted_credentials])
cached_schema_json = None
try:
cached_schema_json = redis_client.get(cache_key)
except (RedisError, RuntimeError) as exc:
logger.warning(
"Failed to read plugin model schema cache for model %s: %s",
model,
str(exc),
exc_info=True,
)
if cached_schema_json:
try:
return AIModelEntity.model_validate_json(cached_schema_json)
except ValidationError:
logger.warning(
"Failed to validate cached plugin model schema for model %s",
model,
exc_info=True,
)
try:
redis_client.delete(cache_key)
except (RedisError, RuntimeError) as exc:
logger.warning(
"Failed to delete invalid plugin model schema cache for model %s: %s",
model,
str(exc),
exc_info=True,
)
schema = plugin_model_manager.get_model_schema(
tenant_id=self.tenant_id,
user_id="unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
model_type=self.model_type.value,
return self.model_runtime.get_model_schema(
provider=self.provider,
model_type=self.model_type,
model=model,
credentials=credentials or {},
)
if schema:
try:
redis_client.setex(cache_key, dify_config.PLUGIN_MODEL_SCHEMA_CACHE_TTL, schema.model_dump_json())
except (RedisError, RuntimeError) as exc:
logger.warning(
"Failed to write plugin model schema cache for model %s: %s",
model,
str(exc),
exc_info=True,
)
return schema
def get_customizable_model_schema_from_credentials(self, model: str, credentials: dict) -> AIModelEntity | None:
"""
Get customizable model schema from credentials

66
api/dify_graph/model_runtime/model_providers/__base/large_language_model.py
View File

@@ -4,9 +4,6 @@ import uuid
from collections.abc import Callable, Generator, Iterator, Sequence
from typing import Union
from pydantic import ConfigDict
from configs import dify_config
from dify_graph.model_runtime.callbacks.base_callback import Callback
from dify_graph.model_runtime.callbacks.logging_callback import LoggingCallback
from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMUsage
@@ -140,11 +137,9 @@ def _build_llm_result_from_chunks(
)
def _invoke_llm_via_plugin(
def _invoke_llm_via_runtime(
*,
tenant_id: str,
user_id: str,
plugin_id: str,
llm_model: "LargeLanguageModel",
provider: str,
model: str,
credentials: dict,
@@ -154,25 +149,19 @@ def _invoke_llm_via_plugin(
stop: Sequence[str] | None,
stream: bool,
) -> Union[LLMResult, Generator[LLMResultChunk, None, None]]:
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.invoke_llm(
tenant_id=tenant_id,
user_id=user_id,
plugin_id=plugin_id,
return llm_model.model_runtime.invoke_llm(
provider=provider,
model=model,
credentials=credentials,
model_parameters=model_parameters,
prompt_messages=list(prompt_messages),
tools=tools,
stop=list(stop) if stop else None,
stop=stop,
stream=stream,
)
def _normalize_non_stream_plugin_result(
def _normalize_non_stream_runtime_result(
model: str,
prompt_messages: Sequence[PromptMessage],
result: Union[LLMResult, Iterator[LLMResultChunk]],
@@ -208,9 +197,6 @@ class LargeLanguageModel(AIModel):
model_type: ModelType = ModelType.LLM
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
def invoke(
self,
model: str,
@@ -220,7 +206,6 @@ class LargeLanguageModel(AIModel):
tools: list[PromptMessageTool] | None = None,
stop: list[str] | None = None,
stream: bool = True,
user: str | None = None,
callbacks: list[Callback] | None = None,
) -> Union[LLMResult, Generator[LLMResultChunk, None, None]]:
"""
@@ -233,7 +218,6 @@ class LargeLanguageModel(AIModel):
:param tools: tools for tool calling
:param stop: stop words
:param stream: is stream response
:param user: unique user id
:param callbacks: callbacks
:return: full response or stream response chunk generator result
"""
@@ -245,7 +229,7 @@ class LargeLanguageModel(AIModel):
callbacks = callbacks or []
if dify_config.DEBUG:
if logger.isEnabledFor(logging.DEBUG):
callbacks.append(LoggingCallback())
# trigger before invoke callbacks
@@ -257,18 +241,15 @@ class LargeLanguageModel(AIModel):
tools=tools,
stop=stop,
stream=stream,
user=user,
callbacks=callbacks,
)
result: Union[LLMResult, Generator[LLMResultChunk, None, None]]
try:
result = _invoke_llm_via_plugin(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
result = _invoke_llm_via_runtime(
llm_model=self,
provider=self.provider,
model=model,
credentials=credentials,
model_parameters=model_parameters,
@@ -279,7 +260,7 @@ class LargeLanguageModel(AIModel):
)
if not stream:
result = _normalize_non_stream_plugin_result(
result = _normalize_non_stream_runtime_result(
model=model, prompt_messages=prompt_messages, result=result
)
except Exception as e:
@@ -292,7 +273,6 @@ class LargeLanguageModel(AIModel):
tools=tools,
stop=stop,
stream=stream,
user=user,
callbacks=callbacks,
)
@@ -309,7 +289,6 @@ class LargeLanguageModel(AIModel):
tools=tools,
stop=stop,
stream=stream,
user=user,
callbacks=callbacks,
)
elif isinstance(result, LLMResult):
@@ -322,7 +301,6 @@ class LargeLanguageModel(AIModel):
tools=tools,
stop=stop,
stream=stream,
user=user,
callbacks=callbacks,
)
# Following https://github.com/langgenius/dify/issues/17799,
@@ -435,22 +413,14 @@ class LargeLanguageModel(AIModel):
:param tools: tools for tool calling
:return:
"""
if dify_config.PLUGIN_BASED_TOKEN_COUNTING_ENABLED:
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.get_llm_num_tokens(
tenant_id=self.tenant_id,
user_id="unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
model_type=self.model_type.value,
model=model,
credentials=credentials,
prompt_messages=prompt_messages,
tools=tools,
)
return 0
return self.model_runtime.get_llm_num_tokens(
provider=self.provider,
model_type=self.model_type,
model=model,
credentials=credentials,
prompt_messages=prompt_messages,
tools=tools,
)
def calc_response_usage(
self, model: str, credentials: dict, prompt_tokens: int, completion_tokens: int

18
api/dify_graph/model_runtime/model_providers/__base/moderation_model.py
View File

@@ -1,7 +1,5 @@
import time
from pydantic import ConfigDict
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.model_providers.__base.ai_model import AIModel
@@ -13,30 +11,20 @@ class ModerationModel(AIModel):
model_type: ModelType = ModelType.MODERATION
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
def invoke(self, model: str, credentials: dict, text: str, user: str | None = None) -> bool:
def invoke(self, model: str, credentials: dict, text: str) -> bool:
"""
Invoke moderation model
:param model: model name
:param credentials: model credentials
:param text: text to moderate
:param user: unique user id
:return: false if text is safe, true otherwise
"""
self.started_at = time.perf_counter()
try:
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.invoke_moderation(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.invoke_moderation(
provider=self.provider,
model=model,
credentials=credentials,
text=text,

30
api/dify_graph/model_runtime/model_providers/__base/rerank_model.py
View File

@@ -1,5 +1,5 @@
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.entities.rerank_entities import RerankResult
from dify_graph.model_runtime.entities.rerank_entities import MultimodalRerankInput, RerankResult
from dify_graph.model_runtime.model_providers.__base.ai_model import AIModel
@@ -18,7 +18,6 @@ class RerankModel(AIModel):
docs: list[str],
score_threshold: float | None = None,
top_n: int | None = None,
user: str | None = None,
) -> RerankResult:
"""
Invoke rerank model
@@ -29,18 +28,11 @@ class RerankModel(AIModel):
:param docs: docs for reranking
:param score_threshold: score threshold
:param top_n: top n
:param user: unique user id
:return: rerank result
"""
try:
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.invoke_rerank(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.invoke_rerank(
provider=self.provider,
model=model,
credentials=credentials,
query=query,
@@ -55,11 +47,10 @@ class RerankModel(AIModel):
self,
model: str,
credentials: dict,
query: dict,
docs: list[dict],
query: MultimodalRerankInput,
docs: list[MultimodalRerankInput],
score_threshold: float | None = None,
top_n: int | None = None,
user: str | None = None,
) -> RerankResult:
"""
Invoke multimodal rerank model
@@ -69,18 +60,11 @@ class RerankModel(AIModel):
:param docs: docs for reranking
:param score_threshold: score threshold
:param top_n: top n
:param user: unique user id
:return: rerank result
"""
try:
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.invoke_multimodal_rerank(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.invoke_multimodal_rerank(
provider=self.provider,
model=model,
credentials=credentials,
query=query,

18
api/dify_graph/model_runtime/model_providers/__base/speech2text_model.py
View File

@@ -1,7 +1,5 @@
from typing import IO
from pydantic import ConfigDict
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.model_providers.__base.ai_model import AIModel
@@ -13,28 +11,18 @@ class Speech2TextModel(AIModel):
model_type: ModelType = ModelType.SPEECH2TEXT
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
def invoke(self, model: str, credentials: dict, file: IO[bytes], user: str | None = None) -> str:
def invoke(self, model: str, credentials: dict, file: IO[bytes]) -> str:
"""
Invoke speech to text model
:param model: model name
:param credentials: model credentials
:param file: audio file
:param user: unique user id
:return: text for given audio file
"""
try:
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.invoke_speech_to_text(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.invoke_speech_to_text(
provider=self.provider,
model=model,
credentials=credentials,
file=file,

37
api/dify_graph/model_runtime/model_providers/__base/text_embedding_model.py
View File

@@ -1,8 +1,5 @@
from pydantic import ConfigDict
from core.entities.embedding_type import EmbeddingInputType
from dify_graph.model_runtime.entities.model_entities import ModelPropertyKey, ModelType
from dify_graph.model_runtime.entities.text_embedding_entities import EmbeddingResult
from dify_graph.model_runtime.entities.text_embedding_entities import EmbeddingInputType, EmbeddingResult
from dify_graph.model_runtime.model_providers.__base.ai_model import AIModel
@@ -13,16 +10,12 @@ class TextEmbeddingModel(AIModel):
model_type: ModelType = ModelType.TEXT_EMBEDDING
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
def invoke(
self,
model: str,
credentials: dict,
texts: list[str] | None = None,
multimodel_documents: list[dict] | None = None,
user: str | None = None,
input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
) -> EmbeddingResult:
"""
@@ -32,31 +25,21 @@ class TextEmbeddingModel(AIModel):
:param credentials: model credentials
:param texts: texts to embed
:param files: files to embed
:param user: unique user id
:param input_type: input type
:return: embeddings result
"""
from core.plugin.impl.model import PluginModelClient
try:
plugin_model_manager = PluginModelClient()
if texts:
return plugin_model_manager.invoke_text_embedding(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.invoke_text_embedding(
provider=self.provider,
model=model,
credentials=credentials,
texts=texts,
input_type=input_type,
)
if multimodel_documents:
return plugin_model_manager.invoke_multimodal_embedding(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.invoke_multimodal_embedding(
provider=self.provider,
model=model,
credentials=credentials,
documents=multimodel_documents,
@@ -75,14 +58,8 @@ class TextEmbeddingModel(AIModel):
:param texts: texts to embed
:return:
"""
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.get_text_embedding_num_tokens(
tenant_id=self.tenant_id,
user_id="unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.get_text_embedding_num_tokens(
provider=self.provider,
model=model,
credentials=credentials,
texts=texts,

29
api/dify_graph/model_runtime/model_providers/__base/tts_model.py
View File

@@ -1,8 +1,6 @@
import logging
from collections.abc import Iterable
from pydantic import ConfigDict
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.model_providers.__base.ai_model import AIModel
@@ -16,38 +14,25 @@ class TTSModel(AIModel):
model_type: ModelType = ModelType.TTS
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
def invoke(
self,
model: str,
tenant_id: str,
credentials: dict,
content_text: str,
voice: str,
user: str | None = None,
) -> Iterable[bytes]:
"""
Invoke large language model
:param model: model name
:param tenant_id: user tenant id
:param credentials: model credentials
:param voice: model timbre
:param content_text: text content to be translated
:param user: unique user id
:return: translated audio file
"""
try:
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.invoke_tts(
tenant_id=self.tenant_id,
user_id=user or "unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.invoke_tts(
provider=self.provider,
model=model,
credentials=credentials,
content_text=content_text,
@@ -65,14 +50,8 @@ class TTSModel(AIModel):
:param credentials: The credentials required to access the TTS model.
:return: A list of voices supported by the TTS model.
"""
from core.plugin.impl.model import PluginModelClient
plugin_model_manager = PluginModelClient()
return plugin_model_manager.get_tts_model_voices(
tenant_id=self.tenant_id,
user_id="unknown",
plugin_id=self.plugin_id,
provider=self.provider_name,
return self.model_runtime.get_tts_model_voices(
provider=self.provider,
model=model,
credentials=credentials,
language=language,

344
api/dify_graph/model_runtime/model_providers/model_provider_factory.py
View File

@@ -1,16 +1,7 @@
from __future__ import annotations
import hashlib
import logging
from collections.abc import Sequence
from threading import Lock
from pydantic import ValidationError
from redis import RedisError
import contexts
from configs import dify_config
from core.plugin.entities.plugin_daemon import PluginModelProviderEntity
from dify_graph.model_runtime.entities.model_entities import AIModelEntity, ModelType
from dify_graph.model_runtime.entities.provider_entities import ProviderConfig, ProviderEntity, SimpleProviderEntity
from dify_graph.model_runtime.model_providers.__base.ai_model import AIModel
@@ -20,120 +11,64 @@ from dify_graph.model_runtime.model_providers.__base.rerank_model import RerankM
from dify_graph.model_runtime.model_providers.__base.speech2text_model import Speech2TextModel
from dify_graph.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
from dify_graph.model_runtime.model_providers.__base.tts_model import TTSModel
from dify_graph.model_runtime.runtime import ModelRuntime
from dify_graph.model_runtime.schema_validators.model_credential_schema_validator import ModelCredentialSchemaValidator
from dify_graph.model_runtime.schema_validators.provider_credential_schema_validator import (
ProviderCredentialSchemaValidator,
)
from extensions.ext_redis import redis_client
from models.provider_ids import ModelProviderID
logger = logging.getLogger(__name__)
class ModelProviderFactory:
def __init__(self, tenant_id: str):
from core.plugin.impl.model import PluginModelClient
"""Factory for provider schemas and model-type instances backed by a runtime adapter."""
self.tenant_id = tenant_id
self.plugin_model_manager = PluginModelClient()
def __init__(self, model_runtime: ModelRuntime):
if model_runtime is None:
raise ValueError("model_runtime is required.")
self.model_runtime = model_runtime
def get_providers(self) -> Sequence[ProviderEntity]:
"""
Get all providers
:return: list of providers
Get all providers.
"""
# FIXME(-LAN-): Removed position map sorting since providers are fetched from plugin server
# The plugin server should return providers in the desired order
plugin_providers = self.get_plugin_model_providers()
return [provider.declaration for provider in plugin_providers]
return list(self.get_model_providers())
def get_plugin_model_providers(self) -> Sequence[PluginModelProviderEntity]:
def get_model_providers(self) -> Sequence[ProviderEntity]:
"""
Get all plugin model providers
:return: list of plugin model providers
Get all model providers exposed by the runtime adapter.
"""
# check if context is set
try:
contexts.plugin_model_providers.get()
except LookupError:
contexts.plugin_model_providers.set(None)
contexts.plugin_model_providers_lock.set(Lock())
with contexts.plugin_model_providers_lock.get():
plugin_model_providers = contexts.plugin_model_providers.get()
if plugin_model_providers is not None:
return plugin_model_providers
plugin_model_providers = []
contexts.plugin_model_providers.set(plugin_model_providers)
# Fetch plugin model providers
plugin_providers = self.plugin_model_manager.fetch_model_providers(self.tenant_id)
for provider in plugin_providers:
provider.declaration.provider = provider.plugin_id + "/" + provider.declaration.provider
plugin_model_providers.append(provider)
return plugin_model_providers
return self.model_runtime.fetch_model_providers()
def get_provider_schema(self, provider: str) -> ProviderEntity:
"""
Get provider schema
:param provider: provider name
:return: provider schema
Get provider schema.
"""
plugin_model_provider_entity = self.get_plugin_model_provider(provider=provider)
return plugin_model_provider_entity.declaration
return self.get_model_provider(provider=provider)
def get_plugin_model_provider(self, provider: str) -> PluginModelProviderEntity:
def get_model_provider(self, provider: str) -> ProviderEntity:
"""
Get plugin model provider
:param provider: provider name
:return: provider schema
Get provider schema.
"""
if "/" not in provider:
provider = str(ModelProviderID(provider))
# fetch plugin model providers
plugin_model_provider_entities = self.get_plugin_model_providers()
# get the provider
plugin_model_provider_entity = next(
(p for p in plugin_model_provider_entities if p.declaration.provider == provider),
None,
)
if not plugin_model_provider_entity:
provider_entity = self._resolve_provider(provider)
if provider_entity is None:
raise ValueError(f"Invalid provider: {provider}")
return plugin_model_provider_entity
return provider_entity
def provider_credentials_validate(self, *, provider: str, credentials: dict):
"""
Validate provider credentials
:param provider: provider name
:param credentials: provider credentials, credentials form defined in `provider_credential_schema`.
:return:
Validate provider credentials.
"""
# fetch plugin model provider
plugin_model_provider_entity = self.get_plugin_model_provider(provider=provider)
provider_entity = self.get_model_provider(provider=provider)
# get provider_credential_schema and validate credentials according to the rules
provider_credential_schema = plugin_model_provider_entity.declaration.provider_credential_schema
provider_credential_schema = provider_entity.provider_credential_schema
if not provider_credential_schema:
raise ValueError(f"Provider {provider} does not have provider_credential_schema")
# validate provider credential schema
validator = ProviderCredentialSchemaValidator(provider_credential_schema)
filtered_credentials = validator.validate_and_filter(credentials)
# validate the credentials, raise exception if validation failed
self.plugin_model_manager.validate_provider_credentials(
tenant_id=self.tenant_id,
user_id="unknown",
plugin_id=plugin_model_provider_entity.plugin_id,
provider=plugin_model_provider_entity.provider,
self.model_runtime.validate_provider_credentials(
provider=provider_entity.provider,
credentials=filtered_credentials,
)
@@ -141,33 +76,20 @@ class ModelProviderFactory:
def model_credentials_validate(self, *, provider: str, model_type: ModelType, model: str, credentials: dict):
"""
Validate model credentials
:param provider: provider name
:param model_type: model type
:param model: model name
:param credentials: model credentials, credentials form defined in `model_credential_schema`.
:return:
Validate model credentials.
"""
# fetch plugin model provider
plugin_model_provider_entity = self.get_plugin_model_provider(provider=provider)
provider_entity = self.get_model_provider(provider=provider)
# get model_credential_schema and validate credentials according to the rules
model_credential_schema = plugin_model_provider_entity.declaration.model_credential_schema
model_credential_schema = provider_entity.model_credential_schema
if not model_credential_schema:
raise ValueError(f"Provider {provider} does not have model_credential_schema")
# validate model credential schema
validator = ModelCredentialSchemaValidator(model_type, model_credential_schema)
filtered_credentials = validator.validate_and_filter(credentials)
# call validate_credentials method of model type to validate credentials, raise exception if validation failed
self.plugin_model_manager.validate_model_credentials(
tenant_id=self.tenant_id,
user_id="unknown",
plugin_id=plugin_model_provider_entity.plugin_id,
provider=plugin_model_provider_entity.provider,
model_type=model_type.value,
self.model_runtime.validate_model_credentials(
provider=provider_entity.provider,
model_type=model_type,
model=model,
credentials=filtered_credentials,
)
@@ -178,65 +100,16 @@ class ModelProviderFactory:
self, *, provider: str, model_type: ModelType, model: str, credentials: dict | None
) -> AIModelEntity | None:
"""
Get model schema
Get model schema.
"""
plugin_id, provider_name = self.get_plugin_id_and_provider_name_from_provider(provider)
cache_key = f"{self.tenant_id}:{plugin_id}:{provider_name}:{model_type.value}:{model}"
sorted_credentials = sorted(credentials.items()) if credentials else []
cache_key += ":".join([hashlib.md5(f"{k}:{v}".encode()).hexdigest() for k, v in sorted_credentials])
cached_schema_json = None
try:
cached_schema_json = redis_client.get(cache_key)
except (RedisError, RuntimeError) as exc:
logger.warning(
"Failed to read plugin model schema cache for model %s: %s",
model,
str(exc),
exc_info=True,
)
if cached_schema_json:
try:
return AIModelEntity.model_validate_json(cached_schema_json)
except ValidationError:
logger.warning(
"Failed to validate cached plugin model schema for model %s",
model,
exc_info=True,
)
try:
redis_client.delete(cache_key)
except (RedisError, RuntimeError) as exc:
logger.warning(
"Failed to delete invalid plugin model schema cache for model %s: %s",
model,
str(exc),
exc_info=True,
)
schema = self.plugin_model_manager.get_model_schema(
tenant_id=self.tenant_id,
user_id="unknown",
plugin_id=plugin_id,
provider=provider_name,
model_type=model_type.value,
provider_entity = self.get_model_provider(provider)
return self.model_runtime.get_model_schema(
provider=provider_entity.provider,
model_type=model_type,
model=model,
credentials=credentials or {},
)
if schema:
try:
redis_client.setex(cache_key, dify_config.PLUGIN_MODEL_SCHEMA_CACHE_TTL, schema.model_dump_json())
except (RedisError, RuntimeError) as exc:
logger.warning(
"Failed to write plugin model schema cache for model %s: %s",
model,
str(exc),
exc_info=True,
)
return schema
def get_models(
self,
*,
@@ -245,143 +118,56 @@ class ModelProviderFactory:
provider_configs: list[ProviderConfig] | None = None,
) -> list[SimpleProviderEntity]:
"""
Get all models for given model type
:param provider: provider name
:param model_type: model type
:param provider_configs: list of provider configs
:return: list of models
Get all models for given model type.
"""
provider_configs = provider_configs or []
# scan all providers
plugin_model_provider_entities = self.get_plugin_model_providers()
# traverse all model_provider_extensions
providers = []
for plugin_model_provider_entity in plugin_model_provider_entities:
# filter by provider if provider is present
if provider and plugin_model_provider_entity.declaration.provider != provider:
for provider_entity in self.get_model_providers():
if provider and not self._matches_provider(provider_entity, provider):
continue
# get provider schema
provider_schema = plugin_model_provider_entity.declaration
model_types = provider_schema.supported_model_types
if model_type:
if model_type not in model_types:
continue
model_types = [model_type]
all_model_type_models = []
for model_schema in provider_schema.models:
if model_schema.model_type != model_type:
continue
all_model_type_models.append(model_schema)
simple_provider_schema = provider_schema.to_simple_provider()
if model_type:
simple_provider_schema.models = all_model_type_models
if model_type and model_type not in provider_entity.supported_model_types:
continue
simple_provider_schema = provider_entity.to_simple_provider()
if model_type is not None:
simple_provider_schema.models = [
model_schema for model_schema in provider_entity.models if model_schema.model_type == model_type
]
providers.append(simple_provider_schema)
return providers
def get_model_type_instance(self, provider: str, model_type: ModelType) -> AIModel:
"""
Get model type instance by provider name and model type
:param provider: provider name
:param model_type: model type
:return: model type instance
Get model type instance by provider name and model type.
"""
plugin_id, provider_name = self.get_plugin_id_and_provider_name_from_provider(provider)
init_params = {
"tenant_id": self.tenant_id,
"plugin_id": plugin_id,
"provider_name": provider_name,
"plugin_model_provider": self.get_plugin_model_provider(provider),
}
provider_schema = self.get_model_provider(provider)
if model_type == ModelType.LLM:
return LargeLanguageModel.model_validate(init_params)
elif model_type == ModelType.TEXT_EMBEDDING:
return TextEmbeddingModel.model_validate(init_params)
elif model_type == ModelType.RERANK:
return RerankModel.model_validate(init_params)
elif model_type == ModelType.SPEECH2TEXT:
return Speech2TextModel.model_validate(init_params)
elif model_type == ModelType.MODERATION:
return ModerationModel.model_validate(init_params)
elif model_type == ModelType.TTS:
return TTSModel.model_validate(init_params)
return LargeLanguageModel(provider_schema=provider_schema, model_runtime=self.model_runtime)
if model_type == ModelType.TEXT_EMBEDDING:
return TextEmbeddingModel(provider_schema=provider_schema, model_runtime=self.model_runtime)
if model_type == ModelType.RERANK:
return RerankModel(provider_schema=provider_schema, model_runtime=self.model_runtime)
if model_type == ModelType.SPEECH2TEXT:
return Speech2TextModel(provider_schema=provider_schema, model_runtime=self.model_runtime)
if model_type == ModelType.MODERATION:
return ModerationModel(provider_schema=provider_schema, model_runtime=self.model_runtime)
if model_type == ModelType.TTS:
return TTSModel(provider_schema=provider_schema, model_runtime=self.model_runtime)
raise ValueError(f"Unsupported model type: {model_type}")
def get_provider_icon(self, provider: str, icon_type: str, lang: str) -> tuple[bytes, str]:
"""
Get provider icon
:param provider: provider name
:param icon_type: icon type (icon_small or icon_small_dark)
:param lang: language (zh_Hans or en_US)
:return: provider icon
Get provider icon.
"""
# get the provider schema
provider_schema = self.get_provider_schema(provider)
provider_entity = self.get_model_provider(provider)
return self.model_runtime.get_provider_icon(provider=provider_entity.provider, icon_type=icon_type, lang=lang)
if icon_type.lower() == "icon_small":
if not provider_schema.icon_small:
raise ValueError(f"Provider {provider} does not have small icon.")
def _resolve_provider(self, provider: str) -> ProviderEntity | None:
return next((item for item in self.get_model_providers() if self._matches_provider(item, provider)), None)
if lang.lower() == "zh_hans":
file_name = provider_schema.icon_small.zh_Hans
else:
file_name = provider_schema.icon_small.en_US
elif icon_type.lower() == "icon_small_dark":
if not provider_schema.icon_small_dark:
raise ValueError(f"Provider {provider} does not have small dark icon.")
if lang.lower() == "zh_hans":
file_name = provider_schema.icon_small_dark.zh_Hans
else:
file_name = provider_schema.icon_small_dark.en_US
else:
raise ValueError(f"Unsupported icon type: {icon_type}.")
if not file_name:
raise ValueError(f"Provider {provider} does not have icon.")
image_mime_types = {
"jpg": "image/jpeg",
"jpeg": "image/jpeg",
"png": "image/png",
"gif": "image/gif",
"bmp": "image/bmp",
"tiff": "image/tiff",
"tif": "image/tiff",
"webp": "image/webp",
"svg": "image/svg+xml",
"ico": "image/vnd.microsoft.icon",
"heif": "image/heif",
"heic": "image/heic",
}
extension = file_name.split(".")[-1]
mime_type = image_mime_types.get(extension, "image/png")
# get icon bytes from plugin asset manager
from core.plugin.impl.asset import PluginAssetManager
plugin_asset_manager = PluginAssetManager()
return plugin_asset_manager.fetch_asset(tenant_id=self.tenant_id, id=file_name), mime_type
def get_plugin_id_and_provider_name_from_provider(self, provider: str) -> tuple[str, str]:
"""
Get plugin id and provider name from provider name
:param provider: provider name
:return: plugin id and provider name
"""
provider_id = ModelProviderID(provider)
return provider_id.plugin_id, provider_id.provider_name
@staticmethod
def _matches_provider(provider_entity: ProviderEntity, provider: str) -> bool:
return provider in (provider_entity.provider, provider_entity.provider_name)

159
api/dify_graph/model_runtime/runtime.py Normal file
View File

@@ -0,0 +1,159 @@
from __future__ import annotations
from collections.abc import Generator, Iterable, Sequence
from typing import IO, Any, Protocol, Union, runtime_checkable
from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk
from dify_graph.model_runtime.entities.message_entities import PromptMessage, PromptMessageTool
from dify_graph.model_runtime.entities.model_entities import AIModelEntity, ModelType
from dify_graph.model_runtime.entities.provider_entities import ProviderEntity
from dify_graph.model_runtime.entities.rerank_entities import MultimodalRerankInput, RerankResult
from dify_graph.model_runtime.entities.text_embedding_entities import EmbeddingInputType, EmbeddingResult
@runtime_checkable
class ModelRuntime(Protocol):
"""Port for provider discovery, schema lookup, and model execution.
`provider` is the model runtime's canonical provider identifier. Adapters may
derive transport-specific details from it, but those details stay outside
this boundary.
"""
def fetch_model_providers(self) -> Sequence[ProviderEntity]: ...
def get_provider_icon(self, *, provider: str, icon_type: str, lang: str) -> tuple[bytes, str]: ...
def validate_provider_credentials(self, *, provider: str, credentials: dict[str, Any]) -> None: ...
def validate_model_credentials(
self,
*,
provider: str,
model_type: ModelType,
model: str,
credentials: dict[str, Any],
) -> None: ...
def get_model_schema(
self,
*,
provider: str,
model_type: ModelType,
model: str,
credentials: dict[str, Any],
) -> AIModelEntity | None: ...
def invoke_llm(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
model_parameters: dict[str, Any],
prompt_messages: Sequence[PromptMessage],
tools: list[PromptMessageTool] | None,
stop: Sequence[str] | None,
stream: bool,
) -> Union[LLMResult, Generator[LLMResultChunk, None, None]]: ...
def get_llm_num_tokens(
self,
*,
provider: str,
model_type: ModelType,
model: str,
credentials: dict[str, Any],
prompt_messages: Sequence[PromptMessage],
tools: Sequence[PromptMessageTool] | None,
) -> int: ...
def invoke_text_embedding(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
texts: list[str],
input_type: EmbeddingInputType,
) -> EmbeddingResult: ...
def invoke_multimodal_embedding(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
documents: list[dict[str, Any]],
input_type: EmbeddingInputType,
) -> EmbeddingResult: ...
def get_text_embedding_num_tokens(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
texts: list[str],
) -> list[int]: ...
def invoke_rerank(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
query: str,
docs: list[str],
score_threshold: float | None,
top_n: int | None,
) -> RerankResult: ...
def invoke_multimodal_rerank(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
query: MultimodalRerankInput,
docs: list[MultimodalRerankInput],
score_threshold: float | None,
top_n: int | None,
) -> RerankResult: ...
def invoke_tts(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
content_text: str,
voice: str,
) -> Iterable[bytes]: ...
def get_tts_model_voices(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
language: str | None,
) -> Any: ...
def invoke_speech_to_text(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
file: IO[bytes],
) -> str: ...
def invoke_moderation(
self,
*,
provider: str,
model: str,
credentials: dict[str, Any],
text: str,
) -> bool: ...

55
api/dify_graph/nodes/base/node.py
View File

@@ -6,13 +6,12 @@ from abc import abstractmethod
from collections.abc import Generator, Mapping, Sequence
from functools import singledispatchmethod
from types import MappingProxyType
from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, get_args, get_origin
from typing import Any, ClassVar, Generic, TypeVar, cast, get_args, get_origin
from uuid import uuid4
from dify_graph.entities import GraphInitParams
from dify_graph.entities.base_node_data import BaseNodeData, RetryConfig
from dify_graph.entities.graph_config import NodeConfigDict
from dify_graph.entities.graph_init_params import DIFY_RUN_CONTEXT_KEY
from dify_graph.enums import (
ErrorStrategy,
NodeExecutionType,
@@ -60,7 +59,7 @@ from dify_graph.node_events import (
StreamCompletedEvent,
)
from dify_graph.runtime import GraphRuntimeState
from libs.datetime_utils import naive_utc_now
from dify_graph.utils.datetime_utils import naive_utc_now
NodeDataT = TypeVar("NodeDataT", bound=BaseNodeData)
_MISSING_RUN_CONTEXT_VALUE = object()
@@ -68,23 +67,6 @@ _MISSING_RUN_CONTEXT_VALUE = object()
logger = logging.getLogger(__name__)
class DifyRunContextProtocol(Protocol):
tenant_id: str
app_id: str
user_id: str
user_from: Any
invoke_from: Any
class _MappingDifyRunContext:
def __init__(self, mapping: Mapping[str, Any]) -> None:
self.tenant_id = str(mapping["tenant_id"])
self.app_id = str(mapping["app_id"])
self.user_id = str(mapping["user_id"])
self.user_from = mapping["user_from"]
self.invoke_from = mapping["invoke_from"]
class Node(Generic[NodeDataT]):
"""BaseNode serves as the foundational class for all node implementations.
@@ -177,8 +159,9 @@ class Node(Generic[NodeDataT]):
# Skip base class itself
if cls is Node:
return
# Only register production node implementations defined under the
# canonical workflow namespaces.
# Only treat nodes from the base dify_graph package as production
# registrations. Higher-layer packages may still register subclasses,
# but dify_graph itself should not know their module identities.
# This prevents test helper subclasses from polluting the global registry and
# accidentally overriding real node types (e.g., a test Answer node).
module_name = getattr(cls, "__module__", "")
@@ -186,7 +169,7 @@ class Node(Generic[NodeDataT]):
node_type = cls.node_type
version = cls.version()
bucket = Node._registry.setdefault(node_type, {})
if module_name.startswith(("dify_graph.nodes.", "core.workflow.nodes.")):
if module_name.startswith("dify_graph.nodes."):
# Production node definitions take precedence and may override
bucket[version] = cls # type: ignore[index]
else:
@@ -299,25 +282,6 @@ class Node(Generic[NodeDataT]):
raise ValueError(f"run_context missing required key: {key}")
return value
def require_dify_context(self) -> DifyRunContextProtocol:
raw_ctx = self.require_run_context_value(DIFY_RUN_CONTEXT_KEY)
if raw_ctx is None:
raise ValueError(f"run_context missing required key: {DIFY_RUN_CONTEXT_KEY}")
if isinstance(raw_ctx, Mapping):
missing_keys = [
key for key in ("tenant_id", "app_id", "user_id", "user_from", "invoke_from") if key not in raw_ctx
]
if missing_keys:
raise ValueError(f"dify context missing required keys: {', '.join(missing_keys)}")
return _MappingDifyRunContext(raw_ctx)
for attr in ("tenant_id", "app_id", "user_id", "user_from", "invoke_from"):
if not hasattr(raw_ctx, attr):
raise TypeError(f"invalid dify context object, missing attribute: {attr}")
return cast(DifyRunContextProtocol, raw_ctx)
@property
def execution_id(self) -> str:
return self._node_execution_id
@@ -793,16 +757,11 @@ class Node(Generic[NodeDataT]):
@_dispatch.register
def _(self, event: RunRetrieverResourceEvent) -> NodeRunRetrieverResourceEvent:
from core.rag.entities.citation_metadata import RetrievalSourceMetadata
retriever_resources = [
RetrievalSourceMetadata.model_validate(resource) for resource in event.retriever_resources
]
return NodeRunRetrieverResourceEvent(
id=self.execution_id,
node_id=self._node_id,
node_type=self.node_type,
retriever_resources=retriever_resources,
retriever_resources=event.retriever_resources,
context=event.context,
node_version=self.version(),
)

25
api/dify_graph/nodes/http_request/node.py
View File

@@ -11,9 +11,13 @@ from dify_graph.nodes.base import variable_template_parser
from dify_graph.nodes.base.entities import VariableSelector
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.http_request.executor import Executor
from dify_graph.nodes.protocols import FileManagerProtocol, HttpClientProtocol, ToolFileManagerProtocol
from dify_graph.nodes.protocols import (
FileManagerProtocol,
FileReferenceFactoryProtocol,
HttpClientProtocol,
ToolFileManagerProtocol,
)
from dify_graph.variables.segments import ArrayFileSegment
from factories import file_factory
from .config import build_http_request_config, resolve_http_request_config
from .entities import (
@@ -46,6 +50,7 @@ class HttpRequestNode(Node[HttpRequestNodeData]):
http_client: HttpClientProtocol,
tool_file_manager_factory: Callable[[], ToolFileManagerProtocol],
file_manager: FileManagerProtocol,
file_reference_factory: FileReferenceFactoryProtocol,
) -> None:
super().__init__(
id=id,
@@ -58,6 +63,7 @@ class HttpRequestNode(Node[HttpRequestNodeData]):
self._http_client = http_client
self._tool_file_manager_factory = tool_file_manager_factory
self._file_manager = file_manager
self._file_reference_factory = file_reference_factory
@classmethod
def get_default_config(cls, filters: Mapping[str, object] | None = None) -> Mapping[str, object]:
@@ -212,7 +218,6 @@ class HttpRequestNode(Node[HttpRequestNodeData]):
"""
Extract files from response by checking both Content-Type header and URL
"""
dify_ctx = self.require_dify_context()
files: list[File] = []
is_file = response.is_file
content_type = response.content_type
@@ -237,20 +242,16 @@ class HttpRequestNode(Node[HttpRequestNodeData]):
tool_file_manager = self._tool_file_manager_factory()
tool_file = tool_file_manager.create_file_by_raw(
user_id=dify_ctx.user_id,
tenant_id=dify_ctx.tenant_id,
conversation_id=None,
file_binary=content,
mimetype=mime_type,
)
mapping = {
"tool_file_id": tool_file.id,
"transfer_method": FileTransferMethod.TOOL_FILE,
}
file = file_factory.build_from_mapping(
mapping=mapping,
tenant_id=dify_ctx.tenant_id,
file = self._file_reference_factory.build_from_mapping(
mapping={
"tool_file_id": tool_file.id,
"transfer_method": FileTransferMethod.TOOL_FILE,
}
)
files.append(file)

28
api/dify_graph/nodes/human_input/human_input_node.py
View File

@@ -15,15 +15,16 @@ from dify_graph.node_events import (
from dify_graph.node_events.base import NodeEventBase
from dify_graph.node_events.node import StreamCompletedEvent
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.runtime import HumanInputNodeRuntimeProtocol
from dify_graph.repositories.human_input_form_repository import (
FormCreateParams,
HumanInputFormEntity,
HumanInputFormRepository,
)
from dify_graph.utils.datetime_utils import naive_utc_now
from dify_graph.workflow_type_encoder import WorkflowRuntimeTypeConverter
from libs.datetime_utils import naive_utc_now
from .entities import DeliveryChannelConfig, HumanInputNodeData, apply_debug_email_recipient
from .entities import DeliveryChannelConfig, HumanInputNodeData
from .enums import DeliveryMethodType, HumanInputFormStatus, PlaceholderType
if TYPE_CHECKING:
@@ -68,6 +69,7 @@ class HumanInputNode(Node[HumanInputNodeData]):
graph_init_params: "GraphInitParams",
graph_runtime_state: "GraphRuntimeState",
form_repository: HumanInputFormRepository,
runtime: HumanInputNodeRuntimeProtocol | None = None,
) -> None:
super().__init__(
id=id,
@@ -76,6 +78,9 @@ class HumanInputNode(Node[HumanInputNodeData]):
graph_runtime_state=graph_runtime_state,
)
self._form_repository = form_repository
if runtime is None:
raise ValueError("runtime is required")
self._runtime = runtime
@classmethod
def version(cls) -> str:
@@ -171,25 +176,14 @@ class HumanInputNode(Node[HumanInputNodeData]):
return self._node_data.is_webapp_enabled()
def _effective_delivery_methods(self) -> Sequence[DeliveryChannelConfig]:
dify_ctx = self.require_dify_context()
invoke_from = self._invoke_from_value()
enabled_methods = [method for method in self._node_data.delivery_methods if method.enabled]
if invoke_from in {_INVOKE_FROM_DEBUGGER, _INVOKE_FROM_EXPLORE}:
enabled_methods = [method for method in enabled_methods if method.type != DeliveryMethodType.WEBAPP]
return [
apply_debug_email_recipient(
method,
enabled=invoke_from == _INVOKE_FROM_DEBUGGER,
user_id=dify_ctx.user_id,
)
for method in enabled_methods
]
return self._runtime.apply_delivery_runtime(methods=enabled_methods)
def _invoke_from_value(self) -> str:
invoke_from = self.require_dify_context().invoke_from
if isinstance(invoke_from, str):
return invoke_from
return str(getattr(invoke_from, "value", invoke_from))
return self._runtime.invoke_source()
def _human_input_required_event(self, form_entity: HumanInputFormEntity) -> HumanInputRequired:
node_data = self._node_data
@@ -224,11 +218,9 @@ class HumanInputNode(Node[HumanInputNodeData]):
"""
repo = self._form_repository
form = repo.get_form(self._workflow_execution_id, self.id)
dify_ctx = self.require_dify_context()
if form is None:
display_in_ui = self._display_in_ui()
params = FormCreateParams(
app_id=dify_ctx.app_id,
workflow_execution_id=self._workflow_execution_id,
node_id=self.id,
form_config=self._node_data,
@@ -238,7 +230,7 @@ class HumanInputNode(Node[HumanInputNodeData]):
resolved_default_values=self.resolve_default_values(),
console_recipient_required=self._should_require_console_recipient(),
console_creator_account_id=(
dify_ctx.user_id
self._runtime.console_actor_id()
if self._invoke_from_value() in {_INVOKE_FROM_DEBUGGER, _INVOKE_FROM_EXPLORE}
else None
),

2
api/dify_graph/nodes/iteration/iteration_node.py
View File

@@ -34,10 +34,10 @@ from dify_graph.nodes.base import LLMUsageTrackingMixin
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.iteration.entities import ErrorHandleMode, IterationNodeData
from dify_graph.runtime import VariablePool
from dify_graph.utils.datetime_utils import naive_utc_now
from dify_graph.variables import IntegerVariable, NoneSegment
from dify_graph.variables.segments import ArrayAnySegment, ArraySegment
from dify_graph.variables.variables import Variable
from libs.datetime_utils import naive_utc_now
from .exc import (
InvalidIteratorValueError,

2
api/dify_graph/nodes/llm/entities.py
View File

@@ -3,11 +3,11 @@ from typing import Any, Literal
from pydantic import BaseModel, Field, field_validator
from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
from dify_graph.entities.base_node_data import BaseNodeData
from dify_graph.enums import BuiltinNodeTypes, NodeType
from dify_graph.model_runtime.entities import ImagePromptMessageContent, LLMMode
from dify_graph.nodes.base.entities import VariableSelector
from dify_graph.prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
class ModelConfig(BaseModel):

53
api/dify_graph/nodes/llm/file_saver.py
View File

@@ -2,10 +2,8 @@ import mimetypes
import typing as tp
from constants.mimetypes import DEFAULT_EXTENSION, DEFAULT_MIME_TYPE
from core.tools.signature import sign_tool_file
from core.tools.tool_file_manager import ToolFileManager
from dify_graph.file import File, FileTransferMethod, FileType
from dify_graph.nodes.protocols import HttpClientProtocol
from dify_graph.nodes.protocols import FileReferenceFactoryProtocol, HttpClientProtocol, ToolFileManagerProtocol
class LLMFileSaver(tp.Protocol):
@@ -57,17 +55,20 @@ class LLMFileSaver(tp.Protocol):
class FileSaverImpl(LLMFileSaver):
_tenant_id: str
_user_id: str
_tool_file_manager: ToolFileManagerProtocol
_file_reference_factory: FileReferenceFactoryProtocol
def __init__(self, user_id: str, tenant_id: str, http_client: HttpClientProtocol):
self._user_id = user_id
self._tenant_id = tenant_id
def __init__(
self,
*,
tool_file_manager: ToolFileManagerProtocol,
file_reference_factory: FileReferenceFactoryProtocol,
http_client: HttpClientProtocol,
):
self._tool_file_manager = tool_file_manager
self._file_reference_factory = file_reference_factory
self._http_client = http_client
def _get_tool_file_manager(self):
return ToolFileManager()
def save_remote_url(self, url: str, file_type: FileType) -> File:
http_response = self._http_client.get(url)
http_response.raise_for_status()
@@ -83,10 +84,7 @@ class FileSaverImpl(LLMFileSaver):
file_type: FileType,
extension_override: str | None = None,
) -> File:
tool_file_manager = self._get_tool_file_manager()
tool_file = tool_file_manager.create_file_by_raw(
user_id=self._user_id,
tenant_id=self._tenant_id,
tool_file = self._tool_file_manager.create_file_by_raw(
# TODO(QuantumGhost): what is conversation id?
conversation_id=None,
file_binary=data,
@@ -94,19 +92,18 @@ class FileSaverImpl(LLMFileSaver):
)
extension_override = _validate_extension_override(extension_override)
extension = _get_extension(mime_type, extension_override)
url = sign_tool_file(tool_file.id, extension)
return File(
tenant_id=self._tenant_id,
type=file_type,
transfer_method=FileTransferMethod.TOOL_FILE,
filename=tool_file.name,
extension=extension,
mime_type=mime_type,
size=len(data),
related_id=tool_file.id,
url=url,
storage_key=tool_file.file_key,
return self._file_reference_factory.build_from_mapping(
mapping={
"type": file_type,
"transfer_method": FileTransferMethod.TOOL_FILE,
"filename": tool_file.name,
"extension": extension,
"mime_type": mime_type,
"size": len(data),
"tool_file_id": tool_file.id,
"related_id": tool_file.id,
"storage_key": tool_file.file_key,
}
)

46
api/dify_graph/nodes/llm/llm_utils.py
View File

@@ -1,9 +1,8 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import Any, cast
from collections.abc import Mapping, Sequence
from typing import Any, Protocol, TypeAlias, cast
from core.model_manager import ModelInstance
from dify_graph.file import FileType, file_manager
from dify_graph.file.models import File
from dify_graph.model_runtime.entities import (
@@ -35,15 +34,36 @@ from .exc import (
TemplateTypeNotSupportError,
)
from .protocols import TemplateRenderer
from .runtime_protocols import PreparedLLMProtocol
def fetch_model_schema(*, model_instance: ModelInstance) -> AIModelEntity:
model_schema = cast(LargeLanguageModel, model_instance.model_type_instance).get_model_schema(
model_instance.model_name,
dict(model_instance.credentials),
)
class _LegacyModelInstance(Protocol):
model_type_instance: object
model_name: str
credentials: object
parameters: Mapping[str, Any]
def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int: ...
PreparedModelInstance: TypeAlias = PreparedLLMProtocol | _LegacyModelInstance
def fetch_model_schema(*, model_instance: PreparedModelInstance) -> AIModelEntity:
get_model_schema = getattr(model_instance, "get_model_schema", None)
if callable(get_model_schema):
model_schema = cast(PreparedLLMProtocol, model_instance).get_model_schema()
else:
legacy_model_instance = cast(_LegacyModelInstance, model_instance)
credentials = legacy_model_instance.credentials
if isinstance(credentials, Mapping):
credentials = dict(credentials)
model_schema = cast(LargeLanguageModel, legacy_model_instance.model_type_instance).get_model_schema(
legacy_model_instance.model_name,
credentials,
)
if not model_schema:
raise ValueError(f"Model schema not found for {model_instance.model_name}")
raise ValueError(f"Model schema not found for {getattr(model_instance, 'model_name', 'unknown model')}")
return model_schema
@@ -116,7 +136,7 @@ def fetch_prompt_messages(
sys_files: Sequence[File],
context: str | None = None,
memory: PromptMessageMemory | None = None,
model_instance: ModelInstance,
model_instance: PreparedModelInstance,
prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate,
stop: Sequence[str] | None = None,
memory_config: MemoryConfig | None = None,
@@ -391,7 +411,7 @@ def combine_message_content_with_role(
raise NotImplementedError(f"Role {role} is not supported")
def calculate_rest_token(*, prompt_messages: list[PromptMessage], model_instance: ModelInstance) -> int:
def calculate_rest_token(*, prompt_messages: list[PromptMessage], model_instance: PreparedModelInstance) -> int:
rest_tokens = 2000
runtime_model_schema = fetch_model_schema(model_instance=model_instance)
runtime_model_parameters = model_instance.parameters
@@ -421,7 +441,7 @@ def handle_memory_chat_mode(
*,
memory: PromptMessageMemory | None,
memory_config: MemoryConfig | None,
model_instance: ModelInstance,
model_instance: PreparedModelInstance,
) -> Sequence[PromptMessage]:
if not memory or not memory_config:
return []
@@ -436,7 +456,7 @@ def handle_memory_completion_mode(
*,
memory: PromptMessageMemory | None,
memory_config: MemoryConfig | None,
model_instance: ModelInstance,
model_instance: PreparedModelInstance,
) -> str:
if not memory or not memory_config:
return ""

547
api/dify_graph/nodes/llm/node.py
View File

@@ -7,16 +7,8 @@ import logging
import re
import time
from collections.abc import Generator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Literal
from typing import TYPE_CHECKING, Any, Literal, cast
from sqlalchemy import select
from core.llm_generator.output_parser.errors import OutputParserError
from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output
from core.model_manager import ModelInstance
from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptTemplate, MemoryConfig
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from core.tools.signature import sign_upload_file
from dify_graph.constants import SYSTEM_VARIABLE_NODE_ID
from dify_graph.entities import GraphInitParams
from dify_graph.entities.graph_config import NodeConfigDict
@@ -27,10 +19,11 @@ from dify_graph.enums import (
WorkflowNodeExecutionMetadataKey,
WorkflowNodeExecutionStatus,
)
from dify_graph.file import File, FileTransferMethod, FileType
from dify_graph.file import File, FileType, file_manager
from dify_graph.model_runtime.entities import (
ImagePromptMessageContent,
PromptMessage,
PromptMessageContentType,
TextPromptMessageContent,
)
from dify_graph.model_runtime.entities.llm_entities import (
@@ -41,7 +34,14 @@ from dify_graph.model_runtime.entities.llm_entities import (
LLMStructuredOutput,
LLMUsage,
)
from dify_graph.model_runtime.entities.message_entities import PromptMessageContentUnionTypes
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessageContentUnionTypes,
PromptMessageRole,
SystemPromptMessage,
UserPromptMessage,
)
from dify_graph.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey
from dify_graph.model_runtime.memory import PromptMessageMemory
from dify_graph.model_runtime.utils.encoders import jsonable_encoder
from dify_graph.node_events import (
@@ -55,19 +55,23 @@ from dify_graph.node_events import (
from dify_graph.nodes.base.entities import VariableSelector
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.base.variable_template_parser import VariableTemplateParser
from dify_graph.nodes.llm.protocols import CredentialsProvider, ModelFactory, TemplateRenderer
from dify_graph.nodes.llm.runtime_protocols import (
PreparedLLMProtocol,
PromptMessageSerializerProtocol,
RetrieverAttachmentLoaderProtocol,
)
from dify_graph.nodes.protocols import HttpClientProtocol
from dify_graph.prompt_entities import CompletionModelPromptTemplate, MemoryConfig
from dify_graph.runtime import VariablePool
from dify_graph.template_rendering import Jinja2TemplateRenderer, TemplateRenderError
from dify_graph.variables import (
ArrayFileSegment,
ArraySegment,
FileSegment,
NoneSegment,
ObjectSegment,
StringSegment,
)
from extensions.ext_database import db
from models.dataset import SegmentAttachmentBinding
from models.model import UploadFile
from . import llm_utils
from .entities import (
@@ -79,9 +83,12 @@ from .exc import (
InvalidContextStructureError,
InvalidVariableTypeError,
LLMNodeError,
MemoryRolePrefixRequiredError,
NoPromptFoundError,
TemplateTypeNotSupportError,
VariableNotFoundError,
)
from .file_saver import FileSaverImpl, LLMFileSaver
from .file_saver import LLMFileSaver
if TYPE_CHECKING:
from dify_graph.file.models import File
@@ -101,11 +108,11 @@ class LLMNode(Node[LLMNodeData]):
_file_outputs: list[File]
_llm_file_saver: LLMFileSaver
_credentials_provider: CredentialsProvider
_model_factory: ModelFactory
_model_instance: ModelInstance
_retriever_attachment_loader: RetrieverAttachmentLoaderProtocol | None
_prompt_message_serializer: PromptMessageSerializerProtocol
_jinja2_template_renderer: Jinja2TemplateRenderer | None
_model_instance: PreparedLLMProtocol
_memory: PromptMessageMemory | None
_template_renderer: TemplateRenderer
def __init__(
self,
@@ -114,13 +121,15 @@ class LLMNode(Node[LLMNodeData]):
graph_init_params: GraphInitParams,
graph_runtime_state: GraphRuntimeState,
*,
credentials_provider: CredentialsProvider,
model_factory: ModelFactory,
model_instance: ModelInstance,
credentials_provider: object | None = None,
model_factory: object | None = None,
model_instance: PreparedLLMProtocol,
http_client: HttpClientProtocol,
template_renderer: TemplateRenderer,
memory: PromptMessageMemory | None = None,
llm_file_saver: LLMFileSaver | None = None,
llm_file_saver: LLMFileSaver,
prompt_message_serializer: PromptMessageSerializerProtocol,
retriever_attachment_loader: RetrieverAttachmentLoaderProtocol | None = None,
jinja2_template_renderer: Jinja2TemplateRenderer | None = None,
):
super().__init__(
id=id,
@@ -131,20 +140,14 @@ class LLMNode(Node[LLMNodeData]):
# LLM file outputs, used for MultiModal outputs.
self._file_outputs = []
self._credentials_provider = credentials_provider
self._model_factory = model_factory
_ = credentials_provider, model_factory, http_client
self._model_instance = model_instance
self._memory = memory
self._template_renderer = template_renderer
if llm_file_saver is None:
dify_ctx = self.require_dify_context()
llm_file_saver = FileSaverImpl(
user_id=dify_ctx.user_id,
tenant_id=dify_ctx.tenant_id,
http_client=http_client,
)
self._llm_file_saver = llm_file_saver
self._prompt_message_serializer = prompt_message_serializer
self._retriever_attachment_loader = retriever_attachment_loader
self._jinja2_template_renderer = jinja2_template_renderer
@classmethod
def version(cls) -> str:
@@ -230,7 +233,7 @@ class LLMNode(Node[LLMNodeData]):
variable_pool=variable_pool,
jinja2_variables=self.node_data.prompt_config.jinja2_variables,
context_files=context_files,
template_renderer=self._template_renderer,
jinja2_template_renderer=self._jinja2_template_renderer,
)
# handle invoke result
@@ -238,7 +241,6 @@ class LLMNode(Node[LLMNodeData]):
model_instance=model_instance,
prompt_messages=prompt_messages,
stop=stop,
user_id=self.require_dify_context().user_id,
structured_output_enabled=self.node_data.structured_output_enabled,
structured_output=self.node_data.structured_output,
file_saver=self._llm_file_saver,
@@ -281,7 +283,7 @@ class LLMNode(Node[LLMNodeData]):
process_data = {
"model_mode": self.node_data.model.mode,
"prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
"prompts": self._prompt_message_serializer.serialize(
model_mode=self.node_data.model.mode, prompt_messages=prompt_messages
),
"usage": jsonable_encoder(usage),
@@ -349,10 +351,9 @@ class LLMNode(Node[LLMNodeData]):
@staticmethod
def invoke_llm(
*,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
prompt_messages: Sequence[PromptMessage],
stop: Sequence[str] | None = None,
user_id: str,
structured_output_enabled: bool,
structured_output: Mapping[str, Any] | None = None,
file_saver: LLMFileSaver,
@@ -363,35 +364,28 @@ class LLMNode(Node[LLMNodeData]):
) -> Generator[NodeEventBase | LLMStructuredOutput, None, None]:
model_parameters = model_instance.parameters
invoke_model_parameters = dict(model_parameters)
model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
if structured_output_enabled:
output_schema = LLMNode.fetch_structured_output_schema(
structured_output=structured_output or {},
)
request_start_time = time.perf_counter()
invoke_result = invoke_llm_with_structured_output(
provider=model_instance.provider,
model_schema=model_schema,
model_instance=model_instance,
invoke_result = model_instance.invoke_llm_with_structured_output(
prompt_messages=prompt_messages,
json_schema=output_schema,
model_parameters=invoke_model_parameters,
stop=list(stop or []),
stop=stop,
stream=True,
user=user_id,
)
else:
request_start_time = time.perf_counter()
invoke_result = model_instance.invoke_llm(
prompt_messages=list(prompt_messages),
prompt_messages=prompt_messages,
model_parameters=invoke_model_parameters,
stop=list(stop or []),
tools=None,
stop=stop,
stream=True,
user=user_id,
)
return LLMNode.handle_invoke_result(
@@ -400,6 +394,7 @@ class LLMNode(Node[LLMNodeData]):
file_outputs=file_outputs,
node_id=node_id,
node_type=node_type,
model_instance=model_instance,
reasoning_format=reasoning_format,
request_start_time=request_start_time,
)
@@ -412,6 +407,7 @@ class LLMNode(Node[LLMNodeData]):
file_outputs: list[File],
node_id: str,
node_type: NodeType,
model_instance: PreparedLLMProtocol | object,
reasoning_format: Literal["separated", "tagged"] = "tagged",
request_start_time: float | None = None,
) -> Generator[NodeEventBase | LLMStructuredOutput, None, None]:
@@ -483,8 +479,14 @@ class LLMNode(Node[LLMNodeData]):
usage = result.delta.usage
if finish_reason is None and result.delta.finish_reason:
finish_reason = result.delta.finish_reason
except OutputParserError as e:
raise LLMNodeError(f"Failed to parse structured output: {e}")
except Exception as e:
if hasattr(model_instance, "is_structured_output_parse_error") and cast(
PreparedLLMProtocol, model_instance
).is_structured_output_parse_error(e):
raise LLMNodeError(f"Failed to parse structured output: {e}") from e
if type(e).__name__ == "OutputParserError":
raise LLMNodeError(f"Failed to parse structured output: {e}") from e
raise
# Extract reasoning content from <think> tags in the main text
full_text = full_text_buffer.getvalue()
@@ -687,30 +689,8 @@ class LLMNode(Node[LLMNodeData]):
segment_id = retriever_resource.get("segment_id")
if not segment_id:
continue
attachments_with_bindings = db.session.execute(
select(SegmentAttachmentBinding, UploadFile)
.join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id)
.where(
SegmentAttachmentBinding.segment_id == segment_id,
)
).all()
if attachments_with_bindings:
for _, upload_file in attachments_with_bindings:
attachment_info = File(
id=upload_file.id,
filename=upload_file.name,
extension="." + upload_file.extension,
mime_type=upload_file.mime_type,
tenant_id=self.require_dify_context().tenant_id,
type=FileType.IMAGE,
transfer_method=FileTransferMethod.LOCAL_FILE,
remote_url=upload_file.source_url,
related_id=upload_file.id,
size=upload_file.size,
storage_key=upload_file.key,
url=sign_upload_file(upload_file.id, upload_file.extension),
)
context_files.append(attachment_info)
if self._retriever_attachment_loader is not None:
context_files.extend(self._retriever_attachment_loader.load(segment_id=segment_id))
yield RunRetrieverResourceEvent(
retriever_resources=original_retriever_resource,
context=context_str.strip(),
@@ -755,7 +735,7 @@ class LLMNode(Node[LLMNodeData]):
sys_files: Sequence[File],
context: str | None = None,
memory: PromptMessageMemory | None = None,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate,
stop: Sequence[str] | None = None,
memory_config: MemoryConfig | None = None,
@@ -764,24 +744,186 @@ class LLMNode(Node[LLMNodeData]):
variable_pool: VariablePool,
jinja2_variables: Sequence[VariableSelector],
context_files: list[File] | None = None,
template_renderer: TemplateRenderer | None = None,
jinja2_template_renderer: Jinja2TemplateRenderer | None = None,
) -> tuple[Sequence[PromptMessage], Sequence[str] | None]:
return llm_utils.fetch_prompt_messages(
sys_query=sys_query,
sys_files=sys_files,
context=context,
memory=memory,
model_instance=model_instance,
prompt_template=prompt_template,
stop=stop,
memory_config=memory_config,
vision_enabled=vision_enabled,
vision_detail=vision_detail,
variable_pool=variable_pool,
jinja2_variables=jinja2_variables,
context_files=context_files,
template_renderer=template_renderer,
)
prompt_messages: list[PromptMessage] = []
model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
if isinstance(prompt_template, list):
# For chat model
prompt_messages.extend(
LLMNode.handle_list_messages(
messages=prompt_template,
context=context,
jinja2_variables=jinja2_variables,
variable_pool=variable_pool,
vision_detail_config=vision_detail,
jinja2_template_renderer=jinja2_template_renderer,
)
)
# Get memory messages for chat mode
memory_messages = _handle_memory_chat_mode(
memory=memory,
memory_config=memory_config,
model_instance=model_instance,
)
# Extend prompt_messages with memory messages
prompt_messages.extend(memory_messages)
# Add current query to the prompt messages
if sys_query:
message = LLMNodeChatModelMessage(
text=sys_query,
role=PromptMessageRole.USER,
edition_type="basic",
)
prompt_messages.extend(
LLMNode.handle_list_messages(
messages=[message],
context="",
jinja2_variables=[],
variable_pool=variable_pool,
vision_detail_config=vision_detail,
jinja2_template_renderer=jinja2_template_renderer,
)
)
elif isinstance(prompt_template, LLMNodeCompletionModelPromptTemplate):
# For completion model
prompt_messages.extend(
_handle_completion_template(
template=prompt_template,
context=context,
jinja2_variables=jinja2_variables,
variable_pool=variable_pool,
jinja2_template_renderer=jinja2_template_renderer,
)
)
# Get memory text for completion model
memory_text = _handle_memory_completion_mode(
memory=memory,
memory_config=memory_config,
model_instance=model_instance,
)
# Insert histories into the prompt
prompt_content = prompt_messages[0].content
# For issue #11247 - Check if prompt content is a string or a list
if isinstance(prompt_content, str):
prompt_content = str(prompt_content)
if "#histories#" in prompt_content:
prompt_content = prompt_content.replace("#histories#", memory_text)
else:
prompt_content = memory_text + "\n" + prompt_content
prompt_messages[0].content = prompt_content
elif isinstance(prompt_content, list):
for content_item in prompt_content:
if isinstance(content_item, TextPromptMessageContent):
if "#histories#" in content_item.data:
content_item.data = content_item.data.replace("#histories#", memory_text)
else:
content_item.data = memory_text + "\n" + content_item.data
else:
raise ValueError("Invalid prompt content type")
# Add current query to the prompt message
if sys_query:
if isinstance(prompt_content, str):
prompt_content = str(prompt_messages[0].content).replace("#sys.query#", sys_query)
prompt_messages[0].content = prompt_content
elif isinstance(prompt_content, list):
for content_item in prompt_content:
if isinstance(content_item, TextPromptMessageContent):
content_item.data = sys_query + "\n" + content_item.data
else:
raise ValueError("Invalid prompt content type")
else:
raise TemplateTypeNotSupportError(type_name=str(type(prompt_template)))
# The sys_files will be deprecated later
if vision_enabled and sys_files:
file_prompts = []
for file in sys_files:
file_prompt = file_manager.to_prompt_message_content(file, image_detail_config=vision_detail)
file_prompts.append(file_prompt)
# If last prompt is a user prompt, add files into its contents,
# otherwise append a new user prompt
if (
len(prompt_messages) > 0
and isinstance(prompt_messages[-1], UserPromptMessage)
and isinstance(prompt_messages[-1].content, list)
):
prompt_messages[-1] = UserPromptMessage(content=file_prompts + prompt_messages[-1].content)
else:
prompt_messages.append(UserPromptMessage(content=file_prompts))
# The context_files
if vision_enabled and context_files:
file_prompts = []
for file in context_files:
file_prompt = file_manager.to_prompt_message_content(file, image_detail_config=vision_detail)
file_prompts.append(file_prompt)
# If last prompt is a user prompt, add files into its contents,
# otherwise append a new user prompt
if (
len(prompt_messages) > 0
and isinstance(prompt_messages[-1], UserPromptMessage)
and isinstance(prompt_messages[-1].content, list)
):
prompt_messages[-1] = UserPromptMessage(content=file_prompts + prompt_messages[-1].content)
else:
prompt_messages.append(UserPromptMessage(content=file_prompts))
# Remove empty messages and filter unsupported content
filtered_prompt_messages = []
for prompt_message in prompt_messages:
if isinstance(prompt_message.content, list):
prompt_message_content: list[PromptMessageContentUnionTypes] = []
for content_item in prompt_message.content:
# Skip content if features are not defined
if not model_schema.features:
if content_item.type != PromptMessageContentType.TEXT:
continue
prompt_message_content.append(content_item)
continue
# Skip content if corresponding feature is not supported
if (
(
content_item.type == PromptMessageContentType.IMAGE
and ModelFeature.VISION not in model_schema.features
)
or (
content_item.type == PromptMessageContentType.DOCUMENT
and ModelFeature.DOCUMENT not in model_schema.features
)
or (
content_item.type == PromptMessageContentType.VIDEO
and ModelFeature.VIDEO not in model_schema.features
)
or (
content_item.type == PromptMessageContentType.AUDIO
and ModelFeature.AUDIO not in model_schema.features
)
):
continue
prompt_message_content.append(content_item)
if len(prompt_message_content) == 1 and prompt_message_content[0].type == PromptMessageContentType.TEXT:
prompt_message.content = prompt_message_content[0].data
else:
prompt_message.content = prompt_message_content
if prompt_message.is_empty():
continue
filtered_prompt_messages.append(prompt_message)
if len(filtered_prompt_messages) == 0:
raise NoPromptFoundError(
"No prompt found in the LLM configuration. "
"Please ensure a prompt is properly configured before proceeding."
)
return filtered_prompt_messages, stop
@classmethod
def _extract_variable_selector_to_variable_mapping(
@@ -881,16 +1023,61 @@ class LLMNode(Node[LLMNodeData]):
jinja2_variables: Sequence[VariableSelector],
variable_pool: VariablePool,
vision_detail_config: ImagePromptMessageContent.DETAIL,
template_renderer: TemplateRenderer | None = None,
jinja2_template_renderer: Jinja2TemplateRenderer | None = None,
) -> Sequence[PromptMessage]:
return llm_utils.handle_list_messages(
messages=messages,
context=context,
jinja2_variables=jinja2_variables,
variable_pool=variable_pool,
vision_detail_config=vision_detail_config,
template_renderer=template_renderer,
)
prompt_messages: list[PromptMessage] = []
for message in messages:
if message.edition_type == "jinja2":
result_text = _render_jinja2_message(
template=message.jinja2_text or "",
jinja2_variables=jinja2_variables,
variable_pool=variable_pool,
jinja2_template_renderer=jinja2_template_renderer,
)
prompt_message = _combine_message_content_with_role(
contents=[TextPromptMessageContent(data=result_text)], role=message.role
)
prompt_messages.append(prompt_message)
else:
# Get segment group from basic message
if context:
template = message.text.replace("{#context#}", context)
else:
template = message.text
segment_group = variable_pool.convert_template(template)
# Process segments for images
file_contents = []
for segment in segment_group.value:
if isinstance(segment, ArrayFileSegment):
for file in segment.value:
if file.type in {FileType.IMAGE, FileType.VIDEO, FileType.AUDIO, FileType.DOCUMENT}:
file_content = file_manager.to_prompt_message_content(
file, image_detail_config=vision_detail_config
)
file_contents.append(file_content)
elif isinstance(segment, FileSegment):
file = segment.value
if file.type in {FileType.IMAGE, FileType.VIDEO, FileType.AUDIO, FileType.DOCUMENT}:
file_content = file_manager.to_prompt_message_content(
file, image_detail_config=vision_detail_config
)
file_contents.append(file_content)
# Create message with text from all segments
plain_text = segment_group.text
if plain_text:
prompt_message = _combine_message_content_with_role(
contents=[TextPromptMessageContent(data=plain_text)], role=message.role
)
prompt_messages.append(prompt_message)
if file_contents:
# Create message with image contents
prompt_message = _combine_message_content_with_role(contents=file_contents, role=message.role)
prompt_messages.append(prompt_message)
return prompt_messages
@staticmethod
def handle_blocking_result(
@@ -1027,5 +1214,153 @@ class LLMNode(Node[LLMNodeData]):
return self.node_data.retry_config.retry_enabled
@property
def model_instance(self) -> ModelInstance:
def model_instance(self) -> PreparedLLMProtocol:
return self._model_instance
def _combine_message_content_with_role(
*, contents: str | list[PromptMessageContentUnionTypes] | None = None, role: PromptMessageRole
):
match role:
case PromptMessageRole.USER:
return UserPromptMessage(content=contents)
case PromptMessageRole.ASSISTANT:
return AssistantPromptMessage(content=contents)
case PromptMessageRole.SYSTEM:
return SystemPromptMessage(content=contents)
case _:
raise NotImplementedError(f"Role {role} is not supported")
def _render_jinja2_message(
*,
template: str,
jinja2_variables: Sequence[VariableSelector],
variable_pool: VariablePool,
jinja2_template_renderer: Jinja2TemplateRenderer | None,
):
if not template:
return ""
jinja2_inputs = {}
for jinja2_variable in jinja2_variables:
variable = variable_pool.get(jinja2_variable.value_selector)
jinja2_inputs[jinja2_variable.variable] = variable.to_object() if variable else ""
if jinja2_template_renderer is None:
raise TemplateRenderError("LLMNode requires an injected jinja2_template_renderer for jinja2 prompts.")
return jinja2_template_renderer.render_template(template, jinja2_inputs)
def _calculate_rest_token(
*,
prompt_messages: list[PromptMessage],
model_instance: PreparedLLMProtocol,
) -> int:
rest_tokens = 2000
runtime_model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
runtime_model_parameters = model_instance.parameters
model_context_tokens = runtime_model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
if model_context_tokens:
curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
max_tokens = 0
for parameter_rule in runtime_model_schema.parameter_rules:
if parameter_rule.name == "max_tokens" or (
parameter_rule.use_template and parameter_rule.use_template == "max_tokens"
):
max_tokens = (
runtime_model_parameters.get(parameter_rule.name)
or runtime_model_parameters.get(str(parameter_rule.use_template))
or 0
)
rest_tokens = model_context_tokens - max_tokens - curr_message_tokens
rest_tokens = max(rest_tokens, 0)
return rest_tokens
def _handle_memory_chat_mode(
*,
memory: PromptMessageMemory | None,
memory_config: MemoryConfig | None,
model_instance: PreparedLLMProtocol,
) -> Sequence[PromptMessage]:
memory_messages: Sequence[PromptMessage] = []
# Get messages from memory for chat model
if memory and memory_config:
rest_tokens = _calculate_rest_token(
prompt_messages=[],
model_instance=model_instance,
)
memory_messages = memory.get_history_prompt_messages(
max_token_limit=rest_tokens,
message_limit=memory_config.window.size if memory_config.window.enabled else None,
)
return memory_messages
def _handle_memory_completion_mode(
*,
memory: PromptMessageMemory | None,
memory_config: MemoryConfig | None,
model_instance: PreparedLLMProtocol,
) -> str:
memory_text = ""
# Get history text from memory for completion model
if memory and memory_config:
rest_tokens = _calculate_rest_token(
prompt_messages=[],
model_instance=model_instance,
)
if not memory_config.role_prefix:
raise MemoryRolePrefixRequiredError("Memory role prefix is required for completion model.")
memory_text = llm_utils.fetch_memory_text(
memory=memory,
max_token_limit=rest_tokens,
message_limit=memory_config.window.size if memory_config.window.enabled else None,
human_prefix=memory_config.role_prefix.user,
ai_prefix=memory_config.role_prefix.assistant,
)
return memory_text
def _handle_completion_template(
*,
template: LLMNodeCompletionModelPromptTemplate,
context: str | None,
jinja2_variables: Sequence[VariableSelector],
variable_pool: VariablePool,
jinja2_template_renderer: Jinja2TemplateRenderer | None = None,
) -> Sequence[PromptMessage]:
"""Handle completion template processing outside of LLMNode class.
Args:
template: The completion model prompt template
context: Optional context string
jinja2_variables: Variables for jinja2 template rendering
variable_pool: Variable pool for template conversion
Returns:
Sequence of prompt messages
"""
prompt_messages = []
if template.edition_type == "jinja2":
result_text = _render_jinja2_message(
template=template.jinja2_text or "",
jinja2_variables=jinja2_variables,
variable_pool=variable_pool,
jinja2_template_renderer=jinja2_template_renderer,
)
else:
if context:
template_text = template.text.replace("{#context#}", context)
else:
template_text = template.text
result_text = variable_pool.convert_template(template_text).text
prompt_message = _combine_message_content_with_role(
contents=[TextPromptMessageContent(data=result_text)], role=PromptMessageRole.USER
)
prompt_messages.append(prompt_message)
return prompt_messages

8
api/dify_graph/nodes/llm/protocols.py
View File

@@ -3,7 +3,7 @@ from __future__ import annotations
from collections.abc import Mapping
from typing import Any, Protocol
from core.model_manager import ModelInstance
from dify_graph.nodes.llm.runtime_protocols import PreparedLLMProtocol
class CredentialsProvider(Protocol):
@@ -15,10 +15,10 @@ class CredentialsProvider(Protocol):
class ModelFactory(Protocol):
"""Port for creating initialized LLM model instances for execution."""
"""Port for creating prepared graph-facing LLM runtimes for execution."""
def init_model_instance(self, provider_name: str, model_name: str) -> ModelInstance:
"""Create a model instance that is ready for schema lookup and invocation."""
def init_model_instance(self, provider_name: str, model_name: str) -> PreparedLLMProtocol:
"""Create a prepared LLM runtime that is ready for graph execution."""
...

74
api/dify_graph/nodes/llm/runtime_protocols.py Normal file
View File

@@ -0,0 +1,74 @@
from __future__ import annotations
from collections.abc import Generator, Mapping, Sequence
from typing import Any, Protocol
from dify_graph.file import File
from dify_graph.model_runtime.entities import LLMMode, PromptMessage
from dify_graph.model_runtime.entities.llm_entities import (
LLMResult,
LLMResultChunk,
LLMResultChunkWithStructuredOutput,
LLMResultWithStructuredOutput,
)
from dify_graph.model_runtime.entities.message_entities import PromptMessageTool
from dify_graph.model_runtime.entities.model_entities import AIModelEntity
class PreparedLLMProtocol(Protocol):
"""A graph-facing LLM runtime with provider-specific setup already applied."""
@property
def provider(self) -> str: ...
@property
def model_name(self) -> str: ...
@property
def parameters(self) -> Mapping[str, Any]: ...
@property
def stop(self) -> Sequence[str] | None: ...
def get_model_schema(self) -> AIModelEntity: ...
def get_llm_num_tokens(self, prompt_messages: Sequence[PromptMessage]) -> int: ...
def invoke_llm(
self,
*,
prompt_messages: Sequence[PromptMessage],
model_parameters: Mapping[str, Any],
tools: Sequence[PromptMessageTool] | None,
stop: Sequence[str] | None,
stream: bool,
) -> LLMResult | Generator[LLMResultChunk, None, None]: ...
def invoke_llm_with_structured_output(
self,
*,
prompt_messages: Sequence[PromptMessage],
json_schema: Mapping[str, Any],
model_parameters: Mapping[str, Any],
stop: Sequence[str] | None,
stream: bool,
) -> LLMResultWithStructuredOutput | Generator[LLMResultChunkWithStructuredOutput, None, None]: ...
def is_structured_output_parse_error(self, error: Exception) -> bool: ...
class PromptMessageSerializerProtocol(Protocol):
"""Port for converting compiled prompt messages into persisted process data."""
def serialize(
self,
*,
model_mode: LLMMode,
prompt_messages: Sequence[PromptMessage],
) -> Any: ...
class RetrieverAttachmentLoaderProtocol(Protocol):
"""Port for resolving retriever segment attachments into graph file references."""
def load(self, *, segment_id: str) -> Sequence[File]: ...

2
api/dify_graph/nodes/loop/loop_node.py
View File

@@ -31,9 +31,9 @@ from dify_graph.nodes.base import LLMUsageTrackingMixin
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.loop.entities import LoopCompletedReason, LoopNodeData, LoopVariableData
from dify_graph.utils.condition.processor import ConditionProcessor
from dify_graph.utils.datetime_utils import naive_utc_now
from dify_graph.variables import Segment, SegmentType
from factories.variable_factory import TypeMismatchError, build_segment_with_type, segment_to_variable
from libs.datetime_utils import naive_utc_now
if TYPE_CHECKING:
from dify_graph.graph_engine import GraphEngine

2
api/dify_graph/nodes/parameter_extractor/entities.py
View File

@@ -7,10 +7,10 @@ from pydantic import (
field_validator,
)
from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from dify_graph.entities.base_node_data import BaseNodeData
from dify_graph.enums import BuiltinNodeTypes, NodeType
from dify_graph.nodes.llm.entities import ModelConfig, VisionConfig
from dify_graph.prompt_entities import MemoryConfig
from dify_graph.variables.types import SegmentType
_OLD_BOOL_TYPE_NAME = "bool"

204
api/dify_graph/nodes/parameter_extractor/parameter_extractor_node.py
View File

@@ -5,11 +5,6 @@ import uuid
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, cast
from core.model_manager import ModelInstance
from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate
from core.prompt.simple_prompt_transform import ModelMode
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from dify_graph.entities.graph_config import NodeConfigDict
from dify_graph.enums import (
BuiltinNodeTypes,
@@ -17,8 +12,8 @@ from dify_graph.enums import (
WorkflowNodeExecutionStatus,
)
from dify_graph.file import File
from dify_graph.model_runtime.entities import ImagePromptMessageContent
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.model_runtime.entities import ImagePromptMessageContent, LLMMode
from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessage,
@@ -27,14 +22,15 @@ from dify_graph.model_runtime.entities.message_entities import (
ToolPromptMessage,
UserPromptMessage,
)
from dify_graph.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey
from dify_graph.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey, ModelType
from dify_graph.model_runtime.memory import PromptMessageMemory
from dify_graph.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from dify_graph.model_runtime.utils.encoders import jsonable_encoder
from dify_graph.node_events import NodeRunResult
from dify_graph.nodes.base import variable_template_parser
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.llm import llm_utils
from dify_graph.nodes.llm import LLMNode, llm_utils
from dify_graph.nodes.llm.entities import LLMNodeChatModelMessage, LLMNodeCompletionModelPromptTemplate
from dify_graph.nodes.llm.runtime_protocols import PreparedLLMProtocol, PromptMessageSerializerProtocol
from dify_graph.runtime import VariablePool
from dify_graph.variables.types import ArrayValidation, SegmentType
from factories.variable_factory import build_segment_with_type
@@ -66,7 +62,6 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from dify_graph.entities import GraphInitParams
from dify_graph.nodes.llm.protocols import CredentialsProvider, ModelFactory
from dify_graph.runtime import GraphRuntimeState
@@ -99,9 +94,8 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
node_type = BuiltinNodeTypes.PARAMETER_EXTRACTOR
_model_instance: ModelInstance
_credentials_provider: "CredentialsProvider"
_model_factory: "ModelFactory"
_model_instance: PreparedLLMProtocol
_prompt_message_serializer: PromptMessageSerializerProtocol
_memory: PromptMessageMemory | None
def __init__(
@@ -111,10 +105,11 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
graph_init_params: "GraphInitParams",
graph_runtime_state: "GraphRuntimeState",
*,
credentials_provider: "CredentialsProvider",
model_factory: "ModelFactory",
model_instance: ModelInstance,
credentials_provider: object | None = None,
model_factory: object | None = None,
model_instance: PreparedLLMProtocol,
memory: PromptMessageMemory | None = None,
prompt_message_serializer: PromptMessageSerializerProtocol,
) -> None:
super().__init__(
id=id,
@@ -122,9 +117,9 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
graph_init_params=graph_init_params,
graph_runtime_state=graph_runtime_state,
)
self._credentials_provider = credentials_provider
self._model_factory = model_factory
_ = credentials_provider, model_factory
self._model_instance = model_instance
self._prompt_message_serializer = prompt_message_serializer
self._memory = memory
@classmethod
@@ -164,13 +159,12 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
)
model_instance = self._model_instance
if not isinstance(model_instance.model_type_instance, LargeLanguageModel):
raise InvalidModelTypeError("Model is not a Large Language Model")
try:
model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
except ValueError as exc:
raise ModelSchemaNotFoundError("Model schema not found") from exc
if model_schema.model_type != ModelType.LLM:
raise InvalidModelTypeError("Model is not a Large Language Model")
memory = self._memory
if (
@@ -210,8 +204,9 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
process_data = {
"model_mode": node_data.model.mode,
"prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
model_mode=node_data.model.mode, prompt_messages=prompt_messages
"prompts": self._prompt_message_serializer.serialize(
model_mode=node_data.model.mode,
prompt_messages=prompt_messages,
),
"usage": None,
"function": {} if not prompt_message_tools else jsonable_encoder(prompt_message_tools[0]),
@@ -287,18 +282,20 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
def _invoke(
self,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
prompt_messages: list[PromptMessage],
tools: list[PromptMessageTool],
stop: Sequence[str],
stop: Sequence[str] | None,
) -> tuple[str, LLMUsage, AssistantPromptMessage.ToolCall | None]:
invoke_result = model_instance.invoke_llm(
prompt_messages=prompt_messages,
model_parameters=dict(model_instance.parameters),
tools=tools,
stop=list(stop),
stream=False,
user=self.require_dify_context().user_id,
invoke_result = cast(
LLMResult,
model_instance.invoke_llm(
prompt_messages=prompt_messages,
model_parameters=dict(model_instance.parameters),
tools=tools or None,
stop=stop,
stream=False,
),
)
# handle invoke result
@@ -317,7 +314,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
node_data: ParameterExtractorNodeData,
query: str,
variable_pool: VariablePool,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
memory: PromptMessageMemory | None,
files: Sequence[File],
vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -329,7 +326,6 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
content=query, structure=json.dumps(node_data.get_parameter_json_schema())
)
prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
rest_token = self._calculate_rest_token(
node_data=node_data,
query=query,
@@ -340,15 +336,11 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
prompt_template = self._get_function_calling_prompt_template(
node_data, query, variable_pool, memory, rest_token
)
prompt_messages = prompt_transform.get_prompt(
prompt_template=prompt_template,
inputs={},
query="",
files=files,
context="",
memory_config=node_data.memory,
memory=None,
prompt_messages = self._compile_prompt_messages(
model_instance=model_instance,
prompt_template=prompt_template,
files=files,
vision_enabled=node_data.vision.enabled,
image_detail_config=vision_detail,
)
@@ -405,7 +397,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
data: ParameterExtractorNodeData,
query: str,
variable_pool: VariablePool,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
memory: PromptMessageMemory | None,
files: Sequence[File],
vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -413,9 +405,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
"""
Generate prompt engineering prompt.
"""
model_mode = ModelMode(data.model.mode)
if model_mode == ModelMode.COMPLETION:
if data.model.mode == LLMMode.COMPLETION:
return self._generate_prompt_engineering_completion_prompt(
node_data=data,
query=query,
@@ -425,7 +415,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
files=files,
vision_detail=vision_detail,
)
elif model_mode == ModelMode.CHAT:
if data.model.mode == LLMMode.CHAT:
return self._generate_prompt_engineering_chat_prompt(
node_data=data,
query=query,
@@ -435,15 +425,14 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
files=files,
vision_detail=vision_detail,
)
else:
raise InvalidModelModeError(f"Invalid model mode: {model_mode}")
raise InvalidModelModeError(f"Invalid model mode: {data.model.mode}")
def _generate_prompt_engineering_completion_prompt(
self,
node_data: ParameterExtractorNodeData,
query: str,
variable_pool: VariablePool,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
memory: PromptMessageMemory | None,
files: Sequence[File],
vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -451,7 +440,6 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
"""
Generate completion prompt.
"""
prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
rest_token = self._calculate_rest_token(
node_data=node_data,
query=query,
@@ -462,27 +450,20 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
prompt_template = self._get_prompt_engineering_prompt_template(
node_data=node_data, query=query, variable_pool=variable_pool, memory=memory, max_token_limit=rest_token
)
prompt_messages = prompt_transform.get_prompt(
prompt_template=prompt_template,
inputs={"structure": json.dumps(node_data.get_parameter_json_schema())},
query="",
files=files,
context="",
memory_config=node_data.memory,
# AdvancedPromptTransform is still typed against TokenBufferMemory.
memory=cast(Any, memory),
return self._compile_prompt_messages(
model_instance=model_instance,
prompt_template=prompt_template,
files=files,
vision_enabled=node_data.vision.enabled,
image_detail_config=vision_detail,
)
return prompt_messages
def _generate_prompt_engineering_chat_prompt(
self,
node_data: ParameterExtractorNodeData,
query: str,
variable_pool: VariablePool,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
memory: PromptMessageMemory | None,
files: Sequence[File],
vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -490,7 +471,6 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
"""
Generate chat prompt.
"""
prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
rest_token = self._calculate_rest_token(
node_data=node_data,
query=query,
@@ -508,15 +488,11 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
max_token_limit=rest_token,
)
prompt_messages = prompt_transform.get_prompt(
prompt_template=prompt_template,
inputs={},
query="",
files=files,
context="",
memory_config=node_data.memory,
memory=None,
prompt_messages = self._compile_prompt_messages(
model_instance=model_instance,
prompt_template=prompt_template,
files=files,
vision_enabled=node_data.vision.enabled,
image_detail_config=vision_detail,
)
@@ -717,8 +693,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
variable_pool: VariablePool,
memory: PromptMessageMemory | None,
max_token_limit: int = 2000,
) -> list[ChatModelMessage]:
model_mode = ModelMode(node_data.model.mode)
) -> list[LLMNodeChatModelMessage]:
input_text = query
memory_str = ""
instruction = variable_pool.convert_template(node_data.instruction or "").text
@@ -727,15 +702,14 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
memory_str = llm_utils.fetch_memory_text(
memory=memory, max_token_limit=max_token_limit, message_limit=node_data.memory.window.size
)
if model_mode == ModelMode.CHAT:
system_prompt_messages = ChatModelMessage(
if node_data.model.mode == LLMMode.CHAT:
system_prompt_messages = LLMNodeChatModelMessage(
role=PromptMessageRole.SYSTEM,
text=FUNCTION_CALLING_EXTRACTOR_SYSTEM_PROMPT.format(histories=memory_str, instruction=instruction),
)
user_prompt_message = ChatModelMessage(role=PromptMessageRole.USER, text=input_text)
user_prompt_message = LLMNodeChatModelMessage(role=PromptMessageRole.USER, text=input_text)
return [system_prompt_messages, user_prompt_message]
else:
raise InvalidModelModeError(f"Model mode {model_mode} not support.")
raise InvalidModelModeError(f"Model mode {node_data.model.mode} not support.")
def _get_prompt_engineering_prompt_template(
self,
@@ -744,8 +718,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
variable_pool: VariablePool,
memory: PromptMessageMemory | None,
max_token_limit: int = 2000,
):
model_mode = ModelMode(node_data.model.mode)
) -> list[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate:
input_text = query
memory_str = ""
instruction = variable_pool.convert_template(node_data.instruction or "").text
@@ -754,64 +727,53 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
memory_str = llm_utils.fetch_memory_text(
memory=memory, max_token_limit=max_token_limit, message_limit=node_data.memory.window.size
)
if model_mode == ModelMode.CHAT:
system_prompt_messages = ChatModelMessage(
if node_data.model.mode == LLMMode.CHAT:
system_prompt_messages = LLMNodeChatModelMessage(
role=PromptMessageRole.SYSTEM,
text=CHAT_GENERATE_JSON_PROMPT.format(histories=memory_str, instructions=instruction),
)
user_prompt_message = ChatModelMessage(role=PromptMessageRole.USER, text=input_text)
user_prompt_message = LLMNodeChatModelMessage(role=PromptMessageRole.USER, text=input_text)
return [system_prompt_messages, user_prompt_message]
elif model_mode == ModelMode.COMPLETION:
return CompletionModelPromptTemplate(
if node_data.model.mode == LLMMode.COMPLETION:
return LLMNodeCompletionModelPromptTemplate(
text=COMPLETION_GENERATE_JSON_PROMPT.format(
histories=memory_str, text=input_text, instruction=instruction
)
.replace("{γγγ", "")
.replace("}γγγ", "")
.replace("{ structure }", json.dumps(node_data.get_parameter_json_schema())),
)
else:
raise InvalidModelModeError(f"Model mode {model_mode} not support.")
raise InvalidModelModeError(f"Model mode {node_data.model.mode} not support.")
def _calculate_rest_token(
self,
node_data: ParameterExtractorNodeData,
query: str,
variable_pool: VariablePool,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
context: str | None,
) -> int:
try:
model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
except ValueError as exc:
raise ModelSchemaNotFoundError("Model schema not found") from exc
prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
if set(model_schema.features or []) & {ModelFeature.TOOL_CALL, ModelFeature.MULTI_TOOL_CALL}:
prompt_template = self._get_function_calling_prompt_template(node_data, query, variable_pool, None, 2000)
else:
prompt_template = self._get_prompt_engineering_prompt_template(node_data, query, variable_pool, None, 2000)
prompt_messages = prompt_transform.get_prompt(
prompt_template=prompt_template,
inputs={},
query="",
files=[],
context=context,
memory_config=node_data.memory,
memory=None,
prompt_messages = self._compile_prompt_messages(
model_instance=model_instance,
prompt_template=prompt_template,
files=[],
vision_enabled=False,
context=context,
)
rest_tokens = 2000
model_context_tokens = model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
if model_context_tokens:
model_type_instance = cast(LargeLanguageModel, model_instance.model_type_instance)
curr_message_tokens = (
model_type_instance.get_num_tokens(
model_instance.model_name, model_instance.credentials, prompt_messages
)
+ 1000
) # add 1000 to ensure tool call messages
curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages) + 1000
max_tokens = 0
for parameter_rule in model_schema.parameter_rules:
@@ -828,8 +790,34 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
return rest_tokens
def _compile_prompt_messages(
self,
*,
model_instance: PreparedLLMProtocol,
prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate,
files: Sequence[File],
vision_enabled: bool,
context: str | None = "",
image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
) -> list[PromptMessage]:
prompt_messages, _ = LLMNode.fetch_prompt_messages(
sys_query="",
sys_files=files,
context=context,
memory=None,
model_instance=model_instance,
prompt_template=prompt_template,
stop=model_instance.stop,
memory_config=None,
vision_enabled=vision_enabled,
vision_detail=image_detail_config or ImagePromptMessageContent.DETAIL.HIGH,
variable_pool=self.graph_runtime_state.variable_pool,
jinja2_variables=[],
)
return list(prompt_messages)
@property
def model_instance(self) -> ModelInstance:
def model_instance(self) -> PreparedLLMProtocol:
return self._model_instance
@classmethod

8
api/dify_graph/nodes/protocols.py
View File

@@ -1,4 +1,4 @@
from collections.abc import Generator
from collections.abc import Generator, Mapping
from typing import Any, Protocol
import httpx
@@ -35,8 +35,6 @@ class ToolFileManagerProtocol(Protocol):
def create_file_by_raw(
self,
*,
user_id: str,
tenant_id: str,
conversation_id: str | None,
file_binary: bytes,
mimetype: str,
@@ -44,3 +42,7 @@ class ToolFileManagerProtocol(Protocol):
) -> Any: ...
def get_file_generator_by_tool_file_id(self, tool_file_id: str) -> tuple[Generator | None, ToolFile | None]: ...
class FileReferenceFactoryProtocol(Protocol):
def build_from_mapping(self, *, mapping: Mapping[str, Any]) -> File: ...

2
api/dify_graph/nodes/question_classifier/entities.py
View File

@@ -1,9 +1,9 @@
from pydantic import BaseModel, Field
from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from dify_graph.entities.base_node_data import BaseNodeData
from dify_graph.enums import BuiltinNodeTypes, NodeType
from dify_graph.nodes.llm import ModelConfig, VisionConfig
from dify_graph.prompt_entities import MemoryConfig
class ClassConfig(BaseModel):

56
api/dify_graph/nodes/question_classifier/question_classifier_node.py
View File

@@ -3,9 +3,6 @@ import re
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any
from core.model_manager import ModelInstance
from core.prompt.simple_prompt_transform import ModelMode
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from dify_graph.entities import GraphInitParams
from dify_graph.entities.graph_config import NodeConfigDict
from dify_graph.enums import (
@@ -14,7 +11,7 @@ from dify_graph.enums import (
WorkflowNodeExecutionMetadataKey,
WorkflowNodeExecutionStatus,
)
from dify_graph.model_runtime.entities import LLMUsage, ModelPropertyKey, PromptMessageRole
from dify_graph.model_runtime.entities import LLMMode, LLMUsage, ModelPropertyKey, PromptMessageRole
from dify_graph.model_runtime.memory import PromptMessageMemory
from dify_graph.model_runtime.utils.encoders import jsonable_encoder
from dify_graph.node_events import ModelInvokeCompletedEvent, NodeRunResult
@@ -27,10 +24,11 @@ from dify_graph.nodes.llm import (
LLMNodeCompletionModelPromptTemplate,
llm_utils,
)
from dify_graph.nodes.llm.file_saver import FileSaverImpl, LLMFileSaver
from dify_graph.nodes.llm.protocols import CredentialsProvider, ModelFactory, TemplateRenderer
from dify_graph.nodes.llm.file_saver import LLMFileSaver
from dify_graph.nodes.llm.protocols import TemplateRenderer
from dify_graph.nodes.llm.runtime_protocols import PreparedLLMProtocol, PromptMessageSerializerProtocol
from dify_graph.nodes.protocols import HttpClientProtocol
from libs.json_in_md_parser import parse_and_check_json_markdown
from dify_graph.utils.json_in_md_parser import parse_and_check_json_markdown
from .entities import QuestionClassifierNodeData
from .exc import InvalidModelTypeError
@@ -49,15 +47,20 @@ if TYPE_CHECKING:
from dify_graph.runtime import GraphRuntimeState
class _PassthroughPromptMessageSerializer:
def serialize(self, *, model_mode: Any, prompt_messages: Sequence[Any]) -> Any:
_ = model_mode
return list(prompt_messages)
class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
node_type = BuiltinNodeTypes.QUESTION_CLASSIFIER
execution_type = NodeExecutionType.BRANCH
_file_outputs: list["File"]
_llm_file_saver: LLMFileSaver
_credentials_provider: "CredentialsProvider"
_model_factory: "ModelFactory"
_model_instance: ModelInstance
_prompt_message_serializer: PromptMessageSerializerProtocol
_model_instance: PreparedLLMProtocol
_memory: PromptMessageMemory | None
_template_renderer: TemplateRenderer
@@ -68,13 +71,14 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
graph_init_params: "GraphInitParams",
graph_runtime_state: "GraphRuntimeState",
*,
credentials_provider: "CredentialsProvider",
model_factory: "ModelFactory",
model_instance: ModelInstance,
credentials_provider: object | None = None,
model_factory: object | None = None,
model_instance: PreparedLLMProtocol,
http_client: HttpClientProtocol,
template_renderer: TemplateRenderer,
memory: PromptMessageMemory | None = None,
llm_file_saver: LLMFileSaver | None = None,
llm_file_saver: LLMFileSaver,
prompt_message_serializer: PromptMessageSerializerProtocol | None = None,
):
super().__init__(
id=id,
@@ -85,20 +89,13 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
# LLM file outputs, used for MultiModal outputs.
self._file_outputs = []
self._credentials_provider = credentials_provider
self._model_factory = model_factory
_ = credentials_provider, model_factory, http_client
self._model_instance = model_instance
self._memory = memory
self._template_renderer = template_renderer
if llm_file_saver is None:
dify_ctx = self.require_dify_context()
llm_file_saver = FileSaverImpl(
user_id=dify_ctx.user_id,
tenant_id=dify_ctx.tenant_id,
http_client=http_client,
)
self._llm_file_saver = llm_file_saver
self._prompt_message_serializer = prompt_message_serializer or _PassthroughPromptMessageSerializer()
@classmethod
def version(cls):
@@ -169,7 +166,6 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
model_instance=model_instance,
prompt_messages=prompt_messages,
stop=stop,
user_id=self.require_dify_context().user_id,
structured_output_enabled=False,
structured_output=None,
file_saver=self._llm_file_saver,
@@ -205,7 +201,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
category_id = category_id_result
process_data = {
"model_mode": node_data.model.mode,
"prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
"prompts": self._prompt_message_serializer.serialize(
model_mode=node_data.model.mode, prompt_messages=prompt_messages
),
"usage": jsonable_encoder(usage),
@@ -247,7 +243,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
)
@property
def model_instance(self) -> ModelInstance:
def model_instance(self) -> PreparedLLMProtocol:
return self._model_instance
@classmethod
@@ -285,7 +281,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
self,
node_data: QuestionClassifierNodeData,
query: str,
model_instance: ModelInstance,
model_instance: PreparedLLMProtocol,
context: str | None,
) -> int:
model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
@@ -334,7 +330,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
memory: PromptMessageMemory | None,
max_token_limit: int = 2000,
):
model_mode = ModelMode(node_data.model.mode)
model_mode = LLMMode(node_data.model.mode)
classes = node_data.classes
categories = []
for class_ in classes:
@@ -350,7 +346,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
message_limit=node_data.memory.window.size if node_data.memory and node_data.memory.window else None,
)
prompt_messages: list[LLMNodeChatModelMessage] = []
if model_mode == ModelMode.CHAT:
if model_mode == LLMMode.CHAT:
system_prompt_messages = LLMNodeChatModelMessage(
role=PromptMessageRole.SYSTEM, text=QUESTION_CLASSIFIER_SYSTEM_PROMPT.format(histories=memory_str)
)
@@ -381,7 +377,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
)
prompt_messages.append(user_prompt_message_3)
return prompt_messages
elif model_mode == ModelMode.COMPLETION:
elif model_mode == LLMMode.COMPLETION:
return LLMNodeCompletionModelPromptTemplate(
text=QUESTION_CLASSIFIER_COMPLETION_PROMPT.format(
histories=memory_str,

75
api/dify_graph/nodes/runtime.py Normal file
View File

@@ -0,0 +1,75 @@
from __future__ import annotations
from collections.abc import Generator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Protocol
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.nodes.tool_runtime_entities import (
ToolRuntimeHandle,
ToolRuntimeMessage,
ToolRuntimeParameter,
)
if TYPE_CHECKING:
from dify_graph.nodes.human_input.entities import DeliveryChannelConfig
from dify_graph.nodes.tool.entities import ToolNodeData
from dify_graph.runtime import VariablePool
class ToolNodeRuntimeProtocol(Protocol):
"""Workflow-layer adapter owned by `core.workflow` and consumed by `dify_graph`.
The graph package depends only on these DTOs and lets the workflow layer
translate between graph-owned abstractions and `core.tools` internals.
"""
def get_runtime(
self,
*,
node_id: str,
node_data: ToolNodeData,
variable_pool: VariablePool | None,
) -> ToolRuntimeHandle: ...
def get_runtime_parameters(
self,
*,
tool_runtime: ToolRuntimeHandle,
) -> Sequence[ToolRuntimeParameter]: ...
def invoke(
self,
*,
tool_runtime: ToolRuntimeHandle,
tool_parameters: Mapping[str, Any],
workflow_call_depth: int,
conversation_id: str | None,
provider_name: str,
) -> Generator[ToolRuntimeMessage, None, None]: ...
def get_usage(
self,
*,
tool_runtime: ToolRuntimeHandle,
) -> LLMUsage: ...
def build_file_reference(self, *, mapping: Mapping[str, Any]) -> Any: ...
def resolve_provider_icons(
self,
*,
provider_name: str,
default_icon: str | None = None,
) -> tuple[str | None, str | None]: ...
class HumanInputNodeRuntimeProtocol(Protocol):
def invoke_source(self) -> str: ...
def apply_delivery_runtime(
self,
*,
methods: Sequence[DeliveryChannelConfig],
) -> Sequence[DeliveryChannelConfig]: ...
def console_actor_id(self) -> str | None: ...

44
api/dify_graph/nodes/template_transform/template_transform_node.py
View File

@@ -4,9 +4,10 @@ from typing import TYPE_CHECKING, Any
from dify_graph.entities.graph_config import NodeConfigDict
from dify_graph.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus
from dify_graph.node_events import NodeRunResult
from dify_graph.nodes.base.entities import VariableSelector
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.template_transform.entities import TemplateTransformNodeData
from dify_graph.nodes.template_transform.template_renderer import (
from dify_graph.template_rendering import (
Jinja2TemplateRenderer,
TemplateRenderError,
)
@@ -20,7 +21,7 @@ DEFAULT_TEMPLATE_TRANSFORM_MAX_OUTPUT_LENGTH = 400_000
class TemplateTransformNode(Node[TemplateTransformNodeData]):
node_type = BuiltinNodeTypes.TEMPLATE_TRANSFORM
_template_renderer: Jinja2TemplateRenderer
_jinja2_template_renderer: Jinja2TemplateRenderer
_max_output_length: int
def __init__(
@@ -30,7 +31,7 @@ class TemplateTransformNode(Node[TemplateTransformNodeData]):
graph_init_params: "GraphInitParams",
graph_runtime_state: "GraphRuntimeState",
*,
template_renderer: Jinja2TemplateRenderer,
jinja2_template_renderer: Jinja2TemplateRenderer,
max_output_length: int | None = None,
) -> None:
super().__init__(
@@ -39,7 +40,7 @@ class TemplateTransformNode(Node[TemplateTransformNodeData]):
graph_init_params=graph_init_params,
graph_runtime_state=graph_runtime_state,
)
self._template_renderer = template_renderer
self._jinja2_template_renderer = jinja2_template_renderer
if max_output_length is not None and max_output_length <= 0:
raise ValueError("max_output_length must be a positive integer")
@@ -70,7 +71,7 @@ class TemplateTransformNode(Node[TemplateTransformNodeData]):
variables[variable_name] = value.to_object() if value else None
# Run code
try:
rendered = self._template_renderer.render_template(self.node_data.template, variables)
rendered = self._jinja2_template_renderer.render_template(self.node_data.template, variables)
except TemplateRenderError as e:
return NodeRunResult(inputs=variables, status=WorkflowNodeExecutionStatus.FAILED, error=str(e))
@@ -87,9 +88,32 @@ class TemplateTransformNode(Node[TemplateTransformNodeData]):
@classmethod
def _extract_variable_selector_to_variable_mapping(
cls, *, graph_config: Mapping[str, Any], node_id: str, node_data: TemplateTransformNodeData
cls,
*,
graph_config: Mapping[str, Any],
node_id: str,
node_data: TemplateTransformNodeData | Mapping[str, Any],
) -> Mapping[str, Sequence[str]]:
return {
node_id + "." + variable_selector.variable: variable_selector.value_selector
for variable_selector in node_data.variables
}
_ = graph_config
raw_variables = (
node_data.variables if isinstance(node_data, TemplateTransformNodeData) else node_data.get("variables", [])
)
variable_mapping: dict[str, Sequence[str]] = {}
for variable_selector in raw_variables:
if isinstance(variable_selector, VariableSelector):
variable_mapping[node_id + "." + variable_selector.variable] = variable_selector.value_selector
continue
if not isinstance(variable_selector, Mapping):
continue
variable = variable_selector.get("variable")
value_selector = variable_selector.get("value_selector")
if (
isinstance(variable, str)
and isinstance(value_selector, Sequence)
and all(isinstance(selector_part, str) for selector_part in value_selector)
):
variable_mapping[node_id + "." + variable] = list(value_selector)
return variable_mapping

16
api/dify_graph/nodes/tool/entities.py
View File

@@ -1,13 +1,27 @@
from enum import StrEnum, auto
from typing import Any, Literal, Union
from pydantic import BaseModel, field_validator
from pydantic_core.core_schema import ValidationInfo
from core.tools.entities.tool_entities import ToolProviderType
from dify_graph.entities.base_node_data import BaseNodeData
from dify_graph.enums import BuiltinNodeTypes, NodeType
class ToolProviderType(StrEnum):
"""
Graph-owned enum for persisted tool provider kinds.
"""
PLUGIN = auto()
BUILT_IN = "builtin"
WORKFLOW = auto()
API = auto()
APP = auto()
DATASET_RETRIEVAL = "dataset-retrieval"
MCP = auto()
class ToolEntity(BaseModel):
provider_id: str
provider_type: ToolProviderType

12
api/dify_graph/nodes/tool/exc.py
View File

@@ -4,6 +4,18 @@ class ToolNodeError(ValueError):
pass
class ToolRuntimeResolutionError(ToolNodeError):
"""Raised when the workflow layer cannot construct a tool runtime."""
pass
class ToolRuntimeInvocationError(ToolNodeError):
"""Raised when the workflow layer fails while invoking a tool runtime."""
pass
class ToolParameterError(ToolNodeError):
"""Exception raised for errors in tool parameters."""

189
api/dify_graph/nodes/tool/tool_node.py
View File

@@ -1,12 +1,6 @@
from collections.abc import Generator, Mapping, Sequence
from typing import TYPE_CHECKING, Any
from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
from core.tools.__base.tool import Tool
from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
from core.tools.errors import ToolInvokeError
from core.tools.tool_engine import ToolEngine
from core.tools.utils.message_transformer import ToolFileMessageTransformer
from dify_graph.entities.graph_config import NodeConfigDict
from dify_graph.enums import (
BuiltinNodeTypes,
@@ -20,10 +14,14 @@ from dify_graph.node_events import NodeEventBase, NodeRunResult, StreamChunkEven
from dify_graph.nodes.base.node import Node
from dify_graph.nodes.base.variable_template_parser import VariableTemplateParser
from dify_graph.nodes.protocols import ToolFileManagerProtocol
from dify_graph.nodes.runtime import ToolNodeRuntimeProtocol
from dify_graph.nodes.tool_runtime_entities import (
ToolRuntimeHandle,
ToolRuntimeMessage,
ToolRuntimeParameter,
)
from dify_graph.variables.segments import ArrayAnySegment, ArrayFileSegment
from dify_graph.variables.variables import ArrayAnyVariable
from factories import file_factory
from services.tools.builtin_tools_manage_service import BuiltinToolManageService
from .entities import ToolNodeData
from .exc import (
@@ -52,6 +50,7 @@ class ToolNode(Node[ToolNodeData]):
graph_runtime_state: "GraphRuntimeState",
*,
tool_file_manager_factory: ToolFileManagerProtocol,
runtime: ToolNodeRuntimeProtocol | None = None,
):
super().__init__(
id=id,
@@ -60,6 +59,9 @@ class ToolNode(Node[ToolNodeData]):
graph_runtime_state=graph_runtime_state,
)
self._tool_file_manager_factory = tool_file_manager_factory
if runtime is None:
raise ValueError("runtime is required")
self._runtime = runtime
@classmethod
def version(cls) -> str:
@@ -73,10 +75,6 @@ class ToolNode(Node[ToolNodeData]):
"""
Run the tool node
"""
from core.plugin.impl.exc import PluginDaemonClientSideError, PluginInvokeError
dify_ctx = self.require_dify_context()
# fetch tool icon
tool_info = {
"provider_type": self.node_data.provider_type.value,
@@ -86,8 +84,6 @@ class ToolNode(Node[ToolNodeData]):
# get tool runtime
try:
from core.tools.tool_manager import ToolManager
# This is an issue that caused problems before.
# Logically, we shouldn't use the node_data.version field for judgment
# But for backward compatibility with historical data
@@ -95,13 +91,10 @@ class ToolNode(Node[ToolNodeData]):
variable_pool: VariablePool | None = None
if self.node_data.version != "1" or self.node_data.tool_node_version is not None:
variable_pool = self.graph_runtime_state.variable_pool
tool_runtime = ToolManager.get_workflow_tool_runtime(
dify_ctx.tenant_id,
dify_ctx.app_id,
self._node_id,
self.node_data,
dify_ctx.invoke_from,
variable_pool,
tool_runtime = self._runtime.get_runtime(
node_id=self._node_id,
node_data=self.node_data,
variable_pool=variable_pool,
)
except ToolNodeError as e:
yield StreamCompletedEvent(
@@ -116,7 +109,7 @@ class ToolNode(Node[ToolNodeData]):
return
# get parameters
tool_parameters = tool_runtime.get_merged_runtime_parameters() or []
tool_parameters = self._runtime.get_runtime_parameters(tool_runtime=tool_runtime)
parameters = self._generate_parameters(
tool_parameters=tool_parameters,
variable_pool=self.graph_runtime_state.variable_pool,
@@ -132,14 +125,12 @@ class ToolNode(Node[ToolNodeData]):
conversation_id = self.graph_runtime_state.variable_pool.get(["sys", SystemVariableKey.CONVERSATION_ID])
try:
message_stream = ToolEngine.generic_invoke(
tool=tool_runtime,
message_stream = self._runtime.invoke(
tool_runtime=tool_runtime,
tool_parameters=parameters,
user_id=dify_ctx.user_id,
workflow_tool_callback=DifyWorkflowCallbackHandler(),
workflow_call_depth=self.workflow_call_depth,
app_id=dify_ctx.app_id,
conversation_id=conversation_id.text if conversation_id else None,
provider_name=self.node_data.provider_name,
)
except ToolNodeError as e:
yield StreamCompletedEvent(
@@ -159,38 +150,16 @@ class ToolNode(Node[ToolNodeData]):
messages=message_stream,
tool_info=tool_info,
parameters_for_log=parameters_for_log,
user_id=dify_ctx.user_id,
tenant_id=dify_ctx.tenant_id,
node_id=self._node_id,
tool_runtime=tool_runtime,
)
except ToolInvokeError as e:
except ToolNodeError as e:
yield StreamCompletedEvent(
node_run_result=NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
inputs=parameters_for_log,
metadata={WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info},
error=f"Failed to invoke tool {self.node_data.provider_name}: {str(e)}",
error_type=type(e).__name__,
)
)
except PluginInvokeError as e:
yield StreamCompletedEvent(
node_run_result=NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
inputs=parameters_for_log,
metadata={WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info},
error=e.to_user_friendly_error(plugin_name=self.node_data.provider_name),
error_type=type(e).__name__,
)
)
except PluginDaemonClientSideError as e:
yield StreamCompletedEvent(
node_run_result=NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
inputs=parameters_for_log,
metadata={WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info},
error=f"Failed to invoke tool, error: {e.description}",
error=str(e),
error_type=type(e).__name__,
)
)
@@ -198,7 +167,7 @@ class ToolNode(Node[ToolNodeData]):
def _generate_parameters(
self,
*,
tool_parameters: Sequence[ToolParameter],
tool_parameters: Sequence[ToolRuntimeParameter],
variable_pool: "VariablePool",
node_data: ToolNodeData,
for_log: bool = False,
@@ -207,7 +176,7 @@ class ToolNode(Node[ToolNodeData]):
Generate parameters based on the given tool parameters, variable pool, and node data.
Args:
tool_parameters (Sequence[ToolParameter]): The list of tool parameters.
tool_parameters (Sequence[ToolRuntimeParameter]): The list of tool parameters.
variable_pool (VariablePool): The variable pool containing the variables.
node_data (ToolNodeData): The data associated with the tool node.
@@ -247,40 +216,29 @@ class ToolNode(Node[ToolNodeData]):
def _transform_message(
self,
messages: Generator[ToolInvokeMessage, None, None],
messages: Generator[ToolRuntimeMessage, None, None],
tool_info: Mapping[str, Any],
parameters_for_log: dict[str, Any],
user_id: str,
tenant_id: str,
node_id: str,
tool_runtime: Tool,
tool_runtime: ToolRuntimeHandle,
**_: Any,
) -> Generator[NodeEventBase, None, LLMUsage]:
"""
Convert ToolInvokeMessages into tuple[plain_text, files]
Convert graph-owned tool runtime messages into node outputs.
"""
# transform message and handle file storage
from core.plugin.impl.plugin import PluginInstaller
message_stream = ToolFileMessageTransformer.transform_tool_invoke_messages(
messages=messages,
user_id=user_id,
tenant_id=tenant_id,
conversation_id=None,
)
text = ""
files: list[File] = []
json: list[dict | list] = []
variables: dict[str, Any] = {}
for message in message_stream:
for message in messages:
if message.type in {
ToolInvokeMessage.MessageType.IMAGE_LINK,
ToolInvokeMessage.MessageType.BINARY_LINK,
ToolInvokeMessage.MessageType.IMAGE,
ToolRuntimeMessage.MessageType.IMAGE_LINK,
ToolRuntimeMessage.MessageType.BINARY_LINK,
ToolRuntimeMessage.MessageType.IMAGE,
}:
assert isinstance(message.message, ToolInvokeMessage.TextMessage)
assert isinstance(message.message, ToolRuntimeMessage.TextMessage)
url = message.message.text
if message.meta:
@@ -300,14 +258,11 @@ class ToolNode(Node[ToolNodeData]):
"transfer_method": transfer_method,
"url": url,
}
file = file_factory.build_from_mapping(
mapping=mapping,
tenant_id=tenant_id,
)
file = self._runtime.build_file_reference(mapping=mapping)
files.append(file)
elif message.type == ToolInvokeMessage.MessageType.BLOB:
elif message.type == ToolRuntimeMessage.MessageType.BLOB:
# get tool file id
assert isinstance(message.message, ToolInvokeMessage.TextMessage)
assert isinstance(message.message, ToolRuntimeMessage.TextMessage)
assert message.meta
tool_file_id = message.message.text.split("/")[-1].split(".")[0]
@@ -320,27 +275,22 @@ class ToolNode(Node[ToolNodeData]):
"transfer_method": FileTransferMethod.TOOL_FILE,
}
files.append(
file_factory.build_from_mapping(
mapping=mapping,
tenant_id=tenant_id,
)
)
elif message.type == ToolInvokeMessage.MessageType.TEXT:
assert isinstance(message.message, ToolInvokeMessage.TextMessage)
files.append(self._runtime.build_file_reference(mapping=mapping))
elif message.type == ToolRuntimeMessage.MessageType.TEXT:
assert isinstance(message.message, ToolRuntimeMessage.TextMessage)
text += message.message.text
yield StreamChunkEvent(
selector=[node_id, "text"],
chunk=message.message.text,
is_final=False,
)
elif message.type == ToolInvokeMessage.MessageType.JSON:
assert isinstance(message.message, ToolInvokeMessage.JsonMessage)
elif message.type == ToolRuntimeMessage.MessageType.JSON:
assert isinstance(message.message, ToolRuntimeMessage.JsonMessage)
# JSON message handling for tool node
if message.message.json_object:
json.append(message.message.json_object)
elif message.type == ToolInvokeMessage.MessageType.LINK:
assert isinstance(message.message, ToolInvokeMessage.TextMessage)
elif message.type == ToolRuntimeMessage.MessageType.LINK:
assert isinstance(message.message, ToolRuntimeMessage.TextMessage)
# Check if this LINK message is a file link
file_obj = (message.meta or {}).get("file")
@@ -356,8 +306,8 @@ class ToolNode(Node[ToolNodeData]):
chunk=stream_text,
is_final=False,
)
elif message.type == ToolInvokeMessage.MessageType.VARIABLE:
assert isinstance(message.message, ToolInvokeMessage.VariableMessage)
elif message.type == ToolRuntimeMessage.MessageType.VARIABLE:
assert isinstance(message.message, ToolRuntimeMessage.VariableMessage)
variable_name = message.message.variable_name
variable_value = message.message.variable_value
if message.message.stream:
@@ -374,7 +324,7 @@ class ToolNode(Node[ToolNodeData]):
)
else:
variables[variable_name] = variable_value
elif message.type == ToolInvokeMessage.MessageType.FILE:
elif message.type == ToolRuntimeMessage.MessageType.FILE:
assert message.meta is not None
assert isinstance(message.meta, dict)
# Validate that meta contains a 'file' key
@@ -385,38 +335,16 @@ class ToolNode(Node[ToolNodeData]):
if not isinstance(message.meta["file"], File):
raise ToolNodeError(f"Expected File object but got {type(message.meta['file']).__name__}")
files.append(message.meta["file"])
elif message.type == ToolInvokeMessage.MessageType.LOG:
assert isinstance(message.message, ToolInvokeMessage.LogMessage)
elif message.type == ToolRuntimeMessage.MessageType.LOG:
assert isinstance(message.message, ToolRuntimeMessage.LogMessage)
if message.message.metadata:
icon = tool_info.get("icon", "")
dict_metadata = dict(message.message.metadata)
if dict_metadata.get("provider"):
manager = PluginInstaller()
plugins = manager.list_plugins(tenant_id)
try:
current_plugin = next(
plugin
for plugin in plugins
if f"{plugin.plugin_id}/{plugin.name}" == dict_metadata["provider"]
)
icon = current_plugin.declaration.icon
except StopIteration:
pass
icon_dark = None
try:
builtin_tool = next(
provider
for provider in BuiltinToolManageService.list_builtin_tools(
user_id,
tenant_id,
)
if provider.name == dict_metadata["provider"]
)
icon = builtin_tool.icon
icon_dark = builtin_tool.icon_dark
except StopIteration:
pass
icon, icon_dark = self._runtime.resolve_provider_icons(
provider_name=dict_metadata["provider"],
default_icon=icon,
)
dict_metadata["icon"] = icon
dict_metadata["icon_dark"] = icon_dark
message.message.metadata = dict_metadata
@@ -446,7 +374,7 @@ class ToolNode(Node[ToolNodeData]):
is_final=True,
)
usage = self._extract_tool_usage(tool_runtime)
usage = self._runtime.get_usage(tool_runtime=tool_runtime)
metadata: dict[WorkflowNodeExecutionMetadataKey, Any] = {
WorkflowNodeExecutionMetadataKey.TOOL_INFO: tool_info,
@@ -468,21 +396,6 @@ class ToolNode(Node[ToolNodeData]):
return usage
@staticmethod
def _extract_tool_usage(tool_runtime: Tool) -> LLMUsage:
# Avoid importing WorkflowTool at module import time; rely on duck typing
# Some runtimes expose `latest_usage`; mocks may synthesize arbitrary attributes.
latest = getattr(tool_runtime, "latest_usage", None)
# Normalize into a concrete LLMUsage. MagicMock returns truthy attribute objects
# for any name, so we must type-check here.
if isinstance(latest, LLMUsage):
return latest
if isinstance(latest, dict):
# Allow dict payloads from external runtimes
return LLMUsage.model_validate(latest)
# Fallback to empty usage when attribute is missing or not a valid payload
return LLMUsage.empty_usage()
@classmethod
def _extract_variable_selector_to_variable_mapping(
cls,

101
api/dify_graph/nodes/tool_runtime_entities.py Normal file
View File

@@ -0,0 +1,101 @@
from __future__ import annotations
from dataclasses import dataclass
from enum import StrEnum, auto
from typing import Any
from pydantic import BaseModel, ConfigDict, Field
class _ToolRuntimeModel(BaseModel):
model_config = ConfigDict(extra="forbid")
@dataclass(frozen=True, slots=True)
class ToolRuntimeHandle:
"""Opaque graph-owned handle for a workflow-layer tool runtime."""
raw: object
@dataclass(frozen=True, slots=True)
class ToolRuntimeParameter:
"""Graph-owned parameter shape used by tool nodes."""
name: str
required: bool = False
class ToolRuntimeMessage(_ToolRuntimeModel):
"""Graph-owned tool invocation message DTO."""
class TextMessage(_ToolRuntimeModel):
text: str
class JsonMessage(_ToolRuntimeModel):
json_object: dict[str, Any] | list[Any]
suppress_output: bool = Field(default=False)
class BlobMessage(_ToolRuntimeModel):
blob: bytes
class BlobChunkMessage(_ToolRuntimeModel):
id: str
sequence: int
total_length: int
blob: bytes
end: bool
class FileMessage(_ToolRuntimeModel):
file_marker: str = Field(default="file_marker")
class VariableMessage(_ToolRuntimeModel):
variable_name: str
variable_value: dict[str, Any] | list[Any] | str | int | float | bool | None
stream: bool = Field(default=False)
class LogMessage(_ToolRuntimeModel):
class LogStatus(StrEnum):
START = auto()
ERROR = auto()
SUCCESS = auto()
id: str
label: str
parent_id: str | None = None
error: str | None = None
status: LogStatus
data: dict[str, Any]
metadata: dict[str, Any] = Field(default_factory=dict)
class RetrieverResourceMessage(_ToolRuntimeModel):
retriever_resources: list[dict[str, Any]]
context: str
class MessageType(StrEnum):
TEXT = auto()
IMAGE = auto()
LINK = auto()
BLOB = auto()
JSON = auto()
IMAGE_LINK = auto()
BINARY_LINK = auto()
VARIABLE = auto()
FILE = auto()
LOG = auto()
BLOB_CHUNK = auto()
RETRIEVER_RESOURCES = auto()
type: MessageType = MessageType.TEXT
message: (
JsonMessage
| TextMessage
| BlobChunkMessage
| BlobMessage
| LogMessage
| FileMessage
| None
| VariableMessage
| RetrieverResourceMessage
)
meta: dict[str, Any] | None = None

47
api/dify_graph/prompt_entities.py Normal file
View File

@@ -0,0 +1,47 @@
from typing import Literal
from pydantic import BaseModel
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
class ChatModelMessage(BaseModel):
"""Graph-owned chat prompt template message."""
text: str
role: PromptMessageRole
edition_type: Literal["basic", "jinja2"] | None = None
class CompletionModelPromptTemplate(BaseModel):
"""Graph-owned completion prompt template."""
text: str
edition_type: Literal["basic", "jinja2"] | None = None
class MemoryConfig(BaseModel):
"""Graph-owned memory configuration for prompt assembly."""
class RolePrefix(BaseModel):
"""Role labels used when serializing completion-model histories."""
user: str
assistant: str
class WindowConfig(BaseModel):
"""History windowing controls."""
enabled: bool
size: int | None = None
role_prefix: RolePrefix | None = None
window: WindowConfig
query_prompt_template: str | None = None
__all__ = [
"ChatModelMessage",
"CompletionModelPromptTemplate",
"MemoryConfig",
]

6
api/dify_graph/repositories/human_input_form_repository.py
View File

@@ -18,9 +18,6 @@ class FormNotFoundError(HumanInputError):
@dataclasses.dataclass
class FormCreateParams:
# app_id is the identifier for the app that the form belongs to.
# It is a string with uuid format.
app_id: str
# None when creating a delivery test form; set for runtime forms.
workflow_execution_id: str | None
@@ -45,6 +42,9 @@ class FormCreateParams:
resolved_default_values: Mapping[str, Any]
form_kind: HumanInputFormKind = HumanInputFormKind.RUNTIME
# Optional application identifier. Implementations may bind this at construction time.
app_id: str | None = None
# Force creating a console-only recipient for submission in Console.
console_recipient_required: bool = False
console_creator_account_id: str | None = None

10
api/dify_graph/nodes/template_transform/template_renderer.py → api/dify_graph/template_rendering.py
View File

@@ -8,19 +8,17 @@ from dify_graph.nodes.code.entities import CodeLanguage
class TemplateRenderError(ValueError):
"""Raised when rendering a Jinja2 template fails."""
"""Raised when rendering a template fails."""
class Jinja2TemplateRenderer(Protocol):
"""Render Jinja2 templates for template transform nodes."""
"""Shared contract for rendering Jinja2 templates in graph nodes."""
def render_template(self, template: str, variables: Mapping[str, Any]) -> str:
"""Render a Jinja2 template with provided variables."""
raise NotImplementedError
def render_template(self, template: str, variables: Mapping[str, Any]) -> str: ...
class CodeExecutorJinja2TemplateRenderer(Jinja2TemplateRenderer):
"""Adapter that renders Jinja2 templates via CodeExecutor."""
"""Adapter that renders Jinja2 templates via the workflow code executor."""
_code_executor: WorkflowCodeExecutor

20
api/dify_graph/utils/datetime_utils.py Normal file
View File

@@ -0,0 +1,20 @@
from __future__ import annotations
import abc
import datetime
from typing import Protocol
class _NowFunction(Protocol):
@abc.abstractmethod
def __call__(self, tz: datetime.timezone | None) -> datetime.datetime:
"""Return the current time for the requested timezone."""
...
_now_func: _NowFunction = datetime.datetime.now
def naive_utc_now() -> datetime.datetime:
"""Return the current UTC time as a naive datetime."""
return _now_func(datetime.UTC).replace(tzinfo=None)

58
api/dify_graph/utils/json_in_md_parser.py Normal file
View File

@@ -0,0 +1,58 @@
from __future__ import annotations
import json
class OutputParserError(ValueError):
"""Raised when a markdown-wrapped JSON payload cannot be parsed or validated."""
def parse_json_markdown(json_string: str) -> dict | list:
"""Extract and parse the first JSON object or array embedded in markdown text."""
json_string = json_string.strip()
starts = ["```json", "```", "``", "`", "{", "["]
ends = ["```", "``", "`", "}", "]"]
end_index = -1
start_index = 0
for start_marker in starts:
start_index = json_string.find(start_marker)
if start_index != -1:
if json_string[start_index] not in ("{", "["):
start_index += len(start_marker)
break
if start_index != -1:
for end_marker in ends:
end_index = json_string.rfind(end_marker, start_index)
if end_index != -1:
if json_string[end_index] in ("}", "]"):
end_index += 1
break
if start_index == -1 or end_index == -1 or start_index >= end_index:
raise ValueError("could not find json block in the output.")
extracted_content = json_string[start_index:end_index].strip()
return json.loads(extracted_content)
def parse_and_check_json_markdown(text: str, expected_keys: list[str]) -> dict:
try:
json_obj = parse_json_markdown(text)
except json.JSONDecodeError as exc:
raise OutputParserError(f"got invalid json object. error: {exc}") from exc
if isinstance(json_obj, list):
if len(json_obj) == 1 and isinstance(json_obj[0], dict):
json_obj = json_obj[0]
else:
raise OutputParserError(f"got invalid return object. obj:{json_obj}")
for key in expected_keys:
if key not in json_obj:
raise OutputParserError(
f"got invalid return object. expected key `{key}` to be present, but got {json_obj}"
)
return json_obj