fischer-agentkit/tests/unit/test_prompt_template.py

"""Tests for PromptTemplate - Prompt 模板渲染"""

import pytest

from agentkit.prompts.section import PromptSection
from agentkit.prompts.template import PromptTemplate


class TestPromptTemplateInit:
    """PromptTemplate 初始化测试"""

    def test_default_name_and_version(self):
        section = PromptSection(identity="I am a bot")
        tpl = PromptTemplate(sections=section)
        assert tpl.name == ""
        assert tpl.version == "1.0.0"

    def test_custom_name_and_version(self):
        section = PromptSection()
        tpl = PromptTemplate(sections=section, name="my_template", version="2.0")
        assert tpl.name == "my_template"
        assert tpl.version == "2.0"

    def test_sections_property(self):
        section = PromptSection(identity="Bot")
        tpl = PromptTemplate(sections=section)
        assert tpl.sections is section


class TestPromptTemplateRender:
    """PromptTemplate.render 渲染测试"""

    def test_render_empty_sections(self):
        section = PromptSection()
        tpl = PromptTemplate(sections=section)
        messages = tpl.render()
        assert messages == []

    def test_render_system_parts(self):
        section = PromptSection(
            identity="You are an assistant.",
            context="Context info here.",
            constraints="Do not lie.",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render()

        assert len(messages) == 1
        assert messages[0]["role"] == "system"
        assert "You are an assistant." in messages[0]["content"]
        assert "Context info here." in messages[0]["content"]
        assert "Do not lie." in messages[0]["content"]

    def test_render_user_parts(self):
        section = PromptSection(
            instructions="Answer the question.",
            output_format="JSON format.",
            examples="Q: 1+1? A: 2",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render()

        assert len(messages) == 1
        assert messages[0]["role"] == "user"
        assert "Answer the question." in messages[0]["content"]
        assert "JSON format." in messages[0]["content"]
        assert "Q: 1+1? A: 2" in messages[0]["content"]

    def test_render_system_and_user(self):
        section = PromptSection(
            identity="Bot",
            instructions="Do stuff",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render()

        assert len(messages) == 2
        assert messages[0]["role"] == "system"
        assert messages[1]["role"] == "user"

    def test_render_variable_substitution_in_context(self):
        section = PromptSection(
            context="Hello ${name}, welcome to ${place}.",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render(variables={"name": "Alice", "place": "Wonderland"})

        assert len(messages) == 1
        assert "Hello Alice, welcome to Wonderland." in messages[0]["content"]

    def test_render_variable_substitution_in_instructions(self):
        section = PromptSection(
            instructions="Process ${item} with ${method}.",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render(variables={"item": "data", "method": "AI"})

        assert len(messages) == 1
        assert "Process data with AI." in messages[0]["content"]

    def test_render_unsubstituted_variables_remain(self):
        section = PromptSection(
            context="Hello ${name}, ${unknown} stays.",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render(variables={"name": "Bob"})

        assert "Hello Bob, ${unknown} stays." in messages[0]["content"]

    def test_render_no_variables(self):
        section = PromptSection(
            identity="Bot",
            context="No vars here.",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render()
        assert "No vars here." in messages[0]["content"]

    def test_render_system_parts_joined_by_double_newline(self):
        section = PromptSection(
            identity="Part1",
            context="Part2",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render()
        assert messages[0]["content"] == "Part1\n\nPart2"

    def test_render_user_parts_joined_by_double_newline(self):
        section = PromptSection(
            instructions="Step1",
            output_format="Step2",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render()
        assert messages[0]["content"] == "Step1\n\nStep2"

    def test_render_identity_and_constraints_not_substituted(self):
        """identity 和 constraints 不做变量替换"""
        section = PromptSection(
            identity="I am ${name}",
            constraints="Never say ${word}",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render(variables={"name": "Bot", "word": "hello"})

        assert "I am ${name}" in messages[0]["content"]
        assert "Never say ${word}" in messages[0]["content"]

    def test_render_output_format_and_examples_not_substituted(self):
        """output_format 和 examples 不做变量替换"""
        section = PromptSection(
            output_format="Return ${format}",
            examples="Example: ${example}",
        )
        tpl = PromptTemplate(sections=section)
        messages = tpl.render(variables={"format": "JSON", "example": "test"})

        assert "Return ${format}" in messages[0]["content"]
        assert "Example: ${example}" in messages[0]["content"]

    def test_render_context_budget_parameter_accepted(self):
        """context_budget 参数被接受（当前实现未使用）"""
        section = PromptSection(identity="Bot")
        tpl = PromptTemplate(sections=section)
        messages = tpl.render(context_budget=5000)
        assert len(messages) == 1