llm-quant/tests/test_llm_templates.py

"""Test cases for LLM template management."""
import pytest

from app.llm.templates import PromptTemplate, TemplateRegistry


def test_prompt_template_validation():
    """Test template validation logic."""
    # Valid template
    template = PromptTemplate(
        id="test",
        name="Test Template",
        description="A test template",
        template="Hello {name}!",
        variables=["name"]
    )
    assert not template.validate()

    # Missing variable
    template = PromptTemplate(
        id="test",
        name="Test Template",
        description="A test template",
        template="Hello {name}!",
        variables=["name", "missing"]
    )
    errors = template.validate()
    assert len(errors) == 1
    assert "missing" in errors[0]

    # Empty required context
    template = PromptTemplate(
        id="test",
        name="Test Template",
        description="A test template",
        template="Hello {name}!",
        variables=["name"],
        required_context=["", "name"]
    )
    errors = template.validate()
    assert len(errors) == 1
    assert "Empty required context" in errors[0]

    # Empty validation rule
    template = PromptTemplate(
        id="test",
        name="Test Template",
        description="A test template",
        template="Hello {name}!",
        variables=["name"],
        validation_rules=["len(name) > 0", ""]
    )
    errors = template.validate()
    assert len(errors) == 1
    assert "Empty validation rule" in errors[0]


def test_prompt_template_format():
    """Test template formatting."""
    template = PromptTemplate(
        id="test",
        name="Test Template",
        description="A test template",
        template="Hello {name}!",
        variables=["name"],
        required_context=["name"],
        max_length=10
    )

    # Valid context
    result = template.format({"name": "World"})
    assert result == "Hello W..."

    # Missing required context
    with pytest.raises(ValueError) as exc:
        template.format({})
    assert "Missing required context" in str(exc.value)

    # Missing variable
    template_no_required = PromptTemplate(
        id="test2",
        name="Test Template",
        description="A test template",
        template="Hello {name}!",
        variables=["name"],
    )
    with pytest.raises(ValueError) as exc:
        template_no_required.format({"wrong": "value"})
    assert "Missing template variable" in str(exc.value)


def test_template_registry():
    """Test template registry operations."""
    TemplateRegistry.clear()

    # Register valid template
    template = PromptTemplate(
        id="test",
        name="Test Template",
        description="A test template",
        template="Hello {name}!",
        variables=["name"]
    )
    TemplateRegistry.register(template)
    assert TemplateRegistry.get("test") is not None
    assert TemplateRegistry.get_active_version("test") == "1.0.0"

    # Register invalid template
    invalid = PromptTemplate(
        id="invalid",
        name="Invalid Template",
        description="An invalid template",
        template="Hello {name}!",
        variables=["wrong"]
    )
    with pytest.raises(ValueError) as exc:
        TemplateRegistry.register(invalid)
    assert "Invalid template" in str(exc.value)

    # List templates
    templates = TemplateRegistry.list()
    assert len(templates) == 1
    assert templates[0].id == "test"

    # Load from JSON
    json_str = '''
    {
        "json_test": {
            "name": "JSON Test",
            "description": "Test template from JSON",
            "template": "Hello {name}!",
            "variables": ["name"],
            "version": "2024.10",
            "metadata": {"author": "qa"},
            "activate": true
        }
    }
    '''
    TemplateRegistry.load_from_json(json_str)
    loaded = TemplateRegistry.get("json_test")
    assert loaded is not None
    assert TemplateRegistry.get_active_version("json_test") == "2024.10"

    # Invalid JSON
    with pytest.raises(ValueError) as exc:
        TemplateRegistry.load_from_json("invalid json")
    assert "Invalid JSON" in str(exc.value)

    # Non-object JSON
    with pytest.raises(ValueError) as exc:
        TemplateRegistry.load_from_json("[1, 2, 3]")
    assert "JSON root must be an object" in str(exc.value)


def test_default_templates():
    """Test default template registration."""
    TemplateRegistry.clear(reload_defaults=True)
    from app.llm.templates import DEFAULT_TEMPLATES

    # Verify default templates are loaded
    assert len(TemplateRegistry.list()) > 0

    # Check specific templates
    dept_base = TemplateRegistry.get("department_base")
    assert dept_base is not None
    assert "部门基础模板" in dept_base.name

    momentum = TemplateRegistry.get("momentum_dept")
    assert momentum is not None
    assert "动量研究部门" in momentum.name

    # Validate template content
    assert all("{" + var + "}" in dept_base.template for var in dept_base.variables)
    assert all("{" + var + "}" in momentum.template for var in momentum.variables)

    # Test template usage
    context = {
        "title": "测试部门",
        "ts_code": "000001.SZ",
        "trade_date": "20251005",
        "description": "测试描述",
        "instruction": "测试指令",
        "data_scope": "daily,daily_basic",
        "features": "特征1,特征2",
        "market_snapshot": "市场数据1,市场数据2",
        "supplements": "补充数据"
    }
    result = dept_base.format(context)
    assert "测试部门" in result
    assert "000001.SZ" in result
    assert "20251005" in result