Source files / scripts · raw
longctx-test.py
#!/usr/bin/env python3 """Long-context agent session: recall and speed past 131k tokens, on any OpenAI-compatible server. longctx-test.py [--port 8080] [--label NAME] [--first 120000] [--total 200000] Builds an agent-style history (tool outputs = source files) with three planted facts (codenames at ~10 %, ~50 % and ~90 % of the final length), then: a) cold prompt of --first tokens, asks for fact 1 (prefill from zero) b) same history + more tool output up to --total, fact 3 (prefill of the new part, deep in) c) one more question appended, fact 2 (cache reuse, recall from the middle) d) a 512-token answer at full depth (decode speed at --total) Thinking off; client-side timing (streamed). Appends JSON lines to --out. """ import argparse, glob, json, os, random, time, urllib.request ap = argparse.ArgumentParser() ap.add_argument("--port", type=int, default=8080) ap.add_argument("--label", default="") ap.add_argument("--first", type=int, default=120000) ap.add_argument("--total", type=int, default=200000) ap.add_argument("--src", default="/opt/tersler/src/llama.cpp") ap.add_argument("--out", default="/opt/tersler/bench/longctx.jsonl") a = ap.parse_args() URL = f"http://127.0.0.1:{a.port}/v1/chat/completions" def ntok(text): """Token count from the server (Anthropic-style count_tokens, which Strata offers); None if unsupported.""" body = {"messages": [{"role": "user", "content": text}]} req = urllib.request.Request(f"http://127.0.0.1:{a.port}/v1/messages/count_tokens", json.dumps(body).encode(), {"Content-Type": "application/json"}) try: with urllib.request.urlopen(req, timeout=60) as r: return json.load(r)["input_tokens"] except Exception: return None files = sorted(f for f in glob.glob(os.path.join(a.src, "**/*.c*"), recursive=True) + glob.glob(os.path.join(a.src, "**/*.h"), recursive=True) if os.path.getsize(f) > 4000) random.seed(7) random.shuffle(files) words = random.sample(["AMBER FALCON", "COPPER LANTERN", "SILENT ORCHARD", "VELVET GLACIER", "IRON MEADOW", "PAPER COMET"], 3) facts = [f"NOTE FROM THE TEAM: the release codename for milestone {i + 1} is {w}." for i, w in enumerate(words)] def turns(tok_target, start_idx, plant): """Tool-output turns totalling ~tok_target tokens; plant = {fraction_of_this_block: fact}. Token-dense files (tables, embedded data: < 2.5 characters per token) are skipped: an agent rarely reads those.""" out, n, i = [], 0, start_idx pending = sorted(plant.items()) while n < tok_target and i < len(files): body = open(files[i], errors="replace").read()[:24000] i += 1 t = ntok(body) or len(body) // 3 if len(body) / t < 2.5: continue while pending and n >= pending[0][0] * tok_target: body = pending.pop(0)[1] + "\n\n" + body rel = os.path.relpath(files[i - 1], a.src) out += [{"role": "assistant", "content": f"Reading {rel}."}, {"role": "user", "content": f"Tool output (read_file {rel}):\n{body}"}] n += t + 30 return out, i def ask(messages, max_tokens=64): body = {"messages": messages, "max_tokens": max_tokens, "stream": True, "temperature": 0.6, "top_p": 0.95, "stream_options": {"include_usage": True}, "chat_template_kwargs": {"enable_thinking": False}} req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"}) t0 = time.time(); t1 = None; text = []; usage = {} with urllib.request.urlopen(req, timeout=7200) as r: for line in r: line = line.strip() if not line.startswith(b"data:") or line == b"data: [DONE]": continue ev = json.loads(line[5:]) usage = ev.get("usage") or usage for ch in ev.get("choices", []): c = ch.get("delta", {}).get("content") if c: t1 = t1 or time.time() text.append(c) t2 = time.time() pt = usage.get("prompt_tokens", 0) cached = (usage.get("prompt_tokens_details") or {}).get("cached_tokens") or 0 ct = usage.get("completion_tokens", 0) return {"prompt_tokens": pt, "cached": cached, "new": pt - cached, "ttft_s": round(t1 - t0, 1), "pp_new_tok_s": round((pt - cached) / (t1 - t0), 1) if t1 > t0 else None, "completion_tokens": ct, "tg_tok_s": round((ct - 1) / (t2 - t1), 1) if ct > 1 and t2 > t1 else None, "answer": "".join(text).strip()[:200]} def log(step, rec, expect=None): if expect: rec["recall_ok"] = expect.lower() in rec["answer"].lower() rec.update(step=step, label=a.label, date=time.strftime("%Y-%m-%dT%H:%M:%S")) print(json.dumps(rec), flush=True) with open(a.out, "a") as f: f.write(json.dumps(rec) + "\n") system = {"role": "system", "content": "You are a coding agent working in the llama.cpp repository. Tool outputs follow."} start = {"role": "user", "content": "Read through the code base; I will ask questions afterwards."} f1 = a.first / a.total part1, idx = turns(a.first, 0, {0.10 / f1: facts[0], 0.50 / f1: facts[1]}) part2, idx = turns(a.total - a.first, idx, {(0.90 - f1) / (1 - f1): facts[2]}) Q = "Answer with the codename only: what is the release codename for milestone %d?" h = [system, start] + part1 log("a_cold_first", ask(h + [{"role": "user", "content": Q % 1}]), words[0]) h = h + [{"role": "user", "content": Q % 1}, {"role": "assistant", "content": words[0]}] + part2 log("b_extend_to_total", ask(h + [{"role": "user", "content": Q % 3}]), words[2]) h = h + [{"role": "user", "content": Q % 3}, {"role": "assistant", "content": words[2]}] log("c_cached_recall_middle", ask(h + [{"role": "user", "content": Q % 2}]), words[1]) h = h + [{"role": "user", "content": Q % 2}, {"role": "assistant", "content": words[1]}] log("d_decode_at_depth", ask(h + [{"role": "user", "content": "Summarise in about 400 words what the files you read do."}], 512))