llama.py (20845B)
1 """Semantic decision probe using an existing llama-server; no local weights needed. 2 3 Uses SemIf's messages, but the GGUF's own template/tokenizer. This is an 4 independent serving backend, not a numerically equivalent torch replacement. 5 """ 6 from __future__ import annotations 7 8 import argparse 9 import json 10 import math 11 import os 12 import sys 13 import time 14 from urllib.error import HTTPError, URLError 15 from urllib.request import Request, urlopen 16 17 from semif_phase1.core import LETTERS, digest, direct_messages, softmax 18 19 # Planner framing for /plan: ONE universal system prompt for every simulation. 20 # The caller composes the user prompt from facts it owns — the simulation's 21 # one-line goal statement, a note on what triggered planning, and the rules 22 # previously in effect — so no game needs its own prompt. The output 23 # discipline matters because the consumer is a logit-probed decision pass: 24 # one forward pass, no deliberation, so the rule set must be facts and 25 # imperatives the model can weigh in a single read — never if-state-then- 26 # action branches it would have to execute. Length is bounded by spirit, not 27 # numbers: a handful of brief lines, no counting. Everything uncertain gets 28 # litigated inside the thinking; the visible reply is committal rules only. 29 # The style example is placeholder-only on purpose: weaker models copy 30 # concrete example vocabulary into their plans, so the example carries no 31 # content to copy — only shape. 32 PLAN_SYSTEM = ( 33 "You are writing the prompt for a decision-making model that runs a " 34 "single forward pass — it cannot deliberate, only decide. Your " 35 "entire response is that prompt: it is prepended to every game state " 36 "the model reads, and it is the only guidance the model ever gets. " 37 "Write it addressed to the model in its own voice: you are a " 38 "decision maker, this is what your input looks like, this is how " 39 "to respond.\n\n" 40 "Structure the prompt as:\n" 41 "1. What the input is — a plain-text state (what the actor currently " 42 "sees and knows) followed by a question and a fixed list of possible " 43 "actions.\n" 44 "2. How to read the state — what each standing part of it means for " 45 "the decision at hand.\n" 46 "3. How to decide — the objective, the mechanics that bear on it, " 47 "stated as facts, and the priorities for choosing among the " 48 "actions.\n" 49 "4. The limits that end an attempt, stated as plain facts.\n\n" 50 "Never write if-state-then-action branches; the model picks actions, " 51 "you decide what it should know. It is a capable decision maker " 52 "whose only limitation is that it cannot deliberate.\n\n" 53 "Keep the prompt short: a handful of brief lines, each a single " 54 "imperative or fact. Nothing else — no headings, no preamble, no " 55 "explanation.\n\n" 56 "If previous rules are attached, the actor failed while following " 57 "them: that failure indicts the rules — their facts, priorities, or " 58 "framing — not the actor's comprehension of them. Never resubmit a " 59 "reworded or lightly edited version; change the substance, or " 60 "discard the set and write a fresh one.\n\n" 61 "The environment resets to its initial state after planning: the " 62 "rules must hold from the very first state, not from the situation " 63 "that triggered the plan. Resolve uncertainty inside your thinking; " 64 "the rules themselves must be plain and committal.\n\n" 65 "Style example — the shape, not the words:\n" 66 "GOOD:\n" 67 "<you are a decision maker, and this is the input you receive>\n" 68 "<one part of the state and what it means for the decision>\n" 69 "<the objective and one priority, stated as plain facts>\n" 70 "<one limit that ends an attempt, stated as a plain fact>\n\n" 71 "BAD:\n" 72 "<if this state, then that action; if that state, then this action; " 73 "and so on — a branching chain of situations and prescriptions>" 74 ) 75 PLAN_ROLES = {"system", "user", "assistant"} 76 # Sampling pinned to the model card's thinking-mode recommendation: the 77 # planner must explore while it reasons, so no per-request overrides. All six 78 # are accepted by llama-server's OpenAI-compatible endpoint (top_k, min_p, 79 # repetition_penalty are llama.cpp extensions); min_p/presence_penalty/ 80 # repetition_penalty are sent explicitly to pin them against server defaults. 81 PLAN_SAMPLING = {"temperature": 1.0, "top_p": 0.95, "top_k": 20, 82 "min_p": 0.0, "presence_penalty": 0.0, "repetition_penalty": 1.0} 83 84 # Default llama-server model alias. Any served GGUF works; the torch backend 85 # defaults to the HF repo form (Qwen/Qwen3.5-4B) instead. 86 DEFAULT_MODEL = "Qwen3.5-4B" 87 88 89 class BackendError(RuntimeError): 90 """Backend transport or readout failure (never an invented score).""" 91 92 93 class MissingOptions(BackendError): 94 """Valid readout, but the candidate list omitted declared options.""" 95 96 97 def option_logprobs(response: dict, slots: list[int]) -> list[float]: 98 """Read pre-sampling logprobs from modern or legacy native responses.""" 99 entries = response.get("completion_probabilities", response.get("probs")) 100 if not isinstance(entries, list) or len(entries) != 1: 101 raise BackendError("Expected exactly one token's probability readout") 102 entry = entries[0] 103 if not isinstance(entry, dict): 104 raise BackendError("Malformed probability readout") 105 candidates = entry.get("top_logprobs", entry.get("probs", [])) 106 if not isinstance(candidates, list) or not all(isinstance(c, dict) for c in candidates): 107 raise BackendError("Malformed candidate list") 108 scores = {} 109 for candidate in candidates: 110 token = candidate.get("id") 111 if token not in slots: 112 continue 113 if token in scores: 114 raise BackendError(f"Duplicate candidate token ID: {token}") 115 if "logprob" in candidate: 116 value = candidate["logprob"] 117 else: 118 probability = candidate.get("prob") 119 if not isinstance(probability, (int, float)) or not 0 < probability <= 1: 120 raise BackendError(f"Invalid probability for token {token}") 121 value = math.log(probability) 122 if not isinstance(value, (int, float)) or not math.isfinite(value): 123 raise BackendError(f"Non-finite logprob for token {token}") 124 scores[token] = value 125 missing = [token for token in slots if token not in scores] 126 if missing: 127 raise MissingOptions( 128 f"Missing option token IDs {missing}; increase the candidate limit " 129 "(--max-n-probs / SEMIF_LLAMA_MAX_N_PROBS). No scores were fabricated." 130 ) 131 return [scores[token] for token in slots] 132 133 134 class LlamaBackend: 135 def __init__(self, url: str, model: str = DEFAULT_MODEL, *, timeout: float = 180, 136 n_probs: int = 1024, max_n_probs: int = 16384, max_tokens: int = 4096, cache_prompt: bool = True): 137 if max_n_probs < n_probs or n_probs < 16 or max_tokens < 1 or timeout <= 0: 138 raise ValueError("Require max_n_probs >= n_probs >= 16, max_tokens >= 1, timeout > 0") 139 self.url = url.rstrip("/") 140 self.model = model 141 self.timeout = timeout 142 self.n_probs = n_probs 143 self.max_n_probs = max_n_probs 144 self.max_tokens = max_tokens 145 self.cache_prompt = cache_prompt 146 self._slots: dict[str, int] = {} 147 self.metadata = {"source": self.model, "backend": "llama", "url": self.url} 148 149 def request(self, method: str, path: str, payload: dict | None = None) -> dict: 150 body = None 151 if payload is not None: 152 body = json.dumps({**payload, "model": self.model}, allow_nan=False).encode() 153 req = Request(self.url + path, data=body, headers={"Content-Type": "application/json"}, 154 method=method) 155 try: 156 with urlopen(req, timeout=self.timeout) as response: 157 raw = response.read() 158 except HTTPError as error: 159 detail = error.read(2048).decode(errors="replace") 160 raise BackendError(f"{method} {path}: HTTP {error.code}: {detail}") from error 161 except (URLError, TimeoutError, OSError, ValueError) as error: 162 raise BackendError(f"{method} {path}: {error}") from error 163 try: 164 result = json.loads(raw) if raw.strip() else {} 165 except ValueError as error: 166 raise BackendError(f"{method} {path}: response was not JSON: {str(error)}") from error 167 if not isinstance(result, dict) or "error" in result: 168 raise BackendError(f"{method} {path}: unexpected response: {str(result)[:500]}") 169 return result 170 171 def post(self, path: str, payload: dict) -> dict: 172 return self.request("POST", path, payload) 173 174 def get(self, path: str) -> dict: 175 return self.request("GET", path) 176 177 def list_models(self) -> list[str]: 178 """Model aliases known to the server (models dir + currently loaded).""" 179 data = self.get("/v1/models").get("data") 180 if not isinstance(data, list): 181 raise BackendError("/v1/models did not return a data list") 182 return sorted({entry["id"] for entry in data 183 if isinstance(entry, dict) and isinstance(entry.get("id"), str)}) 184 185 def unload(self) -> None: 186 """Unload the selected model via the server's model-management API. 187 188 POST /models/unload with the model name in the body (llama-server's 189 router-mode model-management API). Raises on any non-success response; whether 190 that blocks a switch is the caller's decision (see switch). 191 """ 192 result = self.post("/models/unload", {}) 193 if result.get("success") is not True: 194 raise BackendError(f"/models/unload: unexpected response: {str(result)[:500]}") 195 196 def switch(self, model: str) -> dict: 197 """Select a different served model: unload the current one, lazy-load on first request. 198 199 The new model is validated against the server's own list BEFORE the current 200 model is unloaded, so a typo cannot leave the server with nothing loaded. 201 The server loads the new weights lazily, on the next request that names it. 202 203 Unloading the previous model is best-effort: a failed unload must not 204 block the selection, or the API would be stuck expecting a model the 205 caller no longer wants (recoverable only by scoring it). The old model 206 may stay resident until the server evicts it; that is reported as a 207 warning, not a refusal. 208 """ 209 model = model.strip() 210 if not model: 211 raise ValueError("model must be a nonempty string") 212 if model == self.model: 213 return {"selected": model, "unchanged": True, "warning": None} 214 known = self.list_models() 215 if model not in known: 216 raise ValueError(f"model {model!r} is not served (known: {', '.join(known) or 'none'})") 217 warning = None 218 try: 219 self.unload() 220 except BackendError as error: 221 warning = f"previous model may still be loaded: {error}" 222 self.model = model 223 self.metadata = {"source": self.model, "backend": "llama", "url": self.url} 224 # Option-letter token IDs are model-specific; force re-probing on next score. 225 self._slots.clear() 226 return {"selected": model, "unchanged": False, "warning": warning} 227 228 def tokenize(self, text: str) -> list[int]: 229 ids = self.post("/tokenize", { 230 "content": text, "add_special": False, "parse_special": True, 231 }).get("tokens") 232 if not isinstance(ids, list) or not all(type(token) is int for token in ids): 233 raise BackendError("/tokenize did not return integer token IDs") 234 return ids 235 236 def score(self, row: dict) -> dict: 237 started = time.perf_counter() 238 messages = direct_messages(row) # upstream validation and decision format 239 prompt = self.post("/apply-template", { 240 "messages": messages, "add_generation_prompt": True, 241 "chat_template_kwargs": {"enable_thinking": False}, 242 }).get("prompt") 243 if not isinstance(prompt, str) or not prompt: 244 raise BackendError("/apply-template did not return a nonempty prompt") 245 ids = self.tokenize(prompt) 246 if not ids or len(ids) > self.max_tokens: 247 raise ValueError(f"Prompt has {len(ids)} tokens; limit is {self.max_tokens}") 248 slots = [] 249 for letter in LETTERS[:len(row["options"])]: 250 if letter not in self._slots: 251 encoded = self.tokenize(letter) 252 if len(encoded) != 1 or self.post("/detokenize", {"tokens": encoded}).get("content") != letter: 253 raise BackendError(f"Option letter {letter} is not one round-trip token") 254 self._slots[letter] = encoded[0] 255 token = self._slots[letter] 256 if self.tokenize(prompt + letter) != ids + [token]: 257 raise BackendError(f"Answer boundary changes tokenization for {letter}") 258 slots.append(token) 259 if len(set(slots)) != len(slots): 260 raise BackendError("Option token IDs collide") 261 forward_start = time.perf_counter() 262 n_probs = self.n_probs 263 attempts = 0 264 while True: 265 attempts += 1 266 response = self.post("/completion", { 267 "prompt": ids, "n_predict": 1, "n_probs": n_probs, 268 "post_sampling_probs": False, "temperature": 1.0, 269 "samplers": [], "seed": 0, "stream": False, 270 "cache_prompt": self.cache_prompt, "return_tokens": True, 271 }) 272 if response.get("truncated"): 273 raise BackendError("llama-server truncated the prompt") 274 try: 275 selected = option_logprobs(response, slots) 276 break 277 except MissingOptions: 278 if n_probs >= self.max_n_probs: 279 raise 280 n_probs = min(n_probs * 4, self.max_n_probs) 281 forward_seconds = time.perf_counter() - forward_start 282 probabilities = softmax(selected) 283 option_ids = [option["id"] for option in row["options"]] 284 return { 285 "id": row["id"], "choice": option_ids[max(range(len(slots)), key=probabilities.__getitem__)], 286 "option_ids": option_ids, "probabilities": probabilities, 287 "option_logits": selected, 288 "option_logits_kind": "full-vocabulary log probabilities; logits up to an additive constant", 289 "input_tokens": len(ids), "forward_seconds": forward_seconds, 290 "total_seconds": time.perf_counter() - started, 291 "prompt_sha256": digest(prompt), "prompt_version": "direct-options-v1-gguf-template", 292 "model": self.metadata, 293 "readout": "pre-sampling next-token scores restricted to declared answer slots", 294 "probability_status": "conditional option score; uncalibrated as decision confidence", 295 "llama": {"timings": response.get("timings"), "tokens_cached": response.get("tokens_cached"), 296 "slot_id": response.get("id_slot"), "n_probs": n_probs, "attempts": attempts, 297 "cache_n": (response.get("timings") or {}).get("cache_n"), 298 "cache_prompt": self.cache_prompt}, 299 } 300 301 def plan(self, plan_id: str, prompt: str, transcript: list[dict]) -> dict: 302 """Reasoning chat completion that derives environment rules from a transcript. 303 304 Unlike score(), this generates text: a regular /v1/chat/completions with 305 thinking enabled (enable_thinking template kwarg), so the model can reason 306 about how previous actions went wrong before committing to rules. The 307 caller (the /plan endpoint) supplies the game-specific instruction as the 308 user prompt; completed actions arrive as extra transcript turns. No 309 max_tokens is sent — it would cap the reasoning trace, not just the 310 reply, and the server applies its own generation limit. Sampling is 311 pinned to PLAN_SAMPLING (the model card's thinking-mode settings). 312 Returns the generated rules, the thinking trace when the server 313 surfaces one, and usage/timing. Raises ValueError on contract 314 violations before any HTTP. 315 """ 316 started = time.perf_counter() 317 if not isinstance(plan_id, str) or not plan_id: 318 raise ValueError("plan id must be a nonempty string") 319 if not isinstance(prompt, str) or not prompt.strip(): 320 raise ValueError("prompt must be a nonempty string") 321 messages = [{"role": "system", "content": PLAN_SYSTEM}, 322 {"role": "user", "content": prompt}] 323 for index, turn in enumerate(transcript): 324 if not isinstance(turn, dict) or turn.get("role") not in PLAN_ROLES: 325 raise ValueError(f"transcript[{index}] needs role system|user|assistant") 326 if not isinstance(turn.get("content"), str) or not turn["content"].strip(): 327 raise ValueError(f"transcript[{index}] needs nonempty content") 328 messages.append({"role": turn["role"], "content": turn["content"]}) 329 response = self.post("/v1/chat/completions", { 330 "messages": messages, 331 "stream": False, 332 "chat_template_kwargs": {"enable_thinking": True}, 333 **PLAN_SAMPLING, 334 }) 335 choices = response.get("choices") 336 if not isinstance(choices, list) or len(choices) != 1: 337 raise BackendError("/v1/chat/completions: expected exactly one choice") 338 message = choices[0].get("message") if isinstance(choices[0], dict) else None 339 if not isinstance(message, dict): 340 raise BackendError("/v1/chat/completions: malformed message") 341 content = message.get("content") 342 if isinstance(content, list): # content-parts form: keep the text pieces 343 content = "".join(part.get("text", "") for part in content if isinstance(part, dict)) 344 if not isinstance(content, str) or not content.strip(): 345 raise BackendError("Planner produced no rules content") 346 reasoning = message.get("reasoning_content") 347 if not isinstance(reasoning, str) or not reasoning.strip(): 348 reasoning = None # template/server may not split thinking out 349 usage = response.get("usage") 350 return { 351 "id": plan_id, 352 "rules": content, 353 "reasoning": reasoning, 354 "truncated": choices[0].get("finish_reason") == "length", 355 "model": self.metadata, 356 "usage": usage if isinstance(usage, dict) else {}, 357 "total_seconds": time.perf_counter() - started, 358 "prompt_version": "plan-transcript-v1-chat-thinking", 359 } 360 361 362 EXAMPLE = { 363 "id": "interrupt-1", 364 "state": "Alex is in a meeting. A production service is down and customers cannot sign in.", 365 "question": "Should this notification interrupt Alex now?", 366 "options": [ 367 {"id": "interrupt", "description": "Interrupt now: urgent action is needed."}, 368 {"id": "later", "description": "Queue for after the meeting."}, 369 {"id": "ignore", "description": "No notification is necessary."}, 370 ], 371 } 372 373 374 def main() -> None: 375 parser = argparse.ArgumentParser(description=__doc__) 376 parser.add_argument("--url", default=os.environ.get("SEMIF_LLAMA_URL", "http://127.0.0.1:8080")) 377 parser.add_argument("--model", default=os.environ.get("SEMIF_LLAMA_MODEL", DEFAULT_MODEL)) 378 parser.add_argument("--input", help="JSONL decisions, or - for stdin; omitted: built-in notification example") 379 parser.add_argument("--n-probs", type=int, default=1024) 380 parser.add_argument("--max-n-probs", type=int, default=16384) 381 parser.add_argument("--max-tokens", type=int, default=4096) 382 parser.add_argument("--timeout", type=float, default=180) 383 parser.add_argument("--no-cache", action="store_true") 384 args = parser.parse_args() 385 try: 386 backend = LlamaBackend(args.url, args.model, timeout=args.timeout, n_probs=args.n_probs, 387 max_n_probs=args.max_n_probs, max_tokens=args.max_tokens, cache_prompt=not args.no_cache) 388 if args.input: 389 if args.input == "-": 390 rows = [json.loads(line) for line in sys.stdin if line.strip()] 391 else: 392 with open(args.input) as source: 393 rows = [json.loads(line) for line in source if line.strip()] 394 else: 395 rows = [EXAMPLE] 396 for row in rows: 397 print(json.dumps(backend.score(row), allow_nan=False), flush=True) 398 except (BackendError, ValueError, OSError) as error: 399 parser.exit(1, f"semif llama probe: {error}\n") 400 401 402 if __name__ == "__main__": 403 main()