tests/test_prompting.py

from __future__ import annotations
from typing import Any
import pytest
from app.generation.prompting import render_prompt, SYSTEM_PROMPT

class TokenizerWithTemplate:
    def __init__(self, template: str | None = "some template"):
        self.chat_template = template
    
    def apply_chat_template(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
        if self.chat_template is None:
            raise ValueError("No template set")
        return f"TEMPLATED: {messages[-1]['content']}"

class TokenizerWithoutTemplate:
    pass

def test_render_prompt_with_valid_template():
    tokenizer = TokenizerWithTemplate()
    prompt = render_prompt(tokenizer, "Hello")
    assert prompt == "TEMPLATED: Hello"

def test_render_prompt_with_none_template():
    # This simulates the failure case reported by the user
    tokenizer = TokenizerWithTemplate(template=None)
    prompt = render_prompt(tokenizer, "Hello")
    assert "User: Hello" in prompt
    assert "Assistant:" in prompt
    assert SYSTEM_PROMPT in prompt

def test_render_prompt_with_exception_in_apply():
    tokenizer = TokenizerWithTemplate()
    # Mock apply_chat_template to raise exception
    def broken_apply(*args, **kwargs):
        raise RuntimeError("broken")
    tokenizer.apply_chat_template = broken_apply
    
    prompt = render_prompt(tokenizer, "Hello")
    assert "User: Hello" in prompt

def test_render_prompt_without_apply_method():
    tokenizer = TokenizerWithoutTemplate()
    prompt = render_prompt(tokenizer, "Hello")
    assert "User: Hello" in prompt

def test_render_prompt_hrm_text_fallback():
    tokenizer = TokenizerWithoutTemplate()
    setattr(tokenizer, "name_or_path", "sapientinc/HRM-Text-1B")
    prompt = render_prompt(tokenizer, "Hello")
    assert "<|im_start|><|quad_end|><|object_ref_end|><|im_end|>" in prompt
    assert "User: Hello" in prompt