Tools

The functions behind the agent's tool calls. See Agent tools for what each one is for.

Parameter search

helioai.tools.rag

ChromaDB semantic search over the speasy catalog.

The index is built by indexer.py (run: helioai index). This module provides the read-only search path used at agent query time. Models load lazily and are cached at module scope.

search

search(query: str, top_k: int = 5, *, provider: str | None = None, region: str | None = None, measurement_type: str | None = None) -> list[dict]

Semantic search over speasy catalog (single query).

Optional metadata filters narrow the search at query time (the metadata is already indexed by indexer.py): provider (amda/cda/csa/ssc) is the most useful — it counters CDA's dominance of the catalog.

With hybrid search enabled (default), a BM25 lexical channel is fused with the dense channel via RRF — this recovers exact id/code matches (e.g. BGSEc) that dense embeddings miss. When the cross-encoder reranker is off, score is a RELATIVE confidence in [0,1] (top≈1.0); with the reranker on it is the absolute sigmoid score.

For several queries at once use search_batch (one embedding pass + one Chroma call). Returns a list of dicts: {id, name, description, score}.

Parameters:

Name	Type	Description	Default
query	str	Free-text English description of ONE parameter.	required
top_k	int	Number of results.	5
provider	str | None	Restrict to one provider (amda/cda/csa/ssc).	None
region	str | None	SPASE region filter (exact indexed string).	None
measurement_type	str | None	Measurement-type filter (exact indexed string).	None

Example

search("ACE solar wind proton density", top_k=2)[0] {'id': 'cda/AC_H2_SWE/Np', 'name': 'Proton No. density', 'description': 'Proton No. density. Solar Wind Proton Number Density, ' 'scalar. ACE/SWEPAM ... Units: #/cc. ...', 'score': 1.0, ...}

Source code in helioai/tools/rag.py

def search(
    query: str,
    top_k: int = 5,
    *,
    provider: str | None = None,
    region: str | None = None,
    measurement_type: str | None = None,
) -> list[dict]:
    """Semantic search over speasy catalog (single query).

    Optional metadata filters narrow the search at query time (the metadata is
    already indexed by indexer.py): `provider` (amda/cda/csa/ssc) is the most
    useful — it counters CDA's dominance of the catalog.

    With hybrid search enabled (default), a BM25 lexical channel is fused with
    the dense channel via RRF — this recovers exact id/code matches (e.g.
    `BGSEc`) that dense embeddings miss. When the cross-encoder reranker is off,
    `score` is a RELATIVE confidence in [0,1] (top≈1.0); with the reranker on it
    is the absolute sigmoid score.

    For several queries at once use `search_batch` (one embedding pass + one
    Chroma call). Returns a list of dicts: {id, name, description, score}.

    Args:
        query: Free-text English description of ONE parameter.
        top_k: Number of results.
        provider: Restrict to one provider (amda/cda/csa/ssc).
        region: SPASE region filter (exact indexed string).
        measurement_type: Measurement-type filter (exact indexed string).

    Example:
        >>> search("ACE solar wind proton density", top_k=2)[0]
        {'id': 'cda/AC_H2_SWE/Np', 'name': 'Proton No. density',
         'description': 'Proton No. density. Solar Wind Proton Number Density, '
                        'scalar. ACE/SWEPAM ... Units: #/cc. ...', 'score': 1.0, ...}
    """
    if not query or not query.strip():
        return []
    return search_batch(
        [query], top_k, provider=provider, region=region, measurement_type=measurement_type
    )[0]

search_batch

search_batch(queries: list[str], top_k: int = 5, *, provider: str | None = None, region: str | None = None, measurement_type: str | None = None) -> list[list[dict]]

Resolve several queries in ONE pass — the 'composed RAG'.

Encodes all queries in a single embedding pass and issues a single multi-vector ChromaDB query (Chroma returns one result set per query natively), then fuses each independently. Returns one result list per input query, aligned by index (blank queries map to []).

Parameters:

Name	Type	Description	Default
queries	list[str]	One free-text query per parameter to resolve.	required
top_k	int	Results per query.	5
provider	str | None	Same filter as search().	None
region	str | None	Same filter as search().	None
measurement_type	str | None	Same filter as search().	None

Example

ace, wind = search_batch(["ACE solar wind proton density", ... "Wind magnetic field GSE"], top_k=3) ace[0].get("id"), wind[0].get("id") ('cda/AC_H2_SWE/Np', 'cda/WI_H0_MFI/BGSEa')

Source code in helioai/tools/rag.py

def search_batch(
    queries: list[str],
    top_k: int = 5,
    *,
    provider: str | None = None,
    region: str | None = None,
    measurement_type: str | None = None,
) -> list[list[dict]]:
    """Resolve several queries in ONE pass — the 'composed RAG'.

    Encodes all queries in a single embedding pass and issues a single
    multi-vector ChromaDB query (Chroma returns one result set per query
    natively), then fuses each independently. Returns one result list per input
    query, aligned by index (blank queries map to []).

    Args:
        queries: One free-text query per parameter to resolve.
        top_k: Results per query.
        provider: Same filter as `search()`.
        region: Same filter as `search()`.
        measurement_type: Same filter as `search()`.

    Example:
        >>> ace, wind = search_batch(["ACE solar wind proton density",
        ...                           "Wind magnetic field GSE"], top_k=3)
        >>> ace[0].get("id"), wind[0].get("id")
        ('cda/AC_H2_SWE/Np', 'cda/WI_H0_MFI/BGSEa')
    """
    results: list[list[dict]] = [[] for _ in queries]
    active = [(i, q) for i, q in enumerate(queries) if q and q.strip()]
    if not active:
        return results

    # Check per-query cache; only encode/query Chroma for cache misses.
    def _cache_key(q: str) -> tuple:
        return (q, provider, region, measurement_type, top_k)

    uncached = []
    for i, q in active:
        hit = _search_cache.get(_cache_key(q))
        if hit is not None:
            results[i] = hit
        else:
            uncached.append((i, q))

    if not uncached:
        return results

    active = uncached
    model, collection = _load()
    hybrid = settings.rag.hybrid_enabled
    if hybrid:
        dense_k = settings.rag.hybrid_fetch_k
    elif settings.rag.rerank_enabled:
        dense_k = max(top_k, settings.rag.rerank_fetch_k)
    else:
        dense_k = top_k

    vecs = model.encode(
        [q for _, q in active],
        normalize_embeddings=True,
        convert_to_numpy=True,
    ).tolist()

    res = collection.query(
        query_embeddings=vecs,
        n_results=dense_k,
        where=_build_where(provider, region, measurement_type),
        include=["documents", "metadatas", "distances"],
    )
    all_ids = res.get("ids", []) or []
    all_docs = res.get("documents", []) or []
    all_metas = res.get("metadatas", []) or []
    all_dists = res.get("distances", []) or []

    # Second pass without the provider constraint, so an over-eager filter cannot
    # hide a parameter entirely. The embeddings are already computed, so this
    # costs one extra ANN search per batch and no re-encoding.
    open_hits: list[tuple] | None = None
    if provider:
        open_res = collection.query(
            query_embeddings=vecs,
            n_results=dense_k,
            where=_build_where(None, region, measurement_type),
            include=["documents", "metadatas", "distances"],
        )
        open_hits = (
            open_res.get("ids", []) or [],
            open_res.get("documents", []) or [],
            open_res.get("metadatas", []) or [],
            open_res.get("distances", []) or [],
        )

    for j, (i, q) in enumerate(active):
        dense_hit = (
            all_ids[j] if j < len(all_ids) else [],
            all_docs[j] if j < len(all_docs) else [],
            all_metas[j] if j < len(all_metas) else [],
            all_dists[j] if j < len(all_dists) else [],
        )
        res = _fuse_query(
            q,
            dense_hit,
            top_k,
            provider=provider,
            region=region,
            measurement_type=measurement_type,
            hybrid=hybrid,
        )
        if open_hits is not None:
            unfiltered = _fuse_query(
                q,
                tuple(col[j] if j < len(col) else [] for col in open_hits),
                top_k,
                provider=None,
                region=region,
                measurement_type=measurement_type,
                hybrid=hybrid,
            )
            res = _append_cross_provider(res, unfiltered, provider, _CROSS_PROVIDER_EXTRA)
        results[i] = res
        key = _cache_key(q)
        if len(_search_cache) >= _SEARCH_CACHE_MAX:
            _search_cache.pop(next(iter(_search_cache)))
        _search_cache[key] = res
    return results

search_catalogs

search_catalogs(query: str, top_k: int = 5, *, product_type: str | None = None) -> list[dict]

Semantic search over the AMDA catalog/timetable index.

product_type can be 'catalog', 'timetable', or None (both). Returns {id, name, description, score, nb_events, product_type}. Requires helioai index to have been run at least once. Falls back to an empty list if the catalog collection is absent.

Source code in helioai/tools/rag.py

def search_catalogs(
    query: str,
    top_k: int = 5,
    *,
    product_type: str | None = None,
) -> list[dict]:
    """Semantic search over the AMDA catalog/timetable index.

    `product_type` can be 'catalog', 'timetable', or None (both).
    Returns {id, name, description, score, nb_events, product_type}.
    Requires `helioai index` to have been run at least once.
    Falls back to an empty list if the catalog collection is absent.
    """
    if not query or not query.strip():
        return []
    try:
        global _catalog_collection
        model, _ = _load()  # reuse the already-loaded embedding model
        if _catalog_collection is None:
            with _lock:
                if _catalog_collection is None:
                    import chromadb

                    client = chromadb.PersistentClient(path=str(settings.rag.chroma_dir))
                    _catalog_collection = client.get_collection(
                        name=settings.rag.catalogs_collection_name
                    )
        col = _catalog_collection
    except Exception as e:
        log.debug("catalog collection unavailable (%s)", e)
        return []

    try:
        vec = model.encode([query], normalize_embeddings=True, convert_to_numpy=True).tolist()
        where: dict | None = {"product_type": product_type} if product_type else None
        res = col.query(
            query_embeddings=vec,
            n_results=min(top_k, col.count() or 1),
            where=where,
            include=["documents", "metadatas", "distances"],
        )
        results: list[dict] = []
        for pid, doc, meta, dist in zip(
            res["ids"][0],
            res["documents"][0],
            res["metadatas"][0],
            res["distances"][0],
            strict=False,
        ):
            score = round(max(0.0, min(1.0, (1.0 - float(dist) + 1.0) / 2.0)), 4)
            results.append(
                {
                    "id": pid,
                    "name": (meta or {}).get("name", pid),
                    "description": _truncate(doc or ""),
                    "score": score,
                    "nb_events": (meta or {}).get("nb_events", 0),
                    "product_type": (meta or {}).get("product_type", ""),
                }
            )
        return results
    except Exception as e:
        log.warning("catalog search failed: %s", e)
        return []

Data access

helioai.tools.speasy_tools

speasy tools: data access layer wrapping the speasy library.

speasy provides unified access to 70+ missions and 65k+ products from CDAWeb, AMDA, CSA, SSC and others. These tools are the helioai equivalent of AMDA's download_timeseries and list_parameters.

get_timeseries async

get_timeseries(param_id: str, start: str, stop: str, max_points: int = 5000) -> dict

Download a time series from any speasy provider.

Parameters:

Name	Type	Description	Default
param_id	str	speasy parameter id (e.g. 'amda/imf', 'cdaweb/AC_H0_MFI/BGSEc')	required
start	str	ISO 8601 start time (e.g. '2024-01-01T00:00:00')	required
stop	str	ISO 8601 stop time	required
max_points	int	max samples to return (downsampled if needed)	5000

The param_id should be in speasy format: "{provider}/{xmlid}" e.g. "amda/ace_epam_ca60_he", "cda/ACE_H0_MFI/BGSEc" (returned by search_parameters).

Returns dict with: param_id, start, stop, units, shape, n_points, preview (first 10 rows as CSV)

Example

await get_timeseries("amda/imf", "2010-01-01T00:00:00", "2010-01-01T06:00:00") {'dataset': 'imf', 'param_id': 'amda/imf', 'units': 'nT', 'components': ['bx', 'by', 'bz'], 'cadence': '16 s', 'shape': [1350, 3], 'n_points': 1350, 'n_valid': 1350, 'quality': {'missing_pct': 0.0, ...}, 'preview': '2010-01-01T00:00:09.000000000 -1.839, 2.308, 0.108\n...', ...}

Source code in helioai/tools/speasy_tools.py

async def get_timeseries(
    param_id: str,
    start: str,
    stop: str,
    max_points: int = 5000,
) -> dict:
    """Download a time series from any speasy provider.

    Args:
        param_id: speasy parameter id (e.g. 'amda/imf', 'cdaweb/AC_H0_MFI/BGSEc')
        start: ISO 8601 start time (e.g. '2024-01-01T00:00:00')
        stop:  ISO 8601 stop time
        max_points: max samples to return (downsampled if needed)

    The param_id should be in speasy format: "{provider}/{xmlid}"
    e.g. "amda/ace_epam_ca60_he", "cda/ACE_H0_MFI/BGSEc"
    (returned by search_parameters).

    Returns dict with: param_id, start, stop, units, shape, n_points, preview (first 10 rows as CSV)

    Example:
        >>> await get_timeseries("amda/imf", "2010-01-01T00:00:00", "2010-01-01T06:00:00")
        {'dataset': 'imf', 'param_id': 'amda/imf', 'units': 'nT',
         'components': ['bx', 'by', 'bz'], 'cadence': '16 s', 'shape': [1350, 3],
         'n_points': 1350, 'n_valid': 1350, 'quality': {'missing_pct': 0.0, ...},
         'preview': '2010-01-01T00:00:09.000000000  -1.839, 2.308, 0.108\\n...', ...}
    """
    try:
        import numpy as np
        import speasy as spz
    except ImportError:
        return {"error": "speasy is not installed. Run: pip install speasy"}

    from helioai.datastore import find_existing

    cached = find_existing(param_id, start, stop)
    if cached is not None:
        return {
            "dataset": cached,
            "dataset_note": f"use load_data({cached!r}) in run_python — never spz.get_data",
            "param_id": param_id,
            "start": start,
            "stop": stop,
            "already_downloaded": True,
            "note": "This session already holds this parameter for this exact interval. "
            "Nothing was re-fetched; read it with load_data().",
        }

    blocking, coverage_note = _coverage_check(spz, np, param_id, start, stop)
    if blocking is not None:
        return blocking

    try:
        var = spz.get_data(param_id, start, stop)
    except Exception as e:
        log.warning("speasy.get_data failed: %s", e)
        # `str(e)` alone can be a bare "tuple index out of range", which tells the agent
        # nothing about what to do next. Name the exception, and recognise the shape a
        # non-data variable (a CDF time axis) produces so the agent looks for a real
        # product instead of retrying the same id.
        detail = f"{type(e).__name__}: {e}"
        if isinstance(e, IndexError | TypeError):
            detail += (
                " — this often means the id is not a plottable data variable "
                "(a CDF time axis or support variable). Search for the measured "
                "quantity itself and use that id."
            )
        return {"error": f"Failed to retrieve {param_id!r}: {detail}"}

    if var is None:
        return {"error": f"No data returned for {param_id!r} between {start} and {stop}"}

    times = var.time
    values = var.values

    n_points = len(times)
    if n_points == 0:
        return {"error": f"Empty dataset for {param_id!r}"}

    # Replace fill values with NaN HERE, once, before anything else sees the array.
    # Everything downstream — the preview shown to the agent, the .npz the sandbox
    # loads, the exported notebook — then works on data where "no measurement" is
    # NaN rather than -1e31 or 99999.9. Leaving it to the caller meant the sandbox
    # had to remember to call clean(), which cannot see FILLVAL anyway, so a
    # forgotten call put a 99999.9 "speed" straight into a plot and an average.
    from helioai.datastore import blank_fill

    values, fill_mask = blank_fill(values, (getattr(var, "meta", {}) or {}).get("FILLVAL"))
    quality = _data_quality(times, values, np, fill_mask=fill_mask)

    # An all-fill series is functionally as empty as a zero-length one: the rows
    # exist but not one of them holds a measurement. Reporting it as a successful
    # download of n points is a lie the agent then plots — the ACE/SWEPAM
    # saturation during the 2003 Halloween storm returns 1913 such rows. Fail
    # here instead, and do not persist the garbage.
    if quality.get("missing_pct") == 100.0:
        return {
            "error": (
                f"{param_id!r} returned {n_points} rows between {start} and {stop} but "
                f"every value is a fill value — the instrument has no valid data for "
                f"this window."
            ),
            "suggestion": (
                "Try another instrument or spacecraft for this quantity; the parameter "
                "id is fine, the coverage is not."
            ),
            "n_points": n_points,
            "missing_pct": 100.0,
        }

    # Persist full-resolution data before downsampling
    from helioai.datastore import save_timeseries

    saved = save_timeseries(
        param_id,
        time=times,
        values=values,
        param_id=param_id,
        units=str(getattr(var, "unit", "") or ""),
        start=start,
        stop=stop,
        columns=list(getattr(var, "columns", None) or []),
        source="get_timeseries",
    )

    # Cadence and valid-sample count describe the PERSISTED series, so they are taken
    # before the preview is thinned: `load_data()` hands back the full resolution, and
    # a cadence measured on the thinned copy would describe something the agent never
    # works with.
    cadence, n_valid = _sample_cadence(times, values)

    # Downsample if needed
    if n_points > max_points:
        step = n_points // max_points
        times = times[::step]
        values = values[::step]
        n_points = len(times)

    # Build a brief preview (first 10 rows as CSV text). Nothing here may raise: a
    # non-numeric variable (datetime64 values, string labels) used to kill the whole
    # call from inside this loop, with no try/except above it.
    preview_lines: list[str] = []
    for i in range(min(10, n_points)):
        v = values[i]
        try:
            if hasattr(v, "__len__"):
                v_str = ", ".join(f"{x:.4g}" for x in v)
            else:
                v_str = f"{float(v):.6g}"
        except (TypeError, ValueError):
            v_str = str(v)[:60]
        preview_lines.append(f"{times[i]}  {v_str}")
    preview = "\n".join(preview_lines)

    shape = list(values.shape)
    units = getattr(var, "unit", "") or ""
    name = getattr(var, "name", "") or ""
    components = list(getattr(var, "columns", None) or [])

    # Mission / instrument: best-effort from param_id prefix + var.meta
    mission = ""
    instrument = ""
    try:
        parts = param_id.split("/")
        dataset = parts[1] if len(parts) > 1 else parts[0]
        # First token before '_'/'-' is the mission (cda/ACE_H0_MFI/.. → ACE, amda/ace_imf_all → ace)
        mission = dataset.split("_")[0].split("-")[0]
    except Exception:
        pass
    try:
        meta = getattr(var, "meta", {}) or {}
        instrument = str(meta.get("FIELDNAM", "") or meta.get("VAR_NOTES", "") or "")[:80]
    except Exception:
        pass

    # `dataset` leads the payload on purpose. Stale tool results are summarised by
    # serialising this dict and cutting it to a fixed length, so trailing keys are the
    # ones that vanish — and losing the handle is what makes the agent re-download a
    # parameter it already has instead of calling load_data().
    result: dict = {}
    if saved:
        ds_name = saved["dataset"]
        result["dataset"] = ds_name
        result["dataset_note"] = f"use load_data({ds_name!r}) in run_python — never spz.get_data"
    result |= {
        "param_id": param_id,
        "name": name,
        "start": start,
        "stop": stop,
        "units": str(units),
        "components": components,
        "cadence": cadence,
        "mission": mission,
        "instrument": instrument,
        "shape": shape,
        "n_points": n_points,
        "n_valid": n_valid,
        "preview": preview,
    }
    if quality:
        result["quality"] = quality
    if coverage_note:
        result["coverage_note"] = coverage_note
    return result

list_missions async

list_missions() -> dict

List available speasy data providers and their top-level missions.

Returns a summary dict with provider names and approximate product counts.

Example

await list_missions() {'providers': ['amda', 'archive', 'cda', 'csa', 'ssc', 'uiowaephtool'], 'note': 'Use search_parameters to find specific parameters. ...'}

Source code in helioai/tools/speasy_tools.py

async def list_missions() -> dict:
    """List available speasy data providers and their top-level missions.

    Returns a summary dict with provider names and approximate product counts.

    Example:
        >>> await list_missions()
        {'providers': ['amda', 'archive', 'cda', 'csa', 'ssc', 'uiowaephtool'],
         'note': 'Use search_parameters to find specific parameters. ...'}
    """
    try:
        import speasy as spz
    except ImportError:
        return {"error": "speasy is not installed. Run: pip install speasy"}

    providers: dict[str, int] = {}
    try:
        tree = spz.inventories.tree
        for attr in dir(tree):
            if attr.startswith("_"):
                continue
            node = getattr(tree, attr, None)
            if node is not None:
                providers[attr] = "available"
    except Exception as e:
        log.warning("Failed to walk speasy inventory: %s", e)
        return {"error": str(e)}

    return {
        "providers": list(providers.keys()),
        "note": (
            "Use search_parameters to find specific parameters. "
            "Common provider prefixes: amda/, cdaweb/, csa/, ssc/"
        ),
    }

search_parameters async

search_parameters(query: str | None = None, top_k: int = 5, provider: str | None = None, queries: list[str] | None = None, start: str | None = None, stop: str | None = None) -> dict

Semantic search over the speasy catalog (83k+ products).

Requires the index to be built first (run: helioai index). Falls back to a direct speasy text match if no index is found.

Every result carries the product's published coverage. Passing the window you intend to download sorts products that cannot cover it to the bottom and flags them, which is cheaper than discovering it one failed download at a time — resolving a 2015 ephemeris used to cost four turns against products that stop in 1997.

Parameters:

Name	Type	Description	Default
query	str | None	free-text English query for a SINGLE parameter.	None
queries	list[str] | None	list of queries to resolve SEVERAL parameters in ONE call (preferred when 2+ parameters are needed — much cheaper).	None
top_k	int	number of results per query.	5
provider	str | None	optional — restrict to one provider (amda/cda/csa/ssc).	None
start	str | None	optional ISO start of the interval you intend to download.	None
stop	str | None	optional ISO stop. Both are needed for the filter to apply.	None

Returns either {query, provider, results} (single) or {provider, groups: [{query, results}]} (batch).

Example

await search_parameters(query="ACE solar wind proton density", top_k=3) {'query': 'ACE solar wind proton density', 'provider': None, 'results': [ {'id': 'cda/AC_H2_SWE/Np', 'description': 'Proton No. density. Solar Wind Proton Number Density, scalar. ' 'ACE/SWEPAM ... 1-Hour Level 2 Data ... Units: #/cc. ...', 'coverage': '1998-02-04 → 2024-07-09', 'score': 1.0}, ...]}

Source code in helioai/tools/speasy_tools.py

async def search_parameters(
    query: str | None = None,
    top_k: int = 5,
    provider: str | None = None,
    queries: list[str] | None = None,
    start: str | None = None,
    stop: str | None = None,
) -> dict:
    """Semantic search over the speasy catalog (83k+ products).

    Requires the index to be built first (run: helioai index).
    Falls back to a direct speasy text match if no index is found.

    Every result carries the product's published `coverage`. Passing the window you
    intend to download sorts products that cannot cover it to the bottom and flags them,
    which is cheaper than discovering it one failed download at a time — resolving a
    2015 ephemeris used to cost four turns against products that stop in 1997.

    Args:
        query: free-text English query for a SINGLE parameter.
        queries: list of queries to resolve SEVERAL parameters in ONE call
                 (preferred when 2+ parameters are needed — much cheaper).
        top_k: number of results per query.
        provider: optional — restrict to one provider (amda/cda/csa/ssc).
        start: optional ISO start of the interval you intend to download.
        stop: optional ISO stop. Both are needed for the filter to apply.

    Returns either {query, provider, results} (single) or
    {provider, groups: [{query, results}]} (batch).

    Example:
        >>> await search_parameters(query="ACE solar wind proton density", top_k=3)
        {'query': 'ACE solar wind proton density', 'provider': None, 'results': [
         {'id': 'cda/AC_H2_SWE/Np',
          'description': 'Proton No. density. Solar Wind Proton Number Density, scalar. '
                         'ACE/SWEPAM ... 1-Hour Level 2 Data ... Units: #/cc. ...',
          'coverage': '1998-02-04 → 2024-07-09', 'score': 1.0}, ...]}
    """
    window = (start, stop) if start and stop else None
    if queries:
        try:
            from helioai.tools.rag import search_batch as rag_search_batch

            batch = rag_search_batch(queries, top_k=top_k, provider=provider)
            return {
                "provider": provider,
                "groups": [
                    {"query": q, "results": _apply_window(r, window)}
                    for q, r in zip(queries, batch, strict=False)
                ],
            }
        except Exception as e:
            log.warning("RAG batch search failed (%s), falling back to text scan", e)
        try:
            import speasy as spz

            groups = [{"query": q, "results": _fallback_search(spz, q, top_k)} for q in queries]
            return {
                "provider": provider,
                "groups": groups,
                "note": "RAG index not built — using text fallback (provider filter ignored)",
            }
        except Exception as e2:
            return {"error": f"Search failed: {e2}"}

    if not query:
        return {"error": "provide query (string) or queries (list of strings)"}

    try:
        from helioai.tools.rag import search as rag_search

        results = rag_search(query, top_k=top_k, provider=provider)
        return {"query": query, "provider": provider, "results": _apply_window(results, window)}
    except Exception as e:
        log.warning("RAG search failed (%s), falling back to speasy inventory scan", e)

    # Fallback: naive text search on speasy inventory (provider filter ignored — best effort)
    try:
        import speasy as spz

        results = _fallback_search(spz, query, top_k)
        return {
            "query": query,
            "results": results,
            "note": "RAG index not built — using text fallback (provider filter ignored)",
        }
    except Exception as e2:
        return {"error": f"Search failed: {e2}"}

Event catalogs

helioai.tools.catalog_tools

Catalog and timetable tools for HelioAI.

Exposes the 29 CatalogIndex + 188 TimetableIndex from the AMDA speasy inventory as first-class agent tools. The key capability is get_events_timeseries: download a parameter for every event in a catalog in one speasy call, opening the door to superposed epoch analysis.

list_catalogs async

list_catalogs(type: str = 'all', region: str | None = None) -> dict

List available AMDA event catalogs and timetables.

Parameters:

Name	Type	Description	Default
type	str	'catalog', 'timetable', or 'all' (default).	'all'
region	str | None	optional keyword filter on name/description (e.g. 'ICME', 'bow shock', 'MMS').	None

Returns a list of entries with id, name, type, nb_events, survey range and description. Use the id field with get_catalog() and get_events_timeseries().

Example

await list_catalogs(type="catalog", region="ICME") {'total': 3, 'type_filter': 'catalog', 'region_filter': 'ICME', 'catalogs': [ {'id': 'amda/sharedcatalog_41', 'name': 'ICME_multi-catalog', 'type': 'catalog', 'nb_events': 2003, 'survey_start': '1975-01-08', 'survey_stop': '2022-10-21', 'description': '...'}, ...]}

Source code in helioai/tools/catalog_tools.py

async def list_catalogs(
    type: str = "all",
    region: str | None = None,
) -> dict:
    """List available AMDA event catalogs and timetables.

    Args:
        type: 'catalog', 'timetable', or 'all' (default).
        region: optional keyword filter on name/description (e.g. 'ICME', 'bow shock', 'MMS').

    Returns a list of entries with id, name, type, nb_events, survey range and description.
    Use the `id` field with get_catalog() and get_events_timeseries().

    Example:
        >>> await list_catalogs(type="catalog", region="ICME")
        {'total': 3, 'type_filter': 'catalog', 'region_filter': 'ICME', 'catalogs': [
         {'id': 'amda/sharedcatalog_41', 'name': 'ICME_multi-catalog', 'type': 'catalog',
          'nb_events': 2003, 'survey_start': '1975-01-08', 'survey_stop': '2022-10-21',
          'description': '...'}, ...]}
    """
    spz = _get_spz()
    if spz is None:
        return {"error": "speasy is not installed"}

    entries = _walk_catalogs(spz)

    # Append local/ catalogs (direct disk read, bypasses TTL cache)
    try:
        for p in sorted(_catalogs_dir().glob("*.json")):
            data = json.loads(p.read_text(encoding="utf-8"))
            nb = len(data.get("events", []))
            entries.append(
                {
                    "id": f"local/{p.stem}",
                    "name": data.get("name", p.stem),
                    "type": "catalog",
                    "nb_events": nb,
                    "survey_start": "",
                    "survey_stop": "",
                    "description": data.get("description", "")[:200],
                }
            )
    except Exception as e:
        log.warning("list_catalogs: local catalog scan failed: %s", e)

    for key, spec in _HELIO4CAST.items():
        nb = spec["approx_events"]
        cache = _helio4cast_cache_path(key)
        if cache.exists():
            try:
                with cache.open(encoding="utf-8") as f:
                    nb = max(sum(1 for _ in f) - 1, 0)
            except OSError:
                pass
        entries.append(
            {
                "id": f"helio4cast/{key}",
                "name": spec["name"],
                "type": "catalog",
                "nb_events": nb,
                "survey_start": spec["survey"][0],
                "survey_stop": spec["survey"][1],
                "description": spec["description"][:200],
            }
        )

    if type in ("catalog", "timetable"):
        entries = [e for e in entries if e["type"] == type]

    if region:
        kw = region.lower()
        entries = [e for e in entries if kw in e["name"].lower() or kw in e["description"].lower()]

    entries.sort(key=lambda e: e["nb_events"], reverse=True)

    return {
        "total": len(entries),
        "type_filter": type,
        "region_filter": region,
        "catalogs": entries,
    }

get_catalog async

get_catalog(catalog_id: str, start: str | None = None, stop: str | None = None, max_events: int = 10, columns: list[str] | None = None, where: dict | None = None, sort_by: str | None = None, descending: bool = False, offset: int = 0) -> dict

Download and summarize an AMDA event catalog or timetable.

Parameters:

Name	Type	Description	Default
catalog_id	str	speasy uid from list_catalogs (e.g. 'amda/sharedcatalog_41').	required
start	str | None	optional ISO 8601 start — filter events beginning after this time.	None
stop	str | None	optional ISO 8601 stop — filter events beginning before this time.	None
max_events	int	maximum events to include in the sample (default 10).	10
columns	list[str] | None	restrict the metadata columns returned per event.	None
where	dict | None	server-side row filter — {"column": str, "op": "eq|ne|gt|gte|lt|lte|contains", "value": any}.	None
sort_by	str | None	column name to sort events by before slicing.	None
descending	bool	sort direction (default ascending).	False
offset	int	pagination offset into the filtered+sorted events.	0

Returns catalog metadata + a sample of events (start, stop, key columns). Use get_events_timeseries() to download a parameter over all events.

Example

await get_catalog("amda/sharedcatalog_41", start="2015-01-01", stop="2016-01-01", ... max_events=5, sort_by="start") {'_kind': 'catalog_preview', 'catalog_id': 'amda/sharedcatalog_41', 'name': 'ICME_multi-catalog', 'nb_events_total': 2003, 'nb_events_filtered': ..., 'returned': 5, 'columns': [...], 'sample': [{'start': ..., 'stop': ..., ...}, ...], 'survey_start': '1975-01-08', 'survey_stop': '2022-10-21'}

Source code in helioai/tools/catalog_tools.py

async def get_catalog(
    catalog_id: str,
    start: str | None = None,
    stop: str | None = None,
    max_events: int = 10,
    columns: list[str] | None = None,
    where: dict | None = None,
    sort_by: str | None = None,
    descending: bool = False,
    offset: int = 0,
) -> dict:
    """Download and summarize an AMDA event catalog or timetable.

    Args:
        catalog_id: speasy uid from list_catalogs (e.g. 'amda/sharedcatalog_41').
        start:      optional ISO 8601 start — filter events beginning after this time.
        stop:       optional ISO 8601 stop  — filter events beginning before this time.
        max_events: maximum events to include in the sample (default 10).
        columns:    restrict the metadata columns returned per event.
        where:      server-side row filter — {"column": str, "op": "eq|ne|gt|gte|lt|lte|contains", "value": any}.
        sort_by:    column name to sort events by before slicing.
        descending: sort direction (default ascending).
        offset:     pagination offset into the filtered+sorted events.

    Returns catalog metadata + a sample of events (start, stop, key columns).
    Use get_events_timeseries() to download a parameter over all events.

    Example:
        >>> await get_catalog("amda/sharedcatalog_41", start="2015-01-01", stop="2016-01-01",
        ...                   max_events=5, sort_by="start")
        {'_kind': 'catalog_preview', 'catalog_id': 'amda/sharedcatalog_41',
         'name': 'ICME_multi-catalog', 'nb_events_total': 2003, 'nb_events_filtered': ...,
         'returned': 5, 'columns': [...], 'sample': [{'start': ..., 'stop': ..., ...}, ...],
         'survey_start': '1975-01-08', 'survey_stop': '2022-10-21'}
    """
    spz = _get_spz()
    if spz is None:
        return {"error": "speasy is not installed"}

    try:
        cat, index = _resolve_catalog(catalog_id, spz)
    except Exception as e:
        return {"error": f"Failed to download catalog {catalog_id!r}: {e}"}

    if cat is None:
        return {"error": f"Catalog {catalog_id!r} not found"}

    try:
        events = list(cat)
    except Exception as e:
        return {"error": f"Failed to iterate catalog events: {e}"}

    nb_total = len(events)

    # 1. Time window filter
    if start or stop:
        filtered: list = []
        for ev in events:
            ev_start = _ev_iso(ev, "start_time", "start")
            if start and ev_start < start:
                continue
            if stop and ev_start > stop:
                continue
            filtered.append(ev)
        events = filtered

    # 2. where filter
    if where and isinstance(where, dict):
        col = where.get("column", "")
        op = where.get("op", "eq")
        val = where.get("value")
        if col and op and val is not None:
            events = [ev for ev in events if _match(op, _event_value(ev, col), val)]

    nb_filtered = len(events)

    # 3. sort
    if sort_by:

        def _sort_key(ev):
            v = _event_value(ev, sort_by)
            try:
                return (0, float(v))
            except (TypeError, ValueError):
                return (1, str(v) if v is not None else "")

        events = sorted(events, key=_sort_key, reverse=descending)

    # 4. pagination + slice
    offset = max(0, offset)
    page = events[offset : offset + max_events]

    # 5. Build event rows
    rows: list[dict] = []
    for ev in page:
        ev_start = _ev_iso(ev, "start_time", "start")
        ev_stop = _ev_iso(ev, "stop_time", "stop")
        row: dict[str, Any] = {"start": ev_start, "stop": ev_stop}
        meta = getattr(ev, "meta", None)
        if meta and isinstance(meta, dict):
            if columns:
                for k in columns:
                    row[k] = meta.get(k)
            else:
                for k, v in list(meta.items())[:8]:
                    row[k] = v
        if sort_by and sort_by not in row:
            row[sort_by] = _event_value(ev, sort_by)
        rows.append(row)

    all_columns = list(rows[0].keys()) if rows else (["start", "stop"] + (columns or []))

    survey_start, survey_stop = _survey(index) if index is not None else ("", "")
    cat_name = _name(index) if index is not None else catalog_id.split("/")[-1]
    cat_type = _spz_type(index) if index is not None else "catalog"
    returned = len(rows)
    return {
        "_kind": "catalog_preview",
        "catalog_id": catalog_id,
        "name": cat_name,
        "type": cat_type,
        "nb_events_total": nb_total,
        "nb_events_filtered": nb_filtered,
        "offset": offset,
        "returned": returned,
        "columns": all_columns,
        "sample": rows,
        "survey_start": survey_start,
        "survey_stop": survey_stop,
        "note": (
            f"Showing rows {offset}–{offset + returned} of {nb_filtered} filtered "
            f"({nb_total} total). "
            + (
                f"Use offset={offset + returned} for the next page. "
                if offset + returned < nb_filtered
                else ""
            )
            + "Use get_events_timeseries() to download a parameter over all events."
        ),
    }

get_events_timeseries async

get_events_timeseries(catalog_id: str, param_id: str, start: str, stop: str, max_events: int = 50) -> dict

Download a parameter for every event in a catalog window (superposed epoch).

This is the core catalog tool: it fetches N time series in a SINGLE speasy call using the native multi-interval API. Use it for: - Superposed epoch analysis (stack-plot across events) - Statistical summaries per event (min/max/mean) - Comparing a parameter across e.g. all ICME crossings in a year

Parameters:

Name	Type	Description	Default
catalog_id	str	speasy uid from list_catalogs (e.g. 'amda/sharedcatalog_41').	required
param_id	str	speasy parameter id (e.g. 'amda/imf_gsm') — resolve via search_parameters first.	required
start	str	ISO 8601 start — restrict to events beginning after this time.	required
stop	str	ISO 8601 stop — restrict to events beginning before this time.	required
max_events	int	cap on events to download (default 20 — each is one speasy call slot).	50

Returns per-event statistics and saves the raw data to the workspace for run_python.

Example

await get_events_timeseries("amda/sharedcatalog_41", "amda/imf", ... "2015-01-01", "2016-01-01", max_events=10) {'catalog_id': 'amda/sharedcatalog_41', 'param_id': 'amda/imf', 'stats': [ {'event': 0, 'start': '2015-01-03T...', 'stop': '2015-01-04T...', 'n_points': ..., ...}, ...], ...}

Source code in helioai/tools/catalog_tools.py

async def get_events_timeseries(
    catalog_id: str,
    param_id: str,
    start: str,
    stop: str,
    max_events: int = 50,
) -> dict:
    """Download a parameter for every event in a catalog window (superposed epoch).

    This is the core catalog tool: it fetches N time series in a SINGLE speasy call
    using the native multi-interval API.  Use it for:
    - Superposed epoch analysis (stack-plot across events)
    - Statistical summaries per event (min/max/mean)
    - Comparing a parameter across e.g. all ICME crossings in a year

    Args:
        catalog_id: speasy uid from list_catalogs (e.g. 'amda/sharedcatalog_41').
        param_id:   speasy parameter id (e.g. 'amda/imf_gsm') — resolve via search_parameters first.
        start:      ISO 8601 start — restrict to events beginning after this time.
        stop:       ISO 8601 stop  — restrict to events beginning before this time.
        max_events: cap on events to download (default 20 — each is one speasy call slot).

    Returns per-event statistics and saves the raw data to the workspace for run_python.

    Example:
        >>> await get_events_timeseries("amda/sharedcatalog_41", "amda/imf",
        ...                             "2015-01-01", "2016-01-01", max_events=10)
        {'catalog_id': 'amda/sharedcatalog_41', 'param_id': 'amda/imf', 'stats': [
         {'event': 0, 'start': '2015-01-03T...', 'stop': '2015-01-04T...',
          'n_points': ..., ...}, ...], ...}
    """
    spz = _get_spz()
    if spz is None:
        return {"error": "speasy is not installed"}

    # --- resolve catalog ---
    try:
        cat, _ = _resolve_catalog(catalog_id, spz)
    except Exception as e:
        return {"error": f"Failed to download catalog {catalog_id!r}: {e}"}

    if cat is None:
        return {"error": f"Catalog {catalog_id!r} not found"}

    # --- filter events ---
    try:
        events = list(cat)
    except Exception as e:
        return {"error": f"Cannot iterate catalog: {e}"}

    filtered = []
    for ev in events:
        ev_start = _ev_iso(ev, "start_time", "start")
        if ev_start < start or ev_start > stop:
            continue
        filtered.append(ev)

    if not filtered:
        return {
            "warning": f"No events found in [{start}, {stop}] for {catalog_id!r}.",
            "suggestion": "Widen the time window or use get_catalog() to inspect the survey range.",
        }

    selected = filtered[:max_events]
    cap_warning = (
        f"Showing first {max_events}/{len(filtered)} events. "
        f"Pass max_events={len(filtered)} for full SEA."
        if len(filtered) > max_events
        else None
    )

    # --- batch download: ONE speasy call for all events ---
    try:
        timeseries_list = spz.get_data(param_id, selected)
    except Exception as e:
        return {"error": f"speasy.get_data({param_id!r}, events) failed: {e}"}

    if timeseries_list is None:
        return {"error": f"No data returned for {param_id!r} over {len(selected)} events"}

    if not isinstance(timeseries_list, list):
        timeseries_list = [timeseries_list]

    # --- per-event statistics ---
    import numpy as np

    stats: list[dict] = []
    for i, (ev, ts) in enumerate(zip(selected, timeseries_list, strict=False)):
        ev_start = _ev_iso(ev, "start_time", "start")
        ev_stop = _ev_iso(ev, "stop_time", "stop")
        if ts is None or len(ts.time) == 0:
            stats.append({"event": i, "start": ev_start, "stop": ev_stop, "status": "no_data"})
            continue
        from helioai.datastore import blank_fill

        vals, _ = blank_fill(ts.values, (getattr(ts, "meta", {}) or {}).get("FILLVAL"))
        with np.errstate(all="ignore"):
            entry: dict = {
                "event": i,
                "start": ev_start,
                "stop": ev_stop,
                "n_points": int(len(ts.time)),
            }
            if vals.ndim == 1 or vals.shape[1] == 1:
                flat = vals.ravel()
                entry.update(
                    mean=_fmt(np.nanmean(flat)),
                    std=_fmt(np.nanstd(flat)),
                    min=_fmt(np.nanmin(flat)),
                    max=_fmt(np.nanmax(flat)),
                )
            else:
                col_names = list(getattr(ts, "columns", None) or [])
                if len(col_names) != vals.shape[1]:
                    col_names = [f"c{j}" for j in range(vals.shape[1])]
                components = {}
                for j, cname in enumerate(col_names):
                    col = vals[:, j]
                    components[cname] = {
                        "mean": _fmt(np.nanmean(col)),
                        "std": _fmt(np.nanstd(col)),
                        "min": _fmt(np.nanmin(col)),
                        "max": _fmt(np.nanmax(col)),
                    }
                entry["components"] = components
                n_mag = min(vals.shape[1], 3)
                mag = np.linalg.norm(vals[:, :n_mag], axis=1)
                mag[~np.isfinite(mag)] = np.nan
                entry["magnitude"] = {
                    "mean": _fmt(np.nanmean(mag)),
                    "std": _fmt(np.nanstd(mag)),
                    "min": _fmt(np.nanmin(mag)),
                    "max": _fmt(np.nanmax(mag)),
                }
            stats.append(entry)

    good = [s for s in stats if s.get("status") != "no_data"]
    units = str(getattr(timeseries_list[0], "unit", "") or "") if timeseries_list else ""

    if len(stats) > 10:
        per_event_stats = stats[:5] + stats[-5:]
        stats_note = (
            f"per_event_stats shows first 5 + last 5 of {len(stats)} events. "
            "Full data available via load_data()."
        )
    else:
        per_event_stats = stats
        stats_note = None

    # Persist event collection for reuse in run_python via load_data()
    from helioai.datastore import save_event_collection

    series = []
    for ev, ts in zip(selected, timeseries_list, strict=False):
        ev_start = _ev_iso(ev, "start_time", "start")
        ev_stop = _ev_iso(ev, "stop_time", "stop")
        series.append((ev_start, ev_stop, ts if (ts is not None and len(ts.time) > 0) else None))

    saved = save_event_collection(
        param_id,
        series=series,
        param_id=param_id,
        units=units,
        source="get_events_timeseries",
    )

    result: dict = {
        "catalog_id": catalog_id,
        "param_id": param_id,
        "time_window": [start, stop],
        "n_events_found": len(filtered),
        "n_events_downloaded": len(selected),
        "n_events_with_data": len(good),
        "units": units,
        "per_event_stats": per_event_stats,
    }
    if cap_warning:
        result["cap_warning"] = cap_warning
    if stats_note:
        result["stats_note"] = stats_note
    if saved:
        ds_name = saved["dataset"]
        result["dataset"] = ds_name
        result["note"] = (
            f"In run_python: events = load_data({ds_name!r}) — "
            "a list of objects with .time, .values, .start, .stop, .units per event. "
            "Use the superposed_epoch recipe: load_recipe('superposed_epoch')."
        )
    else:
        result["note"] = (
            "Use run_python with spz.get_data(param_id, events) for custom plots. "
            "The catalog events are speasy DateTimeRange objects iterable from the catalog."
        )
    return result

save_catalog async

save_catalog(name: str, events: list[dict], description: str = '') -> dict

Save a list of events as a local catalog under the local/ prefix.

Parameters:

Name	Type	Description	Default
name	str	Catalog name — lowercase letters, digits, hyphens, underscores (1-40 chars).	required
events	list[dict]	List of dicts with 'start' and 'stop' ISO 8601 strings plus optional extra keys.	required
description	str	Short description (optional).	''

Returns {"catalog_id": "local/", "nb_events": N, "note": "..."}. Overwrites an existing catalog with the same name. Use list_catalogs() then get_catalog("local/") to inspect it.

Example

await save_catalog("my-shocks", ... [{"start": "2015-03-17T04:01:00", "stop": "2015-03-17T05:00:00", ... "note": "St. Patrick's Day storm shock"}])

Source code in helioai/tools/catalog_tools.py

async def save_catalog(
    name: str,
    events: list[dict],
    description: str = "",
) -> dict:
    """Save a list of events as a local catalog under the local/<name> prefix.

    Args:
        name:        Catalog name — lowercase letters, digits, hyphens, underscores (1-40 chars).
        events:      List of dicts with 'start' and 'stop' ISO 8601 strings plus optional extra keys.
        description: Short description (optional).

    Returns {"catalog_id": "local/<name>", "nb_events": N, "note": "..."}.
    Overwrites an existing catalog with the same name.
    Use list_catalogs() then get_catalog("local/<name>") to inspect it.

    Example:
        >>> await save_catalog("my-shocks",
        ...                    [{"start": "2015-03-17T04:01:00", "stop": "2015-03-17T05:00:00",
        ...                      "note": "St. Patrick's Day storm shock"}])
        {'catalog_id': 'local/my-shocks', 'nb_events': 1, 'overwritten': False, 'note': '...'}
    """
    if not _LOCAL_NAME_RE.fullmatch(name):
        return {
            "error": (
                f"Invalid catalog name {name!r} — "
                "use 1-40 lowercase letters, digits, hyphens or underscores"
            )
        }
    if not events:
        return {"error": "events list is empty — provide at least one event"}
    if len(events) > _MAX_EVENTS_LOCAL:
        return {"error": f"Too many events ({len(events)} > {_MAX_EVENTS_LOCAL} cap)"}

    validated: list[dict] = []
    for i, ev in enumerate(events):
        s = str(ev.get("start", "")).strip()
        e = str(ev.get("stop", "")).strip()
        if not s or not e:
            return {"error": f"Event {i}: 'start' and 'stop' are required"}
        if s >= e:
            return {"error": f"Event {i}: start >= stop ({s!r} >= {e!r})"}
        meta = {k: v for k, v in ev.items() if k not in ("start", "stop")}
        validated.append({"start": s, "stop": e, "meta": meta})

    import datetime

    payload = {
        "name": name,
        "description": description,
        "created": datetime.datetime.now(datetime.UTC).isoformat(timespec="seconds"),
        "events": validated,
    }
    path = _catalogs_dir() / f"{name}.json"
    overwritten = path.exists()
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")

    return {
        "catalog_id": f"local/{name}",
        "nb_events": len(validated),
        "overwritten": overwritten,
        "note": (
            f"Saved {len(validated)} events as local/{name}. "
            "Use get_catalog('local/" + name + "') to inspect or "
            "get_events_timeseries('local/" + name + "', param_id, ...) to analyse."
        ),
    }

Plasma physics

helioai.tools.plasmapy_tools

PlasmaPy-based plasma physics calculations exposed as agent tools.

Each function accepts plain SI-ish numbers (nT, cm⁻³, eV) and returns a dict with value, unit, and a brief physical context — ready for LLM consumption.

SPASE ParticleQuantity / FieldQuantity mapping: plasma_beta → PlasmaBeta (ActivityIndex) gyrofrequency → Gyrofrequency (FieldQuantity + ParticleQuantity) debye_length → (ParticleQuantity implied) alfven_speed → AlfvenVelocity (ParticleQuantity) inertial_length → (ParticleQuantity implied) power_spectrum → Spectrum (MeasurementType)

plasma_beta async

plasma_beta(B_nT: float, n_cm3: float, T_eV: float) -> dict

Compute plasma beta — ratio of thermal pressure to magnetic pressure.

Parameters:

Name	Type	Description	Default
B_nT	float	Magnetic field magnitude in nT	required
n_cm3	float	Number density in cm⁻³	required
T_eV	float	Temperature in eV	required

Returns dict with beta (dimensionless) and regime interpretation.

Example

await plasma_beta(B_nT=5.0, n_cm3=10.0, T_eV=20.0) {'beta': 3.221367, 'unit': 'dimensionless', 'regime': 'high-β plasma (β ~ 1-10) — typical magnetosheath / plasma sheet', ...}

Source code in helioai/tools/plasmapy_tools.py

async def plasma_beta(B_nT: float, n_cm3: float, T_eV: float) -> dict:
    """Compute plasma beta — ratio of thermal pressure to magnetic pressure.

    Args:
        B_nT:  Magnetic field magnitude in nT
        n_cm3: Number density in cm⁻³
        T_eV:  Temperature in eV

    Returns dict with beta (dimensionless) and regime interpretation.

    Example:
        >>> await plasma_beta(B_nT=5.0, n_cm3=10.0, T_eV=20.0)
        {'beta': 3.221367, 'unit': 'dimensionless',
         'regime': 'high-β plasma (β ~ 1-10) — typical magnetosheath / plasma sheet', ...}
    """
    try:
        import astropy.units as u
        import plasmapy.formulary as pf

        B = B_nT * u.nT
        n = n_cm3 * u.cm**-3
        T = T_eV * u.eV

        beta_val = float(pf.beta(T, n, B).value)

        if beta_val < 0.01:
            regime = "magnetically dominated (β ≪ 1) — typical inner magnetosphere / coronal loop"
        elif beta_val < 1.0:
            regime = "low-β plasma (β < 1) — typical solar wind / outer magnetosphere"
        elif beta_val < 10.0:
            regime = "high-β plasma (β ~ 1-10) — typical magnetosheath / plasma sheet"
        else:
            regime = "pressure-dominated (β ≫ 1) — typical ionosphere / dense plasma"

        return {
            "beta": round(beta_val, 6),
            "unit": "dimensionless",
            "regime": regime,
            "inputs": {"B_nT": B_nT, "n_cm3": n_cm3, "T_eV": T_eV},
        }
    except Exception as e:
        return {"error": str(e)}

gyrofrequency async

gyrofrequency(B_nT: float, particle: str = 'proton') -> dict

Compute particle gyrofrequency (cyclotron frequency).

Parameters:

Name	Type	Description	Default
B_nT	float	Magnetic field magnitude in nT	required
particle	str	'proton', 'electron', 'alpha' (default: proton)	'proton'

Returns dict with frequency in Hz and angular frequency in rad/s.

Example

await gyrofrequency(B_nT=5.0) {'frequency_Hz': 0.0762, 'angular_frequency_rad_s': 0.4789, 'particle': 'proton', 'B_nT': 5.0, 'period_s': 13.118895} await gyrofrequency(B_nT=5.0, particle="electron")

Source code in helioai/tools/plasmapy_tools.py

async def gyrofrequency(B_nT: float, particle: str = "proton") -> dict:
    """Compute particle gyrofrequency (cyclotron frequency).

    Args:
        B_nT:    Magnetic field magnitude in nT
        particle: 'proton', 'electron', 'alpha' (default: proton)

    Returns dict with frequency in Hz and angular frequency in rad/s.

    Example:
        >>> await gyrofrequency(B_nT=5.0)
        {'frequency_Hz': 0.0762, 'angular_frequency_rad_s': 0.4789, 'particle': 'proton',
         'B_nT': 5.0, 'period_s': 13.118895}
        >>> await gyrofrequency(B_nT=5.0, particle="electron")
        {'frequency_Hz': 139.9624, ...}
    """
    try:
        import astropy.units as u
        import plasmapy.formulary as pf

        B = B_nT * u.nT
        p = _parse_particle(particle)

        omega = pf.gyrofrequency(B, particle=p, signed=False)
        f_hz = float((omega / (2 * math.pi * u.rad)).to(u.Hz).value)
        omega_rad_s = float(omega.to(u.rad / u.s).value)

        return {
            "frequency_Hz": round(f_hz, 4),
            "angular_frequency_rad_s": round(omega_rad_s, 4),
            "particle": particle,
            "B_nT": B_nT,
            "period_s": round(1.0 / f_hz, 6) if f_hz > 0 else None,
        }
    except Exception as e:
        return {"error": str(e)}

debye_length async

debye_length(n_cm3: float, T_eV: float) -> dict

Compute electron Debye length.

Parameters:

Name	Type	Description	Default
n_cm3	float	Electron number density in cm⁻³	required
T_eV	float	Electron temperature in eV	required

Returns dict with Debye length in km and meters.

Example

await debye_length(n_cm3=10.0, T_eV=12.0) {'debye_length_m': 8.143475, 'debye_length_km': 0.008143475, 'inputs': {'n_cm3': 10.0, 'T_eV': 12.0}}

Source code in helioai/tools/plasmapy_tools.py

async def debye_length(n_cm3: float, T_eV: float) -> dict:
    """Compute electron Debye length.

    Args:
        n_cm3: Electron number density in cm⁻³
        T_eV:  Electron temperature in eV

    Returns dict with Debye length in km and meters.

    Example:
        >>> await debye_length(n_cm3=10.0, T_eV=12.0)
        {'debye_length_m': 8.143475, 'debye_length_km': 0.008143475,
         'inputs': {'n_cm3': 10.0, 'T_eV': 12.0}}
    """
    try:
        import astropy.units as u
        import plasmapy.formulary as pf

        n = n_cm3 * u.cm**-3
        T = T_eV * u.eV

        lambda_D = pf.Debye_length(T, n)
        lambda_m = float(lambda_D.to(u.m).value)
        lambda_km = float(lambda_D.to(u.km).value)

        return {
            "debye_length_m": round(lambda_m, 6),
            "debye_length_km": round(lambda_km, 9),
            "inputs": {"n_cm3": n_cm3, "T_eV": T_eV},
        }
    except Exception as e:
        return {"error": str(e)}

alfven_speed async

alfven_speed(B_nT: float, n_cm3: float, mass_amu: float = 1.0) -> dict

Compute Alfvén speed.

Parameters:

Name	Type	Description	Default
B_nT	float	Magnetic field magnitude in nT	required
n_cm3	float	Ion number density in cm⁻³	required
mass_amu	float	Ion mass in atomic mass units (default 1.0 = proton)	1.0

Returns dict with Alfvén speed in km/s.

Example

await alfven_speed(B_nT=5.0, n_cm3=5.0) {'alfven_speed_km_s': 48.937, 'alfven_speed_m_s': 48936.9, 'inputs': {'B_nT': 5.0, 'n_cm3': 5.0, 'mass_amu': 1.0}, ...}

Source code in helioai/tools/plasmapy_tools.py

async def alfven_speed(B_nT: float, n_cm3: float, mass_amu: float = 1.0) -> dict:
    """Compute Alfvén speed.

    Args:
        B_nT:      Magnetic field magnitude in nT
        n_cm3:     Ion number density in cm⁻³
        mass_amu:  Ion mass in atomic mass units (default 1.0 = proton)

    Returns dict with Alfvén speed in km/s.

    Example:
        >>> await alfven_speed(B_nT=5.0, n_cm3=5.0)
        {'alfven_speed_km_s': 48.937, 'alfven_speed_m_s': 48936.9,
         'inputs': {'B_nT': 5.0, 'n_cm3': 5.0, 'mass_amu': 1.0}, ...}
    """
    try:
        import astropy.constants as const
        import astropy.units as u
        import plasmapy.formulary as pf

        B = B_nT * u.nT
        n = n_cm3 * u.cm**-3
        from plasmapy.particles import CustomParticle

        ion = CustomParticle(mass=mass_amu * const.u, charge=1 * const.e.si)

        V_A = pf.Alfven_speed(B, n, ion=ion)
        va_km_s = float(V_A.to(u.km / u.s).value)

        return {
            "alfven_speed_km_s": round(va_km_s, 3),
            "alfven_speed_m_s": round(va_km_s * 1000, 1),
            "inputs": {"B_nT": B_nT, "n_cm3": n_cm3, "mass_amu": mass_amu},
            "note": "Typical solar wind: 40-80 km/s. Magnetosphere: 100-1000 km/s.",
        }
    except Exception as e:
        return {"error": str(e)}

inertial_length async

inertial_length(n_cm3: float, particle: str = 'proton') -> dict

Compute ion or electron inertial length (skin depth).

Parameters:

Name	Type	Description	Default
n_cm3	float	Number density in cm⁻³	required
particle	str	'proton' or 'electron' (default: proton)	'proton'

Returns dict with inertial length in km and meters.

Example

await inertial_length(n_cm3=5.0) {'inertial_length_km': 101.8354, 'inertial_length_m': 101835.35, 'particle': 'proton', 'inputs': {'n_cm3': 5.0}}

Source code in helioai/tools/plasmapy_tools.py

async def inertial_length(n_cm3: float, particle: str = "proton") -> dict:
    """Compute ion or electron inertial length (skin depth).

    Args:
        n_cm3:    Number density in cm⁻³
        particle: 'proton' or 'electron' (default: proton)

    Returns dict with inertial length in km and meters.

    Example:
        >>> await inertial_length(n_cm3=5.0)
        {'inertial_length_km': 101.8354, 'inertial_length_m': 101835.35,
         'particle': 'proton', 'inputs': {'n_cm3': 5.0}}
    """
    try:
        import astropy.units as u
        import plasmapy.formulary as pf

        n = n_cm3 * u.cm**-3
        p = _parse_particle(particle)

        d = pf.inertial_length(n, particle=p)
        d_km = float(d.to(u.km).value)
        d_m = float(d.to(u.m).value)

        return {
            "inertial_length_km": round(d_km, 4),
            "inertial_length_m": round(d_m, 2),
            "particle": particle,
            "inputs": {"n_cm3": n_cm3},
        }
    except Exception as e:
        return {"error": str(e)}

power_spectrum async

power_spectrum(values: list[float], dt_s: float, nperseg: int | None = None) -> dict

Compute power spectral density using Welch's method.

Parameters:

Name	Type	Description	Default
values	list[float]	Time series as a list of floats	required
dt_s	float	Sampling interval in seconds	required
nperseg	int | None	Samples per FFT segment (default: min(256, len(values)//4))	None

Returns dict with frequencies (Hz), PSD values, peak frequency, and export-ready summary for LLM interpretation.

Example

import math wave = [math.sin(2 * math.pi * 0.1 * i) for i in range(512)] psd = await power_spectrum(wave, dt_s=1.0) psd["peak_frequency_Hz"], psd["peak_period_s"] (0.101562, 9.846)

Source code in helioai/tools/plasmapy_tools.py

async def power_spectrum(
    values: list[float],
    dt_s: float,
    nperseg: int | None = None,
) -> dict:
    """Compute power spectral density using Welch's method.

    Args:
        values:  Time series as a list of floats
        dt_s:    Sampling interval in seconds
        nperseg: Samples per FFT segment (default: min(256, len(values)//4))

    Returns dict with frequencies (Hz), PSD values, peak frequency,
    and export-ready summary for LLM interpretation.

    Example:
        >>> import math
        >>> wave = [math.sin(2 * math.pi * 0.1 * i) for i in range(512)]
        >>> psd = await power_spectrum(wave, dt_s=1.0)
        >>> psd["peak_frequency_Hz"], psd["peak_period_s"]
        (0.101562, 9.846)
    """
    try:
        import numpy as np
        from scipy import signal

        arr = np.asarray(values, dtype=float)
        arr[~np.isfinite(arr)] = np.nan
        arr[np.abs(arr) >= 1e30] = np.nan
        n_original = len(arr)
        valid_mask = ~np.isnan(arr)
        arr = arr[valid_mask]
        n_dropped = n_original - len(arr)
        if len(arr) < 8:
            return {"error": f"Need at least 8 finite samples, got {len(arr)}"}

        fs = 1.0 / dt_s
        seg = nperseg or min(256, len(arr) // 4)
        seg = max(seg, 8)

        freqs, psd = signal.welch(arr, fs=fs, nperseg=seg)

        peak_idx = int(np.argmax(psd[1:])) + 1
        peak_freq = float(freqs[peak_idx])
        peak_power = float(psd[peak_idx])

        return {
            "frequencies_Hz": [round(f, 6) for f in freqs.tolist()],
            "psd": [round(p, 8) for p in psd.tolist()],
            "peak_frequency_Hz": round(peak_freq, 6),
            "peak_period_s": round(1.0 / peak_freq, 3) if peak_freq > 0 else None,
            "peak_power": round(peak_power, 8),
            "n_points": len(arr),
            "n_original": n_original,
            "n_dropped": n_dropped,
            "gap_fraction": round(n_dropped / max(n_original, 1), 4),
            "fs_Hz": round(fs, 6),
            "freq_resolution_Hz": round(freqs[1] - freqs[0], 8) if len(freqs) > 1 else None,
        }
    except Exception as e:
        return {"error": str(e)}

Sandbox

helioai.tools.sandbox

Python sandbox: execute user/LLM-generated code in an isolated subprocess.

Security model

• Runs in a fresh subprocess (separate memory, no shared globals)
• Hard timeout (default 30s) — kills the process if exceeded
• stdout/stderr captured and returned
• No network isolation (speasy needs network access) — trust LLM-generated code

Pre-imports available in sandbox: speasy, plasmapy, numpy, scipy, matplotlib, astropy Figures are saved to a temp directory; paths are returned (not base64). Use export(name, array) to share numerical data with the LLM.

run_python async

run_python(code: str, timeout: float = 60.0, _plot_dir: str | None = None, _run_idx: int | None = None) -> dict

Execute Python code in an isolated subprocess.

Parameters:

Name	Type	Description	Default
code	str	Python source code to execute. Has access to speasy (spz), plasmapy (pf), numpy (np), scipy, matplotlib (Agg — plt.show() saves to disk), astropy units (u). Call export(name, array) to share numerical results with the LLM.	required
timeout	float	maximum execution time in seconds — clamped to _MAX_TIMEOUT_S	60.0
_plot_dir	str | None	injected by the agent loop — workspace dir for this run. Not exposed in the LLM tool schema.	None

Returns dict with

• stdout: captured text output
• stderr: captured errors/warnings
• figure_paths: list of absolute paths to saved PNG files
• exports: dict of named numerical summaries (from export() calls)
• error: error message if execution failed

Example

await run_python( ... "import numpy as np\n" ... "export('rms', np.sqrt(np.mean(np.arange(8) ** 2)))\n" ... "print('done')" ... ) {'stdout': 'done', 'exports': {'rms': {'mean': 4.1833..., 'min': 4.1833..., 'max': 4.1833..., 'std': 0.0, 'n_finite': 1, 'n_nan': 0, ...}}, 'figure_paths': [], 'error': None, ...}

Source code in helioai/tools/sandbox.py

async def run_python(
    code: str, timeout: float = 60.0, _plot_dir: str | None = None, _run_idx: int | None = None
) -> dict:
    """Execute Python code in an isolated subprocess.

    Args:
        code: Python source code to execute. Has access to speasy (spz), plasmapy (pf),
              numpy (np), scipy, matplotlib (Agg — plt.show() saves to disk),
              astropy units (u).
              Call export(name, array) to share numerical results with the LLM.
        timeout: maximum execution time in seconds — clamped to _MAX_TIMEOUT_S
        _plot_dir: injected by the agent loop — workspace dir for this run.
                   Not exposed in the LLM tool schema.

    Returns dict with:
        - stdout: captured text output
        - stderr: captured errors/warnings
        - figure_paths: list of absolute paths to saved PNG files
        - exports: dict of named numerical summaries (from export() calls)
        - error: error message if execution failed

    Example:
        >>> await run_python(
        ...     "import numpy as np\\n"
        ...     "export('rms', np.sqrt(np.mean(np.arange(8) ** 2)))\\n"
        ...     "print('done')"
        ... )
        {'stdout': 'done', 'exports': {'rms': {'mean': 4.1833..., 'min': 4.1833...,
         'max': 4.1833..., 'std': 0.0, 'n_finite': 1, 'n_nan': 0, ...}},
         'figure_paths': [], 'error': None, ...}
    """
    timeout = min(timeout, _MAX_TIMEOUT_S)
    if _plot_dir is None:
        from helioai.workspace import get_run_dir_for_sandbox

        _plot_dir = get_run_dir_for_sandbox()
    run_idx = _run_idx if _run_idx is not None else 0
    plot_dir = _plot_dir
    from helioai.logging_config import get_logger as _get_logger

    _get_logger(__name__).info("sandbox_plot_dir", plot_dir=plot_dir, run_idx=run_idx)
    code_file = Path(plot_dir, f"code_{run_idx}.py")
    dedented_code = textwrap.dedent(code)
    code_file.write_text(dedented_code, encoding="utf-8")
    n_lines = len(dedented_code.splitlines())
    plot_dir_line = f"__sandbox_plot_dir = {plot_dir!r}\n__sandbox_run_idx = {run_idx!r}\n"
    full_code = (
        plot_dir_line + _SANDBOX_PREAMBLE + textwrap.dedent(code) + "\n" + _SANDBOX_POSTAMBLE
    )

    cmd = _build_sandbox_cmd(plot_dir, full_code)
    using_bwrap = cmd[0].endswith("bwrap") if cmd else False

    try:
        if using_bwrap:
            _seed_speasy_inventory(plot_dir)
            sandbox_env = _sandbox_env(home=plot_dir)
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=sandbox_env,
                start_new_session=True,
            )
        else:
            _warn_if_not_isolated()
            sandbox_env = _sandbox_env()
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=sandbox_env,
                start_new_session=True,
                preexec_fn=_preexec_fn(),
                cwd=plot_dir,  # same working directory as the bwrap path's --chdir
            )
        try:
            stdout_bytes, stderr_bytes = await asyncio.wait_for(proc.communicate(), timeout=timeout)
        except TimeoutError:
            _kill_proc_tree(proc)
            stdout_bytes, stderr_bytes = await proc.communicate()
            return {
                "error": f"Execution timed out after {timeout}s",
                "stdout": stdout_bytes.decode("utf-8", errors="replace")[-2000:],
                "stderr": stderr_bytes.decode("utf-8", errors="replace")[-2000:],
            }

        stdout = stdout_bytes.decode("utf-8", errors="replace")
        stderr = stderr_bytes.decode("utf-8", errors="replace")

        figure_paths: list[str] = []
        exports: dict = {}
        cards: list[dict] = []
        clean_stdout_lines: list[str] = []
        for line in stdout.splitlines():
            if line.startswith("__HELIOAI_RESULT__"):
                try:
                    payload = json.loads(line[len("__HELIOAI_RESULT__") :])
                    figure_paths = payload.get("figure_paths", [])
                    exports = payload.get("exports", {})
                    cards = payload.get("cards", [])
                except json.JSONDecodeError:
                    pass
            else:
                clean_stdout_lines.append(line)

        _MAX_STDOUT = 4000
        clean_stdout = "\n".join(clean_stdout_lines).strip()
        if len(clean_stdout) > _MAX_STDOUT:
            clean_stdout = (
                clean_stdout[:_MAX_STDOUT]
                + f"\n[stdout truncated — {len(clean_stdout)} chars total; use export() for numerical data]"
            )

        if proc.returncode != 0:
            agent_stderr = _rewrite_traceback(stderr.strip())
            return {
                "error": _error_summary(agent_stderr, proc.returncode),
                "stdout": clean_stdout,
                "stderr": agent_stderr,
                "figure_paths": figure_paths,
                "exports": exports,
                "cards": cards,
                "code_path": str(code_file),
                "n_lines": n_lines,
            }

        return {
            "stdout": clean_stdout,
            "stderr": stderr.strip() if stderr.strip() else None,
            "figure_paths": figure_paths,
            "n_figures": len(figure_paths),
            "exports": exports,
            "cards": cards,
            "code_path": str(code_file),
            "n_lines": n_lines,
        }

    except Exception as e:
        return {"error": f"Sandbox error: {e}"}

Sandbox helpers

Available inside run_python, and re-emitted into exported notebooks.

helioai.tools.sandbox_helpers

Standalone physics helpers importable inside the sandbox.

MUST NOT import anything from helioai.* — the sandbox masks .env and strips the environment, so helioai.config would fail fast at import time.

Boundary models are clean-room implementations from the published papers: - Shue et al. (1998), JGR 103, 17691, doi:10.1029/98JA01103 - Jelinek et al. (2012), JGR 117, A05208, doi:10.1029/2011JA017252 Coordinate transforms wrap geopack (MIT, Tsyganenko models port).

transform_coords

transform_coords(time, vectors, frm: str = 'gse', to: str = 'gsm') -> np.ndarray

Rotate vectors between geocentric frames: gse, gsm, sm, geo, mag, gei.

time: ISO string(s), datetime(s), numpy datetime64 or epoch seconds (UTC); vectors: shape (3,) or (N, 3). Returns the same shape. Per-point geopack.recalc — fine up to ~1e4 points.

Example

t = np.array(["2015-03-17T04:00:00"], dtype="datetime64[s]") transform_coords(t, np.array([[10.0, 0.0, 0.0]]), "gse", "gsm") array([[10., 0., 0.]]) # the X axis is shared by GSE and GSM

Source code in helioai/tools/sandbox_helpers.py

def transform_coords(time, vectors, frm: str = "gse", to: str = "gsm") -> np.ndarray:
    """Rotate vectors between geocentric frames: gse, gsm, sm, geo, mag, gei.

    time: ISO string(s), datetime(s), numpy datetime64 or epoch seconds (UTC);
    vectors: shape (3,) or (N, 3). Returns the same shape.
    Per-point geopack.recalc — fine up to ~1e4 points.

    Example:
        >>> t = np.array(["2015-03-17T04:00:00"], dtype="datetime64[s]")
        >>> transform_coords(t, np.array([[10.0, 0.0, 0.0]]), "gse", "gsm")
        array([[10., 0., 0.]])   # the X axis is shared by GSE and GSM
    """
    from geopack import geopack as gp

    frm, to = frm.lower(), to.lower()
    to_gsm = {
        "gsm": lambda x, y, z: (x, y, z),
        "gse": lambda x, y, z: gp.gsmgse(x, y, z, -1),
        "sm": lambda x, y, z: gp.smgsm(x, y, z, 1),
        "geo": lambda x, y, z: gp.geogsm(x, y, z, 1),
        "mag": lambda x, y, z: gp.geogsm(*gp.geomag(x, y, z, -1), 1),
        "gei": lambda x, y, z: gp.geogsm(*gp.geigeo(x, y, z, 1), 1),
    }
    from_gsm = {
        "gsm": lambda x, y, z: (x, y, z),
        "gse": lambda x, y, z: gp.gsmgse(x, y, z, 1),
        "sm": lambda x, y, z: gp.smgsm(x, y, z, -1),
        "geo": lambda x, y, z: gp.geogsm(x, y, z, -1),
        "mag": lambda x, y, z: gp.geomag(*gp.geogsm(x, y, z, -1), 1),
        "gei": lambda x, y, z: gp.geigeo(*gp.geogsm(x, y, z, -1), -1),
    }
    if frm not in to_gsm or to not in from_gsm:
        raise ValueError(f"unknown frame: {frm!r} or {to!r} — use one of {sorted(to_gsm)}")

    vec = np.asarray(vectors, dtype=float)
    single = vec.ndim == 1
    vec = np.atleast_2d(vec)
    if vec.shape[1] != 3:
        raise ValueError(f"vectors must be (N, 3), got {vec.shape}")
    ts = _epoch_seconds(time)
    if ts.size == 1:
        ts = np.full(vec.shape[0], ts[0])
    if ts.size != vec.shape[0]:
        raise ValueError(f"{ts.size} times for {vec.shape[0]} vectors")

    out = np.empty_like(vec)
    for i in range(vec.shape[0]):
        gp.recalc(ts[i])
        out[i] = from_gsm[to](*to_gsm[frm](*vec[i]))
    return out[0] if single else out

mp_shue1998

mp_shue1998(pdyn_nPa: float, bz_nT: float, theta_deg=None)

Shue et al. (1998) magnetopause: r = r0 * (2 / (1 + cos(theta)))**alpha.

r0 = (10.22 + 1.29tanh(0.184(Bz + 8.14))) * Pdyn(-1/6.6) alpha = (0.58 - 0.007Bz) * (1 + 0.024ln(Pdyn)) Returns (theta_deg, r_RE); theta defaults to 0..170 deg. theta is the angle from the Earth-Sun line, r in Earth radii (aberrated GSE). Reference: Shue et al. (1998), JGR 103, 17691, doi:10.1029/98JA01103.

Example

theta, r = mp_shue1998(2.0, -5.0) # Pdyn=2 nPa, Bz=-5 nT round(float(r[0]), 2), round(float(r[90]), 2) (9.81, 15.13) # standoff and flank, in R_E

Source code in helioai/tools/sandbox_helpers.py

def mp_shue1998(pdyn_nPa: float, bz_nT: float, theta_deg=None):
    """Shue et al. (1998) magnetopause: r = r0 * (2 / (1 + cos(theta)))**alpha.

    r0 = (10.22 + 1.29*tanh(0.184*(Bz + 8.14))) * Pdyn**(-1/6.6)
    alpha = (0.58 - 0.007*Bz) * (1 + 0.024*ln(Pdyn))
    Returns (theta_deg, r_RE); theta defaults to 0..170 deg. theta is the
    angle from the Earth-Sun line, r in Earth radii (aberrated GSE).
    Reference: Shue et al. (1998), JGR 103, 17691, doi:10.1029/98JA01103.

    Example:
        >>> theta, r = mp_shue1998(2.0, -5.0)   # Pdyn=2 nPa, Bz=-5 nT
        >>> round(float(r[0]), 2), round(float(r[90]), 2)
        (9.81, 15.13)                            # standoff and flank, in R_E
    """
    theta = (
        np.linspace(0.0, 170.0, 171)
        if theta_deg is None
        else np.atleast_1d(np.asarray(theta_deg, dtype=float))
    )
    r0 = (10.22 + 1.29 * np.tanh(0.184 * (bz_nT + 8.14))) * pdyn_nPa ** (-1.0 / 6.6)
    alpha = (0.58 - 0.007 * bz_nT) * (1.0 + 0.024 * np.log(pdyn_nPa))
    r = r0 * (2.0 / (1.0 + np.cos(np.radians(theta)))) ** alpha
    return theta, r

bs_jelinek2012

bs_jelinek2012(pdyn_nPa: float, theta_deg=None)

Jelinek et al. (2012) bow shock: parabola rho^2 = 4R(R - x) / lam^2.

R = 15.02 * Pdyn**(-1/6.55) is the subsolar standoff (RE), lam = 1.17. Solved in polar form r(theta) from the Earth-Sun line; returns (theta_deg, r_RE) with NaN where the parabola has no solution. theta defaults to 0..120 deg. Reference: Jelinek et al. (2012), JGR 117, A05208, doi:10.1029/2011JA017252.

Example

theta, r = bs_jelinek2012(2.0) # Pdyn=2 nPa round(float(r[0]), 2) 13.51 # subsolar standoff, in R_E

Source code in helioai/tools/sandbox_helpers.py

def bs_jelinek2012(pdyn_nPa: float, theta_deg=None):
    """Jelinek et al. (2012) bow shock: parabola rho^2 = 4*R*(R - x) / lam^2.

    R = 15.02 * Pdyn**(-1/6.55) is the subsolar standoff (RE), lam = 1.17.
    Solved in polar form r(theta) from the Earth-Sun line; returns
    (theta_deg, r_RE) with NaN where the parabola has no solution.
    theta defaults to 0..120 deg.
    Reference: Jelinek et al. (2012), JGR 117, A05208, doi:10.1029/2011JA017252.

    Example:
        >>> theta, r = bs_jelinek2012(2.0)   # Pdyn=2 nPa
        >>> round(float(r[0]), 2)
        13.51                                 # subsolar standoff, in R_E
    """
    theta = (
        np.linspace(0.0, 120.0, 121)
        if theta_deg is None
        else np.atleast_1d(np.asarray(theta_deg, dtype=float))
    )
    lam = 1.17
    big_r = 15.02 * pdyn_nPa ** (-1.0 / 6.55)
    s, c = np.sin(np.radians(theta)), np.cos(np.radians(theta))
    with np.errstate(divide="ignore", invalid="ignore"):
        r = 2.0 * big_r * (np.sqrt(c**2 + lam**2 * s**2) - c) / (lam**2 * s**2)
    on_axis = np.isclose(s, 0.0)
    r = np.where(on_axis, np.where(c > 0, big_r, np.nan), r)
    return theta, r

Recipes

helioai.tools.recipes

Derived-recipe tools — list and load scientific Python recipes.

Recipes live in data/recipes/ as .py files with a YAML comment header: # name: theta_bn # description: Compute the shock normal angle theta_Bn from upstream/downstream B. # inputs: B_up (nT vec), B_dn (nT vec) # outputs: theta_bn_deg

list_recipes() returns the catalogue; load_recipe(name) returns the source code.

list_recipes async

list_recipes() -> dict

List all available derived recipes with their name and description.

Returns dict with 'recipes' list (sorted by name). Each entry has 'name', 'description', 'inputs', 'outputs' (when present in header). Returns {"recipes": []} when the recipes directory does not exist.

Example

await list_recipes() {'recipes': [{'name': 'fill_values', 'description': '...'}, {'name': 'mvab', ...}, {'name': 'rankine_hugoniot', ...}, ...]}

Source code in helioai/tools/recipes.py

async def list_recipes() -> dict:
    """List all available derived recipes with their name and description.

    Returns dict with 'recipes' list (sorted by name). Each entry has
    'name', 'description', 'inputs', 'outputs' (when present in header).
    Returns {"recipes": []} when the recipes directory does not exist.

    Example:
        >>> await list_recipes()
        {'recipes': [{'name': 'fill_values', 'description': '...'},
                     {'name': 'mvab', ...}, {'name': 'rankine_hugoniot', ...}, ...]}
    """
    try:
        recipes_dir = settings.recipes.recipes_dir
        if not recipes_dir.exists():
            return {"recipes": []}
        entries = []
        for path in sorted(recipes_dir.glob("*.py")):
            try:
                text = path.read_text(encoding="utf-8")
                meta = _parse_header(text)
                entry = {"name": meta.get("name", path.stem)}
                for field in ("description", "inputs", "outputs"):
                    if field in meta:
                        entry[field] = meta[field]
                entries.append(entry)
            except OSError as exc:
                log.warning("recipe_read_error", path=str(path), error=str(exc))
        return {"recipes": entries}
    except Exception as e:
        return {"error": str(e)}

load_recipe async

load_recipe(name: str) -> dict

Load the source code of a named recipe.

Parameters:

Name	Type	Description	Default
name	str	Recipe name without .py extension (e.g. 'theta_bn').	required

Returns dict with 'name' and 'code'. Returns {'error': ...} when not found or when the name contains path-traversal characters.

Example

await load_recipe("theta_bn")

Source code in helioai/tools/recipes.py

async def load_recipe(name: str) -> dict:
    """Load the source code of a named recipe.

    Args:
        name: Recipe name without .py extension (e.g. 'theta_bn').

    Returns dict with 'name' and 'code'. Returns {'error': ...} when not found
    or when the name contains path-traversal characters.

    Example:
        >>> await load_recipe("theta_bn")
        {'name': 'theta_bn', 'code': '# name: theta_bn\\n# description: Compute the shock...'}
    """
    try:
        if not name or any(c in name for c in ("/", "\\", "..")):
            return {"error": f"invalid recipe name: {name!r}"}
        recipes_dir = settings.recipes.recipes_dir.resolve()
        candidate = (recipes_dir / f"{name}.py").resolve()
        if not candidate.is_relative_to(recipes_dir):
            return {"error": f"recipe {name!r} not found"}
        if not candidate.is_file():
            return {"error": f"recipe {name!r} not found"}
        code = candidate.read_text(encoding="utf-8")
        meta = _parse_header(code)
        return {
            "name": meta.get("name", name),
            "code": code,
            "metadata": meta,
        }
    except Exception as e:
        return {"error": str(e)}

Literature

helioai.tools.literature

NASA ADS literature search — find_papers tool.

find_papers async

find_papers(query: str, max_results: int = 5, year_start: int | None = None, year_end: int | None = None, sort: str = 'relevance', _transport: AsyncBaseTransport | None = None) -> dict

Search NASA ADS for papers relevant to an event, parameter or method.

Requires ADS_API_TOKEN; without it the tool returns an error rather than raising, so the agent can tell the user what is missing.

Example

await find_papers("interplanetary shock Rankine-Hugoniot multi-spacecraft", ... max_results=2) {'query': '...', 'papers': [ {'title': 'Multiple spacecraft observations of interplanetary shocks: ...', 'authors': 'Russell, C. T. et al.', 'year': '1983', 'bibcode': '1983JGR....88.9941R', 'doi': '10.1029/JA088iA12p09941', 'citations': 72, 'abstract': '...'}, ...], 'note': '...'}

Source code in helioai/tools/literature.py

async def find_papers(
    query: str,
    max_results: int = 5,
    year_start: int | None = None,
    year_end: int | None = None,
    sort: str = "relevance",
    _transport: httpx.AsyncBaseTransport | None = None,
) -> dict:
    """Search NASA ADS for papers relevant to an event, parameter or method.

    Requires `ADS_API_TOKEN`; without it the tool returns an error rather than
    raising, so the agent can tell the user what is missing.

    Example:
        >>> await find_papers("interplanetary shock Rankine-Hugoniot multi-spacecraft",
        ...                   max_results=2)
        {'query': '...', 'papers': [
         {'title': 'Multiple spacecraft observations of interplanetary shocks: ...',
          'authors': 'Russell, C. T. et al.', 'year': '1983',
          'bibcode': '1983JGR....88.9941R', 'doi': '10.1029/JA088iA12p09941',
          'citations': 72, 'abstract': '...'}, ...], 'note': '...'}
    """
    token = settings.literature.ads_token
    if not token:
        return {
            "error": (
                "ADS_API_TOKEN is not set — get a free key at "
                "https://ui.adsabs.harvard.edu/user/settings/token and add it to .env"
            )
        }

    q = query
    if year_start or year_end:
        q = f"{query} year:{year_start or ''}-{year_end or ''}"
    params = {
        "q": q,
        "fl": _FIELDS,
        "rows": min(max(max_results, 1), _MAX_ROWS),
        "sort": _SORTS.get(sort, _SORTS["relevance"]),
    }

    async with httpx.AsyncClient(timeout=15, transport=_transport) as client:
        try:
            resp = await client.get(
                _ADS_URL, params=params, headers={"Authorization": f"Bearer {token}"}
            )
        except httpx.HTTPError as e:
            return {"error": f"ADS request failed: {e}"}

    if resp.status_code != 200:
        return {"error": f"ADS returned HTTP {resp.status_code}: {resp.text[:200]}"}

    docs = resp.json().get("response", {}).get("docs", [])
    papers = [_slim(d) for d in docs]
    log.info("find_papers", query=query, n_results=len(papers))
    return {
        "query": q,
        "papers": papers,
        "note": (
            "Cite as: Authors (year), bibcode. "
            "Full record: https://ui.adsabs.harvard.edu/abs/<bibcode>"
        ),
    }

Registry

helioai.tools.registry

Tool registry: wraps Python async functions with JSON Schema metadata.

The agent loop calls registry.call_tool(name, args) and never imports tool modules directly, which keeps the dependency surface small and makes sub-agent tool whitelisting trivial.

Tool dataclass

A registered tool: an async function plus the schema shown to the model.

Source code in helioai/tools/registry.py

@dataclass
class Tool:
    """A registered tool: an async function plus the schema shown to the model."""

    name: str
    description: str
    parameters: dict  # JSON Schema object
    func: Callable[..., Coroutine[Any, Any, Any]]

ToolRegistry

Maps tool names to async functions and their JSON Schemas.

The agent loop dispatches through here and never imports tool modules, which keeps the dependency surface small and makes per-role whitelisting trivial.

Source code in helioai/tools/registry.py

class ToolRegistry:
    """Maps tool names to async functions and their JSON Schemas.

    The agent loop dispatches through here and never imports tool modules, which
    keeps the dependency surface small and makes per-role whitelisting trivial.
    """

    def __init__(self) -> None:
        self._tools: dict[str, Tool] = {}

    def register(self, name: str, description: str, parameters: dict) -> Callable:
        """Decorator that registers an async function as a tool."""

        def decorator(func: Callable) -> Callable:
            self._tools[name] = Tool(
                name=name,
                description=description,
                parameters=parameters,
                func=func,
            )
            return func

        return decorator

    def list_tool_defs(self, only: set[str] | None = None) -> list[ToolDef]:
        """Return tool definitions for the model.

        Args:
        only: Restrict to these names — how sub-agent whitelists are applied.
        None returns every registered tool.
        """
        tools = self._tools.values()
        if only is not None:
            tools = [t for t in tools if t.name in only]
        return [
            ToolDef(name=t.name, description=t.description, parameters=t.parameters) for t in tools
        ]

    async def call_tool(
        self, name: str, arguments: dict | None, *, trusted: dict | None = None
    ) -> str:
        """Invoke a tool and return its JSON-serialized result string.

        `arguments` is caller-supplied (LLM/MCP) and may not carry private
        `_*` keys. `trusted` is framework-injected (e.g. the sandbox output
        dir) and bypasses that guard.
        """
        if name not in self._tools:
            return json.dumps({"error": f"unknown tool {name!r}"})
        if arguments and any(k.startswith("_") for k in arguments):
            bad = sorted(k for k in arguments if k.startswith("_"))
            return json.dumps({"error": f"rejected private argument(s): {bad}"})
        try:
            result = await self._tools[name].func(**{**(arguments or {}), **(trusted or {})})
            if isinstance(result, str):
                return result
            return json.dumps(result, ensure_ascii=False, default=str)
        except Exception as e:
            return json.dumps({"error": str(e)})

    def __contains__(self, name: str) -> bool:
        return name in self._tools

register

register(name: str, description: str, parameters: dict) -> Callable

Decorator that registers an async function as a tool.

Source code in helioai/tools/registry.py

def register(self, name: str, description: str, parameters: dict) -> Callable:
    """Decorator that registers an async function as a tool."""

    def decorator(func: Callable) -> Callable:
        self._tools[name] = Tool(
            name=name,
            description=description,
            parameters=parameters,
            func=func,
        )
        return func

    return decorator

list_tool_defs

list_tool_defs(only: set[str] | None = None) -> list[ToolDef]

Return tool definitions for the model.

Args: only: Restrict to these names — how sub-agent whitelists are applied. None returns every registered tool.

Source code in helioai/tools/registry.py

def list_tool_defs(self, only: set[str] | None = None) -> list[ToolDef]:
    """Return tool definitions for the model.

    Args:
    only: Restrict to these names — how sub-agent whitelists are applied.
    None returns every registered tool.
    """
    tools = self._tools.values()
    if only is not None:
        tools = [t for t in tools if t.name in only]
    return [
        ToolDef(name=t.name, description=t.description, parameters=t.parameters) for t in tools
    ]

call_tool async

call_tool(name: str, arguments: dict | None, *, trusted: dict | None = None) -> str

Invoke a tool and return its JSON-serialized result string.

arguments is caller-supplied (LLM/MCP) and may not carry private _* keys. trusted is framework-injected (e.g. the sandbox output dir) and bypasses that guard.

Source code in helioai/tools/registry.py

async def call_tool(
    self, name: str, arguments: dict | None, *, trusted: dict | None = None
) -> str:
    """Invoke a tool and return its JSON-serialized result string.

    `arguments` is caller-supplied (LLM/MCP) and may not carry private
    `_*` keys. `trusted` is framework-injected (e.g. the sandbox output
    dir) and bypasses that guard.
    """
    if name not in self._tools:
        return json.dumps({"error": f"unknown tool {name!r}"})
    if arguments and any(k.startswith("_") for k in arguments):
        bad = sorted(k for k in arguments if k.startswith("_"))
        return json.dumps({"error": f"rejected private argument(s): {bad}"})
    try:
        result = await self._tools[name].func(**{**(arguments or {}), **(trusted or {})})
        if isinstance(result, str):
            return result
        return json.dumps(result, ensure_ascii=False, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})

MCP client

helioai.tools.mcp_client

Generic MCP client — mounts remote MCP server tools into the ToolRegistry.

Configured via HELIOAI_MCP_SERVERS, a JSON object keyed by server alias: {"ads": {"url": "https://.../mcp", "headers": {"Authorization": "Bearer x"}}, "alphaxiv": {"command": "npx", "args": ["-y", "mcp-remote", "https://api.alphaxiv.org/mcp/v1"]}}

Remote tools are registered as "_" and proxied with a fresh connection per call: the CLI runs one asyncio.run() per query, so a persistent session would not outlive a single request anyway.

discover_and_register async

discover_and_register() -> list[str]

Connect to every configured MCP server and register its tools.

Idempotent: only the first call does any work. Servers that are unreachable are logged and skipped rather than fatal — a dead remote must not stop HelioAI from starting.

Returns:

Type	Description
list[str]	The names under which remote tools were registered, prefixed by alias.

Source code in helioai/tools/mcp_client.py

async def discover_and_register() -> list[str]:
    """Connect to every configured MCP server and register its tools.

    Idempotent: only the first call does any work. Servers that are unreachable
    are logged and skipped rather than fatal — a dead remote must not stop
    HelioAI from starting.

    Returns:
        The names under which remote tools were registered, prefixed by alias.
    """
    global _discovered
    if _discovered:
        return []
    _discovered = True

    registered: list[str] = []
    for alias, spec in _server_specs().items():
        try:
            # `spec` is bound as a default: the closure would otherwise read the
            # loop variable at await time, so every server after the first would
            # be discovered against the last spec if this were ever deferred.
            async def _list(spec=spec):
                async with _session(spec) as session:
                    return await session.list_tools()

            tools = (await asyncio.wait_for(_list(), timeout=_DISCOVER_TIMEOUT_S)).tools
        except Exception as e:
            log.warning("mcp_server_unreachable", server=alias, error=str(e))
            continue

        for t in tools:
            name = _safe_name(alias, t.name)
            if name in registry:
                log.warning("mcp_tool_collision_skipped", server=alias, tool=name)
                continue
            schema = t.inputSchema or {"type": "object", "properties": {}}
            props = schema.get("properties") or {}
            if any(k.startswith("_") for k in props):
                log.warning("mcp_tool_private_args", server=alias, tool=name)
            registry.register(
                name=name,
                description=f"[{alias} MCP] {t.description or t.name}",
                parameters=schema,
            )(_make_proxy(alias, spec, t.name))
            registered.append(name)
        log.info("mcp_server_mounted", server=alias, n_tools=len(tools))
    return registered