diff --git a/examples/github_trigger/events/code_scanning_alert/code_scanning_alert.py b/examples/github_trigger/events/code_scanning_alert/code_scanning_alert.py index 9927779d..ba17dcc0 100644 --- a/examples/github_trigger/events/code_scanning_alert/code_scanning_alert.py +++ b/examples/github_trigger/events/code_scanning_alert/code_scanning_alert.py @@ -91,7 +91,6 @@ def _check_branch(self, alert: Mapping[str, Any], value: str | None) -> None: or alert.get("ref") or "" ) - # Convert refs/heads/main -> main - branch = ref.split("/", 2)[-1] if ref.startswith("refs/heads/") else ref + branch = ref.removeprefix("refs/heads/") if branch not in branches: raise EventIgnoreError diff --git a/examples/github_trigger/events/push/push.py b/examples/github_trigger/events/push/push.py index b471d13c..7f8326ef 100644 --- a/examples/github_trigger/events/push/push.py +++ b/examples/github_trigger/events/push/push.py @@ -62,11 +62,7 @@ def _check_branch( return current_ref = payload.get("ref") or "" - branch = ( - current_ref.split("/", 2)[-1] - if current_ref.startswith("refs/heads/") - else current_ref - ) + branch = current_ref.removeprefix("refs/heads/") if branch not in allowed_branches: raise EventIgnoreError diff --git a/examples/github_trigger/events/secret_scanning/secret_scanning.py b/examples/github_trigger/events/secret_scanning/secret_scanning.py index 0eff9a19..d17ffbfe 100644 --- a/examples/github_trigger/events/secret_scanning/secret_scanning.py +++ b/examples/github_trigger/events/secret_scanning/secret_scanning.py @@ -87,6 +87,6 @@ def _check_branch(self, payload: Mapping[str, Any], value: str | None) -> None: ref = (payload.get("alert") or {}).get("most_recent_instance", {}).get( "ref" ) or "" - branch = ref.split("/", 2)[-1] if ref.startswith("refs/heads/") else ref + branch = ref.removeprefix("refs/heads/") if branch and branch not in branches: raise EventIgnoreError diff --git a/examples/google_calendar_trigger/provider/google_calendar_trigger.py b/examples/google_calendar_trigger/provider/google_calendar_trigger.py index a7f75929..785c8f80 100644 --- a/examples/google_calendar_trigger/provider/google_calendar_trigger.py +++ b/examples/google_calendar_trigger/provider/google_calendar_trigger.py @@ -35,6 +35,7 @@ FALSY_STRINGS = frozenset({"false", "0", "no", "off"}) WATCH_SUCCESS_STATUSES = frozenset({HTTPStatus.OK, HTTPStatus.CREATED}) CHANNEL_STOP_SUCCESS_STATUSES = frozenset({HTTPStatus.OK, HTTPStatus.NO_CONTENT}) +RFC3339_UTC_SUFFIX = "Z" class SyncTokenExpiredError(TriggerError): @@ -146,8 +147,8 @@ def _parse_rfc3339(value: object) -> datetime.datetime | None: return None try: text = value.strip() - if text.endswith("Z"): - text = text[:-1] + "+00:00" + if text.endswith(RFC3339_UTC_SUFFIX): + text = f"{text.removesuffix(RFC3339_UTC_SUFFIX)}+00:00" return datetime.datetime.fromisoformat(text) except Exception: return None diff --git a/examples/notion_datasource/datasources/utils/notion_client.py b/examples/notion_datasource/datasources/utils/notion_client.py index c8b05ceb..9ebb9953 100644 --- a/examples/notion_datasource/datasources/utils/notion_client.py +++ b/examples/notion_datasource/datasources/utils/notion_client.py @@ -474,6 +474,35 @@ def format_rich_text(self, content: str) -> list[dict[str, Any]]: """ return self.create_rich_text(content) + def _normalize_icon( + self, page_icon: dict[str, Any] | None + ) -> dict[str, str] | None: + if not page_icon: + return None + + icon_type = page_icon["type"] + if icon_type in FILE_ICON_TYPES: + url = page_icon[icon_type]["url"] + return { + "type": "url", + "url": url if url.startswith("http") else f"https://www.notion.so{url}", + } + + return {"type": "emoji", "emoji": page_icon[icon_type]} + + def _resolve_parent_id(self, access_token: str, parent: dict[str, Any]) -> str: + parent_type = parent["type"] + match parent_type: + case "block_id": + return self._resolve_block_parent_page_id( + access_token, + parent[parent_type], + ) + case "workspace": + return "root" + case _: + return parent[parent_type] + def get_authorized_pages(self) -> list[OnlineDocumentPage]: pages: list[OnlineDocumentPage] = [] access_token = self.integration_token @@ -489,33 +518,10 @@ def get_authorized_pages(self) -> list[OnlineDocumentPage]: and page_result["properties"][key]["title"] ): title_list = page_result["properties"][key]["title"] - if len(title_list) > 0 and "plain_text" in title_list[0]: + if title_list and "plain_text" in title_list[0]: page_name = title_list[0]["plain_text"] - page_icon = page_result["icon"] - if page_icon: - icon_type = page_icon["type"] - if icon_type in FILE_ICON_TYPES: - url = page_icon[icon_type]["url"] - icon = { - "type": "url", - "url": url - if url.startswith("http") - else f"https://www.notion.so{url}", - } - else: - icon = {"type": "emoji", "emoji": page_icon[icon_type]} - else: - icon = None - parent = page_result["parent"] - parent_type = parent["type"] - if parent_type == "block_id": - parent_id = self._resolve_block_parent_page_id( - access_token, parent[parent_type] - ) - elif parent_type == "workspace": - parent_id = "root" - else: - parent_id = parent[parent_type] + icon = self._normalize_icon(page_result["icon"]) + parent_id = self._resolve_parent_id(access_token, page_result["parent"]) page = OnlineDocumentPage( page_id=page_id, page_name=page_name, @@ -530,35 +536,12 @@ def get_authorized_pages(self) -> list[OnlineDocumentPage]: page_id = database_result["id"] page_name = ( database_result["title"][0]["plain_text"] - if len(database_result["title"]) > 0 + if database_result["title"] else "Untitled" ) - page_icon = database_result["icon"] - if page_icon: - icon_type = page_icon["type"] - if icon_type in FILE_ICON_TYPES: - url = page_icon[icon_type]["url"] - icon = { - "type": "url", - "url": url - if url.startswith("http") - else f"https://www.notion.so{url}", - } - else: - icon = {"type": "emoji", "emoji": page_icon[icon_type]} - else: - icon = None - parent = database_result["parent"] - parent_type = parent["type"] - if parent_type == "block_id": - parent_id = self._resolve_block_parent_page_id( - access_token, parent[parent_type] - ) - elif parent_type == "workspace": - parent_id = "root" - else: - parent_id = parent[parent_type] + icon = self._normalize_icon(database_result["icon"]) + parent_id = self._resolve_parent_id(access_token, database_result["parent"]) page = OnlineDocumentPage( page_id=page_id, page_name=page_name, diff --git a/examples/notion_datasource/datasources/utils/notion_extractor.py b/examples/notion_datasource/datasources/utils/notion_extractor.py index 7bea0953..0fa16ee1 100644 --- a/examples/notion_datasource/datasources/utils/notion_extractor.py +++ b/examples/notion_datasource/datasources/utils/notion_extractor.py @@ -304,40 +304,35 @@ def _format_page_data(self, page_data: dict[str, Any]) -> dict[str, Any]: prop_type = prop_data.get("type") # Extract value based on property type - if prop_type == "title": - title_content = prop_data.get("title", []) - value = self._client.extract_plain_text(title_content) - if value: - title = value # Save title for the result - elif prop_type == "rich_text": - text_content = prop_data.get("rich_text", []) - value = self._client.extract_plain_text(text_content) - elif prop_type == "number": - value = prop_data.get("number") - elif prop_type == "select": - select_data = prop_data.get("select", {}) - value = select_data.get("name") if select_data else None - elif prop_type == "multi_select": - multi_select = prop_data.get("multi_select", []) - value = ( - [item.get("name") for item in multi_select] if multi_select else [] - ) - elif prop_type == "date": - date_data = prop_data.get("date", {}) - start = date_data.get("start") if date_data else None - end = date_data.get("end") if date_data else None - value = {"start": start, "end": end} if start else None - elif prop_type == "checkbox": - value = prop_data.get("checkbox") - elif prop_type == "url": - value = prop_data.get("url") - elif prop_type == "email": - value = prop_data.get("email") - elif prop_type == "phone_number": - value = prop_data.get("phone_number") - else: - # For other property types, just note the type - value = f"<{prop_type}>" + match prop_type: + case "title": + title_content = prop_data.get("title", []) + value = self._client.extract_plain_text(title_content) + if value: + title = value # Save title for the result + case "rich_text": + text_content = prop_data.get("rich_text", []) + value = self._client.extract_plain_text(text_content) + case "number" | "checkbox" | "url" | "email" | "phone_number": + value = prop_data.get(prop_type) + case "select": + select_data = prop_data.get("select", {}) + value = select_data.get("name") if select_data else None + case "multi_select": + multi_select = prop_data.get("multi_select", []) + value = ( + [item.get("name") for item in multi_select] + if multi_select + else [] + ) + case "date": + date_data = prop_data.get("date", {}) + start = date_data.get("start") if date_data else None + end = date_data.get("end") if date_data else None + value = {"start": start, "end": end} if start else None + case _: + # For other property types, just note the type + value = f"<{prop_type}>" formatted_properties[prop_name] = value diff --git a/src/dify_plugin/core/documentation/generator.py b/src/dify_plugin/core/documentation/generator.py index 58f19757..927d5942 100644 --- a/src/dify_plugin/core/documentation/generator.py +++ b/src/dify_plugin/core/documentation/generator.py @@ -262,12 +262,8 @@ def _extract_referenced_types(self, field_type: object) -> set[type]: if isinstance(field_type, type): if issubclass(field_type, (BaseModel, Enum)): referenced.add(field_type) - # Handle generic types (List, Dict, Union, etc) - elif ( - hasattr(field_type, "__origin__") and field_type.__origin__ == Union - ) or hasattr(field_type, "__args__"): - # Handle Union types - for arg in field_type.__args__: + else: + for arg in get_args(field_type): referenced.update(self._extract_referenced_types(arg)) return referenced @@ -347,22 +343,13 @@ def _is_container_type( container_types: tuple[type, ...] = (list, set), ) -> bool: """Check if a field type is a container type (list, set, etc).""" - try: - return ( - hasattr(field_type, "__origin__") - and isinstance(getattr(field_type, "__origin__", None), type) - and getattr(field_type, "__origin__", None) in container_types - ) - except Exception: - return False + origin = get_origin(field_type) + return isinstance(origin, type) and origin in container_types def _get_container_name(self, field_type: object) -> str: """Get the name of a container type.""" - try: - origin = getattr(field_type, "__origin__", None) - return origin.__name__ if origin else str(field_type) - except Exception: - return str(field_type) + origin = get_origin(field_type) + return origin.__name__ if isinstance(origin, type) else str(field_type) def _write_schema_doc(self, f: TextIO, type_: type) -> None: """Write documentation for a single schema.""" diff --git a/src/dify_plugin/core/entities/invocation.py b/src/dify_plugin/core/entities/invocation.py index 205f436a..8c5857fc 100644 --- a/src/dify_plugin/core/entities/invocation.py +++ b/src/dify_plugin/core/entities/invocation.py @@ -30,12 +30,5 @@ def value_of(cls, value: str) -> "InvokeType": Returns: The return value. - - Raises: - ValueError: If input values are invalid. """ - for mode in cls: - if mode.value == value: - return mode - msg = f"invalid type value {value}" - raise ValueError(msg) + return cls(value) diff --git a/src/dify_plugin/core/entities/plugin/io.py b/src/dify_plugin/core/entities/plugin/io.py index 94e456fe..801d9b4a 100644 --- a/src/dify_plugin/core/entities/plugin/io.py +++ b/src/dify_plugin/core/entities/plugin/io.py @@ -13,11 +13,7 @@ class PluginInStreamEvent(Enum): @classmethod def value_of(cls, v: str) -> "PluginInStreamEvent": - for e in cls: - if e.value == v: - return e - msg = f"Invalid value for PluginInStream.Event: {v}" - raise ValueError(msg) + return cls(v) @dataclass(frozen=True, slots=True) diff --git a/src/dify_plugin/core/entities/plugin/request.py b/src/dify_plugin/core/entities/plugin/request.py index 4f1f9eb4..486e20da 100644 --- a/src/dify_plugin/core/entities/plugin/request.py +++ b/src/dify_plugin/core/entities/plugin/request.py @@ -11,13 +11,9 @@ ) from dify_plugin.entities.model import EmbeddingInputType, ModelType from dify_plugin.entities.model.message import ( - AssistantPromptMessage, PromptMessage, - PromptMessageRole, PromptMessageTool, - SystemPromptMessage, - ToolPromptMessage, - UserPromptMessage, + ensure_prompt_message, ) from dify_plugin.entities.model.text_embedding import MultiModalContent from dify_plugin.entities.provider_config import CredentialType @@ -171,22 +167,7 @@ def convert_prompt_messages(cls, v: list[object]) -> list[PromptMessage]: msg = "prompt_messages must be a list" raise TypeError(msg) - for i in range(len(v)): - if isinstance(v[i], PromptMessage): - continue - - if v[i]["role"] == PromptMessageRole.USER.value: - v[i] = UserPromptMessage(**v[i]) - elif v[i]["role"] == PromptMessageRole.ASSISTANT.value: - v[i] = AssistantPromptMessage(**v[i]) - elif v[i]["role"] == PromptMessageRole.SYSTEM.value: - v[i] = SystemPromptMessage(**v[i]) - elif v[i]["role"] == PromptMessageRole.TOOL.value: - v[i] = ToolPromptMessage(**v[i]) - else: - v[i] = PromptMessage(**v[i]) - - return v + return [ensure_prompt_message(item) for item in v] class ModelInvokeLLMRequest(PluginAccessModelRequest, PromptMessageMixin): diff --git a/src/dify_plugin/entities/model/llm.py b/src/dify_plugin/entities/model/llm.py index 8fca97bd..5acc44b4 100644 --- a/src/dify_plugin/entities/model/llm.py +++ b/src/dify_plugin/entities/model/llm.py @@ -34,15 +34,8 @@ def value_of(cls, value: str) -> "LLMMode": Returns: The return value. - - Raises: - ValueError: If input values are invalid. """ - for mode in cls: - if mode.value == value: - return mode - msg = f"invalid mode value {value}" - raise ValueError(msg) + return cls(value) class LLMPollingStatus(StrEnum): diff --git a/src/dify_plugin/entities/model/message.py b/src/dify_plugin/entities/model/message.py index b28e8c3d..611a142f 100644 --- a/src/dify_plugin/entities/model/message.py +++ b/src/dify_plugin/entities/model/message.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from enum import Enum, StrEnum from typing import Annotated, Literal @@ -25,15 +25,8 @@ def value_of(cls, value: str) -> PromptMessageRole: Returns: The return value. - - Raises: - ValueError: If input values are invalid. """ - for mode in cls: - if mode.value == value: - return mode - msg = f"invalid prompt message type value {value}" - raise ValueError(msg) + return cls(value) class PromptMessageTool(BaseModel): @@ -278,3 +271,28 @@ def is_empty(self) -> bool: The return value. """ return not (not super().is_empty() and not self.tool_call_id) + + +PROMPT_MESSAGE_TYPES_BY_ROLE: dict[str, type[PromptMessage]] = { + PromptMessageRole.USER.value: UserPromptMessage, + PromptMessageRole.ASSISTANT.value: AssistantPromptMessage, + PromptMessageRole.SYSTEM.value: SystemPromptMessage, + PromptMessageRole.DEVELOPER.value: DeveloperPromptMessage, + PromptMessageRole.TOOL.value: ToolPromptMessage, +} + + +def ensure_prompt_message(value: PromptMessage | Mapping[str, object]) -> PromptMessage: + if isinstance(value, PromptMessage): + return value + + if not isinstance(value, Mapping): + msg = "prompt message must be a PromptMessage or a Mapping" + raise TypeError(msg) + + role = value.get("role") + if isinstance(role, PromptMessageRole): + role = role.value + + message_cls = PROMPT_MESSAGE_TYPES_BY_ROLE.get(role, PromptMessage) + return message_cls(**dict(value)) diff --git a/src/dify_plugin/entities/model/schema.py b/src/dify_plugin/entities/model/schema.py index a5b6ce47..fc0256fb 100644 --- a/src/dify_plugin/entities/model/schema.py +++ b/src/dify_plugin/entities/model/schema.py @@ -37,15 +37,8 @@ def value_of(cls, value: object) -> "DefaultParameterName": Returns: The return value. - - Raises: - ValueError: If input values are invalid. """ - for name in cls: - if name.value == value: - return name - msg = f"invalid parameter name {value}" - raise ValueError(msg) + return cls(value) PARAMETER_RULE_TEMPLATE: dict[DefaultParameterName, dict] = { diff --git a/src/dify_plugin/entities/provider_config.py b/src/dify_plugin/entities/provider_config.py index 867fcf22..856885a8 100644 --- a/src/dify_plugin/entities/provider_config.py +++ b/src/dify_plugin/entities/provider_config.py @@ -104,15 +104,8 @@ def value_of(cls, value: str) -> "ProviderConfig.Config": Returns: The return value. - - Raises: - ValueError: If input values are invalid. """ - for mode in cls: - if mode.value == value: - return mode - msg = f"invalid mode value {value}" - raise ValueError(msg) + return cls(value) name: str = Field(..., description="The name of the credentials") type: Config = Field(..., description="The type of the credentials") diff --git a/src/dify_plugin/entities/tool.py b/src/dify_plugin/entities/tool.py index 5fda0019..750d3f95 100644 --- a/src/dify_plugin/entities/tool.py +++ b/src/dify_plugin/entities/tool.py @@ -304,15 +304,8 @@ def value_of(cls, value: str) -> "ToolProviderType": Returns: The return value. - - Raises: - ValueError: If input values are invalid. """ - for mode in cls: - if mode.value == value: - return mode - msg = f"invalid mode value {value}" - raise ValueError(msg) + return cls(value) class ToolSelector(BaseModel): diff --git a/src/dify_plugin/entities/workflow_node.py b/src/dify_plugin/entities/workflow_node.py index b3e0cdeb..a58f28ac 100644 --- a/src/dify_plugin/entities/workflow_node.py +++ b/src/dify_plugin/entities/workflow_node.py @@ -43,15 +43,8 @@ def value_of(cls, value: str) -> "NodeType": Returns: The return value. - - Raises: - ValueError: If input values are invalid. """ - for node_type in cls: - if node_type.value == value: - return node_type - msg = f"invalid node type value {value}" - raise ValueError(msg) + return cls(value) class ModelConfig(BaseModel): diff --git a/src/dify_plugin/file/entities.py b/src/dify_plugin/file/entities.py index 435684ad..24abebd1 100644 --- a/src/dify_plugin/file/entities.py +++ b/src/dify_plugin/file/entities.py @@ -10,8 +10,4 @@ class FileType(StrEnum): @staticmethod def value_of(value: str) -> "FileType": - for member in FileType: - if member.value == value: - return member - msg = f"No such file type: {value}" - raise ValueError(msg) + return FileType(value) diff --git a/src/dify_plugin/interfaces/agent/__init__.py b/src/dify_plugin/interfaces/agent/__init__.py index f2dcd01d..d715b2da 100644 --- a/src/dify_plugin/interfaces/agent/__init__.py +++ b/src/dify_plugin/interfaces/agent/__init__.py @@ -1,3 +1,13 @@ +from dify_plugin.entities.model.message import ( + AssistantPromptMessage, + PromptMessage, + PromptMessageRole, + PromptMessageTool, + SystemPromptMessage, + ToolPromptMessage, + UserPromptMessage, +) + from .strategy import ( FILE_PARAMETER_TYPES, STRING_PARAMETER_TYPES, @@ -10,7 +20,6 @@ AgentToolIdentity, AIModelEntity, Any, - AssistantPromptMessage, BaseModel, ConfigDict, CredentialType, @@ -20,21 +29,15 @@ LLMUsage, Mapping, ModelPropertyKey, - PromptMessage, - PromptMessageRole, - PromptMessageTool, Session, - SystemPromptMessage, ToolDescription, ToolEntity, ToolIdentity, ToolInvokeMeta, ToolLike, ToolParameter, - ToolPromptMessage, ToolProvider, ToolProviderType, - UserPromptMessage, ValidationInfo, abstractmethod, field_validator, diff --git a/src/dify_plugin/interfaces/agent/strategy.py b/src/dify_plugin/interfaces/agent/strategy.py index 9d20bdef..ea6b9cb7 100644 --- a/src/dify_plugin/interfaces/agent/strategy.py +++ b/src/dify_plugin/interfaces/agent/strategy.py @@ -10,13 +10,9 @@ from dify_plugin.entities.model import AIModelEntity, ModelPropertyKey from dify_plugin.entities.model.llm import LLMModelConfig, LLMUsage from dify_plugin.entities.model.message import ( - AssistantPromptMessage, PromptMessage, - PromptMessageRole, PromptMessageTool, - SystemPromptMessage, - ToolPromptMessage, - UserPromptMessage, + ensure_prompt_message, ) from dify_plugin.entities.provider_config import CredentialType from dify_plugin.entities.tool import ( @@ -53,19 +49,7 @@ def convert_prompt_messages(cls, v: list[object]) -> list[PromptMessage]: msg = "prompt_messages must be a list" raise TypeError(msg) - for i in range(len(v)): - if v[i]["role"] == PromptMessageRole.USER.value: - v[i] = UserPromptMessage(**v[i]) - elif v[i]["role"] == PromptMessageRole.ASSISTANT.value: - v[i] = AssistantPromptMessage(**v[i]) - elif v[i]["role"] == PromptMessageRole.SYSTEM.value: - v[i] = SystemPromptMessage(**v[i]) - elif v[i]["role"] == PromptMessageRole.TOOL.value: - v[i] = ToolPromptMessage(**v[i]) - else: - v[i] = PromptMessage(**v[i]) - - return v + return [ensure_prompt_message(item) for item in v] class AgentScratchpadUnit(BaseModel): @@ -295,6 +279,43 @@ def _init_prompt_tools( return prompt_messages_tools + def _set_prompt_tool_parameter( + self, + prompt_tool: PromptMessageTool, + parameter: ToolParameter, + ) -> None: + if parameter.form != ToolParameter.ToolParameterForm.LLM: + return + + parameter_type = parameter.type + if parameter.type in FILE_PARAMETER_TYPES: + return + if parameter.type in STRING_PARAMETER_TYPES: + parameter_type = ToolParameter.ToolParameterType.STRING + + prompt_tool.parameters["properties"][parameter.name] = ( + { + "type": parameter_type, + "description": parameter.llm_description or "", + } + if parameter.input_schema is None + else dict(parameter.input_schema) + ) + + if ( + parameter.type == ToolParameter.ToolParameterType.SELECT + and parameter.options + ): + prompt_tool.parameters["properties"][parameter.name]["enum"] = [ + option.value for option in parameter.options + ] + + if ( + parameter.required + and parameter.name not in prompt_tool.parameters["required"] + ): + prompt_tool.parameters["required"].append(parameter.name) + def _convert_tool_to_prompt_message_tool( self, tool: ToolEntity, @@ -312,36 +333,7 @@ def _convert_tool_to_prompt_message_tool( parameters = tool.parameters for parameter in parameters: - if parameter.form != ToolParameter.ToolParameterForm.LLM: - continue - - parameter_type = parameter.type - if parameter.type in FILE_PARAMETER_TYPES: - continue - if parameter.type in STRING_PARAMETER_TYPES: - parameter_type = ToolParameter.ToolParameterType.STRING - enum = [] - if parameter.type == ToolParameter.ToolParameterType.SELECT: - enum = ( - [option.value for option in parameter.options] - if parameter.options - else [] - ) - - message_tool.parameters["properties"][parameter.name] = ( - { - "type": parameter_type, - "description": parameter.llm_description or "", - } - if parameter.input_schema is None - else parameter.input_schema - ) - - if len(enum) > 0: - message_tool.parameters["properties"][parameter.name]["enum"] = enum - - if parameter.required: - message_tool.parameters["required"].append(parameter.name) + self._set_prompt_tool_parameter(message_tool, parameter) return message_tool @@ -355,38 +347,6 @@ def update_prompt_message_tool( tool_runtime_parameters = tool.parameters for parameter in tool_runtime_parameters: - if parameter.form != ToolParameter.ToolParameterForm.LLM: - continue - - parameter_type = parameter.type - if parameter.type in FILE_PARAMETER_TYPES: - continue - if parameter.type in STRING_PARAMETER_TYPES: - parameter_type = ToolParameter.ToolParameterType.STRING - enum = [] - if parameter.type == ToolParameter.ToolParameterType.SELECT: - enum = ( - [option.value for option in parameter.options] - if parameter.options - else [] - ) - - prompt_tool.parameters["properties"][parameter.name] = ( - { - "type": parameter_type, - "description": parameter.llm_description or "", - } - if parameter.input_schema is None - else parameter.input_schema - ) - - if len(enum) > 0: - prompt_tool.parameters["properties"][parameter.name]["enum"] = enum - - if ( - parameter.required - and parameter.name not in prompt_tool.parameters["required"] - ): - prompt_tool.parameters["required"].append(parameter.name) + self._set_prompt_tool_parameter(prompt_tool, parameter) return prompt_tool diff --git a/src/dify_plugin/interfaces/model/openai_compatible/common.py b/src/dify_plugin/interfaces/model/openai_compatible/common.py index e7426af2..14501ed0 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/common.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/common.py @@ -1,3 +1,5 @@ +from urllib.parse import urljoin + import requests from dify_plugin.errors.model import ( @@ -11,6 +13,9 @@ class _CommonOaiApiCompat: + def _join_endpoint_url(self, endpoint_url: str, path: str) -> str: + return urljoin(f"{endpoint_url.rstrip('/')}/", path) + @property def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]: """ diff --git a/src/dify_plugin/interfaces/model/openai_compatible/llm.py b/src/dify_plugin/interfaces/model/openai_compatible/llm.py index 0d811811..f7b0e4a4 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/llm.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/llm.py @@ -8,7 +8,6 @@ from decimal import Decimal from http import HTTPStatus from typing import Any, cast -from urllib.parse import urljoin import requests from pydantic import TypeAdapter, ValidationError @@ -201,8 +200,6 @@ def validate_credentials(self, model: str, credentials: dict) -> None: headers["Authorization"] = f"Bearer {api_key}" endpoint_url = credentials["endpoint_url"] - if not endpoint_url.endswith("/"): - endpoint_url += "/" # prepare the payload for a simple ping to the model validate_credentials_max_tokens = ( @@ -219,10 +216,13 @@ def validate_credentials(self, model: str, credentials: dict) -> None: data["messages"] = [ {"role": "user", "content": "ping"}, ] - endpoint_url = urljoin(endpoint_url, "chat/completions") + endpoint_url = self._join_endpoint_url( + endpoint_url, + "chat/completions", + ) elif completion_type is LLMMode.COMPLETION: data["prompt"] = "ping" - endpoint_url = urljoin(endpoint_url, "completions") + endpoint_url = self._join_endpoint_url(endpoint_url, "completions") else: msg = "Unsupported completion type for model configuration." raise ValueError(msg) @@ -537,8 +537,6 @@ def _generate( headers["Authorization"] = f"Bearer {api_key}" endpoint_url = credentials["endpoint_url"] - if not endpoint_url.endswith("/"): - endpoint_url += "/" response_format = model_parameters.get("response_format") if response_format: @@ -578,13 +576,13 @@ def _generate( completion_type = LLMMode.value_of(credentials["mode"]) if completion_type is LLMMode.CHAT: - endpoint_url = urljoin(endpoint_url, "chat/completions") + endpoint_url = self._join_endpoint_url(endpoint_url, "chat/completions") data["messages"] = [ self._convert_prompt_message_to_dict(m, credentials) for m in prompt_messages ] elif completion_type is LLMMode.COMPLETION: - endpoint_url = urljoin(endpoint_url, "completions") + endpoint_url = self._join_endpoint_url(endpoint_url, "completions") data["prompt"] = prompt_messages[0].content else: msg = "Unsupported completion type for model configuration." diff --git a/src/dify_plugin/interfaces/model/openai_compatible/speech2text.py b/src/dify_plugin/interfaces/model/openai_compatible/speech2text.py index 5bc5d71e..370c4a8b 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/speech2text.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/speech2text.py @@ -1,7 +1,6 @@ import pathlib from http import HTTPStatus from typing import IO -from urllib.parse import urljoin import requests @@ -44,10 +43,10 @@ def _invoke( if api_key: headers["Authorization"] = f"Bearer {api_key}" - endpoint_url = credentials.get("endpoint_url", "https://api.openai.com/v1/") - if not endpoint_url.endswith("/"): - endpoint_url += "/" - endpoint_url = urljoin(endpoint_url, "audio/transcriptions") + endpoint_url = self._join_endpoint_url( + credentials.get("endpoint_url", "https://api.openai.com/v1/"), + "audio/transcriptions", + ) payload = {"model": credentials.get("endpoint_model_name", model)} files = [("file", file)] diff --git a/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py b/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py index 21b2e046..ab8ac5b9 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py @@ -2,7 +2,6 @@ import time from decimal import Decimal from http import HTTPStatus -from urllib.parse import urljoin import requests @@ -62,11 +61,10 @@ def _invoke( if api_key: headers["Authorization"] = f"Bearer {api_key}" - endpoint_url = credentials.get("endpoint_url", "") - if not endpoint_url.endswith("/"): - endpoint_url += "/" - - endpoint_url = urljoin(endpoint_url, "embeddings") + endpoint_url = self._join_endpoint_url( + credentials.get("endpoint_url", ""), + "embeddings", + ) extra_model_kwargs = {} if user: @@ -170,11 +168,10 @@ def validate_credentials(self, model: str, credentials: dict) -> None: if api_key: headers["Authorization"] = f"Bearer {api_key}" - endpoint_url = credentials.get("endpoint_url", "") - if not endpoint_url.endswith("/"): - endpoint_url += "/" - - endpoint_url = urljoin(endpoint_url, "embeddings") + endpoint_url = self._join_endpoint_url( + credentials.get("endpoint_url", ""), + "embeddings", + ) payload = { "input": ["ping"], diff --git a/src/dify_plugin/interfaces/model/openai_compatible/tts.py b/src/dify_plugin/interfaces/model/openai_compatible/tts.py index b61608c1..f36b350b 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/tts.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/tts.py @@ -1,6 +1,5 @@ from collections.abc import Generator from http import HTTPStatus -from urllib.parse import urljoin import requests @@ -58,10 +57,10 @@ def _invoke( headers["Authorization"] = f"Bearer {api_key}" # Construct endpoint URL - endpoint_url = credentials.get("endpoint_url", "") - if not endpoint_url.endswith("/"): - endpoint_url += "/" - endpoint_url = urljoin(endpoint_url, "audio/speech") + endpoint_url = self._join_endpoint_url( + credentials.get("endpoint_url", ""), + "audio/speech", + ) # Get audio format from model properties audio_format = self._get_model_audio_type(model, credentials) diff --git a/tests/interfaces/agent/test_agent.py b/tests/interfaces/agent/test_agent.py index fb0cc2e6..3e387b55 100644 --- a/tests/interfaces/agent/test_agent.py +++ b/tests/interfaces/agent/test_agent.py @@ -4,9 +4,20 @@ from dify_plugin.core.runtime import Session from dify_plugin.core.server.stdio.request_reader import StdioRequestReader from dify_plugin.core.server.stdio.response_writer import StdioResponseWriter +from dify_plugin.entities import I18nObject from dify_plugin.entities.agent import AgentInvokeMessage, AgentRuntime from dify_plugin.entities.model.message import PromptMessage, PromptMessageRole -from dify_plugin.interfaces.agent import AgentModelConfig, AgentStrategy +from dify_plugin.entities.tool import ( + ToolDescription, + ToolParameter, + ToolParameterOption, +) +from dify_plugin.interfaces.agent import ( + AgentModelConfig, + AgentStrategy, + AgentToolIdentity, + ToolEntity, +) def _make_agent_model_config() -> AgentModelConfig: @@ -17,6 +28,24 @@ def _make_agent_model_config() -> AgentModelConfig: ) +def _make_agent_strategy() -> AgentStrategy: + class AgentStrategyImpl(AgentStrategy): + def _invoke( + self, + parameters: dict, + ) -> Generator[AgentInvokeMessage, None, None]: + del parameters + yield self.create_text_message("Hello, world!") + + session = Session( + session_id="test", + executor=ThreadPoolExecutor(max_workers=1), + reader=StdioRequestReader(), + writer=StdioResponseWriter(), + ) + return AgentStrategyImpl(runtime=AgentRuntime(user_id="test"), session=session) + + def test_agent_model_config_ensure_history_prompt_messages_not_shared() -> None: prompt_message = PromptMessage( role=PromptMessageRole.USER, content="Content", name=None @@ -40,20 +69,54 @@ def test_constructor_of_agent_strategy() -> None: - And ensure a breaking change will be detected by CI. """ - class AgentStrategyImpl(AgentStrategy): - def _invoke( - self, parameters: dict - ) -> Generator[AgentInvokeMessage, None, None]: - del parameters - yield self.create_text_message("Hello, world!") + agent_strategy = _make_agent_strategy() + assert agent_strategy is not None - session = Session( - session_id="test", - executor=ThreadPoolExecutor(max_workers=1), - reader=StdioRequestReader(), - writer=StdioResponseWriter(), - ) - agent_strategy = AgentStrategyImpl( - runtime=AgentRuntime(user_id="test"), session=session + +def test_agent_strategy_converts_tool_parameters_once() -> None: + label = I18nObject(en_US="Search") + input_schema = {"type": "string"} + tool = ToolEntity( + identity=AgentToolIdentity( + author="test", + name="search", + label=label, + provider="test", + ), + description=ToolDescription(human=label, llm="Search documents"), + parameters=[ + ToolParameter( + name="query", + label=label, + human_description=label, + type=ToolParameter.ToolParameterType.SELECT, + form=ToolParameter.ToolParameterForm.LLM, + llm_description="Query", + required=True, + input_schema=input_schema, + options=[ + ToolParameterOption(value="docs", label=label), + ToolParameterOption(value="web", label=label), + ], + ), + ToolParameter( + name="upload", + label=label, + human_description=label, + type=ToolParameter.ToolParameterType.FILE, + form=ToolParameter.ToolParameterForm.LLM, + ), + ], ) - assert agent_strategy is not None + + agent_strategy = _make_agent_strategy() + prompt_tool = agent_strategy._convert_tool_to_prompt_message_tool(tool) + agent_strategy.update_prompt_message_tool(tool, prompt_tool) + + query_schema = prompt_tool.parameters["properties"]["query"] + assert query_schema["type"] == ToolParameter.ToolParameterType.STRING + assert query_schema["enum"] == ["docs", "web"] + assert query_schema is not tool.parameters[0].input_schema + assert tool.parameters[0].input_schema == {"type": "string"} + assert "upload" not in prompt_tool.parameters["properties"] + assert prompt_tool.parameters["required"] == ["query"] diff --git a/tests/interfaces/model/openai_compatible/test_common.py b/tests/interfaces/model/openai_compatible/test_common.py new file mode 100644 index 00000000..6ce8c869 --- /dev/null +++ b/tests/interfaces/model/openai_compatible/test_common.py @@ -0,0 +1,15 @@ +from dify_plugin.interfaces.model.openai_compatible import common + + +def test_join_endpoint_url_normalizes_single_separator() -> None: + compat = common._CommonOaiApiCompat() + + assert ( + compat._join_endpoint_url("https://example.com/v1", "embeddings") + == "https://example.com/v1/embeddings" + ) + assert ( + compat._join_endpoint_url("https://example.com/v1/", "audio/speech") + == "https://example.com/v1/audio/speech" + ) + assert compat._join_endpoint_url("", "embeddings") == "/embeddings" diff --git a/tests/test_documentation_generator.py b/tests/test_documentation_generator.py new file mode 100644 index 00000000..98845ba1 --- /dev/null +++ b/tests/test_documentation_generator.py @@ -0,0 +1,25 @@ +from enum import Enum + +from pydantic import BaseModel + +from dify_plugin.core.documentation.generator import SchemaDocumentationGenerator + + +class ReferencedModel(BaseModel): + value: str + + +class ReferencedEnum(Enum): + VALUE = "value" + + +def test_documentation_generator_handles_generic_type_args() -> None: + generator = SchemaDocumentationGenerator() + + refs = generator._extract_referenced_types( + list[ReferencedModel] | dict[str, ReferencedEnum], + ) + + assert refs == {ReferencedModel, ReferencedEnum} + assert generator._is_container_type(list[ReferencedModel]) + assert generator._get_container_name(list[ReferencedModel]) == "list" diff --git a/tests/test_prompt_message.py b/tests/test_prompt_message.py index ec130efd..52a67457 100644 --- a/tests/test_prompt_message.py +++ b/tests/test_prompt_message.py @@ -1,7 +1,12 @@ +import pytest + from dify_plugin.entities.model.message import ( + AssistantPromptMessage, ImagePromptMessageContent, + PromptMessageRole, TextPromptMessageContent, UserPromptMessage, + ensure_prompt_message, ) @@ -51,3 +56,21 @@ def test_validate_prompt_message() -> None: assert prompt_content[0].data == "Hello, World!" assert isinstance(prompt_content[1], ImagePromptMessageContent) assert prompt_content[1].url == "https://example.com/image.jpg" + + +def test_ensure_prompt_message_uses_role_specific_class() -> None: + message = ensure_prompt_message({"role": "assistant", "content": "ok"}) + assert isinstance(message, AssistantPromptMessage) + + prompt = UserPromptMessage(content="hello") + assert ensure_prompt_message(prompt) is prompt + + enum_message = ensure_prompt_message( + {"role": PromptMessageRole.USER, "content": "hello"}, + ) + assert isinstance(enum_message, UserPromptMessage) + + +def test_ensure_prompt_message_rejects_non_mapping() -> None: + with pytest.raises(TypeError, match="PromptMessage or a Mapping"): + ensure_prompt_message("assistant")