DataOps-platform/tests/data_research/test_ontology_suggestions.py

from __future__ import annotations

import pytest


def test_rule_suggestions_are_deterministic_explainable_and_confidence_sorted():
    from app.core.data_research.ontology.suggestions import OntologySuggestionEngine

    engine = OntologySuggestionEngine()
    suggestions = engine.generate(
        domains=[{"uid": "domain-customer", "name": "客户"}],
        data_elements=[
            {"uid": "element-name", "name": "客户名称", "status": "published", "aliases": ["姓名"]},
            {"uid": "element-id", "name": "客户编号", "status": "published", "aliases": []},
            {"uid": "draft", "name": "未发布", "status": "draft"},
        ],
        foreign_keys=[
            {"from_uid": "customer", "to_uid": "account", "evidence_uid": "evidence-fk"}
        ],
    )

    assert [item.kind for item in suggestions[:3]] == ["class", "property", "property"]
    assert suggestions == engine.generate(
        domains=[{"uid": "domain-customer", "name": "客户"}],
        data_elements=[
            {"uid": "element-name", "name": "客户名称", "status": "published", "aliases": ["姓名"]},
            {"uid": "element-id", "name": "客户编号", "status": "published", "aliases": []},
            {"uid": "draft", "name": "未发布", "status": "draft"},
        ],
        foreign_keys=[{"from_uid": "customer", "to_uid": "account", "evidence_uid": "evidence-fk"}],
    )
    assert all(item.evidence_uids for item in suggestions)
    assert [item.confidence for item in suggestions] == sorted(
        [item.confidence for item in suggestions], reverse=True
    )


def test_ai_suggestion_requires_evidence_and_records_model_versions():
    from app.core.data_research.ontology.suggestions import (
        OntologySuggestionEngine,
        SuggestionInvalid,
    )

    ai = lambda _context: [
        {
            "kind": "relation",
            "payload": {"uid": "rel-ai", "from_class_uid": "a", "to_class_uid": "b"},
            "evidence_uids": [],
            "confidence": 0.7,
        }
    ]
    with pytest.raises(SuggestionInvalid, match="evidence"):
        OntologySuggestionEngine(ai_port=ai, model_version="model-v1", prompt_version="prompt-v2").generate(
            domains=[], data_elements=[], foreign_keys=[], include_ai=True
        )