test_llama.py (11130B)
1 """Offline tests: python -m unittest discover -s . -p test_llama.py""" 2 import copy 3 import math 4 import re 5 import unittest 6 from unittest.mock import patch 7 from urllib.error import URLError 8 9 from semif_api.llama import BackendError, EXAMPLE, LlamaBackend, PLAN_SAMPLING, PLAN_SYSTEM, option_logprobs 10 11 12 def response(values): 13 return {"completion_probabilities": [{"top_logprobs": [ 14 {"id": token, "logprob": value, "token": "deliberately ignored"} 15 for token, value in values 16 ]}]} 17 18 19 class FakeBackend(LlamaBackend): 20 def __init__(self, **kwargs): 21 super().__init__("http://example.invalid", **kwargs) 22 self.calls = [] 23 self.result = response([(67, -4.0), (65, -1.0), (66, -2.0)]) 24 self.chat_result = { 25 "choices": [{"message": {"role": "assistant", "content": "1. Run moves one space.", 26 "reasoning_content": "the transcript shows…"}, 27 "finish_reason": "stop"}], 28 "usage": {"prompt_tokens": 12, "completion_tokens": 7}, 29 } 30 31 def post(self, path, payload): 32 self.calls.append((path, payload)) 33 if path == "/apply-template": 34 return {"prompt": "test prompt\n"} 35 if path == "/tokenize": 36 return {"tokens": [ord(char) for char in payload["content"]]} 37 if path == "/detokenize": 38 return {"content": "".join(chr(token) for token in payload["tokens"])} 39 if path == "/completion": 40 return self.result 41 if path == "/v1/chat/completions": 42 return self.chat_result 43 raise AssertionError(path) 44 45 46 class ReadoutTests(unittest.TestCase): 47 def test_token_ids_not_text_or_candidate_order(self): 48 self.assertEqual(option_logprobs(response([(66, -2), (65, -1)]), [65, 66]), [-1, -2]) 49 50 def test_legacy_probabilities(self): 51 data = {"completion_probabilities": [{"probs": [ 52 {"id": 65, "prob": .25}, {"id": 66, "prob": .5}]}]} 53 self.assertEqual(option_logprobs(data, [65, 66]), [math.log(.25), math.log(.5)]) 54 55 def test_missing_option_fails(self): 56 with self.assertRaisesRegex(BackendError, "Missing option"): 57 option_logprobs(response([(65, -1)]), [65, 66]) 58 59 def test_nonfinite_and_duplicate_fail(self): 60 for values in ([(65, float("nan")), (66, -2)], [(65, -1), (65, -2), (66, -3)]): 61 with self.subTest(values=values), self.assertRaises(BackendError): 62 option_logprobs(response(values), [65, 66]) 63 64 def test_no_readout_fails(self): 65 with self.assertRaises(BackendError): 66 option_logprobs({"content": "A"}, [65, 66]) 67 68 def test_score(self): 69 backend = FakeBackend(cache_prompt=False) 70 result = backend.score(EXAMPLE) 71 self.assertEqual(result["choice"], "interrupt") 72 self.assertEqual(result["option_ids"], ["interrupt", "later", "ignore"]) 73 self.assertAlmostEqual(sum(result["probabilities"]), 1) 74 completion = next(payload for path, payload in backend.calls if path == "/completion") 75 self.assertTrue(all(type(token) is int for token in completion["prompt"])) 76 self.assertFalse(completion["post_sampling_probs"]) 77 self.assertFalse(completion["cache_prompt"]) 78 self.assertEqual(completion["n_predict"], 1) 79 template = backend.calls[0][1] 80 self.assertFalse(template["chat_template_kwargs"]["enable_thinking"]) 81 82 def test_truncation_fails(self): 83 backend = FakeBackend() 84 backend.result["truncated"] = True 85 with self.assertRaisesRegex(BackendError, "truncated"): 86 backend.score(EXAMPLE) 87 88 def test_prompt_limit_before_inference(self): 89 backend = FakeBackend(max_tokens=2) 90 with self.assertRaises(ValueError): 91 backend.score(EXAMPLE) 92 self.assertNotIn("/completion", [path for path, _ in backend.calls]) 93 94 def test_bad_input_before_http(self): 95 backend = FakeBackend() 96 row = copy.deepcopy(EXAMPLE) 97 row["options"] = [] 98 with self.assertRaises(ValueError): 99 backend.score(row) 100 self.assertFalse(backend.calls) 101 102 def test_missing_scores_retry_with_larger_list(self): 103 backend = FakeBackend() 104 original = backend.post 105 106 def post(path, payload): 107 result = original(path, payload) 108 if path == "/completion" and payload["n_probs"] == 1024: 109 return response([(65, -1)]) 110 return result 111 112 with patch.object(backend, "post", side_effect=post): 113 result = backend.score(EXAMPLE) 114 self.assertEqual(result["llama"]["attempts"], 2) 115 self.assertEqual(result["llama"]["n_probs"], 4096) 116 117 def test_retry_is_bounded(self): 118 backend = FakeBackend() 119 backend.result = response([(65, -1)]) 120 with self.assertRaisesRegex(BackendError, "Missing option"): 121 backend.score(EXAMPLE) 122 self.assertEqual([p["n_probs"] for path, p in backend.calls if path == "/completion"], 123 [1024, 4096, 16384]) 124 125 def test_malformed_readout_does_not_retry(self): 126 backend = FakeBackend() 127 backend.result = {"completion_probabilities": [None]} 128 with self.assertRaises(BackendError): 129 backend.score(EXAMPLE) 130 self.assertEqual(sum(path == "/completion" for path, _ in backend.calls), 1) 131 132 def test_transport_failure(self): 133 with patch("semif_api.llama.urlopen", side_effect=URLError("offline")): 134 with self.assertRaisesRegex(BackendError, "offline"): 135 LlamaBackend("http://example.invalid").score(EXAMPLE) 136 137 138 class PlanTests(unittest.TestCase): 139 def test_plan_system_prompt_is_universal_and_committal(self): 140 # One prompt for every simulation: it frames the single-forward-pass 141 # consumer, dictates the prompt-shaped structure the planner writes 142 # (what the input is, how to read it, how to decide, the limits), 143 # bans if-state-then-action branches, bounds length by spirit rather 144 # than numbers the thinker would count, and shows a GOOD/BAD style 145 # pair from a fictional environment (no demo terms). 146 self.assertIn("single forward pass", PLAN_SYSTEM) 147 self.assertIn("You are writing the prompt", PLAN_SYSTEM) 148 self.assertIn("What the input is", PLAN_SYSTEM) 149 self.assertIn("How to read the state", PLAN_SYSTEM) 150 self.assertIn("How to decide", PLAN_SYSTEM) 151 self.assertIn("if-state-then-action", PLAN_SYSTEM) 152 self.assertIn("handful of brief lines", PLAN_SYSTEM) 153 self.assertIn("resets to its initial state", PLAN_SYSTEM) 154 self.assertIn("committal", PLAN_SYSTEM) 155 self.assertIn("indicts the rules", PLAN_SYSTEM) 156 self.assertIn("Never resubmit a reworded", PLAN_SYSTEM) 157 self.assertIn("GOOD:", PLAN_SYSTEM) 158 self.assertIn("BAD:", PLAN_SYSTEM) 159 self.assertNotIn("3–6", PLAN_SYSTEM) # no exact counts to satisfy 160 self.assertNotIn("under 20 words", PLAN_SYSTEM) 161 for game_term in ("platformer", "flag", "jump", "hole", "tile", 162 "key", "door", "exit", "room", "wall"): 163 self.assertNotIn(game_term, PLAN_SYSTEM) 164 # The style example must be placeholder-only: weaker models copy 165 # concrete example vocabulary (gates, lasers, conveyors) into plans. 166 for example_noun in ("loading dock", "battery", "conveyor", "gate"): 167 self.assertNotIn(example_noun, PLAN_SYSTEM) 168 self.assertIsNone(re.search(r"\brun\b", PLAN_SYSTEM)) # never the game action 169 170 def test_plan_posts_thinking_chat_completion(self): 171 backend = FakeBackend() 172 transcript = [{"role": "user", "content": "Observation: …\nChosen action: run\nOutcome: moved 1 space."}] 173 result = backend.plan("plan-1", "Write the rules.", transcript) 174 self.assertEqual(result["id"], "plan-1") 175 self.assertEqual(result["rules"], "1. Run moves one space.") 176 self.assertEqual(result["reasoning"], "the transcript shows…") 177 self.assertFalse(result["truncated"]) 178 self.assertEqual(result["usage"]["completion_tokens"], 7) 179 path, payload = backend.calls[-1] 180 self.assertEqual(path, "/v1/chat/completions") 181 self.assertTrue(payload["chat_template_kwargs"]["enable_thinking"]) 182 self.assertNotIn("max_tokens", payload) # the server caps generation, not us 183 # Sampling is pinned to the model card's thinking-mode settings. 184 for key, value in PLAN_SAMPLING.items(): 185 self.assertEqual(payload[key], value) 186 self.assertEqual(payload["temperature"], 1.0) 187 self.assertEqual(payload["top_p"], 0.95) 188 self.assertEqual(payload["top_k"], 20) 189 roles = [message["role"] for message in payload["messages"]] 190 self.assertEqual(roles, ["system", "user", "user"]) 191 self.assertEqual(payload["messages"][1]["content"], "Write the rules.") 192 self.assertEqual(payload["messages"][2]["content"], transcript[0]["content"]) 193 194 def test_plan_defaults(self): 195 backend = FakeBackend() 196 result = backend.plan("plan-2", "Write the rules.", []) 197 _, payload = backend.calls[-1] 198 self.assertNotIn("max_tokens", payload) # never sent; it would cap thinking 199 self.assertEqual(payload["temperature"], 1.0) 200 self.assertEqual([m["role"] for m in payload["messages"]], ["system", "user"]) 201 202 def test_plan_validates_before_http(self): 203 backend = FakeBackend() 204 for bad_args in ( 205 ("", "prompt", []), # empty id 206 ("id", "", []), # empty prompt 207 ("id", "prompt", [{"role": "tool", "content": "x"}]), # bad role 208 ("id", "prompt", [{"role": "user", "content": " "}]), # empty content 209 ("id", "prompt", [{"role": "user"}]), # missing content 210 ): 211 with self.subTest(bad_args=bad_args), self.assertRaises(ValueError): 212 backend.plan(*bad_args) 213 self.assertFalse(backend.calls) 214 215 def test_plan_rejects_empty_or_malformed_content(self): 216 backend = FakeBackend() 217 for chat_result in ( 218 {"choices": [{"message": {"content": " "}, "finish_reason": "stop"}]}, 219 {"choices": [{"message": {"content": ["1. Run."]}}]}, # parts without dicts 220 {"choices": []}, 221 {"choices": [{"message": None}]}, 222 ): 223 backend.chat_result = chat_result 224 with self.subTest(chat_result=chat_result), self.assertRaises(BackendError): 225 backend.plan("id", "prompt", []) 226 227 def test_plan_truncation_flagged_not_fatal(self): 228 backend = FakeBackend() 229 backend.chat_result["choices"][0]["finish_reason"] = "length" 230 backend.chat_result["choices"][0]["message"]["reasoning_content"] = "" 231 result = backend.plan("id", "prompt", []) 232 self.assertTrue(result["truncated"]) 233 self.assertIsNone(result["reasoning"]) 234 235 236 if __name__ == "__main__": 237 unittest.main()