mirror of
https://github.com/langgenius/dify.git
synced 2026-11-12 06:11:12 +08:00
259 lines
8.9 KiB
Python
259 lines
8.9 KiB
Python
from __future__ import annotations
|
|
|
|
import builtins
|
|
import inspect
|
|
import json
|
|
import uuid
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import ANY, patch
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
from flask.views import MethodView
|
|
|
|
from controllers.common.controller_schemas import MessageListQuery
|
|
from core.app.entities.app_invoke_entities import InvokeFrom
|
|
from core.entities.execution_extra_content import HumanInputContent
|
|
from models.enums import ConversationFromSource, EndUserType, FeedbackFromSource, FeedbackRating
|
|
from models.model import (
|
|
App,
|
|
AppMode,
|
|
Conversation,
|
|
EndUser,
|
|
Message,
|
|
MessageAgentThought,
|
|
MessageFeedback,
|
|
)
|
|
|
|
"""Unit tests for controllers.web.message message list mapping."""
|
|
|
|
|
|
# Ensure flask_restx.api finds MethodView during import.
|
|
if not hasattr(builtins, "MethodView"):
|
|
builtins.MethodView = MethodView # type: ignore[attr-defined]
|
|
|
|
|
|
def _load_controller_module():
|
|
"""Import controllers.web.message using a stub package."""
|
|
|
|
import importlib
|
|
import importlib.util
|
|
import sys
|
|
|
|
parent_module_name = "controllers.web"
|
|
module_name = f"{parent_module_name}.message"
|
|
|
|
if parent_module_name not in sys.modules:
|
|
from flask_restx import Namespace
|
|
|
|
stub = ModuleType(parent_module_name)
|
|
web_controller_dir = Path(__file__).resolve().parents[4] / "controllers" / "web"
|
|
stub.__file__ = str(web_controller_dir / "__init__.py")
|
|
stub.__path__ = [str(web_controller_dir)]
|
|
stub.__package__ = "controllers"
|
|
stub.__spec__ = importlib.util.spec_from_loader(parent_module_name, loader=None, is_package=True)
|
|
stub.web_ns = Namespace("web", description="Web API", path="/")
|
|
sys.modules[parent_module_name] = stub
|
|
|
|
wraps_module_name = f"{parent_module_name}.wraps"
|
|
if wraps_module_name not in sys.modules:
|
|
wraps_stub = ModuleType(wraps_module_name)
|
|
|
|
class WebApiResource:
|
|
pass
|
|
|
|
wraps_stub.WebApiResource = WebApiResource
|
|
sys.modules[wraps_module_name] = wraps_stub
|
|
|
|
return importlib.import_module(module_name)
|
|
|
|
|
|
message_module = _load_controller_module()
|
|
MessageListApi = message_module.MessageListApi
|
|
|
|
|
|
class TestMessageListCases:
|
|
@staticmethod
|
|
@pytest.fixture
|
|
def app() -> Flask:
|
|
app = Flask(__name__)
|
|
app.config["TESTING"] = True
|
|
return app
|
|
|
|
@staticmethod
|
|
@pytest.mark.parametrize(
|
|
"sqlite_session",
|
|
[(App, Conversation, EndUser, Message, MessageAgentThought, MessageFeedback)],
|
|
indirect=True,
|
|
)
|
|
def test_message_list_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session) -> None:
|
|
conversation_id = str(uuid4())
|
|
message_id = str(uuid4())
|
|
|
|
created_at = datetime(2024, 1, 1, 12, 0, 0)
|
|
resource_created_at = datetime(2024, 1, 1, 13, 0, 0)
|
|
thought_created_at = datetime(2024, 1, 1, 14, 0, 0)
|
|
|
|
retriever_resource = {
|
|
"id": "res-obj",
|
|
"message_id": message_id,
|
|
"position": 2,
|
|
"dataset_id": "ds-1",
|
|
"dataset_name": "dataset",
|
|
"document_id": "doc-1",
|
|
"document_name": "document",
|
|
"data_source_type": "file",
|
|
"segment_id": "seg-1",
|
|
"score": 0.9,
|
|
"hit_count": 1,
|
|
"word_count": 10,
|
|
"segment_position": 0,
|
|
"index_node_hash": "hash",
|
|
"content": "content",
|
|
"created_at": int(resource_created_at.timestamp()),
|
|
}
|
|
|
|
agent_thought = MessageAgentThought(
|
|
message_chain_id="chain-1",
|
|
message_id=message_id,
|
|
position=1,
|
|
created_by_role="end_user",
|
|
created_by="end-user-1",
|
|
thought="thinking",
|
|
tool="tool",
|
|
tool_labels_str=json.dumps({"label": "value"}),
|
|
tool_input="{}",
|
|
observation="observed",
|
|
message_files=json.dumps(["file-a"]),
|
|
)
|
|
agent_thought.id = "thought-1"
|
|
agent_thought.created_at = thought_created_at
|
|
|
|
message_files = [
|
|
{"id": "file-dict", "filename": "a.txt", "type": "file", "transfer_method": "local"},
|
|
{"id": "file-obj", "filename": "b.txt", "type": "file", "transfer_method": "local"},
|
|
]
|
|
|
|
app_model = App(
|
|
id="app-1",
|
|
tenant_id="tenant-1",
|
|
name="Chat App",
|
|
mode=AppMode.CHAT,
|
|
enable_site=False,
|
|
enable_api=False,
|
|
)
|
|
end_user = EndUser(
|
|
id="end-user-1",
|
|
tenant_id="tenant-1",
|
|
app_id=app_model.id,
|
|
type=EndUserType.BROWSER,
|
|
name="Web User",
|
|
session_id="session-1",
|
|
)
|
|
conversation = Conversation(
|
|
id=conversation_id,
|
|
app_id=app_model.id,
|
|
mode=AppMode.CHAT,
|
|
name="Conversation",
|
|
_inputs={},
|
|
status="normal",
|
|
from_source=ConversationFromSource.API,
|
|
from_end_user_id=end_user.id,
|
|
)
|
|
message = Message(
|
|
id=message_id,
|
|
app_id=app_model.id,
|
|
conversation_id=conversation_id,
|
|
parent_message_id=None,
|
|
_inputs={"foo": "bar"},
|
|
query="hello",
|
|
message={},
|
|
answer="answer",
|
|
message_unit_price=0,
|
|
message_price_unit=0,
|
|
answer_unit_price=0,
|
|
answer_price_unit=0,
|
|
provider_response_latency=0,
|
|
total_price=0,
|
|
currency="USD",
|
|
invoke_from=InvokeFrom.SERVICE_API,
|
|
from_source=ConversationFromSource.API,
|
|
from_end_user_id=end_user.id,
|
|
app_mode=AppMode.CHAT,
|
|
message_metadata=json.dumps(
|
|
{
|
|
"meta": "value",
|
|
"retriever_resources": [
|
|
{"id": "res-dict", "message_id": message_id, "position": 1},
|
|
retriever_resource,
|
|
],
|
|
}
|
|
),
|
|
status="normal",
|
|
error=None,
|
|
)
|
|
message.created_at = created_at
|
|
message.set_extra_contents(
|
|
[
|
|
HumanInputContent(
|
|
workflow_run_id=str(uuid.uuid4()),
|
|
submitted=True,
|
|
).model_dump(mode="json")
|
|
]
|
|
)
|
|
feedback = MessageFeedback(
|
|
app_id=app_model.id,
|
|
conversation_id=conversation_id,
|
|
message_id=message_id,
|
|
rating=FeedbackRating.LIKE,
|
|
from_source=FeedbackFromSource.USER,
|
|
from_end_user_id=end_user.id,
|
|
)
|
|
sqlite_session.add_all([app_model, end_user, conversation, message, feedback, agent_thought])
|
|
sqlite_session.commit()
|
|
|
|
pagination = SimpleNamespace(limit=20, has_more=False, data=[message])
|
|
|
|
def message_files_with_session(_self, *, session):
|
|
del session
|
|
return message_files
|
|
|
|
monkeypatch.setattr(Message, "message_files_with_session", message_files_with_session)
|
|
|
|
with (
|
|
patch.object(message_module.MessageService, "pagination_by_first_id", return_value=pagination) as mock_page,
|
|
patch.object(message_module.db, "session", return_value=sqlite_session),
|
|
app.test_request_context(f"/messages?conversation_id={conversation_id}&limit=20"),
|
|
):
|
|
query = MessageListQuery.model_validate({"conversation_id": conversation_id, "limit": 20})
|
|
response = inspect.unwrap(MessageListApi.get)(MessageListApi(), query, app_model, end_user)
|
|
|
|
mock_page.assert_called_once_with(app_model, end_user, conversation_id, None, 20, session=ANY)
|
|
assert response["limit"] == 20
|
|
assert response["has_more"] is False
|
|
assert len(response["data"]) == 1
|
|
|
|
item = response["data"][0]
|
|
assert item["id"] == message_id
|
|
assert item["conversation_id"] == conversation_id
|
|
assert item["inputs"] == {"foo": "bar"}
|
|
assert item["answer"] == "answer"
|
|
assert item["feedback"]["rating"] == "like"
|
|
assert item["metadata"]["meta"] == "value"
|
|
assert item["created_at"] == int(created_at.timestamp())
|
|
|
|
assert item["retriever_resources"][0]["id"] == "res-dict"
|
|
assert item["retriever_resources"][1]["id"] == "res-obj"
|
|
assert item["retriever_resources"][1]["created_at"] == int(resource_created_at.timestamp())
|
|
|
|
assert item["agent_thoughts"][0]["chain_id"] == "chain-1"
|
|
assert item["agent_thoughts"][0]["created_at"] == int(thought_created_at.timestamp())
|
|
assert item["extra_contents"][0]["workflow_run_id"] == message.extra_contents[0]["workflow_run_id"]
|
|
assert item["extra_contents"][0]["submitted"] == message.extra_contents[0]["submitted"]
|
|
|
|
assert item["message_files"][0]["id"] == "file-dict"
|
|
assert item["message_files"][1]["id"] == "file-obj"
|