llm-in-text/backend/tests/test_prompt.py

import sys
from pathlib import Path


BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

import prompt  # noqa: E402


def test_prompt_builds_system_and_user():
    system_prompt, user_prompt, prefill = prompt.build_completion_prompts(
        prefix="The result is ",
        suffix="for this dataset.",
        language_id="markdown",
    )

    assert "inline completion engine" in system_prompt
    assert "$...$" in system_prompt
    assert "$$...$$" in system_prompt
    assert "```{language}" in system_prompt
    assert "Mermaid" in system_prompt
    assert "CURSOR_FENCE_LANGUAGE" in system_prompt
    assert "MERMAID_CONTEXT" in system_prompt
    assert "CURSOR_IN_FENCED_CODE_BLOCK" in user_prompt
    assert "CURSOR_FENCE_LANGUAGE" in user_prompt
    assert "MERMAID_CONTEXT" in user_prompt
    assert "PREFIX_ENDS_WITH_NEWLINE" in user_prompt
    assert "SUFFIX_STARTS_WITH_NEWLINE" in user_prompt
    assert "actual line breaks" in system_prompt
    assert "Use real line breaks instead of spelled-out escape sequences" in user_prompt
    assert "Do not explain newline or boundary choices" in user_prompt
    assert "Continue after the PREFILL text" in user_prompt
    assert "Step 1" not in user_prompt
    assert "Does output need" not in user_prompt
    assert "assistant" in system_prompt
    assert "fim_middle" in system_prompt
    assert prefill == ""
    assert "start output with \\n" not in user_prompt
    assert "Use single \\n" not in system_prompt


def test_completion_prefill_appended_to_fim_middle():
    _, user_prompt, prefill = prompt.build_completion_prompts(
        prefix="即时可用的 LLM 系统",
        suffix="",
    )
    assert prefill == "系统"
    assert user_prompt.endswith("<|fim_middle|>系统")


def test_completion_prefill_empty_after_newline():
    _, user_prompt, prefill = prompt.build_completion_prompts(
        prefix="即时可用的 LLM 系统\n",
        suffix="",
    )
    assert prefill == ""
    assert user_prompt.endswith("<|fim_middle|>")


def test_cursor_in_fence_detection():
    assert prompt._cursor_in_fenced_code_block("") is False
    assert prompt._cursor_in_fenced_code_block("```python\nprint('x')\n") is True
    assert prompt._cursor_in_fenced_code_block("```python\nprint('x')\n```\n") is False
    assert prompt._cursor_in_fenced_code_block("text ```not-a-fence``` tail") is False


def test_active_fence_language_detection():
    assert prompt._active_fence_language("") == "none"
    assert prompt._active_fence_language("```mermaid\nflowchart TD\nA-->B\n") == "mermaid"
    assert prompt._active_fence_language("```python\nprint('x')\n") == "python"
    assert prompt._active_fence_language("```\nline\n") == "unknown"
    assert prompt._active_fence_language("```mermaid\nA-->B\n```\n") == "none"


def test_newline_flags():
    _, user_prompt_a, _ = prompt.build_completion_prompts(
        prefix="Hello",
        suffix="World",
    )
    assert "CURSOR_IN_FENCED_CODE_BLOCK: false" in user_prompt_a
    assert "CURSOR_FENCE_LANGUAGE: none" in user_prompt_a
    assert "MERMAID_CONTEXT: false" in user_prompt_a
    assert "PREFIX_ENDS_WITH_NEWLINE: false" in user_prompt_a
    assert "SUFFIX_STARTS_WITH_NEWLINE: false" in user_prompt_a

    _, user_prompt_b, _ = prompt.build_completion_prompts(
        prefix="Hello\n",
        suffix="\nWorld",
    )
    assert "CURSOR_FENCE_LANGUAGE: none" in user_prompt_b
    assert "PREFIX_ENDS_WITH_NEWLINE: true" in user_prompt_b
    assert "SUFFIX_STARTS_WITH_NEWLINE: true" in user_prompt_b


def test_mermaid_context_flags():
    _, prompt_in_mermaid, _ = prompt.build_completion_prompts(
        prefix="```mermaid\nflowchart TD\nA --> ",
        suffix="\n```",
    )
    assert "CURSOR_IN_FENCED_CODE_BLOCK: true" in prompt_in_mermaid
    assert "CURSOR_FENCE_LANGUAGE: mermaid" in prompt_in_mermaid
    assert "MERMAID_CONTEXT: true" in prompt_in_mermaid

    _, prompt_mermaid_keyword, _ = prompt.build_completion_prompts(
        prefix="Please draw a mermaid flowchart for deploy pipeline.",
        suffix="",
    )
    assert "CURSOR_IN_FENCED_CODE_BLOCK: false" in prompt_mermaid_keyword
    assert "CURSOR_FENCE_LANGUAGE: none" in prompt_mermaid_keyword
    assert "MERMAID_CONTEXT: true" in prompt_mermaid_keyword


def test_examples_coverage():
    _, user_prompt, _ = prompt.build_completion_prompts(prefix="", suffix="")
    for ex in range(1, 15):
        assert f"[EX{ex:02d}]" in user_prompt