deer-flow/backend/tests/test_knowledge_scope.py

import json

import pytest
from pydantic import ValidationError

from deerflow.knowledge_scope import (
    KNOWLEDGE_SCOPE_KEY,
    KnowledgeScope,
    canonicalize_knowledge_scope,
    execution_scope,
)


def test_all_and_disabled_canonicalize_without_selection_fields() -> None:
    assert canonicalize_knowledge_scope({"version": 1, "mode": "all"}) == {
        "version": 1,
        "mode": "all",
    }
    assert canonicalize_knowledge_scope({"version": 1, "mode": "disabled"}) == {
        "version": 1,
        "mode": "disabled",
    }


@pytest.mark.parametrize(
    "payload",
    [
        {"version": 2, "mode": "all"},
        {"version": 1, "mode": "all", "future": True},
        {"version": 1, "mode": "all", "dataset_ids": ["dataset-a"]},
        {"version": 1, "mode": "disabled", "display": {"datasets": []}},
        {"version": 1, "mode": "selected", "dataset_ids": []},
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": ["dataset-a"],
            "document_filters": [{"dataset_id": "dataset-a", "document_ids": []}],
        },
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": ["dataset-a"],
            "document_filters": [
                {"dataset_id": "dataset-a", "document_ids": ["doc-a"]},
                {"dataset_id": "dataset-a", "document_ids": ["doc-b"]},
            ],
        },
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": ["dataset-a"],
            "document_filters": [{"dataset_id": "dataset-b", "document_ids": ["doc-a"]}],
        },
    ],
)
def test_invalid_scope_shapes_are_rejected(payload: dict) -> None:
    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(payload)


def test_selected_scope_trims_and_stably_deduplicates_ids() -> None:
    scope = canonicalize_knowledge_scope(
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": [" dataset-a ", "dataset-b", "dataset-a"],
            "document_filters": [
                {
                    "dataset_id": " dataset-b ",
                    "document_ids": [" doc-1 ", "doc-2", "doc-1"],
                }
            ],
        }
    )

    assert scope == {
        "version": 1,
        "mode": "selected",
        "dataset_ids": ["dataset-a", "dataset-b"],
        "document_filters": [{"dataset_id": "dataset-b", "document_ids": ["doc-1", "doc-2"]}],
    }


def test_execution_scope_drops_untrusted_display_snapshot() -> None:
    scope = canonicalize_knowledge_scope(
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": ["dataset-a"],
            "display": {"datasets": [{"id": "dataset-a", "name": "Agriculture", "documents": []}]},
        }
    )

    assert execution_scope(scope) == {
        "version": 1,
        "mode": "selected",
        "dataset_ids": ["dataset-a"],
    }
    assert KNOWLEDGE_SCOPE_KEY == "knowledge_scope"


@pytest.mark.parametrize(
    "payload",
    [
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": ["dataset-a"],
            "display": {"datasets": [{"id": "dataset-b", "name": "Other"}]},
        },
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": ["dataset-a"],
            "document_filters": [{"dataset_id": "dataset-a", "document_ids": ["doc-a"]}],
            "display": {
                "datasets": [
                    {
                        "id": "dataset-a",
                        "name": "Agriculture",
                        "documents": [{"id": "doc-b", "name": "Other.pdf"}],
                    }
                ]
            },
        },
        {
            "version": 1,
            "mode": "selected",
            "dataset_ids": ["dataset-a"],
            "display": {
                "datasets": [
                    {"id": "dataset-a", "name": "First"},
                    {"id": "dataset-a", "name": "Duplicate"},
                ]
            },
        },
    ],
)
def test_display_must_be_unique_and_related_to_execution_selection(payload: dict) -> None:
    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(payload)


def test_scope_capacity_limits_are_rejected_without_truncation() -> None:
    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(
            {
                "version": 1,
                "mode": "selected",
                "dataset_ids": [f"dataset-{index}" for index in range(101)],
            }
        )

    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(
            {
                "version": 1,
                "mode": "selected",
                "dataset_ids": ["dataset-a"],
                "document_filters": [
                    {
                        "dataset_id": "dataset-a",
                        "document_ids": [f"doc-{index}" for index in range(1001)],
                    }
                ],
            }
        )

    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(
            {
                "version": 1,
                "mode": "selected",
                "dataset_ids": ["x" * 257],
            }
        )


def test_display_capacity_and_canonical_json_byte_limit_are_enforced() -> None:
    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(
            {
                "version": 1,
                "mode": "selected",
                "dataset_ids": [f"dataset-{index}" for index in range(21)],
                "display": {"datasets": [{"id": f"dataset-{index}", "name": f"Dataset {index}"} for index in range(21)]},
            }
        )

    oversized = {
        "version": 1,
        "mode": "selected",
        "dataset_ids": ["dataset-a"],
        "document_filters": [
            {
                "dataset_id": "dataset-a",
                "document_ids": [f"doc-{index}-{'x' * 240}" for index in range(300)],
            }
        ],
    }
    assert len(json.dumps(oversized, ensure_ascii=False).encode()) > 64 * 1024
    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(oversized)


def test_model_rejects_more_than_fifty_display_documents_and_long_names() -> None:
    document_ids = [f"doc-{index}" for index in range(51)]
    with pytest.raises(ValidationError):
        KnowledgeScope.model_validate(
            {
                "version": 1,
                "mode": "selected",
                "dataset_ids": ["dataset-a"],
                "document_filters": [{"dataset_id": "dataset-a", "document_ids": document_ids}],
                "display": {
                    "datasets": [
                        {
                            "id": "dataset-a",
                            "name": "Agriculture",
                            "documents": [{"id": document_id, "name": f"Document {index}"} for index, document_id in enumerate(document_ids)],
                        }
                    ]
                },
            }
        )

    with pytest.raises(ValidationError):
        canonicalize_knowledge_scope(
            {
                "version": 1,
                "mode": "selected",
                "dataset_ids": ["dataset-a"],
                "display": {
                    "datasets": [
                        {"id": "dataset-a", "name": "名" * 257},
                    ]
                },
            }
        )