opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-10-16 15:35:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): enable strict typing for prompt transform tests (#44139)

Browse Source
This commit is contained in:
Eddy ZHANG
2026-10-15 04:57:39 +00:00
committed by GitHub
parent d00d0289ef
commit c83f633cab
6 changed files with 61 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/tests/unit_tests/.ruff.toml
-1
View File
@@ -154,7 +154,6 @@ extend-select = ["ANN401", "ARG"]
"core/plugin/test_model_runtime_adapter.py" = ["ARG002"]
"core/plugin/test_plugin_runtime.py" = ["ARG001", "ARG002", "TID251"]
"core/prompt/test_advanced_prompt_transform.py" = ["ARG005"]
"core/prompt/test_prompt_transform.py" = ["ARG005"]
"core/rag/datasource/keyword/jieba/test_jieba.py" = ["ARG001", "TID251"]
"core/rag/datasource/keyword/jieba/test_jieba_keyword_table_handler.py" = ["ARG004"]
"core/rag/datasource/test_datasource_retrieval.py" = ["ARG001", "ARG002", "ARG005", "TID251"]
api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py
+8 -5
View File
@@ -5,6 +5,7 @@ from pytest_mock import MockerFixture
from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform
from graphon.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessage,
SystemPromptMessage,
ToolPromptMessage,
UserPromptMessage,
@@ -12,12 +13,12 @@ from graphon.model_runtime.entities.message_entities import (
from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory
def test_get_prompt(mocker: MockerFixture):
prompt_messages = [
def test_get_prompt(mocker: MockerFixture) -> None:
prompt_messages: list[PromptMessage] = [
SystemPromptMessage(content="System Template"),
UserPromptMessage(content="User Query"),
]
history_messages = [
history_messages: list[PromptMessage] = [
SystemPromptMessage(content="System Prompt 1"),
UserPromptMessage(content="User Prompt 1"),
AssistantPromptMessage(content="Assistant Thought 1"),
@@ -33,8 +34,10 @@ def test_get_prompt(mocker: MockerFixture):
]
# use message number instead of token for testing
def side_effect_get_num_tokens(*args):
return len(args[2])
def side_effect_get_num_tokens(*args: object) -> int:
prompt_messages = args[2]
assert isinstance(prompt_messages, list)
return len(prompt_messages)
model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
mocker.patch.object(
api/tests/unit_tests/core/prompt/test_prompt_message.py
+4 -4
View File
@@ -10,14 +10,14 @@ from graphon.model_runtime.entities.message_entities import (
)
def test_build_prompt_message_with_prompt_message_contents():
def test_build_prompt_message_with_prompt_message_contents() -> None:
prompt = UserPromptMessage(content=[TextPromptMessageContent(data="Hello, World!")])
assert isinstance(prompt.content, list)
assert isinstance(prompt.content[0], TextPromptMessageContent)
assert prompt.content[0].data == "Hello, World!"
def test_dump_prompt_message():
def test_dump_prompt_message() -> None:
example_url = "https://example.com/image.jpg"
prompt = UserPromptMessage(
content=[
@@ -32,7 +32,7 @@ def test_dump_prompt_message():
assert data["content"][0].get("url") == example_url
def test_prompt_messages_to_prompt_for_saving_chat_mode():
def test_prompt_messages_to_prompt_for_saving_chat_mode() -> None:
chat_messages = [
UserPromptMessage(
content=[
@@ -79,7 +79,7 @@ def test_prompt_messages_to_prompt_for_saving_chat_mode():
assert prompts[2]["role"] == "tool"
def test_prompt_messages_to_prompt_for_saving_completion_mode_with_and_without_files():
def test_prompt_messages_to_prompt_for_saving_completion_mode_with_and_without_files() -> None:
completion_message_with_files = UserPromptMessage(
content=[
TextPromptMessageContent(data="first "),
api/tests/unit_tests/core/prompt/test_prompt_transform.py
+27 -21
View File
@@ -1,3 +1,5 @@
from typing import NoReturn
import pytest
from core.entities.provider_entities import CustomProviderConfiguration
@@ -5,19 +7,21 @@ from core.model_manager import ModelInstance
from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from core.prompt.prompt_transform import PromptTransform
from graphon.model_runtime.entities.common_entities import I18nObject
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage
from graphon.model_runtime.entities.model_entities import ModelPropertyKey, ParameterRule, ParameterType
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage, UserPromptMessage
from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelPropertyKey, ParameterRule, ParameterType
from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory
class TestPromptTransform:
def test_resolve_model_runtime_requires_model_config_or_instance(self):
def test_resolve_model_runtime_requires_model_config_or_instance(self) -> None:
transform = PromptTransform()
with pytest.raises(ValueError, match="Either model_config or model_instance must be provided."):
transform._resolve_model_runtime()
def test_resolve_model_runtime_builds_model_instance_from_model_config(self, monkeypatch: pytest.MonkeyPatch):
def test_resolve_model_runtime_builds_model_instance_from_model_config(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="config-model", mode="chat")
model_config.credentials = {"api_key": "secret"}
@@ -30,7 +34,7 @@ class TestPromptTransform:
model_schema = model_config.model_schema.model_copy(deep=True)
schema_calls = []
def get_model_schema(*, model, credentials):
def get_model_schema(*, model: str, credentials: dict[str, object]) -> AIModelEntity:
schema_calls.append((model, credentials))
return model_schema
@@ -47,7 +51,7 @@ class TestPromptTransform:
assert model_instance.stop == ["END"]
assert resolved_schema is model_schema
def test_resolve_model_runtime_uses_model_config_schema_fallback(self, monkeypatch: pytest.MonkeyPatch):
def test_resolve_model_runtime_uses_model_config_schema_fallback(self, monkeypatch: pytest.MonkeyPatch) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="resolved-model", mode="chat")
model_instance = make_token_buffer_memory(model_config).model_instance
@@ -63,7 +67,7 @@ class TestPromptTransform:
def test_resolve_model_runtime_raises_when_schema_missing_without_model_config(
self, monkeypatch: pytest.MonkeyPatch
):
) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="resolved-model", mode="chat")
model_instance = make_token_buffer_memory(model_config).model_instance
@@ -72,7 +76,7 @@ class TestPromptTransform:
with pytest.raises(ValueError, match="Model schema not found for the provided model instance."):
transform._resolve_model_runtime(model_instance=model_instance)
def test_calculate_rest_token_defaults_when_context_size_missing(self, monkeypatch: pytest.MonkeyPatch):
def test_calculate_rest_token_defaults_when_context_size_missing(self, monkeypatch: pytest.MonkeyPatch) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="test-model", mode="chat")
model_instance = make_token_buffer_memory(model_config).model_instance
@@ -80,7 +84,7 @@ class TestPromptTransform:
model_instance.model_type_instance, "get_model_schema", lambda **_kwargs: model_config.model_schema
)
def unexpected_token_count(_messages):
def unexpected_token_count(_messages) -> NoReturn:
pytest.fail("A model without a context size must use the default token budget.")
monkeypatch.setattr(model_instance, "get_llm_num_tokens", unexpected_token_count)
@@ -106,7 +110,7 @@ class TestPromptTransform:
max_tokens: int,
message_tokens: int,
expected: int,
):
) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="test-model", mode="chat")
model_config.model_schema.model_properties = {ModelPropertyKey.CONTEXT_SIZE: context_size}
@@ -120,10 +124,10 @@ class TestPromptTransform:
]
model_instance = make_token_buffer_memory(model_config).model_instance
model_instance.parameters = {"max_tokens": max_tokens}
messages = [UserPromptMessage(content="hello")]
messages: list[PromptMessage] = [UserPromptMessage(content="hello")]
token_calls = []
def get_llm_num_tokens(prompt_messages):
def get_llm_num_tokens(prompt_messages: list[PromptMessage]) -> int:
token_calls.append(prompt_messages)
return message_tokens
@@ -137,13 +141,13 @@ class TestPromptTransform:
assert rest == expected
assert token_calls == [messages]
def test_get_history_messages_from_memory_with_and_without_window(self, monkeypatch: pytest.MonkeyPatch):
def test_get_history_messages_from_memory_with_and_without_window(self, monkeypatch: pytest.MonkeyPatch) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="test-model", mode="chat")
memory = make_token_buffer_memory(model_config)
history_calls = []
def get_history_prompt_text(**kwargs):
def get_history_prompt_text(**kwargs: object) -> str:
history_calls.append(kwargs)
return "history"
@@ -172,14 +176,16 @@ class TestPromptTransform:
assert result == "history"
assert history_calls[1:] == [{"max_token_limit": 50}]
def test_get_history_messages_list_from_memory_with_and_without_window(self, monkeypatch: pytest.MonkeyPatch):
def test_get_history_messages_list_from_memory_with_and_without_window(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="test-model", mode="chat")
memory = make_token_buffer_memory(model_config)
messages = [UserPromptMessage(content="m1"), AssistantPromptMessage(content="m2")]
messages: list[PromptMessage] = [UserPromptMessage(content="m1"), AssistantPromptMessage(content="m2")]
history_calls = []
def get_history_prompt_messages(**kwargs):
def get_history_prompt_messages(**kwargs: object) -> list[PromptMessage]:
history_calls.append(kwargs)
return messages
@@ -195,16 +201,16 @@ class TestPromptTransform:
assert result == messages
assert history_calls[1:] == [{"max_token_limit": 10, "message_limit": None}]
def test_append_chat_histories_extends_prompt_messages(self, monkeypatch: pytest.MonkeyPatch):
def test_append_chat_histories_extends_prompt_messages(self, monkeypatch: pytest.MonkeyPatch) -> None:
transform = PromptTransform()
model_config = make_model_config(provider="openai", model="test-model", mode="chat")
memory = make_token_buffer_memory(model_config)
memory_config = MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False, size=None))
histories = [UserPromptMessage(content="h1"), AssistantPromptMessage(content="h2")]
messages = [UserPromptMessage(content="p1")]
histories: list[PromptMessage] = [UserPromptMessage(content="h1"), AssistantPromptMessage(content="h2")]
messages: list[PromptMessage] = [UserPromptMessage(content="p1")]
history_calls = []
def get_history_prompt_messages(**kwargs):
def get_history_prompt_messages(**kwargs: object) -> list[PromptMessage]:
history_calls.append(kwargs)
return histories
api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py
+22 -22
View File
@@ -22,7 +22,7 @@ from models.model import AppMode
from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory
def test_get_common_chat_app_prompt_template_with_pcqm():
def test_get_common_chat_app_prompt_template_with_pcqm() -> None:
prompt_transform = SimplePromptTransform()
pre_prompt = "You are a helpful assistant."
prompt_template = prompt_transform.get_prompt_template(
@@ -43,7 +43,7 @@ def test_get_common_chat_app_prompt_template_with_pcqm():
assert prompt_template["special_variable_keys"] == ["#context#", "#histories#", "#query#"]
def test_get_common_completion_app_prompt_template_with_pcq():
def test_get_common_completion_app_prompt_template_with_pcq() -> None:
prompt_transform = SimplePromptTransform()
pre_prompt = "You are a helpful assistant."
prompt_template = prompt_transform.get_prompt_template(
@@ -60,7 +60,7 @@ def test_get_common_completion_app_prompt_template_with_pcq():
assert prompt_template["special_variable_keys"] == ["#context#", "#query#"]
def test_get_common_chat_app_prompt_template_with_q():
def test_get_common_chat_app_prompt_template_with_q() -> None:
prompt_transform = SimplePromptTransform()
pre_prompt = ""
prompt_template = prompt_transform.get_prompt_template(
@@ -75,7 +75,7 @@ def test_get_common_chat_app_prompt_template_with_q():
assert prompt_template["special_variable_keys"] == ["#query#"]
def test_get_common_chat_app_prompt_template_with_cq():
def test_get_common_chat_app_prompt_template_with_cq() -> None:
prompt_transform = SimplePromptTransform()
pre_prompt = ""
prompt_template = prompt_transform.get_prompt_template(
@@ -92,7 +92,7 @@ def test_get_common_chat_app_prompt_template_with_cq():
assert prompt_template["special_variable_keys"] == ["#context#", "#query#"]
def test_get_common_chat_app_prompt_template_with_p():
def test_get_common_chat_app_prompt_template_with_p() -> None:
prompt_transform = SimplePromptTransform()
pre_prompt = "you are {{name}}"
prompt_template = prompt_transform.get_prompt_template(
@@ -107,7 +107,7 @@ def test_get_common_chat_app_prompt_template_with_p():
assert prompt_template["special_variable_keys"] == []
def test__get_chat_model_prompt_messages():
def test__get_chat_model_prompt_messages() -> None:
model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
memory = make_token_buffer_memory(model_config)
@@ -150,7 +150,7 @@ def test__get_chat_model_prompt_messages():
assert prompt_messages[3].content == query
def test__get_completion_model_prompt_messages():
def test__get_completion_model_prompt_messages() -> None:
model_config = make_model_config(provider="openai", model="gpt-3.5-turbo-instruct", mode="completion")
memory = make_token_buffer_memory(model_config)
@@ -201,7 +201,7 @@ def test__get_completion_model_prompt_messages():
assert prompt_messages[0].content == real_prompt
def test_get_prompt_dispatches_chat_and_completion():
def test_get_prompt_dispatches_chat_and_completion() -> None:
transform = SimplePromptTransform()
model_config_chat = make_model_config(provider="openai", model="gpt-4", mode="chat")
model_config_completion = make_model_config(provider="openai", model="gpt-3.5-turbo-instruct", mode="completion")
@@ -215,7 +215,7 @@ def test_get_prompt_dispatches_chat_and_completion():
chat_messages, chat_stops = transform.get_prompt(
app_mode=AppMode.CHAT,
prompt_template_entity=prompt_entity,
inputs={"n": 1},
inputs={"n": "1"},
query="q",
files=[],
context=None,
@@ -228,7 +228,7 @@ def test_get_prompt_dispatches_chat_and_completion():
completion_messages, completion_stops = transform.get_prompt(
app_mode=AppMode.CHAT,
prompt_template_entity=prompt_entity,
inputs={"n": 1},
inputs={"n": "1"},
query="q",
files=[],
context=None,
@@ -239,7 +239,7 @@ def test_get_prompt_dispatches_chat_and_completion():
assert completion_stops == ["stop"]
def test_get_prompt_str_and_rules_type_validation_errors():
def test_get_prompt_str_and_rules_type_validation_errors() -> None:
transform = SimplePromptTransform()
valid_prompt_template = SimplePromptTransform().get_prompt_template(AppMode.CHAT, "", False, False)[
"prompt_template"
@@ -248,8 +248,8 @@ def test_get_prompt_str_and_rules_type_validation_errors():
bad_custom_keys = {
"prompt_template": valid_prompt_template,
"custom_variable_keys": "not-list",
"special_variable_keys": [],
"prompt_rules": {},
"special_variable_keys": list[object](),
"prompt_rules": dict[str, object](),
}
transform.get_prompt_template = MagicMock(return_value=bad_custom_keys)
with pytest.raises(TypeError, match="custom_variable_keys"):
@@ -257,7 +257,7 @@ def test_get_prompt_str_and_rules_type_validation_errors():
bad_special_keys = {
**bad_custom_keys,
"custom_variable_keys": [],
"custom_variable_keys": list[object](),
"special_variable_keys": "not-list",
}
transform.get_prompt_template = MagicMock(return_value=bad_special_keys)
@@ -266,8 +266,8 @@ def test_get_prompt_str_and_rules_type_validation_errors():
bad_prompt_template = {
**bad_custom_keys,
"custom_variable_keys": [],
"special_variable_keys": [],
"custom_variable_keys": list[object](),
"special_variable_keys": list[object](),
"prompt_template": 123,
}
transform.get_prompt_template = MagicMock(return_value=bad_prompt_template)
@@ -276,8 +276,8 @@ def test_get_prompt_str_and_rules_type_validation_errors():
bad_prompt_rules = {
**bad_custom_keys,
"custom_variable_keys": [],
"special_variable_keys": [],
"custom_variable_keys": list[object](),
"special_variable_keys": list[object](),
"prompt_template": valid_prompt_template,
"prompt_rules": "not-dict",
}
@@ -286,7 +286,7 @@ def test_get_prompt_str_and_rules_type_validation_errors():
transform._get_prompt_str_and_rules(AppMode.CHAT, "", {}, query=None, context=None)
def test_chat_model_prompt_messages_uses_prompt_when_query_empty():
def test_chat_model_prompt_messages_uses_prompt_when_query_empty() -> None:
transform = SimplePromptTransform()
model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
transform._get_prompt_str_and_rules = MagicMock(return_value=("prompt-text", {}))
@@ -307,7 +307,7 @@ def test_chat_model_prompt_messages_uses_prompt_when_query_empty():
transform._get_last_user_message.assert_called_once_with("prompt-text", [], None, None)
def test_completion_model_prompt_messages_empty_stops_becomes_none():
def test_completion_model_prompt_messages_empty_stops_becomes_none() -> None:
transform = SimplePromptTransform()
model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
transform._get_prompt_str_and_rules = MagicMock(return_value=("prompt", {"stops": []}))
@@ -327,7 +327,7 @@ def test_completion_model_prompt_messages_empty_stops_becomes_none():
assert stops is None
def test_get_last_user_message_with_files_and_context_files():
def test_get_last_user_message_with_files_and_context_files() -> None:
transform = SimplePromptTransform()
file = File(
type=FileType.IMAGE, transfer_method=FileTransferMethod.REMOTE_URL, remote_url="https://example.com/a.jpg"
@@ -355,7 +355,7 @@ def test_get_last_user_message_with_files_and_context_files():
assert message.content[2].data == "hello"
def test_advanced_prompt_templates_constants_are_importable():
def test_advanced_prompt_templates_constants_are_importable() -> None:
assert isinstance(CONTEXT, str)
assert "completion_prompt_config" in CHAT_APP_COMPLETION_PROMPT_CONFIG
assert "chat_prompt_config" in CHAT_APP_CHAT_PROMPT_CONFIG
api/tests/unit_tests/pyrefly.toml
-4
View File
@@ -317,11 +317,7 @@ project-excludes = [
"core/plugin/test_plugin_runtime.py",
"core/plugin/utils/test_chunk_merger.py",
"core/prompt/test_advanced_prompt_transform.py",
"core/prompt/test_agent_history_prompt_transform.py",
"core/prompt/test_extract_thread_messages.py",
"core/prompt/test_prompt_message.py",
"core/prompt/test_prompt_transform.py",
"core/prompt/test_simple_prompt_transform.py",
"core/rag/cleaner/test_clean_processor.py",
"core/rag/data_post_processor/test_data_post_processor.py",
"core/rag/datasource/keyword/jieba/test_jieba.py",