sweep-agent

Offline-LLM natural-language control for the sweep stack.

Goal: say "here is vp_init.npy and obs.segy, run an FWI starting at 10 Hz" and have a local LLM turn that into a validated sweep task and run it — no cloud, no API keys.

How it works

user (natural language + files)
        │
        ▼
┌──────────────────────────┐     tool_call
│  Agent loop (agent.py)   │ ───────────────►  local LLM (OpenAI-compatible: vLLM / Ollama)
│                          │ ◄───────────────  tool result (observation)
└───────────┬──────────────┘
            │  dispatches to one of ~30 registered tools
            ▼
   tools/  ── inspect_file · list_equations · check_parameters · make_synthetic_model
            · build_forward_spec · build_fwi_spec · run_task · plot_* · run_fwi · ...
            │
            ├─ discovery / modelling  ──►  sweep            (core wave-equation solver)
            └─ build + execute + viz  ──►  sweep_tasks.TaskRunner   (production runner)

Tools import the geophysics stack lazily: if a layer is missing, the tool returns a clear {"error": "... not importable"} instead of crashing, so the agent always starts and the tools that don't need that layer always work.

Install

pip install sweep-agent          # the agent + the sweep solver

or get it as part of the whole sweep umbrella:

pip install sweepx               # sweep-solver + sweep-agent (+ future companions)

Either path installs the sweep-agent CLI, the ~30-tool registry, and the core solver (sweep-solver, imports as sweep) — so natural-language forward modelling, wavefields and shot gathers work out of the box. Python 3.9+.

To chat you also need a local LLM — any OpenAI-compatible endpoint:

• Ollama (Mac / CPU): ollama serve then ollama pull qwen2.5:7b — sweep-agent chat auto-detects it.
• vLLM (GPU node): pip install "sweep-agent[vllm]" then sweep-agent serve-llm --model qwen2.5-14b-instruct.

Full FWI / LSRTM additionally needs sweep-tasks (the production runner: spec schemas, losses, optimizers, multi-GPU, IO). It is not on PyPI yet — install it from source for now. Forward modelling and the inspection tools don't need it; an FWI tool called without it just returns a clean {"error": "sweep_tasks is not importable"}.

Extras: pip install "sweep-agent[ui]" (Gradio web UI), [vllm], [animate] (GIF export).

macOS (Apple Silicon)

Runs end-to-end on M-series with MPS acceleration (CPU 26.7 s → MPS 5.5 s on a 256×384 / 8-shot / 1500-step demo). Use Ollama for the LLM. Two traps: say "on the mps device" — not "GPU", which makes the LLM fill device="cuda" and silently fall back to CPU; and do not set SWEEP_BUILD_CUDA (that's the Linux + NVIDIA compiled-binding path). macOS uses sweep's eager torch.

Usage

sweep-agent chat        # interactive; auto-detects Ollama/vLLM, tells you if none is running
sweep-agent ui          # same agent in a browser (needs [ui] + a running LLM), then open :7860
sweep-agent tools       # list the ~30 tools — no LLM/GPU needed; --json emits OpenAI tool specs

$ sweep-agent chat
>>> here is vp_init.npy — run a 2-D acoustic forward and show me the shot gather
>>> :reset              # clear conversation history

chat / ui are zero-config by default — they auto-detect a running Ollama (:11434) or vLLM (:8000/:8001), pick a 7B model, and pull it on first run. To switch to any other OpenAI-compatible backend (a remote vLLM, a hosted endpoint, llama.cpp, LM Studio, …) pass --url / --model / --api-key, or set SWEEP_AGENT_LLM_URL / SWEEP_AGENT_LLM_MODEL / SWEEP_AGENT_LLM_API_KEY. For a fully custom backend, subclass BaseLLM from sweep_agent.llm.

Every tool is also a plain function (.fn, with a pydantic params model) — handy for scripts and tests:

from sweep_agent.tools.inspect import inspect_file, InspectFileParams
print(inspect_file.fn(InspectFileParams(path="vp_init.npy")))

What works at each layer

tools	pip install sweep-agent	+ sweep-tasks (from source)
sweep-agent tools, inspect_file, check_parameters, make_synthetic_model	✅	✅
plot_wavelet, plot_velocity_slice, compare_shot_gathers, list_equations	✅	✅
run_forward_sweep — forward modelling / wavefields / shot gathers	✅	✅
build_*_spec, run_task, other plot_*, run_fwi, run_multiscale_fwi, …	error dict	✅

A tool whose layer is missing returns {"error": "… is not importable"} — the agent stays up. The last column (sweep-tasks) is our production FWI/LSRTM tier, not on PyPI yet.

Tests

pip install "sweep-agent[test]"
pytest                  # tests that need sweep / sweep_tasks auto-skip when the stack is absent

License

MIT © Shaowen Wang.