tests/test_parser_llama_cpp.py

import json
import unittest
from unittest.mock import patch

from training_coach.parser_llama_cpp import (
    DEFAULT_LLAMA_CPP_MAX_TOKENS,
    DEFAULT_LLAMA_CPP_MODEL_FILE,
    DEFAULT_LLAMA_CPP_MODEL_REPO,
    DEFAULT_LLAMA_CPP_N_CTX,
    MINIFIED_JSON_GBNF,
    LlamaCppRuntimeUnavailableError,
    _load_llama_cpp,
    build_completion_prompt,
    generate_parser_response_llama_cpp,
    llama_cpp_json_grammar,
    load_llama_cpp_model,
    parse_check_in_with_llama_cpp,
    warm_up_llama_cpp_parser,
)


class FakeLlama:
    from_pretrained_calls = []
    completion_calls = []

    @classmethod
    def from_pretrained(cls, **kwargs):
        cls.from_pretrained_calls.append(kwargs)
        return cls()

    def create_completion(self, **kwargs):
        self.completion_calls.append(kwargs)
        return {
            "choices": [
                {
                    "text": json.dumps(
                        {
                            "check_in": {
                                "raw_text": "60 min",
                                "time_available_minutes": 60,
                            }
                        }
                    ),
                    "finish_reason": "stop",
                }
            ],
            "usage": {"prompt_tokens": 100, "completion_tokens": 20},
        }


class FakeGrammar:
    grammars = []

    @classmethod
    def from_string(cls, grammar, verbose=True):
        cls.grammars.append((grammar, verbose))
        return {"grammar": grammar, "verbose": verbose}




class LlamaCppParserTest(unittest.TestCase):
    def setUp(self):
        load_llama_cpp_model.cache_clear()
        llama_cpp_json_grammar.cache_clear()
        FakeLlama.from_pretrained_calls = []
        FakeLlama.completion_calls = []
        FakeGrammar.grammars = []

    def tearDown(self):
        load_llama_cpp_model.cache_clear()
        llama_cpp_json_grammar.cache_clear()

    def test_llama_cpp_defaults_are_locked(self):
        self.assertEqual(DEFAULT_LLAMA_CPP_MODEL_REPO, "unsloth/Qwen3-1.7B-GGUF")
        self.assertEqual(DEFAULT_LLAMA_CPP_MODEL_FILE, "Qwen3-1.7B-Q4_K_M.gguf")
        self.assertEqual(DEFAULT_LLAMA_CPP_MAX_TOKENS, 512)
        self.assertEqual(DEFAULT_LLAMA_CPP_N_CTX, 2048)

    def test_runtime_can_import_llama_cpp_or_reports_clear_error(self):
        try:
            loaded = _load_llama_cpp()
        except LlamaCppRuntimeUnavailableError as error:
            self.assertIn("Install llama-cpp-python", str(error))
        else:
            self.assertEqual(len(loaded), 2)

    def test_load_model_uses_repo_file_and_runtime_settings(self):
        with patch(
            "training_coach.parser_llama_cpp._load_llama_cpp",
            return_value=(FakeLlama, FakeGrammar),
        ):
            load_llama_cpp_model(
                repo_id="repo/model",
                filename="model.gguf",
                n_ctx=4096,
                n_threads=2,
                n_threads_batch=4,
            )

        self.assertEqual(
            FakeLlama.from_pretrained_calls,
            [
                {
                    "repo_id": "repo/model",
                    "filename": "model.gguf",
                    "n_ctx": 4096,
                    "verbose": False,
                    "n_threads": 2,
                    "n_threads_batch": 4,
                }
            ],
        )

    def test_build_completion_prompt_prefills_empty_think_block(self):
        prompt = build_completion_prompt(
            [
                {"role": "system", "content": "system text"},
                {"role": "user", "content": "user text"},
            ]
        )

        self.assertIn("<|im_start|>system\nsystem text<|im_end|>\n", prompt)
        self.assertIn("<|im_start|>user\nuser text<|im_end|>\n", prompt)
        self.assertTrue(
            prompt.endswith("<|im_start|>assistant\n<think>\n\n</think>\n\n")
        )

    def test_generate_parser_response_uses_generic_json_grammar(self):
        with patch(
            "training_coach.parser_llama_cpp._load_llama_cpp",
            return_value=(FakeLlama, FakeGrammar),
        ):
            response = generate_parser_response_llama_cpp(
                "60 min",
                repo_id="repo/model",
                filename="model.gguf",
                max_tokens=128,
                n_ctx=2048,
                n_threads=None,
            )

        completion_call = FakeLlama.completion_calls[0]
        self.assertEqual(completion_call["max_tokens"], 128)
        self.assertEqual(completion_call["temperature"], 0)
        self.assertEqual(completion_call["stop"], ["<|im_end|>"])
        self.assertEqual(completion_call["grammar"]["grammar"], MINIFIED_JSON_GBNF)
        self.assertEqual(FakeGrammar.grammars, [(MINIFIED_JSON_GBNF, False)])
        self.assertIn("60 min", completion_call["prompt"])
        self.assertIn("</think>", completion_call["prompt"])
        self.assertIn("check_in", response)

    def test_parse_check_in_with_llama_cpp_reads_env_and_validates_response(self):
        with patch.dict(
            "os.environ",
            {
                "LLAMA_CPP_MODEL_REPO": "repo/model",
                "LLAMA_CPP_MODEL_FILE": "model.gguf",
                "LLAMA_CPP_MAX_TOKENS": "128",
                "LLAMA_CPP_N_CTX": "4096",
                "LLAMA_CPP_N_THREADS": "2",
            },
            clear=True,
        ), patch(
            "training_coach.parser_llama_cpp._load_llama_cpp",
            return_value=(FakeLlama, FakeGrammar),
        ):
            parsed = parse_check_in_with_llama_cpp("60 min")

        self.assertEqual(parsed.check_in.time_available_minutes, 60)
        self.assertEqual(FakeLlama.from_pretrained_calls[0]["repo_id"], "repo/model")
        self.assertEqual(FakeLlama.from_pretrained_calls[0]["filename"], "model.gguf")
        self.assertEqual(FakeLlama.from_pretrained_calls[0]["n_ctx"], 4096)
        self.assertEqual(FakeLlama.from_pretrained_calls[0]["n_threads"], 2)
        self.assertEqual(FakeLlama.from_pretrained_calls[0]["n_threads_batch"], 2)

    def test_threads_batch_env_overrides_decode_thread_default(self):
        with patch.dict(
            "os.environ",
            {
                "LLAMA_CPP_N_THREADS": "2",
                "LLAMA_CPP_N_THREADS_BATCH": "6",
            },
            clear=True,
        ), patch(
            "training_coach.parser_llama_cpp._load_llama_cpp",
            return_value=(FakeLlama, FakeGrammar),
        ):
            parse_check_in_with_llama_cpp("60 min")

        self.assertEqual(FakeLlama.from_pretrained_calls[0]["n_threads"], 2)
        self.assertEqual(FakeLlama.from_pretrained_calls[0]["n_threads_batch"], 6)

    def test_warm_up_runs_single_token_generation(self):
        with patch.dict("os.environ", {}, clear=True), patch(
            "training_coach.parser_llama_cpp._load_llama_cpp",
            return_value=(FakeLlama, FakeGrammar),
        ):
            warm_up_llama_cpp_parser()

        self.assertEqual(len(FakeLlama.completion_calls), 1)
        completion_call = FakeLlama.completion_calls[0]
        self.assertEqual(completion_call["max_tokens"], 1)
        self.assertIn("warmup", completion_call["prompt"])


if __name__ == "__main__":
    unittest.main()