Skip to content

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):

  1. CLI flags: --backend, --model, -n / --max-iterations, --package
  2. Environment variables: FDP_LLM_BACKEND, ANTHROPIC_API_KEY, OPENAI_API_KEY
  3. ~/.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"
    
  4. 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).

See also

  • LLM Tutorial — end-to-end walkthrough including a DIII-D plot.
  • fdp CLI — wraps toksearch chat/query with FDP environment setup for DIII-D users.
  • Claude Agent SDK — underlies the claude-max backend.