ontorefactor-governance/tests/test_llm.py

import json
from dataclasses import replace
from io import BytesIO
from urllib.error import HTTPError

import pytest
from lambdagent import Context

from ontorefactor_governance.agents import execute_governance_agent, normalize_source
from ontorefactor_governance.db import Database
from ontorefactor_governance.llm import (
    LLMSettings,
    LongCatProvider,
    LongCatSemanticAgent,
    PROMPT_VERSION,
    public_llm_status,
    test_llm_connection as check_llm_connection,
    validate_semantic_output,
)
from ontorefactor_governance.service import GovernanceError
from ontorefactor_governance.service import create_project, list_assertions, list_evidence


def settings(api_key="test-key"):
    return LLMSettings(
        deployment_mode="auto",
        provider="longcat",
        model="LongCat-2.0",
        base_url="https://api.longcat.chat/openai",
        api_key=api_key,
        temperature=0.1,
        max_tokens=4096,
        timeout=30,
        max_input_chars=120000,
        max_retries=0,
    )


class FakeResponse:
    def __init__(self, value):
        self.value = value

    def __enter__(self):
        return self

    def __exit__(self, *_):
        return False

    def read(self):
        return json.dumps(self.value).encode()


def test_longcat_provider_uses_official_openai_endpoint(monkeypatch):
    captured = {}

    def fake_urlopen(request, timeout):
        captured["url"] = request.full_url
        captured["authorization"] = request.get_header("Authorization")
        captured["body"] = json.loads(request.data)
        captured["timeout"] = timeout
        return FakeResponse({
            "model": "LongCat-2.0",
            "choices": [{"message": {"content": '{"ok":true}'}}],
            "usage": {"prompt_tokens": 12, "completion_tokens": 4},
        })

    monkeypatch.setattr("urllib.request.urlopen", fake_urlopen)
    provider = LongCatProvider(settings())
    result = provider.chat([{"role": "user", "content": "health"}])

    assert result == '{"ok":true}'
    assert captured["url"] == "https://api.longcat.chat/openai/v1/chat/completions"
    assert captured["authorization"] == "Bearer test-key"
    assert captured["body"]["model"] == "LongCat-2.0"
    assert captured["body"]["response_format"] == {"type": "json_object"}
    assert captured["body"]["thinking"] == {"type": "disabled"}
    assert provider.last_metrics["input_tokens"] == 12
    assert provider.last_metrics["output_tokens"] == 4
    assert provider.last_metrics["request_hash"]
    assert provider.last_metrics["response_hash"]


def test_public_status_never_exposes_api_key():
    status = public_llm_status(settings("super-secret-value"))
    assert status["configured"] is True
    assert status["prompt_version"] == PROMPT_VERSION
    assert "super-secret-value" not in json.dumps(status)


def test_longcat_retries_rate_limits_with_backoff(monkeypatch):
    attempts, delays = [], []

    def fake_urlopen(_request, timeout):
        attempts.append(1)
        if len(attempts) == 1:
            raise HTTPError("https://api.longcat.chat", 429, "rate limited", {"Retry-After": "0"}, BytesIO(b'{"error":"limited"}'))
        return FakeResponse({"choices": [{"message": {"content": '{"ok":true}'}}], "usage": {}})

    monkeypatch.setattr("urllib.request.urlopen", fake_urlopen)
    monkeypatch.setattr("time.sleep", lambda seconds: delays.append(seconds))
    provider = LongCatProvider(replace(settings(), max_retries=1))

    assert provider.chat([{"role": "user", "content": "health"}]) == '{"ok":true}'
    assert len(attempts) == 2
    assert delays == [0.0]


def test_connection_maps_auth_failure_without_leaking_provider_body(monkeypatch):
    secret = "must-not-leak"

    def fake_urlopen(_request, timeout):
        raise HTTPError(
            "https://api.longcat.chat/openai/v1/chat/completions",
            401,
            "unauthorized",
            {},
            BytesIO(f'{{"error":"invalid {secret}"}}'.encode()),
        )

    monkeypatch.setattr("urllib.request.urlopen", fake_urlopen)
    with pytest.raises(GovernanceError) as error:
        check_llm_connection(settings(secret))

    assert error.value.code == "LLM_AUTH_FAILED"
    assert error.value.status == 502
    assert secret not in str(error.value)


def test_model_output_is_validated_capped_and_pending():
    catalog = [{
        "local_id": "m1:table:public:customer", "layer": "M1", "profile": "software",
        "kind": "Table", "name": "CustomerTable", "attributes": {},
    }]
    raw = json.dumps({
        "business_objects": [{
            "source_asset_id": "m1:table:public:customer", "name": "Customer",
            "description": "客户主数据", "domain": "CRM", "confidence": 0.99, "reason": "表名与字段结构",
        }],
        "owners": [{
            "target_asset_id": "m1:table:public:customer", "owner_role": "CRMDataOwner",
            "confidence": 0.8, "reason": "属于客户域",
        }],
        "classifications": [], "quality_rules": [], "issues": [],
    }, ensure_ascii=False)
    result = validate_semantic_output(raw, {}, catalog, {"model": "LongCat-2.0"})

    assert len(result["elements"]) == 2
    assert len(result["assertions"]) == 2
    assert all(item["review_status"] == "pending" for item in result["assertions"])
    assert max(item["confidence"] for item in result["assertions"]) <= 0.85
    assert all(item["generated_by"].startswith("longcat:LongCat-2.0") for item in result["assertions"])


def test_chinese_semantic_names_produce_unique_stable_uris():
    catalog = [{
        "local_id": f"m1:table:public:source_{index}", "layer": "M1", "profile": "software",
        "kind": "Table", "name": f"Source{index}", "attributes": {},
    } for index in (1, 2)]
    raw = json.dumps({
        "business_objects": [{
            "source_asset_id": item["local_id"], "name": "用户档案", "description": "用户数据",
            "domain": "IAM", "confidence": 0.8, "reason": "名称和结构证据",
        } for item in catalog],
        "owners": [], "classifications": [],
        "quality_rules": [{
            "target_asset_id": item["local_id"], "name": "完整性规则", "dimension": "completeness",
            "description": "字段不能为空", "expression": "not_null", "severity": "high",
            "confidence": 0.8, "reason": "必填字段证据",
        } for item in catalog],
        "issues": [], "relationships": [],
    }, ensure_ascii=False)

    result = validate_semantic_output(raw, {}, catalog, {"model": "LongCat-2.0"})
    uris = [item["uri"] for item in result["elements"]]

    assert len(uris) == len(set(uris)) == 4
    assert all("unnamed" not in uri for uri in uris)


def test_model_cannot_reference_unknown_assets():
    raw = json.dumps({
        "business_objects": [{
            "source_asset_id": "invented:database", "name": "Invented", "description": "",
            "domain": "", "confidence": 0.5, "reason": "guess",
        }],
        "owners": [], "classifications": [], "quality_rules": [], "issues": [],
    })
    with pytest.raises(GovernanceError) as error:
        validate_semantic_output(raw, {}, [], {"model": "LongCat-2.0"})
    assert error.value.code == "UNSAFE_LLM_OUTPUT"


def test_model_aliases_are_normalized_then_strictly_validated():
    catalog = [
        {"local_id": "m1:table:default:customer", "layer": "M1", "profile": "software", "kind": "Table", "name": "customer", "attributes": {}},
        {"local_id": "m1:column:default:customer:email", "layer": "M1", "profile": "software", "kind": "Column", "name": "email", "attributes": {}},
    ]
    raw = json.dumps({
        "business_objects": [{"name": "Customer", "description": "客户", "asset_ids": ["m1:table:default:customer"]}],
        "owners": [{"asset_id": "m1:table:default:customer", "role": "CustomerDataOwner", "confidence": 0.8}],
        "classifications": [{"asset_id": "m1:column:default:customer:email", "classification": "confidential", "confidence": 0.8}],
        "quality_rules": [{"asset_id": "m1:column:default:customer:email", "rule": "邮箱格式校验", "expression": "email LIKE '%@%'", "confidence": 0.7}],
        "issues": [{"asset_ids": ["m1:table:default:customer"], "description": "缺少数据保留策略", "confidence": 0.7}],
    }, ensure_ascii=False)

    result = validate_semantic_output(raw, {}, catalog, {"model": "LongCat-2.0"})

    assert len(result["elements"]) == 4
    assert len(result["assertions"]) == 4
    assert result["issues"][0]["code"].startswith("LLM_MODEL_RISK_")
    assert all(item["review_status"] == "pending" for item in result["assertions"])


def test_longcat_is_a_real_lambdagent_term(monkeypatch):
    provider = LongCatProvider(settings())

    def fake_chat(_messages):
        provider.last_metrics = {
            "provider": "longcat", "model": "LongCat-2.0", "prompt_version": PROMPT_VERSION,
            "input_tokens": 30, "output_tokens": 20, "latency_ms": 10,
            "request_hash": "request", "response_hash": "response",
        }
        return json.dumps({
            "business_objects": [{
                "source_asset_id": "m1:table:default:customer", "name": "Customer",
                "description": "客户", "domain": "CRM", "confidence": 0.8, "reason": "名称证据",
            }],
            "owners": [], "classifications": [], "quality_rules": [], "issues": [],
        })

    monkeypatch.setattr(provider, "chat", fake_chat)
    source = normalize_source({
        "source_type": "ddl", "source_name": "customer.sql", "environment": "development",
        "content": "CREATE TABLE customer (id INTEGER PRIMARY KEY);",
    })
    context = Context(run_id="test-longcat")
    result = LongCatSemanticAgent(settings(), provider).apply(source, context)

    assert result["agent"] == "business-semantics-llm"
    assert result["llm"]["input_tokens"] == 30
    assert result["assertions"][0]["review_status"] == "pending"
    assert any(entry.term_name == "longcat-governance-semantics" for entry in context.trace)


def test_semantic_agent_repairs_one_invalid_model_response(monkeypatch):
    provider = LongCatProvider(settings())
    calls = []

    def fake_chat(_messages):
        calls.append(1)
        provider.last_metrics = {
            "provider": "longcat", "model": "LongCat-2.0", "prompt_version": PROMPT_VERSION,
            "input_tokens": 10, "output_tokens": 5, "latency_ms": 2,
            "request_hash": f"request-{len(calls)}", "response_hash": f"response-{len(calls)}",
        }
        if len(calls) == 1:
            return '{"business_objects": ['
        return json.dumps({
            "business_objects": [], "owners": [], "classifications": [], "quality_rules": [], "issues": [],
        })

    monkeypatch.setattr(provider, "chat", fake_chat)
    source = normalize_source({
        "source_type": "ddl", "source_name": "empty.sql", "content": "CREATE TABLE sample (id INTEGER);",
    })
    result = LongCatSemanticAgent(settings(), provider).apply(source, Context(run_id="repair-test"))

    assert len(calls) == 2
    assert result["llm"]["attempts"] == 2
    assert result["llm"]["input_tokens"] == 20
    assert result["llm"]["output_tokens"] == 10


def test_truncated_model_output_keeps_only_complete_semantic_arrays():
    catalog = [{
        "local_id": "m1:table:default:customer", "layer": "M1", "profile": "software",
        "kind": "Table", "name": "customer", "attributes": {},
    }]
    raw = (
        '{"business_objects":[],"owners":[{"target_asset_id":"m1:table:default:customer",'
        '"owner_role":"CustomerDataOwner","confidence":0.8,"reason":"客户域"}],'
        '"classifications":[],"quality_rules":[{"target_asset_id":"m1:table:default:customer"'
    )

    result = validate_semantic_output(raw, {}, catalog, {"model": "LongCat-2.0"})

    assert result["llm"]["partial_output_recovered"] is True
    assert len(result["assertions"]) == 1
    assert result["assertions"][0]["predicate"] == "OWNED_BY"


def test_llm_run_persists_model_evidence_and_audit_metrics(monkeypatch):
    monkeypatch.setenv("LONGCAT_API_KEY", "test-key")
    monkeypatch.setenv("ONTOREFACTOR_LLM_MODE", "auto")

    def fake_chat(provider, _messages):
        provider.last_metrics = {
            "provider": "longcat", "model": "LongCat-2.0", "prompt_version": PROMPT_VERSION,
            "input_tokens": 40, "output_tokens": 25, "latency_ms": 12.5,
            "request_hash": "request-hash", "response_hash": "response-hash",
        }
        return json.dumps({
            "business_objects": [{
                "source_asset_id": "m1:table:default:customer", "name": "Customer",
                "description": "客户主数据", "domain": "CRM", "confidence": 0.8, "reason": "表结构证据",
            }],
            "owners": [], "classifications": [], "quality_rules": [], "issues": [],
        })

    monkeypatch.setattr(LongCatProvider, "chat", fake_chat)
    db = Database("sqlite:///:memory:")
    project = create_project(db, "tenant-a", {"name": "LLM", "namespace": "urn:test:llm"})
    result = execute_governance_agent(db, "tenant-a", project["id"], {
        "source_type": "ddl", "source_name": "customer.sql", "semantic_mode": "llm",
        "content": "CREATE TABLE customer (id INTEGER PRIMARY KEY);",
    })

    assert result["semantic_mode"] == "llm"
    assert result["llm"]["input_tokens"] == 40
    assert any(item["kind"] == "ModelInference" for item in list_evidence(db, "tenant-a", project["id"]))
    inferred = [item for item in list_assertions(db, "tenant-a", project["id"]) if item["generated_by"].startswith("longcat:")]
    assert len(inferred) == 1
    assert inferred[0]["review_status"] == "pending"
    run = db.fetchone("SELECT * FROM governance_agent_runs WHERE id=?", (result["run_id"],))
    assert run["llm_model"] == "LongCat-2.0"
    assert run["prompt_version"] == PROMPT_VERSION
    assert run["input_tokens"] == 40
    assert run["output_tokens"] == 25
    assert run["llm_request_hash"] == "request-hash"
    assert run["llm_response_hash"] == "response-hash"
    db.close()