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\n\n\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("", 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()