1
ANTHROPIC_API_KEY=your-key-here

1
%%capture
2
%pip install anthropic python-dotenv matplotlib

1
import json
2
import os
3
import tempfile
4
from collections import namedtuple
5
from pathlib import Path
6

7
import anthropic
8
import matplotlib.pyplot as plt
9
from dotenv import load_dotenv
10

11
load_dotenv()
12

13
if not os.environ.get("ANTHROPIC_API_KEY"):
14
    raise ValueError("ANTHROPIC_API_KEY not set. Add it to a .env file or export it.")
15

16
CORPUS_PATH = Path("research_corpus.py")
17
assert CORPUS_PATH.exists(), (
18
    f"research_corpus.py not found in {Path.cwd()}. It should be alongside this notebook."
19
)
20

21
client = anthropic.Anthropic()
22
MODEL = "claude-sonnet-4-6"
23
print(f"anthropic SDK {anthropic.__version__}, model {MODEL}")

1
anthropic SDK 0.84.0, model claude-sonnet-4-6

1
# Force-reload the corpus module in case this kernel has a stale cached version
2
import importlib
3

4
import research_corpus
5

6
importlib.reload(research_corpus)
7
from research_corpus import COMPACTION_PROBES, CORPUS
8

9
# Count tokens via the API (cached in-memory so repeat kernel runs are fast).
10
_token_cache: dict[str, int] = {}
11

12

13
def count_tokens(text: str) -> int:
14
    if text not in _token_cache:
15
        _token_cache[text] = client.messages.count_tokens(
16
            model=MODEL, messages=[{"role": "user", "content": text}]
17
        ).input_tokens
18
    return _token_cache[text]
19

20

21
print(f"CORPUS is a dict of {len(CORPUS)} synthetic documents held in Python memory.")
22
print("When the agent calls read_file, the content is served from this dict and")
23
print("lands directly in the agent's context window — no disk I/O involved.
24
")
25
_total_tokens = 0
26
for path, content in CORPUS.items():
27
    n_tok = count_tokens(content)
28
    _total_tokens += n_tok
29
    # Strip the virtual directory prefix for display clarity
30
    display_name = path.removeprefix("/research/")
31
    print(f"  {display_name:<26} ~{n_tok:>6,} tokens")
32
print(f"
33
  Total corpus: ~{_total_tokens:,} tokens")
34

35
# Sanity check: the corpus should be large enough to push context well past
36
# 200K in a single read-pass. If this fires, research_corpus.py may be stale.
37
assert _total_tokens > 250_000, (
38
    f"Corpus is only {_total_tokens:,} tokens; expected >250K. "
39
    "Restart the kernel and re-run, or verify research_corpus.py is current."
40
)

1
CORPUS is a dict of 8 synthetic documents held in Python memory.
2
When the agent calls read_file, the content is served from this dict and
3
lands directly in the agent's context window — no disk I/O involved.
4

5
  celegans_review.md         ~41,260 tokens
6
  drosophila_review.md       ~41,200 tokens
7
  mouse_review.md            ~41,155 tokens
8
  zebrafish_review.md        ~41,095 tokens
9
  killifish_review.md        ~41,115 tokens
10
  yeast_review.md            ~41,101 tokens
11
  nmr_review.md              ~41,017 tokens
12
  rhesus_review.md           ~41,012 tokens
13

14
  Total corpus: ~328,955 tokens

1
# ── Tool schemas ─────────────────────────────────────────────────────────
2
RESEARCH_TASK = """Compare the model organisms in /research/ on three dimensions:
3
1. Lifespan and experimental throughput (how fast can you get a lifespan result?)
4
2. Genetic tractability (what tools exist for manipulation and screening?)
5
3. Translational relevance (how well do findings transfer to human aging?)
6

7
Work through the organisms in two batches, recording a note after each
8
batch before moving on:
9

10
BATCH 1 (high-throughput models): Read celegans_review.md,
11
drosophila_review.md, yeast_review.md, and killifish_review.md. Record a
12
single detailed note covering all four organisms on all three dimensions.
13
Do not begin Batch 2 until this note is recorded.
14

15
BATCH 2 (low-throughput models): Read mouse_review.md,
16
zebrafish_review.md, nmr_review.md, and rhesus_review.md. Record a
17
single detailed note covering all four organisms on all three
18
dimensions.
19

20
After both batches, write a comprehensive comparative synthesis
21
contrasting the high-throughput and low-throughput models."""
22

23
BASE_TOOLS: list[dict] = [
24
    {
25
        "name": "search_files",
26
        "description": "Search document filenames and contents for a keyword. Returns matching paths with a short excerpt.",
27
        "input_schema": {
28
            "type": "object",
29
            "properties": {"query": {"type": "string", "description": "Search term"}},
30
            "required": ["query"],
31
        },
32
    },
33
    {
34
        "name": "read_file",
35
        "description": "Read the full content of a document by path.",
36
        "input_schema": {
37
            "type": "object",
38
            "properties": {
39
                "path": {"type": "string", "description": "e.g. /research/celegans_review.md"}
40
            },
41
            "required": ["path"],
42
        },
43
    },
44
    {
45
        "name": "record_finding",
46
        "description": "Record a research finding (freeform text). Findings are held for the duration of this session only; they are not persisted across sessions.",
47
        "input_schema": {
48
            "type": "object",
49
            "properties": {"finding": {"type": "string"}},
50
            "required": ["finding"],
51
        },
52
    },
53
]
54

55
# The memory tool spec is defined here but only added to the agent's tool
56
# list when a memory_handler is passed to run_research_session. Baseline
57
# and clearing/compaction runs do NOT get the memory tool.
58
MEMORY_TOOL_SPEC = {"type": "memory_20250818", "name": "memory"}
59

60
SYSTEM_PROMPT = """You are a biology research analyst writing a comparative review of model organisms for aging research.
61

62
Work systematically: search and read review documents, take notes on key facts (lifespans, genetic tools, translational strengths and caveats), and build a comparative picture.
63

64
Be concise in your reasoning text; the goal is notes, not essays."""

1
# ── Tool execution ───────────────────────────────────────────────────────
2
def execute_research_tool(name: str, tool_input: dict, notes: list[str]) -> str:
3
    if name == "search_files":
4
        query = tool_input.get("query", "").lower()
5
        if not query:
6
            return "Error: query is required."
7
        hits: list[str] = []
8
        for path, content in CORPUS.items():
9
            if query in path.lower() or query in content.lower():
10
                for line in content.split("
11
"):
12
                    if query in line.lower():
13
                        hits.append(f"  {path}
14
    → {line.strip()[:120]}")
15
                        break
16
                else:
17
                    hits.append(f"  {path}")
18
        return (
19
            f"Found {len(hits)} match(es):
20
" + "
21
".join(hits)
22
            if hits
23
            else f"No matches for '{query}'."
24
        )
25

26
    if name == "read_file":
27
        path = tool_input.get("path", "")
28
        content = CORPUS.get(path)
29
        if content is None:
30
            return f"Error: '{path}' not found. Available: {', '.join(CORPUS.keys())}"
31
        return content
32

33
    if name == "record_finding":
34
        finding = tool_input.get("finding", "")
35
        notes.append(finding)
36
        return f"Finding #{len(notes)} recorded (session-local)."
37

38
    return f"Error: unknown tool '{name}'"

1
# ── Session result container ─────────────────────────────────────────────
2
SessionResult = namedtuple(
3
    "SessionResult",
4
    [
5
        "messages",  # final message list
6
        "notes",  # notes taken this session
7
        "token_trajectory",  # list of (turn, total_context_tokens)
8
        "events",  # list of dicts describing compaction / clearing events
9
        "tool_counts",  # dict of tool_name -> call count
10
        "file_reads",  # list of (turn, path) for each read_file call
11
        "hit_limit",  # True if the session stopped because it hit the context window
12
        "final_text",  # final assistant text (convenience)
13
    ],
14
)
15

16

17
def _format_tool_arg(name: str, tool_input: dict) -> str:
18
    """Human-readable one-line summary of a tool call's arguments for verbose output."""
19
    if name == "search_files":
20
        return repr(tool_input.get("query", ""))
21
    if name == "read_file":
22
        return tool_input.get("path", "?")
23
    if name == "record_finding":
24
        finding = tool_input.get("finding", "")
25
        preview = finding[:50].replace("
26
", " ")
27
        return f'"{preview}{"..." if len(finding) > 50 else ""}"'
28
    if name == "memory":
29
        cmd = tool_input.get("command", "?")
30
        path = tool_input.get("path", tool_input.get("old_path", ""))
31
        return f"{cmd} {path}"
32
    return str(tool_input)[:60]
33

34

35
# ── Agent loop ───────────────────────────────────────────────────────────
36
def run_research_session(
37
    initial_prompt: str,
38
    *,
39
    context_management: dict | None = None,
40
    betas: list[str] | None = None,
41
    memory_handler=None,
42
    max_turns: int = 12,
43
    label: str = "session",
44
    verbose: bool = True,
45
) -> SessionResult:
46
    """Run the research agent. Catches context-window overflow gracefully."""
47
    tools = list(BASE_TOOLS)
48
    if memory_handler is not None:
49
        tools.append(MEMORY_TOOL_SPEC)
50

51
    messages: list[dict] = [{"role": "user", "content": initial_prompt}]
52
    notes: list[str] = []
53
    token_trajectory: list[tuple[int, int]] = []
54
    events: list[dict] = []
55
    tool_counts: dict[str, int] = {}
56
    file_reads: list[tuple[int, str]] = []  # (turn, path) for each read_file
57
    hit_limit = False
58
    final_text = ""
59

60
    if verbose:
61
        print(f"┌─ [{label}]")
62

63
    for turn in range(1, max_turns + 1):
64
        kwargs: dict = dict(
65
            model=MODEL,
66
            max_tokens=4096,
67
            system=SYSTEM_PROMPT,
68
            tools=tools,
69
            messages=messages,
70
        )
71
        if context_management:
72
            kwargs["context_management"] = context_management
73
        if betas:
74
            kwargs["betas"] = betas
75

76
        # Call the API; catch context-window overflow
77
        try:
78
            if betas:
79
                response = client.beta.messages.create(**kwargs)
80
            else:
81
                response = client.messages.create(**kwargs)
82
        except anthropic.BadRequestError as e:
83
            # Context window exceeded (or similar input-too-large error)
84
            hit_limit = True
85
            if verbose:
86
                print(f"│  ⚠ CONTEXT WINDOW LIMIT REACHED at turn {turn} (API rejected)")
87
                print(f"│    {str(e)[:200]}")
88
            break
89

90
        # Track TOTAL context size: uncached + cache-read + cache-created.
91
        # usage.input_tokens alone excludes cached tokens, which makes the
92
        # plot show dips that are just cache hits, not context management.
93
        u = response.usage
94
        total_in = (
95
            u.input_tokens
96
            + (getattr(u, "cache_read_input_tokens", None) or 0)
97
            + (getattr(u, "cache_creation_input_tokens", None) or 0)
98
        )
99
        token_trajectory.append((turn, total_in))
100

101
        # Surface context-management events on their own prominent lines
102
        cm = getattr(response, "context_management", None)
103
        if cm is not None and getattr(cm, "applied_edits", None):
104
            for edit in cm.applied_edits:
105
                cleared_uses = getattr(edit, "cleared_tool_uses", None)
106
                cleared_toks = getattr(edit, "cleared_input_tokens", None)
107
                events.append(
108
                    {
109
                        "turn": turn,
110
                        "kind": "clearing",
111
                        "cleared_tool_uses": cleared_uses,
112
                        "cleared_input_tokens": cleared_toks,
113
                    }
114
                )
115
                if verbose:
116
                    print(
117
                        f"│  ✂ CLEARING (turn {turn}): {cleared_uses or '?'} tool results cleared, "
118
                        f"~{cleared_toks:,} tokens freed"
119
                        if cleared_toks
120
                        else f"│  ✂ CLEARING (turn {turn}): applied"
121
                    )
122

123
        # Serialize assistant content, track tool calls and compaction
124
        serialized: list[dict] = []
125
        tool_calls: list[dict] = []
126
        turn_tool_calls: list[tuple[str, dict]] = []  # (name, input) for verbose display
127
        compaction_this_turn = False
128
        for block in response.content:
129
            if block.type == "text":
130
                serialized.append({"type": "text", "text": block.text})
131
                if block.text.strip():
132
                    final_text = block.text
133
            elif block.type == "tool_use":
134
                serialized.append(
135
                    {
136
                        "type": "tool_use",
137
                        "id": block.id,
138
                        "name": block.name,
139
                        "input": block.input,
140
                    }
141
                )
142
                tool_calls.append({"id": block.id, "name": block.name, "input": block.input})
143
                turn_tool_calls.append((block.name, block.input))
144
                tool_counts[block.name] = tool_counts.get(block.name, 0) + 1
145
                # Track file reads specifically so we can show what clearing drops
146
                if block.name == "read_file":
147
                    file_reads.append((turn, block.input.get("path", "?")))
148
            elif block.type == "thinking":
149
                serialized.append(
150
                    {"type": "thinking", "thinking": block.thinking, "signature": block.signature}
151
                )
152
            elif block.type == "compaction":
153
                serialized.append({"type": "compaction", "content": block.content})
154
                events.append({"turn": turn, "kind": "compaction", "summary": block.content})
155
                compaction_this_turn = True
156
                if verbose:
157
                    print(
158
                        f"│  ⊟ COMPACTION (turn {turn}): "
159
                        f"~{count_tokens(block.content):,}-token summary replaces prior turns"
160
                    )
161
        messages.append({"role": "assistant", "content": serialized})
162

163
        if not tool_calls:
164
            if verbose and not compaction_this_turn:
165
                print(f"│  turn {turn:2d}  ctx={total_in:>7,}  (final answer)")
166
            break
167

168
        # Execute tools, collecting result sizes so the verbose print
169
        # can show how much each call added to context
170
        tool_results: list[dict] = []
171
        result_sizes: list[int] = []
172
        for call in tool_calls:
173
            if call["name"] == "memory" and memory_handler is not None:
174
                result = memory_handler.handle(call["input"])
175
            else:
176
                result = execute_research_tool(call["name"], call["input"], notes)
177
            tool_results.append(
178
                {
179
                    "type": "tool_result",
180
                    "tool_use_id": call["id"],
181
                    "content": result,
182
                }
183
            )
184
            result_sizes.append(len(result) if isinstance(result, str) else 0)
185
        messages.append({"role": "user", "content": tool_results})
186

187
        if verbose and not compaction_this_turn:
188
            # Header with context size, then one line per tool call with
189
            # its argument and approximate result size (so you can see
190
            # which calls are responsible for the next turn's ctx jump)
191
            print(f"│  turn {turn:2d}  ctx={total_in:>7,}")
192
            for (name, tinput), rsize in zip(turn_tool_calls, result_sizes, strict=False):
193
                size_note = f"  → ~{rsize // 4:,} tok" if rsize > 200 else ""
194
                print(f"│            {name:<16} {_format_tool_arg(name, tinput)}{size_note}")
195

196
    if verbose:
197
        if token_trajectory:
198
            peak = max(t for _, t in token_trajectory)
199
            status = "⚠ HIT CONTEXT LIMIT" if hit_limit else "completed"
200
            print(
201
                f"└─ {status}: {len(token_trajectory)} turns, peak ctx {peak:,}, "
202
                f"final ctx {token_trajectory[-1][1]:,}, {len(events)} context event(s)
203
"
204
            )
205
        else:
206
            print("└─ HIT CONTEXT LIMIT on first turn: 0 turns completed
207
")
208

209
    return SessionResult(
210
        messages, notes, token_trajectory, events, tool_counts, file_reads, hit_limit, final_text
211
    )
212

213

214
def show_cleared_reads(result: SessionResult, keep: int):
215
    """Show which file reads are no longer in context after clearing.
216

217
    Clearing replaces tool results older than the last `keep` tool uses
218
    with placeholders. We reconstruct the tool-use order and mark any
219
    read that falls outside the surviving `keep`-window as cleared.
220

221
    Note: if clearing fires multiple times, only the last event's
222
    boundary is considered and earlier cleared-then-re-read files may
223
    be misclassified.
224
    """
225
    if not result.file_reads:
226
        print("No file reads in this session.")
227
        return
228
    clearing_events = [e for e in result.events if e["kind"] == "clearing"]
229
    if not clearing_events:
230
        print("Clearing never fired; all file reads remain in context.")
231
        return
232

233
    # Walk messages in order to reconstruct the sequence of tool_use blocks
234
    # and which turn each one came from. The last `keep` of these survive
235
    # the most recent clearing; earlier ones are cleared.
236
    tool_use_seq: list[tuple[int, str, str]] = []  # (turn, name, path-if-read)
237
    turn = 0
238
    for msg in result.messages:
239
        if msg["role"] == "assistant" and isinstance(msg["content"], list):
240
            turn += 1
241
            for block in msg["content"]:
242
                if block.get("type") == "tool_use":
243
                    name = block.get("name", "?")
244
                    path = block.get("input", {}).get("path", "") if name == "read_file" else ""
245
                    tool_use_seq.append((turn, name, path))
246

247
    last_clear_turn = clearing_events[-1]["turn"]
248
    seq_before = [t for t in tool_use_seq if t[0] < last_clear_turn]
249
    cleared_reads = [(t, p) for (t, n, p) in seq_before[:-keep] if n == "read_file"]
250
    surviving_reads = [
251
        (t, p) for tu in seq_before[-keep:] for (t, n, p) in [tu] if n == "read_file"
252
    ]
253
    # Reads at or after the last clearing turn are untouched by it.
254
    surviving_reads += [
255
        (t, p) for (t, n, p) in tool_use_seq if t >= last_clear_turn and n == "read_file"
256
    ]
257

258
    total = len(result.file_reads)
259
    print(f"Total file reads across session: {total}")
260
    print(f"Last clearing event fired at turn {last_clear_turn} (keep={keep})")
261
    print(f"
262
Reads cleared from context: {len(cleared_reads)}")
263
    for turn, path in cleared_reads[:12]:
264
        print(f"  ✗ turn {turn:2d}: {path}")
265
    if len(cleared_reads) > 12:
266
        print(f"  ... and {len(cleared_reads) - 12} more")
267
    print(
268
        f"
269
Reads still in context (within the keep={keep} window or after "
270
        f"the last clearing): {len(surviving_reads)}"
271
    )
272
    for turn, path in surviving_reads[:6]:
273
        print(f"  ✓ turn {turn:2d}: {path}")
274
    if len(surviving_reads) > 6:
275
        print(f"  ... and {len(surviving_reads) - 6} more")
276

277

278
# ── Plot helpers ─────────────────────────────────────────────────────────
279

280

281
def plot_trajectories(
282
    results: dict[str, SessionResult],
283
    title: str = "Context size per turn",
284
    triggers: dict[str, int] | None = None,
285
    project_growth_for: str | None = None,
286
):
287
    """Line plot of context tokens per turn.
288

289
    Vertical dashed lines mark turns where clearing fired; dash-dot lines
290
    mark compaction. Horizontal dotted lines mark configured trigger
291
    thresholds (pass triggers={"clearing": 20000, "compaction": 50000} etc).
292

293
    project_growth_for: label of a run to extrapolate. Fits a line to the
294
    last 5 points and draws a dotted segment forward ~8 turns to show where
295
    unmanaged growth is heading.
296
    """
297
    fig, ax = plt.subplots(figsize=(10, 4.5))
298
    colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
299

300
    for i, (label, res) in enumerate(results.items()):
301
        turns = [t for t, _ in res.token_trajectory]
302
        tokens = [tok for _, tok in res.token_trajectory]
303
        color = colors[i % len(colors)]
304
        ax.plot(turns, tokens, marker="o", label=label, markersize=4, color=color)
305
        # Mark events on the x-axis
306
        for ev in res.events:
307
            style = "--" if ev["kind"] == "clearing" else "-."
308
            ax.axvline(x=ev["turn"], color=color, linestyle=style, alpha=0.25, linewidth=1)
309
        # Dotted growth projection for the named run, capped at 1M
310
        if label == project_growth_for and len(turns) >= 3:
311
            HARD_LIMIT = 1_000_000
312
            fit_n = min(5, len(turns))
313
            xs, ys = turns[-fit_n:], tokens[-fit_n:]
314
            n = len(xs)
315
            sx, sy = sum(xs), sum(ys)
316
            slope = (n * sum(x * y for x, y in zip(xs, ys, strict=False)) - sx * sy) / (
317
                n * sum(x * x for x in xs) - sx * sx
318
            )
319
            intercept = (sy - slope * sx) / n
320
            proj_x, proj_y = [], []
321
            for x in range(turns[-1], turns[-1] + 9):
322
                y = slope * x + intercept
323
                if y > HARD_LIMIT:
324
                    # Clip the last segment to the 1M ceiling and stop
325
                    if proj_y and slope > 0:
326
                        proj_x.append(proj_x[-1] + (HARD_LIMIT - proj_y[-1]) / slope)
327
                        proj_y.append(HARD_LIMIT)
328
                    break
329
                proj_x.append(x)
330
                proj_y.append(y)
331
            if proj_y:
332
                ax.plot(proj_x, proj_y, linestyle=":", color=color, alpha=0.6, linewidth=1.5)
333

334
    # Horizontal reference lines for trigger thresholds
335
    if triggers:
336
        for name, value in triggers.items():
337
            ax.axhline(y=value, color="gray", linestyle=":", alpha=0.6, linewidth=1)
338
            ax.annotate(
339
                f"{name} trigger: {value:,}",
340
                xy=(ax.get_xlim()[1], value),
341
                xytext=(-5, 3),
342
                textcoords="offset points",
343
                ha="right",
344
                va="bottom",
345
                fontsize=8,
346
                color="gray",
347
            )
348

349
    # 200K reference: earlier models cap here and would hard-stop
350
    if ax.get_ylim()[1] > 30_000:
351
        ax.axhline(y=200_000, color="gray", linestyle="--", alpha=0.5, linewidth=1)
352
        ax.annotate(
353
            "200K: earlier models stop here",
354
            xy=(ax.get_xlim()[0], 200_000),
355
            xytext=(5, -12),
356
            textcoords="offset points",
357
            ha="left",
358
            va="top",
359
            fontsize=8,
360
            color="gray",
361
            alpha=0.8,
362
        )
363

364
    ax.set_xlabel("Turn")
365
    ax.set_ylabel("Context tokens (incl. cached)")
366
    ax.set_title(title)
367
    ax.legend(loc="best")
368
    ax.grid(alpha=0.3)
369
    plt.tight_layout()
370
    plt.show()
371

372

373
def plot_summary_bars(results: dict[str, SessionResult], title: str = "Session outcomes"):
374
    """Side-by-side bars: final context, file reads."""
375
    labels = list(results.keys())
376
    final_ctx = [r.token_trajectory[-1][1] for r in results.values()]
377
    reads = [r.tool_counts.get("read_file", 0) for r in results.values()]
378

379
    fig, axes = plt.subplots(1, 2, figsize=(9, 3.5))
380
    fig.suptitle(title, y=1.02)
381

382
    for ax, values, ylabel in zip(
383
        axes,
384
        [final_ctx, reads],
385
        ["Final context (tokens)", "File reads"],
386
        strict=False,
387
    ):
388
        bars = ax.bar(range(len(labels)), values, color=plt.cm.Set2(range(len(labels))))
389
        ax.set_ylabel(ylabel)
390
        ax.set_xticks(range(len(labels)))
391
        ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=8)
392
        ax.grid(axis="y", alpha=0.3)
393
        for bar, val in zip(bars, values, strict=False):
394
            ax.text(
395
                bar.get_x() + bar.get_width() / 2,
396
                bar.get_height(),
397
                f"{val:,}" if val > 1000 else str(val),
398
                ha="center",
399
                va="bottom",
400
                fontsize=7,
401
            )
402

403
    plt.tight_layout()
404
    plt.show()