dify/api/tests/unit_tests/fields/test_dataset_fields.py

import pytest
from sqlalchemy.orm import Session

from fields.dataset_fields import DatasetDetailResponse, dataset_detail_response_source
from models.account import Account
from models.dataset import AppDatasetJoin, Dataset
from models.model import App, AppMode, IconType


def _dataset_detail_payload(**overrides):
    payload = {
        "id": "ds-1",
        "name": "Dataset",
        "description": "desc",
        "provider": "vendor",
        "permission": "only_me",
        "data_source_type": None,
        "indexing_technique": "economy",
        "app_count": 0,
        "document_count": 0,
        "word_count": 0,
        "created_by": "account-1",
        "author_name": None,
        "created_at": 1704067200,
        "updated_by": None,
        "updated_at": 1704067200,
        "embedding_model": None,
        "embedding_model_provider": None,
        "embedding_available": True,
        "retrieval_model_dict": {
            "search_method": "hybrid_search",
            "reranking_enable": True,
            "reranking_mode": "weighted_score",
            "reranking_model": {
                "reranking_provider_name": "provider",
                "reranking_model_name": "model",
            },
            "weights": {
                "weight_type": "customized",
                "keyword_setting": {"keyword_weight": 0.3},
                "vector_setting": {
                    "vector_weight": 0.7,
                    "embedding_model_name": "embedding",
                    "embedding_provider_name": "provider",
                },
            },
            "top_k": 3,
            "score_threshold_enabled": False,
            "score_threshold": None,
        },
        "summary_index_setting": {
            "enable": False,
            "model_name": None,
            "model_provider_name": None,
            "summary_prompt": None,
        },
        "tags": [],
        "doc_form": None,
        "external_knowledge_info": {
            "external_knowledge_id": "knowledge-id",
            "external_knowledge_api_id": "api-id",
            "external_knowledge_api_name": "api",
            "external_knowledge_api_endpoint": "https://example.com",
        },
        "external_retrieval_model": None,
        "doc_metadata": [],
        "built_in_field_enabled": False,
        "pipeline_id": None,
        "runtime_mode": "general",
        "chunk_structure": None,
        "icon_info": {
            "icon_type": "emoji",
            "icon": "📙",
            "icon_background": None,
            "icon_url": None,
        },
        "is_published": False,
        "total_documents": 0,
        "total_available_documents": 0,
        "enable_api": False,
        "is_multimodal": False,
    }
    payload.update(overrides)
    return payload


def _dump_dataset_detail(payload):
    return DatasetDetailResponse.model_validate(payload).model_dump(mode="json")


def test_dataset_detail_preserves_permission_keys():
    response = _dump_dataset_detail(
        _dataset_detail_payload(permission_keys=["dataset.acl.readonly", "dataset.acl.edit"])
    )

    assert response["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"]


def test_dataset_detail_expands_legacy_null_nested_fields():
    response = _dump_dataset_detail(
        _dataset_detail_payload(
            summary_index_setting=None,
            external_knowledge_info=None,
            icon_info=None,
        )
    )

    assert response["summary_index_setting"] == {
        "enable": None,
        "model_name": None,
        "model_provider_name": None,
        "summary_prompt": None,
    }
    assert response["external_knowledge_info"] == {
        "external_knowledge_id": None,
        "external_knowledge_api_id": None,
        "external_knowledge_api_name": None,
        "external_knowledge_api_endpoint": None,
    }
    assert response["icon_info"] == {
        "icon_type": None,
        "icon": None,
        "icon_background": None,
        "icon_url": None,
    }
    assert response["external_retrieval_model"] is None


def test_dataset_detail_expands_legacy_null_retrieval_nested_fields():
    response = _dump_dataset_detail(
        _dataset_detail_payload(
            retrieval_model_dict={
                "search_method": "hybrid_search",
                "reranking_enable": True,
                "reranking_mode": "weighted_score",
                "reranking_model": None,
                "weights": {
                    "keyword_setting": None,
                    "vector_setting": None,
                },
                "top_k": 3,
                "score_threshold_enabled": False,
                "score_threshold": None,
            }
        )
    )

    assert response["retrieval_model_dict"]["reranking_model"] == {
        "reranking_provider_name": None,
        "reranking_model_name": None,
    }
    assert response["retrieval_model_dict"]["weights"] == {
        "weight_type": None,
        "keyword_setting": {"keyword_weight": None},
        "vector_setting": {
            "vector_weight": None,
            "embedding_model_name": None,
            "embedding_provider_name": None,
        },
    }


def test_dataset_detail_expands_missing_weighted_score_nested_fields():
    response = _dump_dataset_detail(
        _dataset_detail_payload(
            retrieval_model_dict={
                "search_method": "hybrid_search",
                "reranking_enable": True,
                "reranking_mode": "weighted_score",
                "reranking_model": None,
                "weights": {},
                "top_k": 3,
                "score_threshold_enabled": False,
                "score_threshold": None,
            }
        )
    )

    assert response["retrieval_model_dict"]["weights"] == {
        "weight_type": None,
        "keyword_setting": {"keyword_weight": None},
        "vector_setting": {
            "vector_weight": None,
            "embedding_model_name": None,
            "embedding_provider_name": None,
        },
    }


@pytest.mark.parametrize("sqlite_session", [(Dataset, Account, App, AppDatasetJoin)], indirect=True)
def test_dataset_detail_response_source_uses_caller_session_for_database_fields(sqlite_session: Session):
    account = Account(name="Ada", email="ada@example.com")
    account.id = "account-1"
    dataset = Dataset(
        id="ds-1",
        tenant_id="tenant-1",
        name="Dataset",
        description="desc",
        provider="vendor",
        permission="only_me",
        data_source_type=None,
        indexing_technique="economy",
        created_by=account.id,
        retrieval_model=_dataset_detail_payload()["retrieval_model_dict"],
        summary_index_setting=_dataset_detail_payload()["summary_index_setting"],
        built_in_field_enabled=False,
        icon_info=_dataset_detail_payload()["icon_info"],
        runtime_mode="general",
        enable_api=False,
        is_multimodal=False,
    )
    dataset.embedding_available = True
    decoy_app = App(
        id="decoy-app",
        tenant_id="tenant-1",
        name="Decoy app",
        description="",
        mode=AppMode.CHAT,
        icon_type=IconType.EMOJI,
        icon="app",
        icon_background="#FFFFFF",
        enable_site=False,
        enable_api=False,
        max_active_requests=0,
    )
    decoy_join = AppDatasetJoin(app_id=decoy_app.id, dataset_id="other-dataset")
    sqlite_session.add_all([account, dataset, decoy_app, decoy_join])
    sqlite_session.flush()

    response = DatasetDetailResponse.model_validate(
        dataset_detail_response_source(dataset, session=sqlite_session),
        from_attributes=True,
    )

    assert response.app_count == 0
    assert response.document_count == 0
    assert response.word_count == 0
    assert response.author_name == "Ada"
    assert response.tags == []
    assert response.doc_form is None
    assert response.external_knowledge_info.external_knowledge_api_id is None
    assert response.doc_metadata == []
    assert response.is_published is False
    assert response.total_documents == 0
    assert response.total_available_documents == 0