LLM Tutorial¶

TokSearch's toksearch.llm library lets you ask for fusion data in plain English and have an LLM write the pipeline code for you. This notebook walks through the full surface end-to-end:

1. Setup — install + credentials
2. Quickstart for new users — first toksearch chat session
3. Python API — programmatic use from notebooks and scripts
4. The persistent namespace — why multi-turn analysis works
5. A realistic DIII-D workflow — fetch, plot, iterate
6. Swapping backends — Anthropic / OpenAI / Claude Max / AmSC
7. Slash commands — /help, /reset, /quit
8. Where to next

Read-only: all toksearch chat / toksearch query outputs in this notebook are illustrative. The notebook is not auto-executed; copy any block and run it in your environment to reproduce.

Reference: the LLM reference page covers the full API surface, configuration, and the entry-point contributor mechanism.

1. Setup¶

Install¶

# Conda (recommended)
conda install -c ga-fdp toksearch          # backends bundled by default

# Pip
pip install toksearch[llm]

The conda recipe lists the four backend SDKs (anthropic, openai, claude-agent-sdk, mcp) and matplotlib as hard run-deps; pip users get the same surface via the [llm] extra.

Credentials¶

Pick whichever backend you have a key for and set the corresponding env var:

# Anthropic
export ANTHROPIC_API_KEY=sk-ant-...

# OpenAI
export OPENAI_API_KEY=sk-...

# Claude Max plan (uses your existing Claude Code login)
claude login                                # one-time; or set CLAUDE_CODE_OAUTH_TOKEN

# AmSC (available to American Science Cloud users; via toksearch_d3d's `amsc` preset)
echo 'sk-...' > ~/amsc_api_key && chmod 600 ~/amsc_api_key

Verify¶

toksearch query --backend anthropic "Use run_python to compute 2 + 2."

Expected (abbreviated):

[run_python] Compute 2 + 2 and print the result.
  print(2 + 2)
[output] 4

The result is 4.

If you see error: Could not find the 'claude' CLI or similar auth errors, see Reference > Configuration for the resolution order.

2. Quickstart for new users¶

If you've never used TokSearch before, this is the lowest-friction way in. The CLI does everything; you don't need to write any Python.

Interactive REPL¶

toksearch chat --backend anthropic

toksearch chat (backend: anthropic, model: claude-sonnet-4-6)
Type /help for commands. Ctrl-D to exit.

you> What does TokSearch do?

TokSearch is a Python library for fetching fusion-experiment diagnostic
signals (MDSplus, PTData, IMAS, Zarr) across many shots in parallel and
composing analysis pipelines over them. The core abstraction is `Pipeline`,
which lets you chain fetch / map / filter operations and run them serially,
on multiprocessing pools, or distributed via Ray or Spark.

you> Show me a one-shot example using MdsSignal.

[run_python] Demonstrate fetching ipmhd from efit01 for one shot.
  pipeline = toksearch.Pipeline([165920])
  pipeline.fetch('ip', toksearch.MdsSignal(r'\ipmhd', 'efit01'))
  records = pipeline.compute_serial()
  rec = list(records)[0]
  print('shot:', rec.shot)
  print('peak |Ip|:', np.nanmax(np.abs(rec['ip']['data'])) / 1e6, 'MA')
[output] shot: 165920
         peak |Ip|: 1.1325 MA

you> /quit

The agent writes the code, executes it, and shows you both the code and its output. Variables (pipeline, records, rec) stay in scope for follow-up turns — see §4.

One-shot queries from a shell¶

For scripting or quick lookups, toksearch query runs a single turn and prints the result:

toksearch query --backend anthropic \
    "Use run_python to fetch ipmhd from efit01 for shot 165920 and report peak in MA."

[run_python] Fetch ipmhd via MdsSignal and compute peak.
  pipeline = toksearch.Pipeline([165920])
  pipeline.fetch('ip', toksearch.MdsSignal(r'\ipmhd', 'efit01'))
  rec = list(pipeline.compute_serial())[0]
  print(np.nanmax(np.abs(rec['ip']['data'])) / 1e6, 'MA')
[output] 1.1325 MA

The peak |Ip| for shot 165920 is 1.1325 MA.

3. Python API¶

For notebooks and scripts, drive Session directly. The CLI is just a wrapper.

Session owns:

• The backend (which provider you're talking to)
• The conversation history (provider-neutral; can be serialized for future /save support)
• The run_python namespace (a plain Python dict that persists across turns)
• The tool registry (run_python, lookup_docs)
• The skill registry (auto-discovered from installed contributor packages)

Session.send(prompt, ...) is synchronous. It returns when the assistant finishes its turn. Callbacks fire as events happen along the way.

In [ ]:

import os
from toksearch.llm import Session
from toksearch.llm.backends.anthropic import AnthropicBackend

sess = Session(
    backend=AnthropicBackend(api_key=os.environ["ANTHROPIC_API_KEY"]),
    max_iterations=10,
)

result = sess.send(
    "Use run_python to compute 2 + 2.",
    on_tool_call=lambda c: print(f"[{c.name}] {c.thought}"),
    on_tool_result=lambda r: print(f"  -> {r.output.strip()}"),
)

print()
print("final_text:", result.final_text)
print("stop_reason:", result.stop_reason)

import os from toksearch.llm import Session from toksearch.llm.backends.anthropic import AnthropicBackend sess = Session( backend=AnthropicBackend(api_key=os.environ["ANTHROPIC_API_KEY"]), max_iterations=10, ) result = sess.send( "Use run_python to compute 2 + 2.", on_tool_call=lambda c: print(f"[{c.name}] {c.thought}"), on_tool_result=lambda r: print(f" -> {r.output.strip()}"), ) print() print("final_text:", result.final_text) print("stop_reason:", result.stop_reason)

Sample output:

[run_python] Compute 2 + 2 and print the result.
  -> 4

final_text: The result is 4.
stop_reason: end_turn

Callback options¶

Session.send takes any subset of these keyword-only callbacks:

Callback	Fires when
on_text(text: str)	The assistant emits a chunk of natural-language text. In PR 1 there's one call per turn (no streaming yet).
on_tool_call(call: ToolCall)	The assistant decides to call a tool — before the tool runs.
on_tool_result(result: ToolResult)	A tool finishes — after execution, before the result is sent back to the model.
on_event(event)	Catch-all that fires for every event above (use one switch statement instead of four kwargs).
confirm(call: ToolCall) -> bool	Gate every tool call. Returning False aborts the turn with stop_reason="interrupted".

4. The persistent namespace¶

This is the load-bearing feature that makes multi-turn analysis viable.

run_python executes its code against Session.namespace — a plain Python dict that lives for the Session's lifetime. Variables you (or the agent) define in one turn are available in all subsequent turns.

Without this, every follow-up question would have to re-fetch — and TokSearch fetches are slow (an MdsSignal tree open + decode of a few hundred shots can take a minute). With it, the agent fetches once and iterates on the analysis in memory.

In [ ]:

# Turn 1: fetch (slow)
sess.send("Use run_python to fetch ipmhd from efit01 for shots 200000-200005. Store in `records`.")

# Turn 2: analysis on the already-fetched data (fast — no re-fetch)
sess.send("What are the peak |Ip| values in MA for each shot in `records`?")

# Turn 3: plot the result (still no re-fetch)
sess.send("Make a bar chart of those peak values and save to peaks.png.")

# Turn 1: fetch (slow) sess.send("Use run_python to fetch ipmhd from efit01 for shots 200000-200005. Store in `records`.") # Turn 2: analysis on the already-fetched data (fast — no re-fetch) sess.send("What are the peak |Ip| values in MA for each shot in `records`?") # Turn 3: plot the result (still no re-fetch) sess.send("Make a bar chart of those peak values and save to peaks.png.")

You can also inspect or modify the namespace from Python at any point between sends:

>>> sess.namespace.keys()
dict_keys(['toksearch', 'np', 'pd', 'plt', 'toksearch_d3d', 'records'])

>>> sess.namespace["records"]
[Record(shot=200000, ...), Record(shot=200001, ...), ...]

>>> sess.namespace["my_threshold"] = 1.2e6      # inject something for the agent to use
>>> sess.send("Filter `records` to shots whose peak |Ip| exceeds `my_threshold`.")

Session.reset() clears both the history and the namespace (back to the pre-populated defaults: toksearch, np, pd, plt, plus contributed-package modules).

5. A realistic DIII-D workflow¶

Assumes you have toksearch_d3d installed (registers PtDataSignal, ImasSignal, FDP/Pelican access, and the amsc backend preset).

From the shell, fdp wraps toksearch chat with the FDP environment pre-configured (XRootD plugin, MDSplus tree paths, BEARER_TOKEN):

fdp chat                   # defaults to --backend amsc

toksearch chat (backend: anthropic, model: claude-sonnet-4-6)
Type /help for commands. Ctrl-D to exit.

you> Plot Ip for shot 200000.

[run_python] Fetch Ip via PtDataSignal and plot the time series.
  ip = PtDataSignal('ip').fetch(200000)
  plt.figure(figsize=(8, 4))
  plt.plot(ip['times'] / 1000, ip['data'] / 1e6)
  plt.xlabel('time (s)')
  plt.ylabel('Ip (MA)')
  plt.title('Shot 200000')
  plt.savefig('ip_200000.png', dpi=120)
  plt.close()
[output] (no output)

Plotted to `ip_200000.png`. Peak |Ip| is 1.42 MA.

you> Now overlay shot 200001 on the same axes.

[run_python] Fetch the second shot and re-plot with both traces.
  ip2 = PtDataSignal('ip').fetch(200001)
  plt.figure(figsize=(8, 4))
  plt.plot(ip['times'] / 1000, ip['data'] / 1e6, label='200000')
  plt.plot(ip2['times'] / 1000, ip2['data'] / 1e6, label='200001')
  plt.xlabel('time (s)'); plt.ylabel('Ip (MA)'); plt.legend()
  plt.savefig('ip_overlay.png', dpi=120)
  plt.close()
[output] (no output)

Saved overlay plot to `ip_overlay.png`.

you> What's the peak |Ip| ratio between the two shots?

[run_python] Compute peak |Ip| for each and the ratio.
  p1 = float(np.nanmax(np.abs(ip['data'])))
  p2 = float(np.nanmax(np.abs(ip2['data'])))
  print(f'shot 200000: {p1/1e6:.3f} MA')
  print(f'shot 200001: {p2/1e6:.3f} MA')
  print(f'ratio 200001/200000: {p2/p1:.3f}')
[output] shot 200000: 1.419 MA
         shot 200001: 1.387 MA
         ratio 200001/200000: 0.978

Shot 200001 had ~2.2% lower peak |Ip| than shot 200000.

you> /quit

Notice the agent never re-fetched ip or ip2 after the first time it pulled each shot — that's the persistent namespace earning its keep.

Same workflow from Python¶

Equivalent in a notebook (e.g. for use in a Jupyter analysis):

In [ ]:

from toksearch.llm import Session
from toksearch.llm.presets import resolve_preset
from toksearch.llm.config import load_config
from toksearch.llm.backends import get_backend_class
from toksearch.llm.cli import _resolve_api_key   # public-ish helper for now

cfg = load_config()
preset = resolve_preset("amsc", cfg)
backend = get_backend_class(preset.backend)(
    api_key=_resolve_api_key(preset, cfg),
    base_url=preset.base_url,
)
sess = Session(backend=backend, model=preset.model)

sess.send("Plot Ip for shot 200000. Save to ip_200000.png.")
sess.send("Now overlay shot 200001 on the same axes. Save to ip_overlay.png.")
result = sess.send("What's the peak |Ip| ratio between the two shots?")

print(result.final_text)

from toksearch.llm import Session from toksearch.llm.presets import resolve_preset from toksearch.llm.config import load_config from toksearch.llm.backends import get_backend_class from toksearch.llm.cli import _resolve_api_key # public-ish helper for now cfg = load_config() preset = resolve_preset("amsc", cfg) backend = get_backend_class(preset.backend)( api_key=_resolve_api_key(preset, cfg), base_url=preset.base_url, ) sess = Session(backend=backend, model=preset.model) sess.send("Plot Ip for shot 200000. Save to ip_200000.png.") sess.send("Now overlay shot 200001 on the same axes. Save to ip_overlay.png.") result = sess.send("What's the peak |Ip| ratio between the two shots?") print(result.final_text)

6. Swapping backends¶

The same prompt can drive any backend. The Session interface is uniform; only the backend instance changes.

In [ ]:

from toksearch.llm import Session
from toksearch.llm.backends.anthropic import AnthropicBackend
from toksearch.llm.backends.openai import OpenAIBackend
from toksearch.llm.backends.claude_sdk import ClaudeSDKBackend

import os

PROMPT = "Use run_python to fetch ipmhd from efit01 for shot 165920 and report peak in MA."

# Anthropic Messages API
Session(backend=AnthropicBackend(api_key=os.environ['ANTHROPIC_API_KEY'])).send(PROMPT)

# OpenAI Chat Completions
Session(backend=OpenAIBackend(api_key=os.environ['OPENAI_API_KEY'])).send(PROMPT)

# Claude Max plan (uses your `claude` CLI login -- no API key)
Session(backend=ClaudeSDKBackend()).send(PROMPT)

from toksearch.llm import Session from toksearch.llm.backends.anthropic import AnthropicBackend from toksearch.llm.backends.openai import OpenAIBackend from toksearch.llm.backends.claude_sdk import ClaudeSDKBackend import os PROMPT = "Use run_python to fetch ipmhd from efit01 for shot 165920 and report peak in MA." # Anthropic Messages API Session(backend=AnthropicBackend(api_key=os.environ['ANTHROPIC_API_KEY'])).send(PROMPT) # OpenAI Chat Completions Session(backend=OpenAIBackend(api_key=os.environ['OPENAI_API_KEY'])).send(PROMPT) # Claude Max plan (uses your `claude` CLI login -- no API key) Session(backend=ClaudeSDKBackend()).send(PROMPT)

From the CLI, same idea:

toksearch query --backend anthropic  "..."
toksearch query --backend openai     "..."
toksearch query --backend claude-max "..."
toksearch query --backend amsc       "..."     # requires toksearch_d3d installed

When to use which:

Backend	Why
anthropic	Default for users with a direct Anthropic API key.
openai	Use OpenAI's models (e.g. gpt-4o); easy A/B comparison.
claude-max	Bill against your Claude Max subscription instead of API tokens. Best for high-volume interactive use.
amsc	AmSC users with a populated ~/amsc_api_key.

Site-specific backends can be registered as user presets in ~/.fdp/config.toml — see Reference > Configuration.

7. Slash commands in the REPL¶

toksearch chat and fdp chat accept three commands:

Command	Effect
/help	Print the command list.
/reset	Clear the conversation history and reset the namespace to defaults (toksearch, np, pd, plt, plus contributed packages).
/quit	Exit the REPL (ctrl-D also works).

A typical session:

you> /reset
(session cleared)

you> Fetch ip for shot 200000.
[run_python] ...

you> /help
Slash commands:
  /help   show this help
  /reset  clear history and namespace
  /quit   exit the chat (ctrl-D also works)

you> ^D

8. Where to next¶

• LLM reference — full API surface, configuration precedence, contributor entry points, exception hierarchy.
• Working with Signals — the underlying MdsSignal / ZarrSignal machinery the agent uses for run_python calls.
• Aggregating Data — patterns for collapsing many-shot pipelines into pandas / xarray; useful prompts to give the agent.
• fdp CLI source: GA-FDP/toksearch_d3d — DIII-D-specific wrapping (fdp chat, fdp query, fdp run, the amsc preset).
• Claude Agent SDK: anthropics/claude-agent-sdk-python — underlies the claude-max backend.

Found a bug?¶

File at github.com/GA-FDP/toksearch/issues. If the agent does something surprising or the wire-format with a provider drifts, the integration tests (tests/test_llm_integration.py, opt-in via TOKSEARCH_INTEGRATION=yes) are the first place to add a reproducer.