LLM Interface

TokSearch ships with toksearch.llm — a conversational interface that lets you ask for fusion data in plain English and have an LLM write the pipeline code for you. The agent uses a persistent Python namespace so successive turns iterate on cached results instead of re-fetching, which is what makes multi-turn analysis viable when each fetch takes seconds to minutes.

you> Fetch ip for shot 200000.
agent> [run_python]
         ip = PtDataSignal('ip').fetch(200000)
       [output] (no output)
       Done. The result is in `ip`.

you> What's the peak value in MA?
agent> [run_python]
         print(np.nanmax(np.abs(ip['data'])) / 1e6)
       [output] 1.4193  ← did NOT re-fetch
       Peak |Ip| is 1.42 MA.

Backends

Four backend names ship in core TokSearch:

Name	Provider	Credentials
`anthropic`	Anthropic Messages API	`ANTHROPIC_API_KEY`
`openai`	OpenAI Chat Completions	`OPENAI_API_KEY`
`claude-max`	Claude Max plan via the Claude Agent SDK	the `claude` CLI (run `claude login`)
`amsc` (registered by `toksearch_d3d`)	American Science Cloud (AmSC) Anthropic-compatible endpoint at `api.i2-core.american-science-cloud.org`	`~/amsc_api_key`

Additional backends can be registered by any installed package — see Contributors.

Installation

# Conda (recommended; matches FDP-on-prem usage)
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-dependencies; pip users get the same surface via the [llm] optional-dependency extra.

Quickstart

CLI

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

# Interactive REPL
toksearch chat --backend anthropic

The REPL accepts /help, /reset, and /quit; ctrl-D also exits. From a DIII-D environment, the fdp script wraps the same commands with the FDP environment pre-configured (XRootD plugin, MDSplus tree paths, etc.):

fdp query "Fetch ip for shot 200000 and report peak in MA."  # default: --backend amsc
fdp chat                                                       # interactive

Python

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

sess = Session(backend=AnthropicBackend(api_key="sk-..."))
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(r.output),
)
print(result.final_text)

For end-to-end examples including the persistent-namespace pattern and a real DIII-D workflow, see the LLM Tutorial.

Configuration

Resolution precedence (highest first):

CLI flags: --backend, --model, -n / --max-iterations, --package
Environment variables: FDP_LLM_BACKEND, ANTHROPIC_API_KEY, OPENAI_API_KEY

~/.fdp/config.toml:

[llm]
backend = "anthropic"          # or openai, claude-max, amsc, or a user preset name
model = "claude-sonnet-4-6"    # overrides the backend's default
max_iterations = 20
anthropic_api_key = "sk-..."   # only if not using env var

[llm.presets.mysite]
backend = "anthropic"          # the underlying backend class
base_url = "https://llm.mysite.gov"
api_key_env = "MYSITE_KEY"
model = "claude-sonnet-4-6"

Built-in defaults

Tools

The agent has exactly two tools, registered with every Session:

`run_python`

Executes a Python code string in the Session's persistent namespace. The namespace lives for the Session's lifetime — variables defined in one turn are available in all subsequent turns. Pre-populated with:

toksearch (and toksearch_d3d if installed)
np (numpy)
pd (pandas, if available)
plt (matplotlib.pyplot, if available)

The agent must populate a thought field with a one-sentence description of what each code block does and why. This is what the REPL prints before execution — it's the load-bearing transparency mechanism (see Show-then-run).

`lookup_docs`

Returns the body of a registered skill (a SKILL.md file). The Session's system prompt lists the available skill names with one-line descriptions; the agent calls lookup_docs(skill_name=...) when it needs the details.

Core TokSearch ships with skills covering Pipeline basics, MdsSignal, the backends, datasets, and API exploration. Device packages add their own: toksearch_d3d contributes skills for PtDataSignal, ImasSignal, FDP CLI, and a DIII-D quickstart.

Skills are served via a standalone MCP server launched as a subprocess when Session is constructed: python -m toksearch.llm.mcp. Each SKILL.md is exposed as a skill://<name> MCP resource (and via a read_skill tool). The toksearch.llm.skills entry-point group remains the discovery source; extra directories can be added via the TOKSEARCH_SKILL_DIRS env var (os.pathsep-delimited). An external MCP client can also connect to the server directly to browse or read skills independently of Session.

Show-then-run

The REPL prints each run_python block's thought and code before executing it. This is the default UX — auto-approval with transparency, no confirmation prompts. The realistic threat model isn't a malicious agent; it's the agent making an expensive mistake (e.g. firing off a compute_multiprocessing over 10,000 shots when you wanted 100). Surfacing the code before execution lets you ctrl-C out before it commits to anything.

A confirm= callback is exposed on Session.send for callers who want to gate each call programmatically:

def review(call):
    print(call.args.get("code"))
    return input("Run this? [Y/n] ") != "n"

sess.send("...", confirm=review)

Contributors

Any installed package can contribute three things to toksearch.llm via Python entry points:

Entry-point group	Value resolves to	Effect
`toksearch.llm.namespace`	a Python module / object	Bound under the entry-point name in the run_python namespace. A module-level `__llm_description__` populates the system-prompt catalog.
`toksearch.llm.skills`	a `Path` (or callable returning one)	Directory scanned for `SKILL.md` files.
`toksearch.llm.presets`	a `toksearch.llm.presets.Preset` instance	Added to the preset registry under the entry-point name.

Example: how toksearch_d3d plugs in (in its pyproject.toml):

[project.entry-points."toksearch.llm.namespace"]
toksearch_d3d = "toksearch_d3d"

[project.entry-points."toksearch.llm.skills"]
toksearch_d3d = "toksearch_d3d.llm:skills_path"

[project.entry-points."toksearch.llm.presets"]
amsc = "toksearch_d3d.llm:AMSC_PRESET"

Sessions auto-discover all installed contributors on construction. Filter to a subset with Session(packages=["toksearch_d3d"]) or toksearch chat --package toksearch_d3d (repeatable).

API reference

`toksearch.llm.Session`

A conversational LLM session over the run_python persistent namespace.

`reset()`

Clear history and reset the namespace to its initial pre-populated state.

`close()`

Tear down the skills MCP server subprocess.

`toksearch.llm.Preset` `dataclass`

`toksearch.llm.Config` `dataclass`

Events

The four event types dispatched to Session.send() callbacks:

`toksearch.llm.events.TextDelta` `dataclass`

An incremental (or whole, in non-streaming backends) chunk of assistant text.

`toksearch.llm.events.ToolCall` `dataclass`

The assistant is requesting a tool invocation.

thought is populated for run_python (its schema requires a thought field) and is None for tools without one.

`toksearch.llm.events.ToolResult` `dataclass`

The output of a tool invocation, about to be sent back to the model.

`toksearch.llm.events.TurnComplete` `dataclass`

The assistant has finished this turn.

stop_reason: - "end_turn": assistant stopped emitting tool calls. - "max_iterations": hit Session.max_iterations before end_turn. - "interrupted": user aborted (ctrl-C) or confirm() returned False.

Exceptions

`toksearch.llm.errors.LLMError`

Bases: Exception

Base class for all toksearch.llm exceptions.

All exceptions raised by toksearch.llm inherit from LLMError. Specific subclasses: LLMConfigError, LLMAuthError, LLMBackendError, LLMRateLimitError, LLMUserAbort (the last also inherits from KeyboardInterrupt so REPL frames can use a single except KeyboardInterrupt handler).

LLM Interface

Backends

Installation

Quickstart

CLI

Python

Configuration

Tools

run_python

lookup_docs

Show-then-run

Contributors

API reference

toksearch.llm.Session

reset()

close()

toksearch.llm.Preset dataclass

toksearch.llm.Config dataclass