| 1 | """Prompt Stacker — monitor and record prompt composition for each model call.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import json |
| 6 | import threading |
| 7 | from dataclasses import dataclass, field |
| 8 | from datetime import datetime, timezone |
| 9 | from pathlib import Path |
| 10 | from typing import Any |
| 11 | |
| 12 | |
| 13 | _STACKER_DIR = ".prompt_stacker" |
| 14 | |
| 15 | |
| 16 | @dataclass |
| 17 | class _SessionState: |
| 18 | """Per-session recording state, isolated from other concurrent sessions.""" |
| 19 | |
| 20 | base_parts: list[dict[str, Any]] = field(default_factory=list) |
| 21 | iter_parts: list[dict[str, Any]] = field(default_factory=list) |
| 22 | model: str | None = None |
| 23 | phase: str = "context" |
| 24 | |
| 25 | |
| 26 | class PromptStacker: |
| 27 | _instance: PromptStacker | None = None |
| 28 | _lock = threading.Lock() |
| 29 | |
| 30 | def __init__(self, workspace: Path, *, max_traces: int = 100) -> None: |
| 31 | self._workspace = workspace |
| 32 | self._max_traces = max_traces |
| 33 | self._dir = workspace / _STACKER_DIR |
| 34 | self._dir.mkdir(parents=True, exist_ok=True) |
| 35 | self._sessions: dict[str, _SessionState] = {} |
| 36 | self._turn_counter = self._recover_max_turn_id() |
| 37 | self._tls = threading.local() |
| 38 | |
| 39 | @classmethod |
| 40 | def init(cls, workspace: Path, *, enabled: bool = False, max_traces: int = 100) -> None: |
| 41 | with cls._lock: |
| 42 | if enabled: |
| 43 | cls._instance = cls(workspace, max_traces=max_traces) |
| 44 | else: |
| 45 | cls._instance = None |
| 46 | |
| 47 | @classmethod |
| 48 | def begin_turn(cls, session_key: str | None, model: str | None) -> None: |
| 49 | """Called in loop.py BEFORE build_messages(). Resets state for this session's new turn.""" |
| 50 | inst = cls._instance |
| 51 | if inst is None: |
| 52 | return |
| 53 | key = session_key or "unknown" |
| 54 | inst._tls.session_key = key |
| 55 | inst._sessions[key] = _SessionState(model=model, phase="context") |
| 56 | |
| 57 | @classmethod |
| 58 | def log(cls, label: str, content: Any) -> None: |
| 59 | """Log a prompt part for the current session (identified via thread-local).""" |
| 60 | inst = cls._instance |
| 61 | if inst is None: |
| 62 | return |
| 63 | key = getattr(inst._tls, "session_key", None) |
| 64 | if key is None: |
| 65 | return |
| 66 | state = inst._sessions.get(key) |
| 67 | if state is None: |
| 68 | return |
| 69 | text = content if isinstance(content, str) else json.dumps(content, ensure_ascii=False, default=str) |
| 70 | entry = { |
| 71 | "label": label, |
| 72 | "content": text, |
| 73 | "char_count": len(text), |
| 74 | } |
| 75 | if state.phase == "context": |
| 76 | state.base_parts.append(entry) |
| 77 | else: |
| 78 | state.iter_parts.append(entry) |
| 79 | |
| 80 | @classmethod |
| 81 | def begin_iteration(cls, iteration: int, model: str | None = None) -> None: |
| 82 | """Called in runner.py at the start of each model call. Clears iter_parts only.""" |
| 83 | inst = cls._instance |
| 84 | if inst is None: |
| 85 | return |
| 86 | key = getattr(inst._tls, "session_key", None) |
| 87 | if key is None: |
| 88 | return |
| 89 | state = inst._sessions.get(key) |
| 90 | if state is None: |
| 91 | return |
| 92 | state.iter_parts = [] |
| 93 | state.phase = "iteration" |
| 94 | if model: |
| 95 | state.model = model |
| 96 | |
| 97 | @classmethod |
| 98 | def commit( |
| 99 | cls, |
| 100 | *, |
| 101 | messages: list[dict[str, Any]] | None = None, |
| 102 | response: Any = None, |
| 103 | usage: dict[str, Any] | None = None, |
| 104 | iteration: int = 0, |
| 105 | ) -> None: |
| 106 | """Write a complete record for one model call. Preserves base_parts for next iteration.""" |
| 107 | inst = cls._instance |
| 108 | if inst is None: |
| 109 | return |
| 110 | |
| 111 | key = getattr(inst._tls, "session_key", None) |
| 112 | if key is None: |
| 113 | return |
| 114 | state = inst._sessions.get(key) |
| 115 | if state is None: |
| 116 | return |
| 117 | |
| 118 | with inst._lock: |
| 119 | inst._turn_counter += 1 |
| 120 | turn_id = inst._turn_counter |
| 121 | |
| 122 | now = datetime.now(timezone.utc) |
| 123 | safe_key = key.replace(":", "_").replace("/", "_") |
| 124 | |
| 125 | resp_data: dict[str, Any] = {} |
| 126 | if response is not None: |
| 127 | resp_data["content"] = getattr(response, "content", None) or "" |
| 128 | tool_calls = getattr(response, "tool_calls", None) |
| 129 | if tool_calls: |
| 130 | resp_data["tool_calls"] = [ |
| 131 | {"name": getattr(tc, "name", ""), "arguments": getattr(tc, "arguments", "")} |
| 132 | for tc in tool_calls |
| 133 | ] |
| 134 | reasoning = getattr(response, "reasoning_content", None) |
| 135 | if reasoning: |
| 136 | resp_data["reasoning_content"] = reasoning |
| 137 | if usage: |
| 138 | resp_data["usage"] = usage |
| 139 | |
| 140 | all_parts = list(state.base_parts) + list(state.iter_parts) |
| 141 | |
| 142 | record = { |
| 143 | "id": f"turn_{turn_id:04d}", |
| 144 | "timestamp": now.isoformat(), |
| 145 | "session_key": key, |
| 146 | "iteration": iteration, |
| 147 | "model": state.model or "", |
| 148 | "parts": all_parts, |
| 149 | "messages_count": len(messages) if messages else 0, |
| 150 | "messages": _full_messages(messages) if messages else [], |
| 151 | "response": resp_data, |
| 152 | } |
| 153 | |
| 154 | trace_file = inst._dir / f"{safe_key}.jsonl" |
| 155 | with open(trace_file, "a", encoding="utf-8") as f: |
| 156 | f.write(json.dumps(record, ensure_ascii=False, default=str) + "\n") |
| 157 | |
| 158 | state.iter_parts = [] |
| 159 | inst._gc_traces() |
| 160 | |
| 161 | @classmethod |
| 162 | def end_turn(cls, session_key: str | None = None) -> None: |
| 163 | """Clean up session state after a turn completes. Called optionally.""" |
| 164 | inst = cls._instance |
| 165 | if inst is None: |
| 166 | return |
| 167 | key = session_key or getattr(inst._tls, "session_key", None) |
| 168 | if key: |
| 169 | inst._sessions.pop(key, None) |
| 170 | |
| 171 | def _recover_max_turn_id(self) -> int: |
| 172 | """Scan existing JSONL files to recover the highest turn counter.""" |
| 173 | max_id = 0 |
| 174 | try: |
| 175 | for f in self._dir.glob("*.jsonl"): |
| 176 | with open(f, "rb") as fh: |
| 177 | fh.seek(max(0, f.stat().st_size - 4096)) |
| 178 | tail = fh.read().decode("utf-8", errors="replace") |
| 179 | for line in reversed(tail.strip().split("\n")): |
| 180 | line = line.strip() |
| 181 | if not line: |
| 182 | continue |
| 183 | try: |
| 184 | rec = json.loads(line) |
| 185 | tid = rec.get("id", "") |
| 186 | if tid.startswith("turn_"): |
| 187 | num = int(tid.split("_", 1)[1]) |
| 188 | max_id = max(max_id, num) |
| 189 | break |
| 190 | except (json.JSONDecodeError, ValueError): |
| 191 | continue |
| 192 | except OSError: |
| 193 | pass |
| 194 | return max_id |
| 195 | |
| 196 | def _gc_traces(self) -> None: |
| 197 | try: |
| 198 | traces = sorted(self._dir.glob("*.jsonl"), key=lambda p: p.stat().st_mtime) |
| 199 | while len(traces) > self._max_traces: |
| 200 | traces.pop(0).unlink(missing_ok=True) |
| 201 | except OSError: |
| 202 | pass |
| 203 | |
| 204 | @classmethod |
| 205 | def get_sessions(cls) -> list[dict[str, Any]]: |
| 206 | inst = cls._instance |
| 207 | if inst is None: |
| 208 | return [] |
| 209 | results = [] |
| 210 | for f in sorted(inst._dir.glob("*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True): |
| 211 | stat = f.stat() |
| 212 | last_line = "" |
| 213 | try: |
| 214 | with open(f, "rb") as fh: |
| 215 | fh.seek(max(0, stat.st_size - 4096)) |
| 216 | last_line = fh.read().decode("utf-8", errors="replace").strip().rsplit("\n", 1)[-1] |
| 217 | except OSError: |
| 218 | pass |
| 219 | turn_count = 0 |
| 220 | try: |
| 221 | with open(f, "r", encoding="utf-8") as fh: |
| 222 | turn_count = sum(1 for _ in fh) |
| 223 | except OSError: |
| 224 | pass |
| 225 | last_record = {} |
| 226 | if last_line: |
| 227 | try: |
| 228 | last_record = json.loads(last_line) |
| 229 | except json.JSONDecodeError: |
| 230 | pass |
| 231 | results.append({ |
| 232 | "id": f.stem, |
| 233 | "filename": f.name, |
| 234 | "turn_count": turn_count, |
| 235 | "size_bytes": stat.st_size, |
| 236 | "modified": datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc).isoformat(), |
| 237 | "last_model": last_record.get("model", ""), |
| 238 | "last_session_key": last_record.get("session_key", ""), |
| 239 | }) |
| 240 | return results |
| 241 | |
| 242 | @classmethod |
| 243 | def get_trace(cls, session_id: str) -> list[dict[str, Any]]: |
| 244 | inst = cls._instance |
| 245 | if inst is None: |
| 246 | return [] |
| 247 | trace_file = inst._dir / f"{session_id}.jsonl" |
| 248 | if not trace_file.is_file(): |
| 249 | return [] |
| 250 | try: |
| 251 | trace_file.relative_to(inst._dir) |
| 252 | except ValueError: |
| 253 | return [] |
| 254 | records = [] |
| 255 | with open(trace_file, "r", encoding="utf-8") as f: |
| 256 | for line in f: |
| 257 | line = line.strip() |
| 258 | if line: |
| 259 | try: |
| 260 | records.append(json.loads(line)) |
| 261 | except json.JSONDecodeError: |
| 262 | continue |
| 263 | return records |
| 264 | |
| 265 | |
| 266 | def _full_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: |
| 267 | result = [] |
| 268 | for msg in messages: |
| 269 | role = msg.get("role", "unknown") |
| 270 | content = msg.get("content", "") |
| 271 | if isinstance(content, str): |
| 272 | full_text = content |
| 273 | char_count = len(content) |
| 274 | elif isinstance(content, list): |
| 275 | texts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"] |
| 276 | full_text = "\n".join(texts) |
| 277 | char_count = len(full_text) |
| 278 | else: |
| 279 | full_text = str(content) |
| 280 | char_count = len(full_text) |
| 281 | |
| 282 | preview = full_text[:200] + "..." if len(full_text) > 200 else full_text |
| 283 | |
| 284 | entry: dict[str, Any] = { |
| 285 | "role": role, |
| 286 | "preview": preview, |
| 287 | "content": full_text, |
| 288 | "char_count": char_count, |
| 289 | } |
| 290 | tool_calls = msg.get("tool_calls") |
| 291 | if tool_calls: |
| 292 | entry["tool_calls"] = [ |
| 293 | {"name": tc.get("function", {}).get("name", ""), "id": tc.get("id", ""), |
| 294 | "arguments": tc.get("function", {}).get("arguments", "")} |
| 295 | for tc in tool_calls |
| 296 | if isinstance(tc, dict) |
| 297 | ] |
| 298 | if role == "tool": |
| 299 | entry["tool_call_id"] = msg.get("tool_call_id", "") |
| 300 | entry["tool_name"] = msg.get("name", "") |
| 301 | result.append(entry) |
| 302 | return result |
| 303 |