Core

Agent loop

helioai.core.agent_loop

The agent decision loop.

Given a user message and a session id, run the LLM in a tool-using loop: the LLM may emit tool calls, execute them via the ToolRegistry, feed the results back, and iterate until the LLM produces a final text reply (or we hit the safety cap).

Two consumption modes share the same generator core (stream_chat): - chat() → collects all events, returns a single ChatResult - stream_chat() → async generator, yields one event dict per step

Event shapes

tool_call {turn, name, arguments} tool_result {turn, name, summary} sub_agent_start {task_id, role, description} sub_agent_end {task_id, role, summary, n_iterations, error} plan {title, steps} skill_loaded {name} reply {text} provenance {matched, contradicted, derived, unsourced, details} done {n_iterations} error {message}

ChatResult dataclass

Final outcome of a non-streaming chat() call.

Source code in helioai/core/agent_loop.py

@dataclass
class ChatResult:
    """Final outcome of a non-streaming `chat()` call."""

    reply: str
    n_iterations: int
    artifacts: list[dict] = field(default_factory=list)
    events: list[dict] = field(default_factory=list)

build_lead_system_prompt

build_lead_system_prompt(restricted: bool) -> str

Return the lead agent system prompt.

restricted=True (default / public): appends the scope guardrail so the LLM auto-refuses off-topic requests. restricted=False (dev token supplied): base prompt only, full access.

Source code in helioai/core/agent_loop.py

def build_lead_system_prompt(restricted: bool) -> str:
    """Return the lead agent system prompt.

    restricted=True (default / public): appends the scope guardrail so the LLM
    auto-refuses off-topic requests.
    restricted=False (dev token supplied): base prompt only, full access.
    """
    if restricted:
        return SYSTEM_PROMPT + "\n\n" + SCOPE_GUARDRAIL
    return SYSTEM_PROMPT

stream_chat async

stream_chat(llm_client: LLMClient, user_id: str, session_id: str, user_text: str, *, restricted: bool = True) -> AsyncIterator[dict]

Run one conversational turn of the agent and stream its progress as events.

This is the package's central API: every interface (CLI, web SSE, Jupyter, MCP) is a consumer of this generator. History is loaded from and persisted to the session store keyed by (user_id, session_id), so consecutive calls with the same ids continue the same conversation.

Parameters:

Name	Type	Description	Default
llm_client	LLMClient	Provider client from build_llm_client().	required
user_id	str	Storage namespace — workspaces and profiles live under it.	required
session_id	str	Conversation id; reuse it to continue, mint one to start fresh.	required
user_text	str	The user's message for this turn.	required
restricted	bool	True (default) appends the heliophysics scope guardrail; False (dev token) exposes the base prompt only.	True

Yields:

Type	Description
AsyncIterator[dict]	Dicts with an "event" key, one of: reply (streamed answer text),
AsyncIterator[dict]	tool_call / tool_result, artifact (figure, parameter card, code),
AsyncIterator[dict]	plan, sub_agent_start / sub_agent_end, skill_loaded,
AsyncIterator[dict]	figure_review, provenance, invalid_ids, recipe_bypassed,
AsyncIterator[dict]	error, and finally done.

Example

llm = build_llm_client() async for ev in stream_chat(llm, "cli", "my-session", "IMF Bz at L1 today?"): ... if ev["event"] == "reply": ... print(ev["text"], end="")

Source code in helioai/core/agent_loop.py

async def stream_chat(
    llm_client: LLMClient,
    user_id: str,
    session_id: str,
    user_text: str,
    *,
    restricted: bool = True,
) -> AsyncIterator[dict]:
    """Run one conversational turn of the agent and stream its progress as events.

    This is the package's central API: every interface (CLI, web SSE, Jupyter,
    MCP) is a consumer of this generator. History is loaded from and persisted
    to the session store keyed by (user_id, session_id), so consecutive calls
    with the same ids continue the same conversation.

    Args:
        llm_client: Provider client from `build_llm_client()`.
        user_id: Storage namespace — workspaces and profiles live under it.
        session_id: Conversation id; reuse it to continue, mint one to start fresh.
        user_text: The user's message for this turn.
        restricted: True (default) appends the heliophysics scope guardrail;
            False (dev token) exposes the base prompt only.

    Yields:
        Dicts with an `"event"` key, one of: `reply` (streamed answer text),
        `tool_call` / `tool_result`, `artifact` (figure, parameter card, code),
        `plan`, `sub_agent_start` / `sub_agent_end`, `skill_loaded`,
        `figure_review`, `provenance`, `invalid_ids`, `recipe_bypassed`,
        `error`, and finally `done`.

    Example:
        >>> llm = build_llm_client()
        >>> async for ev in stream_chat(llm, "cli", "my-session", "IMF Bz at L1 today?"):
        ...     if ev["event"] == "reply":
        ...         print(ev["text"], end="")
    """
    import helioai.workspace as _ws

    _ws_token = _ws.set_session(session_id)
    _user_token = _ws.set_user(user_id)

    history = store.get_or_create(user_id, session_id)
    history.append(Message(role="user", content=user_text))

    # What this run exported, kept so the answer can be checked against the recipe shelf
    # the same way a sub-agent's is. The lead does its own physics often enough that
    # leaving it unchecked was the hole, not an edge case.
    run_artifacts: list[dict] = []

    existing_dir = store.get_workspace_dir(user_id, session_id)
    if existing_dir:
        _label_token = _ws.set_label(existing_dir)
    else:
        label = _ws.make_session_label(user_text, session_id)
        store.save(user_id, session_id, history)
        store.set_workspace_dir(user_id, session_id, label)
        _label_token = _ws.set_label(label)

    tools = registry.list_tool_defs() + _INTERNAL_TOOLS + [task_tool_def()]
    log.info("agent_tools_listed", count=len(tools), tools=[t.name for t in tools])

    effective_prompt = build_lead_system_prompt(restricted)
    profile = _load_user_profile(user_id)
    if profile:
        effective_prompt = f"{effective_prompt}\n\n## User profile\n{profile}"

    try:
        for i in range(settings.agent.max_iterations):
            turn = i + 1
            log.info("llm_call_start", turn=turn, n_messages=len(history))
            t0 = time.monotonic()
            history[:] = strip_orphan_tool_calls(history)
            response = await llm_client.chat(
                compact_history(history), tools, system_prompt=effective_prompt
            )
            log.info(
                "llm_call_end",
                turn=turn,
                duration_ms=int((time.monotonic() - t0) * 1000),
                has_tool_calls=bool(response.tool_calls),
            )
            history.append(response)

            if not response.tool_calls:
                store.save(user_id, session_id, history)
                # No tool calls AND no text is a failed turn, not an answer. It was
                # being yielded as an empty reply, so the caller saw the request
                # simply produce nothing — silence indistinguishable from success.
                # The usual cause is the output budget: on Azure, reasoning tokens
                # are drawn from the same allowance, so a long generation can spend
                # it entirely on reasoning and emit no content at all.
                if not (response.content or "").strip():
                    provider, cap = _active_output_budget()
                    log.warning("empty_llm_response", turn=turn, provider=provider, cap=cap)
                    yield {
                        "event": "error",
                        "data": {
                            "message": (
                                "the model returned neither text nor a tool call. This is "
                                "usually the output token budget running out — set "
                                f"HELIOAI_MAX_OUTPUT_TOKENS above {cap} (the current "
                                f"{provider} limit) and retry, or ask for a shorter answer."
                            )
                        },
                    }
                    yield {"event": "done", "data": {"n_iterations": turn}}
                    return
                final_text, bogus, bypassed = check_answer(response.content, history, run_artifacts)
                yield {"event": "reply", "data": {"text": final_text}}
                if bogus:
                    log.warning("lead_invented_ids", ids=bogus)
                    yield {"event": "invalid_ids", "data": {"ids": bogus}}
                if bypassed:
                    log.warning("lead_recipe_bypassed", recipes=bypassed)
                    yield {"event": "recipe_bypassed", "data": {"recipes": bypassed}}
                for ev in _provenance_events(final_text):
                    yield ev
                yield {"event": "done", "data": {"n_iterations": turn}}
                return

            if response.content and response.content.strip():
                yield {"event": "reply", "data": {"text": response.content}}
                for ev in _provenance_events(response.content):
                    yield ev

            for tc in response.tool_calls:
                log.info("tool_call_issued", turn=turn, tool=tc.name)
                yield {
                    "event": "tool_call",
                    "data": {"turn": turn, "name": tc.name, "arguments": tc.arguments},
                }

                sub_end_event: dict | None = None

                try:
                    if tc.name == TASK_TOOL_NAME:
                        args = tc.arguments or {}
                        sub_role = args.get("agent_role", "")
                        sub_desc = args.get("description", "")
                        yield {
                            "event": "sub_agent_start",
                            "data": {
                                "task_id": tc.id,
                                "role": sub_role,
                                "description": sub_desc[:200],
                            },
                        }
                        async for sub_ev in stream_subagent(
                            role=sub_role,
                            description=sub_desc,
                            parent_session_id=session_id,
                            user_id=user_id,
                            llm_client=llm_client,
                            task_id=tc.id,
                        ):
                            if sub_ev["event"] == "sub_agent_end":
                                end_data = sub_ev["data"]
                                result = json.dumps(
                                    {
                                        # First, deliberately: keys at the tail are the ones
                                        # _summarize_tool_result drops when a stale result is
                                        # trimmed, and the measured values must outlive the prose.
                                        "findings": end_data.get("findings", {}),
                                        "summary": end_data.get("summary", ""),
                                        "n_iterations": end_data.get("n_iterations", 0),
                                        "artifacts": end_data.get("artifacts", []),
                                        "error": end_data.get("error"),
                                    }
                                )
                                sub_end_event = {
                                    "task_id": tc.id,
                                    "role": sub_role,
                                    "summary": end_data.get("summary", "")[:200],
                                    "n_iterations": end_data.get("n_iterations", 0),
                                    "error": end_data.get("error"),
                                }
                            else:
                                yield sub_ev
                    elif tc.name in _INTERNAL_TOOL_NAMES:
                        result = _dispatch_internal_tool(tc.name, tc.arguments)
                    else:
                        result = await registry.call_tool(
                            tc.name, tc.arguments, trusted=inject_run_python_args(tc.name)
                        )
                except Exception as e:
                    log.exception("tool_call_failed", turn=turn, tool=tc.name)
                    result = json.dumps({"error": str(e)})
                    if tc.name == TASK_TOOL_NAME:
                        sub_end_event = {
                            "task_id": tc.id,
                            "role": sub_role if "sub_role" in locals() else "",
                            "summary": "",
                            "n_iterations": 0,
                            "error": str(e),
                        }

                result, figure_verdict = await maybe_review(tc.name, result)
                if figure_verdict:
                    yield {"event": "figure_review", "data": {"turn": turn, "text": figure_verdict}}

                for ev in emit_post_tool_events(tc.name, result, tool_result_extra={"turn": turn}):
                    if ev["event"] == "artifact":
                        run_artifacts.append(ev["data"])
                    yield ev
                if sub_end_event is not None:
                    yield {"event": "sub_agent_end", "data": sub_end_event}
                if tc.name == "present_plan":
                    try:
                        plan = json.loads(result)
                        yield {
                            "event": "plan",
                            "data": {
                                "title": plan.get("title", ""),
                                "steps": plan.get("steps", []),
                            },
                        }
                    except (ValueError, TypeError):
                        pass

                history.append(
                    Message(
                        role="tool",
                        tool_call_id=tc.id,
                        content=_history_tool_result(tc.name, result),
                    )
                )

        log.warning("agent_loop_capped", max_iterations=settings.agent.max_iterations)
        store.save(user_id, session_id, history)
        yield {
            "event": "error",
            "data": {"message": f"agent loop exceeded {settings.agent.max_iterations} iterations"},
        }

    except asyncio.CancelledError:
        store.save(user_id, session_id, strip_orphan_tool_calls(history))
        raise

    except Exception:
        log.exception("agent_loop_crashed", turn=locals().get("turn"))
        store.save(user_id, session_id, strip_orphan_tool_calls(history))
        raise

    finally:
        _ws.reset_session(_ws_token)
        _ws.reset_label(_label_token)
        _ws.reset_user(_user_token)

chat async

chat(llm_client: LLMClient, user_id: str, session_id: str, user_text: str, *, restricted: bool = True) -> ChatResult

Run one agent turn to completion and return the final result.

Non-streaming wrapper over stream_chat — same arguments, same session semantics — for callers that want the answer, not the progress feed (Jupyter magic, scripts, tests).

Returns:

Type	Description
ChatResult	ChatResult with reply (final text), n_iterations (LLM turns used),
ChatResult	artifacts (figures, parameter cards, code) and events (full trace).

Example

result = await chat(build_llm_client(), "cli", "my-session", ... "Plasma beta for B=5 nT, n=10 cm^-3, T=20 eV?") print(result.reply) # final assistant text (model-dependent) len(result.artifacts) # figures / parameter cards / code produced

Source code in helioai/core/agent_loop.py

async def chat(
    llm_client: LLMClient,
    user_id: str,
    session_id: str,
    user_text: str,
    *,
    restricted: bool = True,
) -> ChatResult:
    """Run one agent turn to completion and return the final result.

    Non-streaming wrapper over `stream_chat` — same arguments, same session
    semantics — for callers that want the answer, not the progress feed
    (Jupyter magic, scripts, tests).

    Returns:
        ChatResult with `reply` (final text), `n_iterations` (LLM turns used),
        `artifacts` (figures, parameter cards, code) and `events` (full trace).

    Example:
        >>> result = await chat(build_llm_client(), "cli", "my-session",
        ...                     "Plasma beta for B=5 nT, n=10 cm^-3, T=20 eV?")
        >>> print(result.reply)        # final assistant text (model-dependent)
        >>> len(result.artifacts)      # figures / parameter cards / code produced
    """
    artifacts: list[dict] = []
    events: list[dict] = []
    reply = ""
    n_iters = 0
    error_msg: str | None = None

    async for ev in stream_chat(llm_client, user_id, session_id, user_text, restricted=restricted):
        name, data = ev["event"], ev["data"]
        if name == "reply":
            reply = data.get("text", "")
        elif name == "done":
            n_iters = data.get("n_iterations", 0)
        elif name == "artifact":
            artifacts.append(data)
        elif name == "tool_call":
            events.append(
                {
                    "turn": data["turn"],
                    "type": "tool_call",
                    "tool": data["name"],
                    "arguments": data.get("arguments", {}),
                }
            )
        elif name == "tool_result":
            events.append(
                {
                    "turn": data["turn"],
                    "type": "tool_result",
                    "tool": data["name"],
                    "summary": data.get("summary", ""),
                }
            )
        elif name == "error":
            error_msg = data.get("message", "unknown agent error")

    if error_msg is not None:
        raise RuntimeError(error_msg)
    return ChatResult(reply=reply, n_iterations=n_iters, artifacts=artifacts, events=events)

Sub-agents

helioai.core.sub_agents

Sub-agents for delegating focused heliophysics subtasks.

Each role declares a tool whitelist, a system addon, and an optional set of skills auto-loaded into the sub's system prompt. Sub-agents run in isolation with a fresh context — the lead's history is invisible to them.

SubAgentRole dataclass

A specialised agent: its prompt, its tool whitelist and its turn budget.

allowed_tools is enforced, not advisory — a role calling outside its set gets an error naming what it may use, and the tool is never dispatched.

Source code in helioai/core/sub_agents.py

@dataclass(frozen=True)
class SubAgentRole:
    """A specialised agent: its prompt, its tool whitelist and its turn budget.

    `allowed_tools` is enforced, not advisory — a role calling outside its set
    gets an error naming what it may use, and the tool is never dispatched.
    """

    name: str
    description: str
    system_addon: str
    allowed_tools: tuple[str, ...]
    max_turns: int = 5
    auto_load_skills: tuple[str, ...] = ()

SubAgentResult dataclass

What a finished sub-agent hands back to the lead agent.

Source code in helioai/core/sub_agents.py

@dataclass
class SubAgentResult:
    """What a finished sub-agent hands back to the lead agent."""

    summary: str = ""
    artifacts: list[dict] = field(default_factory=list)
    n_iterations: int = 0
    error: str | None = None

task_tool_def

task_tool_def() -> ToolDef

Build the task tool definition offered to the lead agent.

Deliberately not registered in the ToolRegistry: the agent loop intercepts task and spawns a sub-agent instead of dispatching a function.

Returns:

Type	Description
ToolDef	A ToolDef whose agent_role enum lists every available role.

Source code in helioai/core/sub_agents.py

def task_tool_def() -> ToolDef:
    """Build the `task` tool definition offered to the lead agent.

    Deliberately not registered in the ToolRegistry: the agent loop intercepts
    `task` and spawns a sub-agent instead of dispatching a function.

    Returns:
        A ToolDef whose `agent_role` enum lists every available role.
    """
    role_lines = "\n".join(f"  - `{r.name}`: {r.description}" for r in AGENT_ROLES.values())
    return ToolDef(
        name=TASK_TOOL_NAME,
        # Routing rules live in the lead's system prompt, not here. Stating them in both
        # places let them drift into contradiction — this text used to call
        # `parameter_hunter` "required" for unknown ids while the prompt said never to run
        # it first, and the model resolved the conflict differently from one run to the
        # next: three delegations, three, then none, on the same notebook.
        description=(
            "Spawn a specialist sub-agent for ONE focused subtask. "
            "The sub runs in isolation (empty context) — pre-resolve every fact "
            "(param ids, ISO times, missions) inside `description`.\n\nRoles:\n" + role_lines
        ),
        parameters={
            "type": "object",
            "properties": {
                "description": {
                    "type": "string",
                    "description": "Self-contained task description (1-3 sentences) with all needed facts.",
                },
                "agent_role": {
                    "type": "string",
                    "enum": sorted(AGENT_ROLES.keys()),
                    "description": "Sub-agent role to spawn.",
                },
            },
            "required": ["description", "agent_role"],
        },
    )

stream_subagent async

stream_subagent(role: str, description: str, *, parent_session_id: str, user_id: str, llm_client: LLMClient, task_id: str | None = None) -> AsyncIterator[dict]

Async generator that runs a sub-agent and yields progress events.

Yields the same event types as stream_chat (tool_call, tool_result, skill_loaded, artifact) enriched with sub_agent_ctx={role, task_id}, then a final sub_agent_end event carrying summary/artifacts/n_iterations/error.

summary is the sub-agent's whole deliverable and is emitted in full: it becomes the lead's tool result, so anything cut here is a measurement the lead can no longer report and will be tempted to invent. Callers that display it truncate on their own side. Reaching the turn cap is reported as an error, not as a result, for the same reason — a lead handed a capped run has nothing to summarise.

Source code in helioai/core/sub_agents.py

async def stream_subagent(
    role: str,
    description: str,
    *,
    parent_session_id: str,
    user_id: str,
    llm_client: LLMClient,
    task_id: str | None = None,
) -> AsyncIterator[dict]:
    """Async generator that runs a sub-agent and yields progress events.

    Yields the same event types as stream_chat (tool_call, tool_result,
    skill_loaded, artifact) enriched with sub_agent_ctx={role, task_id},
    then a final sub_agent_end event carrying summary/artifacts/n_iterations/error.

    `summary` is the sub-agent's whole deliverable and is emitted in full: it becomes
    the lead's tool result, so anything cut here is a measurement the lead can no
    longer report and will be tempted to invent. Callers that display it truncate on
    their own side. Reaching the turn cap is reported as an `error`, not as a result,
    for the same reason — a lead handed a capped run has nothing to summarise.
    """
    import helioai.workspace as _ws

    if task_id is None:
        task_id = uuid.uuid4().hex[:8]
    ctx = {"role": role, "task_id": task_id}

    # The user has to be bound too, not just the session: `_root()` resolves the workspace
    # under `current_user()`, which defaults to "web". It works today only because the lead
    # already bound it, so a sub-agent started out of band wrote to the wrong user's files.
    _ws_token = _ws.set_session(parent_session_id)
    _user_token = _ws.set_user(user_id) if user_id else None

    if role not in AGENT_ROLES:
        known = ", ".join(sorted(AGENT_ROLES))
        yield {
            "event": "sub_agent_end",
            "data": {
                "task_id": task_id,
                "role": role,
                "summary": "",
                "n_iterations": 0,
                "error": f"unknown agent_role {role!r}. Known: {known}",
                "artifacts": [],
            },
        }
        return

    role_cfg = AGENT_ROLES[role]

    structlog.contextvars.bind_contextvars(
        parent_session_id=parent_session_id,
        sub_role=role,
        sub_task_id=task_id,
    )

    try:
        system_prompt, skills_loaded = _build_system_prompt(role_cfg)
        for skill_name in skills_loaded:
            yield {"event": "skill_loaded", "data": {"name": skill_name, "sub_agent_ctx": ctx}}

        allowed = set(role_cfg.allowed_tools)
        tools = registry.list_tool_defs(only=allowed)

        log.info(
            "subagent_start",
            role=role,
            description=description[:200],
            allowed_tools=sorted(allowed),
            max_turns=role_cfg.max_turns,
        )

        history: list[Message] = [Message(role="user", content=_with_inventory(description))]
        artifacts: list[dict] = []
        final_text = ""
        n_iters = 0
        t0 = time.monotonic()
        capped = False

        for i in range(role_cfg.max_turns):
            n_iters = i + 1
            tc_choice = "required" if i == 0 else "auto"
            response = await llm_client.chat(
                compact_history(history), tools, system_prompt=system_prompt, tool_choice=tc_choice
            )
            history.append(response)

            if not response.tool_calls:
                final_text = response.content or ""
                break

            for tc in response.tool_calls:
                log.info("tool_call_issued", turn=n_iters, tool=tc.name, sub_role=role)
                yield {
                    "event": "tool_call",
                    "data": {
                        "turn": n_iters,
                        "name": tc.name,
                        "arguments": tc.arguments,
                        "sub_agent_ctx": ctx,
                    },
                }

                if tc.name not in allowed:
                    log.warning("subagent_tool_denied", role=role, tool=tc.name)
                    result = json.dumps(
                        {
                            "error": f"tool {tc.name!r} not available to {role!r}. Allowed: {sorted(allowed)}"
                        }
                    )
                else:
                    result = await registry.call_tool(
                        tc.name, tc.arguments, trusted=inject_run_python_args(tc.name)
                    )

                result, figure_verdict = await maybe_review(tc.name, result)
                if figure_verdict:
                    yield {
                        "event": "figure_review",
                        "data": {"turn": n_iters, "text": figure_verdict, "sub_agent_ctx": ctx},
                    }

                for ev in emit_post_tool_events(
                    tc.name,
                    result,
                    tool_result_extra={"turn": n_iters, "sub_agent_ctx": ctx},
                    common_extra={"sub_agent_ctx": ctx},
                ):
                    if ev["event"] == "artifact":
                        artifacts.append(
                            {k: v for k, v in ev["data"].items() if k != "sub_agent_ctx"}
                        )
                    yield ev

                history.append(
                    Message(
                        role="tool",
                        tool_call_id=tc.id,
                        content=_history_tool_result(tc.name, result),
                    )
                )
        else:
            capped = True
            final_text = f"(sub-agent {role!r} reached its {role_cfg.max_turns}-turn cap)"

        log.info(
            "subagent_end",
            role=role,
            n_iterations=n_iters,
            duration_ms=int((time.monotonic() - t0) * 1000),
            capped=capped,
            n_artifacts=len(artifacts),
        )

        final_text, bogus, bypassed_recipes = check_answer(final_text, history, artifacts)
        if bogus:
            log.warning("subagent_invented_ids", role=role, ids=bogus)
            yield {
                "event": "invalid_ids",
                "data": {"ids": bogus, "sub_agent_ctx": ctx},
            }

        if bypassed_recipes:
            log.warning("subagent_recipe_bypassed", role=role, recipes=bypassed_recipes)
            yield {
                "event": "recipe_bypassed",
                "data": {"recipes": bypassed_recipes, "sub_agent_ctx": ctx},
            }

        yield {
            "event": "sub_agent_end",
            "data": {
                "task_id": task_id,
                "role": role,
                "findings": _findings(artifacts),
                "summary": final_text,
                "n_iterations": n_iters,
                "error": final_text if capped else None,
                "artifacts": artifacts,
            },
        }

    except Exception as e:
        log.exception("subagent_error", role=role, task_id=task_id)
        yield {
            "event": "sub_agent_end",
            "data": {
                "task_id": task_id,
                "role": role,
                "findings": _findings(artifacts) if "artifacts" in dir() else {},
                "summary": "",
                "n_iterations": n_iters if "n_iters" in dir() else 0,
                "error": str(e),
                "artifacts": [],
            },
        }

    finally:
        _ws.reset_session(_ws_token)
        if _user_token is not None:
            _ws.reset_user(_user_token)
        structlog.contextvars.unbind_contextvars("parent_session_id", "sub_role", "sub_task_id")

run_subagent async

run_subagent(role: str, description: str, *, parent_session_id: str, user_id: str, llm_client: LLMClient, task_id: str | None = None) -> SubAgentResult

Run a sub-agent to completion and return only its outcome.

Non-streaming wrapper over stream_subagent — same arguments.

Parameters:

Name	Type	Description	Default
role	str	One of the whitelisted roles (parameter_hunter, data_analyst, plasma_physicist, librarian).	required
description	str	The task handed to the sub-agent, in natural language.	required
parent_session_id	str	The lead conversation this run belongs to.	required
user_id	str	Storage namespace of that conversation.	required
llm_client	LLMClient	Provider client shared with the lead.	required
task_id	str | None	Optional id echoed in events, for UI correlation.	None

Returns:

Type	Description
SubAgentResult	SubAgentResult — summary (full final text), artifacts,
SubAgentResult	n_iterations, and error (set when the run failed or hit its cap).

Source code in helioai/core/sub_agents.py

async def run_subagent(
    role: str,
    description: str,
    *,
    parent_session_id: str,
    user_id: str,
    llm_client: LLMClient,
    task_id: str | None = None,
) -> SubAgentResult:
    """Run a sub-agent to completion and return only its outcome.

    Non-streaming wrapper over `stream_subagent` — same arguments.

    Args:
        role: One of the whitelisted roles (parameter_hunter, data_analyst,
            plasma_physicist, librarian).
        description: The task handed to the sub-agent, in natural language.
        parent_session_id: The lead conversation this run belongs to.
        user_id: Storage namespace of that conversation.
        llm_client: Provider client shared with the lead.
        task_id: Optional id echoed in events, for UI correlation.

    Returns:
        SubAgentResult — `summary` (full final text), `artifacts`,
        `n_iterations`, and `error` (set when the run failed or hit its cap).
    """
    async for ev in stream_subagent(
        role=role,
        description=description,
        parent_session_id=parent_session_id,
        user_id=user_id,
        llm_client=llm_client,
        task_id=task_id,
    ):
        if ev["event"] == "sub_agent_end":
            d = ev["data"]
            return SubAgentResult(
                summary=d.get("summary", ""),
                artifacts=d.get("artifacts", []),
                n_iterations=d.get("n_iterations", 0),
                error=d.get("error"),
            )
    return SubAgentResult(error="stream_subagent yielded no sub_agent_end")

Tool execution

Shared between the lead loop and the sub-agent loop.

helioai.core.tool_exec

Shared tool-execution helpers for the agent loops.

Both the lead loop (agent_loop.stream_chat) and the sub-agent loop (sub_agents.stream_subagent) run the same tool-call mechanics: inject the sandbox run dir for run_python, summarise the result, detect skill loads, and extract renderable artifacts. Keeping that logic here — imported by both loops — prevents the two copies from drifting apart (a real bug source, see the session-13 _extract_artifact list/dict regression).

This module imports neither agent_loop nor sub_agents, so there is no cycle.

compact_history

compact_history(messages: list, keep_full: int = 2) -> list

Return a copy of messages where tool-result messages older than the last keep_full are summarized. The most recent results stay verbatim (the next LLM call usually needs them); older ones — already consumed — are trimmed so context does not grow unbounded over a long session. Persisted history is left untouched; only the per-call payload shrinks.

ponytail: fixed window N=2; widen keep_full if a case regresses on stale results.

Source code in helioai/core/tool_exec.py

def compact_history(messages: list, keep_full: int = 2) -> list:
    """Return a copy of `messages` where tool-result messages older than the last
    `keep_full` are summarized. The most recent results stay verbatim (the next LLM
    call usually needs them); older ones — already consumed — are trimmed so context
    does not grow unbounded over a long session. Persisted history is left untouched;
    only the per-call payload shrinks.

    ponytail: fixed window N=2; widen keep_full if a case regresses on stale results.
    """
    tool_idx = [i for i, m in enumerate(messages) if getattr(m, "role", None) == "tool"]
    if len(tool_idx) <= keep_full:
        return messages
    stale = set(tool_idx[:-keep_full])
    return [
        replace(m, content=_summarize_tool_result(m.content, max_chars=300))
        if i in stale and m.content
        else m
        for i, m in enumerate(messages)
    ]

inject_run_python_args

inject_run_python_args(name: str) -> dict

Trusted per-run sandbox args (_plot_dir/_run_idx) for run_python.

Passed via call_tool(..., trusted=...) so they bypass the private-arg guard that rejects LLM/MCP-supplied _* overrides. Empty for any other tool.

Source code in helioai/core/tool_exec.py

def inject_run_python_args(name: str) -> dict:
    """Trusted per-run sandbox args (_plot_dir/_run_idx) for run_python.

    Passed via `call_tool(..., trusted=...)` so they bypass the private-arg
    guard that rejects LLM/MCP-supplied `_*` overrides. Empty for any other tool.
    """
    if name != "run_python":
        return {}
    import helioai.workspace as _ws

    sdir = _ws.get_session_dir()
    ridx = _ws.get_next_run_idx(sdir)
    return {"_plot_dir": str(sdir), "_run_idx": ridx}

emit_post_tool_events

emit_post_tool_events(name: str, result: str, *, tool_result_extra: dict | None = None, common_extra: dict | None = None) -> Iterator[dict]

Yield the events that follow a completed tool call.

Order is tool_result → (skill_loaded if load_skill) → artifact(s), matching what both loops emitted before this was factored out.

• tool_result_extra is merged into the tool_result event data (e.g. {turn}).
• common_extra is merged into skill_loaded and artifact event data (e.g. {sub_agent_ctx} for sub-agents).

Source code in helioai/core/tool_exec.py

def emit_post_tool_events(
    name: str,
    result: str,
    *,
    tool_result_extra: dict | None = None,
    common_extra: dict | None = None,
) -> Iterator[dict]:
    """Yield the events that follow a completed tool call.

    Order is `tool_result` → (`skill_loaded` if load_skill) → `artifact`(s),
    matching what both loops emitted before this was factored out.

    - `tool_result_extra` is merged into the tool_result event data (e.g. {turn}).
    - `common_extra` is merged into skill_loaded and artifact event data
      (e.g. {sub_agent_ctx} for sub-agents).
    """
    tool_result_extra = tool_result_extra or {}
    common_extra = common_extra or {}

    yield {
        "event": "tool_result",
        "data": {
            "name": name,
            "summary": _summarize_tool_result(result),
            **tool_result_extra,
        },
    }

    if name == "load_skill":
        try:
            payload = json.loads(result)
            if payload.get("body") and not payload.get("error"):
                yield {
                    "event": "skill_loaded",
                    "data": {
                        "name": payload.get("name", ""),
                        **common_extra,
                    },
                }
        except (ValueError, TypeError):
            pass

    for art in _extract_artifact(name, result):
        if art.get("kind") == "exports":
            ctx = common_extra.get("sub_agent_ctx") or {}
            provenance.record(
                art.get("values") or {},
                code_path=_code_path(result),
                agent=ctx.get("role") or "lead",
                task_id=ctx.get("task_id"),
                turn=tool_result_extra.get("turn"),
            )
        yield {"event": "artifact", "data": {**art, **common_extra}}

check_answer

check_answer(text: str, history: list, artifacts: list[dict]) -> tuple[str, list[str], list[dict]]

Confront a finished answer with the catalogue and the recipe shelf.

Both loops call this, which is the whole point of it living here. Both checks were written inside the sub-agent loop and stayed there, so a lead agent that did the physics itself — Acts III and IV of the showcase notebook, on a run where load_recipe was called zero times all session — was never checked at all. The detectors were not silent because the run was clean; they were silent because nothing called them.

Returns:

Type	Description
tuple[str, list[str], list[dict]]	The text (annotated when a check fires), the unknown ids, and the recipe flags.

Source code in helioai/core/tool_exec.py

def check_answer(
    text: str, history: list, artifacts: list[dict]
) -> tuple[str, list[str], list[dict]]:
    """Confront a finished answer with the catalogue and the recipe shelf.

    Both loops call this, which is the whole point of it living here. Both checks were
    written inside the sub-agent loop and stayed there, so a lead agent that did the
    physics itself — Acts III and IV of the showcase notebook, on a run where
    `load_recipe` was called zero times all session — was never checked at all. The
    detectors were not silent because the run was clean; they were silent because
    nothing called them.

    Returns:
        The text (annotated when a check fires), the unknown ids, and the recipe flags.
    """
    text, bogus = _flag_unknown_ids(text)
    text, bypassed = _flag_recipe_bypass(text, history, artifacts)
    return text, bogus, bypassed

Sessions

helioai.core.session

Conversation history keyed by (user_id, session_id), persisted to SQLite.

SessionStore

Conversation history keyed by (user_id, session_id), persisted to SQLite.

Histories are cached in memory per key and written back whole on save. Tests use a real database on tmp_path rather than a mock: a mocked store passed happily through a schema migration that broke production.

Example

store = SessionStore(tmp_path / "sessions.db") history = store.get_or_create("cli", "sess-1") # [] on first call history.append(Message(role="user", content="hello")) store.save("cli", "sess-1", history) store.get_or_create("cli", "sess-1")[0].role 'user'

Source code in helioai/core/session.py

class SessionStore:
    """Conversation history keyed by (user_id, session_id), persisted to SQLite.

    Histories are cached in memory per key and written back whole on `save`.
    Tests use a real database on `tmp_path` rather than a mock: a mocked store
    passed happily through a schema migration that broke production.

    Example:
        >>> store = SessionStore(tmp_path / "sessions.db")
        >>> history = store.get_or_create("cli", "sess-1")   # [] on first call
        >>> history.append(Message(role="user", content="hello"))
        >>> store.save("cli", "sess-1", history)
        >>> store.get_or_create("cli", "sess-1")[0].role
        'user'
    """

    def __init__(self, db_path: Path = DEFAULT_DB) -> None:
        self._db_path = db_path
        self._cache: dict[SessionKey, list[Message]] = {}
        self._lock = threading.Lock()
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_schema()

    def _init_schema(self) -> None:
        with self._connect() as conn:
            conn.executescript(_SCHEMA)
            try:
                conn.execute(_MIGRATE)
            except Exception:
                pass
            conn.commit()

    @contextmanager
    def _connect(self) -> Iterator[sqlite3.Connection]:
        conn = sqlite3.connect(self._db_path, check_same_thread=False)
        conn.execute("PRAGMA foreign_keys = ON")
        conn.execute("PRAGMA journal_mode = WAL")
        conn.execute("PRAGMA busy_timeout = 5000")
        try:
            yield conn
        finally:
            conn.close()

    def get_or_create(self, user_id: str, session_id: str) -> list[Message]:
        """Return the cached history for a session, loading it from disk if needed."""
        key: SessionKey = (user_id, session_id)
        with self._lock:
            if key in self._cache:
                return self._cache[key]
            history = self._load(user_id, session_id)
            self._cache[key] = history
            return history

    def _load(self, user_id: str, session_id: str) -> list[Message]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT role, content, tool_calls, tool_call_id "
                "FROM messages WHERE user_id = ? AND session_id = ? ORDER BY seq",
                (user_id, session_id),
            ).fetchall()
        return [
            Message(
                role=role,
                content=content or "",
                tool_calls=_load_tool_calls(tool_calls),
                tool_call_id=tool_call_id,
            )
            for role, content, tool_calls, tool_call_id in rows
        ]

    def save(self, user_id: str, session_id: str, history: list[Message]) -> None:
        """Replace a session's stored history.

        Args:
            user_id: Owner of the session.
            session_id: Session identifier.
            history: Full message list; it replaces whatever was stored.
        """
        with self._lock, self._connect() as conn:
            conn.execute(
                "INSERT INTO sessions(user_id, session_id) VALUES(?, ?) "
                "ON CONFLICT(user_id, session_id) DO UPDATE SET updated_at = julianday('now')",
                (user_id, session_id),
            )
            conn.execute(
                "DELETE FROM messages WHERE user_id = ? AND session_id = ?", (user_id, session_id)
            )
            conn.executemany(
                "INSERT INTO messages(user_id, session_id, seq, role, content, tool_calls, tool_call_id) "
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                [
                    (
                        user_id,
                        session_id,
                        i,
                        m.role,
                        m.content or "",
                        _dump_tool_calls(m.tool_calls),
                        m.tool_call_id,
                    )
                    for i, m in enumerate(history)
                ],
            )
            conn.commit()

    def reset(self, user_id: str, session_id: str) -> None:
        """Delete a session and its messages, and drop it from the cache."""
        with self._lock, self._connect() as conn:
            conn.execute(
                "DELETE FROM sessions WHERE user_id = ? AND session_id = ?",
                (user_id, session_id),
            )
            conn.commit()
        self._cache.pop((user_id, session_id), None)

    def set_workspace_dir(self, user_id: str, session_id: str, workspace_dir: str) -> None:
        """Record which workspace directory a session's artifacts live in."""
        with self._lock, self._connect() as conn:
            conn.execute(
                "UPDATE sessions SET workspace_dir = ? WHERE user_id = ? AND session_id = ?",
                (workspace_dir, user_id, session_id),
            )
            conn.commit()

    def get_workspace_dir(self, user_id: str, session_id: str) -> str | None:
        """Return a session's workspace directory label, or None."""
        with self._connect() as conn:
            row = conn.execute(
                "SELECT workspace_dir FROM sessions WHERE user_id = ? AND session_id = ?",
                (user_id, session_id),
            ).fetchone()
        return row[0] if row else None

    def workspace_dirs(self, user_id: str) -> set[str]:
        """All workspace dir labels owned by a user (for path-ownership checks)."""
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT workspace_dir FROM sessions "
                "WHERE user_id = ? AND workspace_dir IS NOT NULL",
                (user_id,),
            ).fetchall()
        return {r[0] for r in rows}

    def all_sessions(self, user_id: str) -> list[str]:
        """Return a user's session ids, most recently updated first."""
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT session_id FROM sessions WHERE user_id = ? "
                "ORDER BY updated_at DESC, rowid DESC",
                (user_id,),
            ).fetchall()
        return [r[0] for r in rows]

    def list_summaries(self, user_id: str, limit: int = 50) -> list[dict]:
        """Summarise a user's recent sessions for the history view.

        Args:
            user_id: Owner of the sessions.
            limit: Maximum number of sessions to return.

        Returns:
            Dicts with session_id, updated_at, first_message, n_messages and
            workspace_dir, most recent first.
        """
        from datetime import datetime

        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT s.session_id, s.updated_at AS jd,
                       (SELECT content FROM messages WHERE user_id = s.user_id
                          AND session_id = s.session_id AND role = 'user'
                          ORDER BY seq LIMIT 1) AS first_user,
                       (SELECT COUNT(*) FROM messages WHERE user_id = s.user_id
                          AND session_id = s.session_id) AS n_messages,
                       s.workspace_dir
                FROM sessions s
                WHERE s.user_id = ?
                ORDER BY s.updated_at DESC, s.rowid DESC LIMIT ?
                """,
                (user_id, limit),
            ).fetchall()
        out: list[dict] = []
        for session_id, jd, first_user, n_messages, workspace_dir in rows:
            preview = (first_user or "").strip().replace("\n", " ")
            if len(preview) > 80:
                preview = preview[:77] + "..."
            unix_ts = (jd - 2440587.5) * 86400
            iso = datetime.fromtimestamp(unix_ts, tz=UTC).isoformat().replace("+00:00", "Z")
            out.append(
                {
                    "session_id": session_id,
                    "first_message": preview,
                    "n_messages": n_messages,
                    "updated_at": iso,
                    "workspace_dir": workspace_dir,
                }
            )
        return out

get_or_create

get_or_create(user_id: str, session_id: str) -> list[Message]

Return the cached history for a session, loading it from disk if needed.

Source code in helioai/core/session.py

def get_or_create(self, user_id: str, session_id: str) -> list[Message]:
    """Return the cached history for a session, loading it from disk if needed."""
    key: SessionKey = (user_id, session_id)
    with self._lock:
        if key in self._cache:
            return self._cache[key]
        history = self._load(user_id, session_id)
        self._cache[key] = history
        return history

save

save(user_id: str, session_id: str, history: list[Message]) -> None

Replace a session's stored history.

Parameters:

Name	Type	Description	Default
user_id	str	Owner of the session.	required
session_id	str	Session identifier.	required
history	list[Message]	Full message list; it replaces whatever was stored.	required

Source code in helioai/core/session.py

def save(self, user_id: str, session_id: str, history: list[Message]) -> None:
    """Replace a session's stored history.

    Args:
        user_id: Owner of the session.
        session_id: Session identifier.
        history: Full message list; it replaces whatever was stored.
    """
    with self._lock, self._connect() as conn:
        conn.execute(
            "INSERT INTO sessions(user_id, session_id) VALUES(?, ?) "
            "ON CONFLICT(user_id, session_id) DO UPDATE SET updated_at = julianday('now')",
            (user_id, session_id),
        )
        conn.execute(
            "DELETE FROM messages WHERE user_id = ? AND session_id = ?", (user_id, session_id)
        )
        conn.executemany(
            "INSERT INTO messages(user_id, session_id, seq, role, content, tool_calls, tool_call_id) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            [
                (
                    user_id,
                    session_id,
                    i,
                    m.role,
                    m.content or "",
                    _dump_tool_calls(m.tool_calls),
                    m.tool_call_id,
                )
                for i, m in enumerate(history)
            ],
        )
        conn.commit()

reset

reset(user_id: str, session_id: str) -> None

Delete a session and its messages, and drop it from the cache.

Source code in helioai/core/session.py

def reset(self, user_id: str, session_id: str) -> None:
    """Delete a session and its messages, and drop it from the cache."""
    with self._lock, self._connect() as conn:
        conn.execute(
            "DELETE FROM sessions WHERE user_id = ? AND session_id = ?",
            (user_id, session_id),
        )
        conn.commit()
    self._cache.pop((user_id, session_id), None)

set_workspace_dir

set_workspace_dir(user_id: str, session_id: str, workspace_dir: str) -> None

Record which workspace directory a session's artifacts live in.

Source code in helioai/core/session.py

def set_workspace_dir(self, user_id: str, session_id: str, workspace_dir: str) -> None:
    """Record which workspace directory a session's artifacts live in."""
    with self._lock, self._connect() as conn:
        conn.execute(
            "UPDATE sessions SET workspace_dir = ? WHERE user_id = ? AND session_id = ?",
            (workspace_dir, user_id, session_id),
        )
        conn.commit()

get_workspace_dir

get_workspace_dir(user_id: str, session_id: str) -> str | None

Return a session's workspace directory label, or None.

Source code in helioai/core/session.py

def get_workspace_dir(self, user_id: str, session_id: str) -> str | None:
    """Return a session's workspace directory label, or None."""
    with self._connect() as conn:
        row = conn.execute(
            "SELECT workspace_dir FROM sessions WHERE user_id = ? AND session_id = ?",
            (user_id, session_id),
        ).fetchone()
    return row[0] if row else None

workspace_dirs

workspace_dirs(user_id: str) -> set[str]

All workspace dir labels owned by a user (for path-ownership checks).

Source code in helioai/core/session.py

def workspace_dirs(self, user_id: str) -> set[str]:
    """All workspace dir labels owned by a user (for path-ownership checks)."""
    with self._connect() as conn:
        rows = conn.execute(
            "SELECT workspace_dir FROM sessions "
            "WHERE user_id = ? AND workspace_dir IS NOT NULL",
            (user_id,),
        ).fetchall()
    return {r[0] for r in rows}

all_sessions

all_sessions(user_id: str) -> list[str]

Return a user's session ids, most recently updated first.

Source code in helioai/core/session.py

def all_sessions(self, user_id: str) -> list[str]:
    """Return a user's session ids, most recently updated first."""
    with self._connect() as conn:
        rows = conn.execute(
            "SELECT session_id FROM sessions WHERE user_id = ? "
            "ORDER BY updated_at DESC, rowid DESC",
            (user_id,),
        ).fetchall()
    return [r[0] for r in rows]

list_summaries

list_summaries(user_id: str, limit: int = 50) -> list[dict]

Summarise a user's recent sessions for the history view.

Parameters:

Name	Type	Description	Default
user_id	str	Owner of the sessions.	required
limit	int	Maximum number of sessions to return.	50

Returns:

Type	Description
list[dict]	Dicts with session_id, updated_at, first_message, n_messages and
list[dict]	workspace_dir, most recent first.

Source code in helioai/core/session.py

def list_summaries(self, user_id: str, limit: int = 50) -> list[dict]:
    """Summarise a user's recent sessions for the history view.

    Args:
        user_id: Owner of the sessions.
        limit: Maximum number of sessions to return.

    Returns:
        Dicts with session_id, updated_at, first_message, n_messages and
        workspace_dir, most recent first.
    """
    from datetime import datetime

    with self._connect() as conn:
        rows = conn.execute(
            """
            SELECT s.session_id, s.updated_at AS jd,
                   (SELECT content FROM messages WHERE user_id = s.user_id
                      AND session_id = s.session_id AND role = 'user'
                      ORDER BY seq LIMIT 1) AS first_user,
                   (SELECT COUNT(*) FROM messages WHERE user_id = s.user_id
                      AND session_id = s.session_id) AS n_messages,
                   s.workspace_dir
            FROM sessions s
            WHERE s.user_id = ?
            ORDER BY s.updated_at DESC, s.rowid DESC LIMIT ?
            """,
            (user_id, limit),
        ).fetchall()
    out: list[dict] = []
    for session_id, jd, first_user, n_messages, workspace_dir in rows:
        preview = (first_user or "").strip().replace("\n", " ")
        if len(preview) > 80:
            preview = preview[:77] + "..."
        unix_ts = (jd - 2440587.5) * 86400
        iso = datetime.fromtimestamp(unix_ts, tz=UTC).isoformat().replace("+00:00", "Z")
        out.append(
            {
                "session_id": session_id,
                "first_message": preview,
                "n_messages": n_messages,
                "updated_at": iso,
                "workspace_dir": workspace_dir,
            }
        )
    return out

strip_orphan_tool_calls

strip_orphan_tool_calls(history: list[Message]) -> list[Message]

Remove assistant tool_calls that have no matching tool response.

An interrupted generation (e.g. client disconnect mid-tool) can leave an assistant message with tool_calls but no corresponding tool messages in the history. Sending such a sequence to the LLM API causes a 400 error.

For each orphaned tool_call id: - If the assistant message has content too, keep the message but drop the orphaned tool_calls list entry (or clear it entirely if all are orphaned). - If the assistant message has no content and all its tool_calls are orphaned, drop the message entirely.

Source code in helioai/core/session.py

def strip_orphan_tool_calls(history: list[Message]) -> list[Message]:
    """Remove assistant tool_calls that have no matching tool response.

    An interrupted generation (e.g. client disconnect mid-tool) can leave an
    assistant message with tool_calls but no corresponding tool messages in the
    history.  Sending such a sequence to the LLM API causes a 400 error.

    For each orphaned tool_call id:
    - If the assistant message has content too, keep the message but drop the
      orphaned tool_calls list entry (or clear it entirely if all are orphaned).
    - If the assistant message has no content and all its tool_calls are
      orphaned, drop the message entirely.
    """
    answered: set[str] = {m.tool_call_id for m in history if m.tool_call_id}
    cleaned: list[Message] = []
    for m in history:
        if not m.tool_calls:
            cleaned.append(m)
            continue
        live_tcs = [tc for tc in m.tool_calls if tc.id in answered]
        if len(live_tcs) == len(m.tool_calls):
            cleaned.append(m)
        elif live_tcs:
            cleaned.append(Message(role=m.role, content=m.content, tool_calls=live_tcs))
        elif m.content:
            cleaned.append(Message(role=m.role, content=m.content))
        # else: drop the message entirely (no content, no answered tool_calls)
    return cleaned

Skills

helioai.core.skills_loader

Discover and serve markdown-defined skills to the agent.

Each skill lives in skills//SKILL.md with YAML frontmatter: --- name: parameter_hunter description: Find speasy parameter ids from a natural language query. when_to_use: User asks about a parameter without giving its exact id. allowed_tools: [search_parameters] --- # body of the procedure

SkillError

Bases: RuntimeError

Raised when a skill is missing, unreadable, or fails its path check.

Source code in helioai/core/skills_loader.py

class SkillError(RuntimeError):
    """Raised when a skill is missing, unreadable, or fails its path check."""

    pass

SkillMeta dataclass

Header of a skill, as listed to the agent before it loads the body.

Source code in helioai/core/skills_loader.py

@dataclass(frozen=True)
class SkillMeta:
    """Header of a skill, as listed to the agent before it loads the body."""

    name: str
    description: str
    when_to_use: str
    allowed_tools: tuple[str, ...]
    path: Path

load_index

load_index() -> str

Return the markdown index of available skills, for the agent to browse.

Source code in helioai/core/skills_loader.py

def load_index() -> str:
    """Return the markdown index of available skills, for the agent to browse."""
    skills = _discover()
    if not skills:
        return "(no skills available)"
    lines = ["| Skill | When to use |", "|---|---|"]
    for meta in skills.values():
        when = meta.when_to_use.replace("|", "/").replace("\n", " ")
        lines.append(f"| `{meta.name}` | {when} |")
    return "\n".join(lines)

load_skill

load_skill(name: str) -> str

Return a skill's full markdown body.

Parameters:

Name	Type	Description	Default
name	str	Skill name as listed by list_skill_names.	required

Returns:

Type	Description
str	The skill body, ready to append to a system prompt.

Raises:

Type	Description
SkillError	If the skill does not exist or the name escapes the skills directory.

Source code in helioai/core/skills_loader.py

def load_skill(name: str) -> str:
    """Return a skill's full markdown body.

    Args:
        name: Skill name as listed by `list_skill_names`.

    Returns:
        The skill body, ready to append to a system prompt.

    Raises:
        SkillError: If the skill does not exist or the name escapes the skills
            directory.
    """
    skills = _discover()
    if name not in skills:
        known = ", ".join(sorted(skills)) or "(none)"
        raise SkillError(f"unknown skill {name!r}. Known: {known}")
    text = skills[name].path.read_text(encoding="utf-8")
    _, body = _split_frontmatter(text)
    return body

list_skill_names

list_skill_names() -> list[str]

Return the names of all available skills.

Source code in helioai/core/skills_loader.py

def list_skill_names() -> list[str]:
    """Return the names of all available skills."""
    return list(_discover().keys())

Figure review

helioai.core.vision

Stateless vision side-call: review sandbox figures after run_python.

The image is sent once, outside the conversation; only the short text verdict enters the tool result (and thus the history), so the cost stays a few hundred tokens per figure instead of re-sending images every turn. Never blocks the loop: any failure logs a warning and returns the result unchanged.

maybe_review async

maybe_review(tool_name: str, result: str) -> tuple[str, str | None]

Attach a vision verdict to a run_python result carrying figures.

No-op unless HELIOAI_VISION_ENABLED is set and the tool is run_python.

Parameters:

Name	Type	Description	Default
tool_name	str	Name of the tool that just ran.	required
result	str	Its JSON result string (read for figure_paths).	required

Returns:

Type	Description
str	(possibly augmented result, verdict text or None) — the verdict is a
str | None	stateless side-call; only its text enters the history, never the image.

Source code in helioai/core/vision.py

async def maybe_review(tool_name: str, result: str) -> tuple[str, str | None]:
    """Attach a vision verdict to a run_python result carrying figures.

    No-op unless `HELIOAI_VISION_ENABLED` is set and the tool is `run_python`.

    Args:
        tool_name: Name of the tool that just ran.
        result: Its JSON result string (read for `figure_paths`).

    Returns:
        (possibly augmented result, verdict text or None) — the verdict is a
        stateless side-call; only its text enters the history, never the image.
    """
    if not settings.vision.enabled or tool_name != "run_python":
        return result, None
    try:
        data = json.loads(result)
    except (ValueError, TypeError):
        return result, None
    if not isinstance(data, dict) or data.get("error") or not data.get("figure_paths"):
        return result, None
    verdict = await _review(data["figure_paths"])
    if not verdict:
        return result, None
    data["figure_review"] = verdict
    return json.dumps(data, ensure_ascii=False, default=str), verdict

LLM clients

helioai.core.llm.base

Provider-neutral message model and LLMClient interface.

ToolCall dataclass

A tool invocation requested by the model.

Attributes:

Name	Type	Description
id	str	Provider-assigned identifier, echoed back on the matching tool result. Gemini has no native ids, so its client synthesises name::hex and parses the name back out.
name	str	Registered tool name.
arguments	dict	Decoded JSON arguments. Empty when the model emitted malformed JSON — a bad tool call must not kill the loop.

Source code in helioai/core/llm/base.py

@dataclass
class ToolCall:
    """A tool invocation requested by the model.

    Attributes:
        id: Provider-assigned identifier, echoed back on the matching tool
            result. Gemini has no native ids, so its client synthesises
            `name::hex` and parses the name back out.
        name: Registered tool name.
        arguments: Decoded JSON arguments. Empty when the model emitted
            malformed JSON — a bad tool call must not kill the loop.
    """

    id: str
    name: str
    arguments: dict

Message dataclass

One turn of conversation, in a provider-neutral form.

Every client converts to and from this shape, so the agent loop, the session store and the interfaces never see a provider's wire format.

Attributes:

Name	Type	Description
role	Literal['system', 'user', 'assistant', 'tool']	Who produced the turn.
content	str	Text content. Present alongside tool_calls when the model narrated what it was about to do.
tool_calls	list[ToolCall] | None	Tools the assistant wants invoked, when it requested any.
tool_call_id	str | None	For tool messages, the ToolCall.id being answered.

Source code in helioai/core/llm/base.py

@dataclass
class Message:
    """One turn of conversation, in a provider-neutral form.

    Every client converts to and from this shape, so the agent loop, the session
    store and the interfaces never see a provider's wire format.

    Attributes:
        role: Who produced the turn.
        content: Text content. Present alongside `tool_calls` when the model
            narrated what it was about to do.
        tool_calls: Tools the assistant wants invoked, when it requested any.
        tool_call_id: For `tool` messages, the `ToolCall.id` being answered.
    """

    role: Literal["system", "user", "assistant", "tool"]
    content: str = ""
    tool_calls: list[ToolCall] | None = None
    tool_call_id: str | None = None

ToolDef dataclass

A tool as advertised to the model.

Attributes:

Name	Type	Description
name	str	Tool name the model will call.
description	str	What the tool does and when to reach for it — the model's only clue about applicability.
parameters	dict	JSON Schema object describing the accepted arguments.

Source code in helioai/core/llm/base.py

@dataclass
class ToolDef:
    """A tool as advertised to the model.

    Attributes:
        name: Tool name the model will call.
        description: What the tool does and when to reach for it — the model's
            only clue about applicability.
        parameters: JSON Schema object describing the accepted arguments.
    """

    name: str
    description: str
    parameters: dict = field(default_factory=dict)

LLMClient

Bases: ABC

Interface every provider client implements.

One method is required, deliberately: the agent loop only ever needs a single completion with optional tool calling. Streaming happens at the loop level.

Source code in helioai/core/llm/base.py

class LLMClient(ABC):
    """Interface every provider client implements.

    One method is required, deliberately: the agent loop only ever needs a single
    completion with optional tool calling. Streaming happens at the loop level.
    """

    async def aclose(self) -> None:
        """Release the underlying HTTP connection pool.

        Callers that build a client per request — the CLI, the Jupyter magic and
        the web endpoints all do — must await this before their event loop ends.
        An async pool binds to the loop that used it, so a client left to the
        garbage collector schedules its own teardown after `asyncio.run` has
        closed that loop, and asyncio reports an unretrieved
        `RuntimeError: Event loop is closed` while the sockets stay open.

        The default is a no-op so a client without a pool needs no override.
        """
        return None

    @abstractmethod
    async def chat(
        self,
        messages: list[Message],
        tools: list[ToolDef],
        system_prompt: str | None = None,
        tool_choice: str = "auto",
    ) -> Message:
        """Send one turn and return the assistant's reply.

        Args:
            messages: Conversation history.
            tools: Tools the model may call this turn.
            system_prompt: Instructions placed before the history.
            tool_choice: `auto` to let the model decide, `required` to force a
                tool call — used on a sub-agent's first turn so it cannot answer
                from memory without looking anything up.

        Returns:
            The assistant reply, carrying `tool_calls` when the model requested any.
        """
        raise NotImplementedError

aclose async

aclose() -> None

Release the underlying HTTP connection pool.

Callers that build a client per request — the CLI, the Jupyter magic and the web endpoints all do — must await this before their event loop ends. An async pool binds to the loop that used it, so a client left to the garbage collector schedules its own teardown after asyncio.run has closed that loop, and asyncio reports an unretrieved RuntimeError: Event loop is closed while the sockets stay open.

The default is a no-op so a client without a pool needs no override.

Source code in helioai/core/llm/base.py

async def aclose(self) -> None:
    """Release the underlying HTTP connection pool.

    Callers that build a client per request — the CLI, the Jupyter magic and
    the web endpoints all do — must await this before their event loop ends.
    An async pool binds to the loop that used it, so a client left to the
    garbage collector schedules its own teardown after `asyncio.run` has
    closed that loop, and asyncio reports an unretrieved
    `RuntimeError: Event loop is closed` while the sockets stay open.

    The default is a no-op so a client without a pool needs no override.
    """
    return None

chat abstractmethod async

chat(messages: list[Message], tools: list[ToolDef], system_prompt: str | None = None, tool_choice: str = 'auto') -> Message

Send one turn and return the assistant's reply.

Parameters:

Name	Type	Description	Default
messages	list[Message]	Conversation history.	required
tools	list[ToolDef]	Tools the model may call this turn.	required
system_prompt	str | None	Instructions placed before the history.	None
tool_choice	str	auto to let the model decide, required to force a tool call — used on a sub-agent's first turn so it cannot answer from memory without looking anything up.	'auto'

Returns:

Type	Description
Message	The assistant reply, carrying tool_calls when the model requested any.

Source code in helioai/core/llm/base.py

@abstractmethod
async def chat(
    self,
    messages: list[Message],
    tools: list[ToolDef],
    system_prompt: str | None = None,
    tool_choice: str = "auto",
) -> Message:
    """Send one turn and return the assistant's reply.

    Args:
        messages: Conversation history.
        tools: Tools the model may call this turn.
        system_prompt: Instructions placed before the history.
        tool_choice: `auto` to let the model decide, `required` to force a
            tool call — used on a sub-agent's first turn so it cannot answer
            from memory without looking anything up.

    Returns:
        The assistant reply, carrying `tool_calls` when the model requested any.
    """
    raise NotImplementedError

call_with_retry async

call_with_retry(fn, *, attempts: int = 4, base_delay: float = 1.0, max_delay: float = 60.0)

Call async fn with backoff on retryable HTTP errors.

A server-supplied Retry-After wins over the exponential backoff: waiting the window a rate limiter asks for costs less than losing a session that is already several minutes and several downloads deep.

But only when the wait is one we would actually sit through. A per-minute limiter asks for seconds; an exhausted daily quota answers Retry-After: 35464 — nearly ten hours. Capping that to max_delay and retrying anyway just buys four minutes of silence before the same failure, so a window longer than max_delay fails at once and says how long it really is.

Non-retryable errors and exhausted attempts are re-raised immediately.

Source code in helioai/core/llm/base.py

async def call_with_retry(
    fn,
    *,
    attempts: int = 4,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
):
    """Call async fn with backoff on retryable HTTP errors.

    A server-supplied Retry-After wins over the exponential backoff: waiting the
    window a rate limiter asks for costs less than losing a session that is already
    several minutes and several downloads deep.

    But only when the wait is one we would actually sit through. A per-minute limiter
    asks for seconds; an exhausted daily quota answers `Retry-After: 35464` — nearly ten
    hours. Capping that to `max_delay` and retrying anyway just buys four minutes of
    silence before the same failure, so a window longer than `max_delay` fails at once
    and says how long it really is.

    Non-retryable errors and exhausted attempts are re-raised immediately.
    """
    for attempt in range(attempts):
        try:
            return await fn()
        except Exception as exc:
            status = _error_status(exc)
            if status not in RETRYABLE_STATUS:
                raise
            hinted = _retry_after(exc)
            if hinted is not None and hinted > max_delay:
                log.warning(
                    "llm_retry_window_too_long",
                    extra={"status": status, "retry_after_s": hinted},
                )
                raise
            if attempt == attempts - 1:
                raise
            delay = min(
                hinted
                if hinted is not None
                else base_delay * (2**attempt) + random.uniform(0, 0.5),
                max_delay,
            )
            log.info(
                "llm_retry",
                extra={
                    "attempt": attempt + 1,
                    "status": status,
                    "delay": round(delay, 2),
                    "server_hinted": hinted is not None,
                },
            )
            await asyncio.sleep(delay)

close_sdk_client async

close_sdk_client(client) -> None

Close an SDK client's connection pool, whether its close() is sync or async.

openai.AsyncOpenAI.close is a coroutine; google.genai.Client.close is not. Failures are swallowed: this only ever runs while tearing down, and a pool that will not close is not worth crashing a finished analysis over.

Source code in helioai/core/llm/base.py

async def close_sdk_client(client) -> None:
    """Close an SDK client's connection pool, whether its close() is sync or async.

    `openai.AsyncOpenAI.close` is a coroutine; `google.genai.Client.close` is not.
    Failures are swallowed: this only ever runs while tearing down, and a pool
    that will not close is not worth crashing a finished analysis over.
    """
    import inspect

    closer = getattr(client, "close", None) or getattr(client, "aclose", None)
    if closer is None:
        return
    try:
        result = closer()
        if inspect.isawaitable(result):
            await result
    except Exception as e:
        log.debug("sdk_client_close_failed: %s", e)

helioai.core.llm.openai_compat

Single client for every provider that speaks the OpenAI chat-completions wire format.

Groq, Ollama, Azure OpenAI and OpenAI itself all accept the same request shape, so they share one implementation here instead of one near-identical class each. A provider is a base_url plus a couple of dialect flags, not a subclass.

Azure is the one exception that still needs its own SDK client object (deployment routing and api-version), so it subclasses this to swap the constructor only.

OpenAICompatClient

Bases: LLMClient

Chat client for any OpenAI-compatible endpoint.

Parameters:

Name	Type	Description	Default
model	str	Model name sent in the request (the deployment name on Azure).	required
api_key	str	Provider API key. Local endpoints such as Ollama ignore it, but the SDK requires a non-empty value.	''
base_url	str | None	Endpoint root. None targets OpenAI itself.	None
system_role	str	Role used for the system prompt — developer on Azure and the o-series, system everywhere else.	'system'
max_output_tokens	int	Cap on generated tokens.	4096
temperature	float | None	Sampling temperature. None omits the field entirely, which reasoning models require.	0.2
provider	str	Name used to label log messages.	'openai'
client	Any	Pre-built SDK client. Injected by tests and by subclasses.	None

Example

client = OpenAICompatClient( ... model="llama-3.3-70b-versatile", ... api_key="gsk_...", ... base_url="https://api.groq.com/openai/v1", ... ) reply = await client.chat([Message(role="user", content="hi")], tools=[])

Source code in helioai/core/llm/openai_compat.py

class OpenAICompatClient(LLMClient):
    """Chat client for any OpenAI-compatible endpoint.

    Args:
        model: Model name sent in the request (the deployment name on Azure).
        api_key: Provider API key. Local endpoints such as Ollama ignore it, but
            the SDK requires a non-empty value.
        base_url: Endpoint root. `None` targets OpenAI itself.
        system_role: Role used for the system prompt — `developer` on Azure and
            the o-series, `system` everywhere else.
        max_output_tokens: Cap on generated tokens.
        temperature: Sampling temperature. `None` omits the field entirely, which
            reasoning models require.
        provider: Name used to label log messages.
        client: Pre-built SDK client. Injected by tests and by subclasses.

    Example:
        >>> client = OpenAICompatClient(
        ...     model="llama-3.3-70b-versatile",
        ...     api_key="gsk_...",
        ...     base_url="https://api.groq.com/openai/v1",
        ... )
        >>> reply = await client.chat([Message(role="user", content="hi")], tools=[])
    """

    def __init__(
        self,
        *,
        model: str,
        api_key: str = "",
        base_url: str | None = None,
        system_role: str = "system",
        max_output_tokens: int = 4096,
        temperature: float | None = 0.2,
        provider: str = "openai",
        client: Any = None,
    ):
        self._client = client or AsyncOpenAI(
            api_key=api_key or "unused", base_url=base_url, max_retries=0
        )
        self._model = model
        self._system_role = system_role
        self._max_output_tokens = max_output_tokens
        self._temperature = temperature
        self._provider = provider

    async def aclose(self) -> None:
        """Close the httpx pool held by the SDK client."""
        await close_sdk_client(self._client)

    async def chat(
        self,
        messages: list[Message],
        tools: list[ToolDef],
        system_prompt: str | None = None,
        tool_choice: str = "auto",
    ) -> Message:
        """Send one chat turn and return the assistant's reply.

        Args:
            messages: Conversation history.
            tools: Tools the model may call.
            system_prompt: Instructions prepended as the first message.
            tool_choice: `auto` to let the model decide, `required` to force a call.

        Returns:
            The assistant reply, carrying `tool_calls` when the model requested any.
        """
        openai_messages: list[dict] = []
        if system_prompt:
            openai_messages.append({"role": self._system_role, "content": system_prompt})
        openai_messages.extend(to_openai_messages(messages))

        kwargs: dict = {
            "model": self._model,
            "messages": openai_messages,
            "max_tokens": self._max_output_tokens,
        }
        if self._temperature is not None:
            kwargs["temperature"] = self._temperature
        if tools:
            kwargs["tools"] = to_openai_tools(tools)
            kwargs["tool_choice"] = tool_choice

        try:
            response = await call_with_retry(lambda: self._client.chat.completions.create(**kwargs))
        except BadRequestError as e:
            # Forcing a tool call is a preference, never worth losing the turn over.
            # DeepSeek v4 in thinking mode rejects `required` outright ("Thinking mode
            # does not support this tool_choice"), which killed a sub-agent on its very
            # first turn. Whether a model accepts it depends on the model and its
            # reasoning mode, not on the provider, so it is asked rather than tabulated.
            if tool_choice == "auto" or "tool_choice" not in str(e):
                raise
            log.warning("tool_choice_rejected_falling_back_to_auto: %s", self._model)
            kwargs["tool_choice"] = "auto"
            response = await call_with_retry(lambda: self._client.chat.completions.create(**kwargs))

        reply = from_openai_response(response, self._provider)
        if not (reply.content or "").strip() and not reply.tool_calls:
            # A turn with neither text nor a tool call is not an answer, and on a
            # reasoning model it is not rare either: the whole output allowance can go
            # into hidden reasoning and leave nothing to emit. It is also transient —
            # the identical request, replayed, came back with two tool calls in half
            # the wall time. The loop above treats this as fatal and abandons the
            # question, so two acts of a six-act notebook were lost to a condition that
            # one more attempt clears. Retried once, not in a loop: if the second is
            # empty too, the caller's error is the honest outcome.
            log.warning("%s empty turn, retrying once: %s", self._provider, self._model)
            response = await call_with_retry(lambda: self._client.chat.completions.create(**kwargs))
            reply = from_openai_response(response, self._provider)
        return reply

aclose async

aclose() -> None

Close the httpx pool held by the SDK client.

Source code in helioai/core/llm/openai_compat.py

async def aclose(self) -> None:
    """Close the httpx pool held by the SDK client."""
    await close_sdk_client(self._client)

chat async

chat(messages: list[Message], tools: list[ToolDef], system_prompt: str | None = None, tool_choice: str = 'auto') -> Message

Send one chat turn and return the assistant's reply.

Parameters:

Name	Type	Description	Default
messages	list[Message]	Conversation history.	required
tools	list[ToolDef]	Tools the model may call.	required
system_prompt	str | None	Instructions prepended as the first message.	None
tool_choice	str	auto to let the model decide, required to force a call.	'auto'

Returns:

Type	Description
Message	The assistant reply, carrying tool_calls when the model requested any.

Source code in helioai/core/llm/openai_compat.py

async def chat(
    self,
    messages: list[Message],
    tools: list[ToolDef],
    system_prompt: str | None = None,
    tool_choice: str = "auto",
) -> Message:
    """Send one chat turn and return the assistant's reply.

    Args:
        messages: Conversation history.
        tools: Tools the model may call.
        system_prompt: Instructions prepended as the first message.
        tool_choice: `auto` to let the model decide, `required` to force a call.

    Returns:
        The assistant reply, carrying `tool_calls` when the model requested any.
    """
    openai_messages: list[dict] = []
    if system_prompt:
        openai_messages.append({"role": self._system_role, "content": system_prompt})
    openai_messages.extend(to_openai_messages(messages))

    kwargs: dict = {
        "model": self._model,
        "messages": openai_messages,
        "max_tokens": self._max_output_tokens,
    }
    if self._temperature is not None:
        kwargs["temperature"] = self._temperature
    if tools:
        kwargs["tools"] = to_openai_tools(tools)
        kwargs["tool_choice"] = tool_choice

    try:
        response = await call_with_retry(lambda: self._client.chat.completions.create(**kwargs))
    except BadRequestError as e:
        # Forcing a tool call is a preference, never worth losing the turn over.
        # DeepSeek v4 in thinking mode rejects `required` outright ("Thinking mode
        # does not support this tool_choice"), which killed a sub-agent on its very
        # first turn. Whether a model accepts it depends on the model and its
        # reasoning mode, not on the provider, so it is asked rather than tabulated.
        if tool_choice == "auto" or "tool_choice" not in str(e):
            raise
        log.warning("tool_choice_rejected_falling_back_to_auto: %s", self._model)
        kwargs["tool_choice"] = "auto"
        response = await call_with_retry(lambda: self._client.chat.completions.create(**kwargs))

    reply = from_openai_response(response, self._provider)
    if not (reply.content or "").strip() and not reply.tool_calls:
        # A turn with neither text nor a tool call is not an answer, and on a
        # reasoning model it is not rare either: the whole output allowance can go
        # into hidden reasoning and leave nothing to emit. It is also transient —
        # the identical request, replayed, came back with two tool calls in half
        # the wall time. The loop above treats this as fatal and abandons the
        # question, so two acts of a six-act notebook were lost to a condition that
        # one more attempt clears. Retried once, not in a loop: if the second is
        # empty too, the caller's error is the honest outcome.
        log.warning("%s empty turn, retrying once: %s", self._provider, self._model)
        response = await call_with_retry(lambda: self._client.chat.completions.create(**kwargs))
        reply = from_openai_response(response, self._provider)
    return reply

to_openai_messages

to_openai_messages(messages: list[Message]) -> list[dict]

Convert neutral messages to the OpenAI wire format.

System messages already in the history are dropped: the system prompt is passed separately by chat() so it always lands first.

Parameters:

Name	Type	Description	Default
messages	list[Message]	Conversation history in HelioAI's provider-neutral form.	required

Returns:

Type	Description
list[dict]	Message dicts ready to send as the messages request field.

Source code in helioai/core/llm/openai_compat.py

def to_openai_messages(messages: list[Message]) -> list[dict]:
    """Convert neutral messages to the OpenAI wire format.

    System messages already in the history are dropped: the system prompt is
    passed separately by `chat()` so it always lands first.

    Args:
        messages: Conversation history in HelioAI's provider-neutral form.

    Returns:
        Message dicts ready to send as the `messages` request field.
    """
    out: list[dict] = []
    for msg in messages:
        if msg.role == "system":
            continue
        if msg.role == "user":
            out.append({"role": "user", "content": msg.content})
        elif msg.role == "assistant":
            if msg.tool_calls:
                out.append(
                    {
                        "role": "assistant",
                        "content": msg.content or None,
                        "tool_calls": [
                            {
                                "id": tc.id,
                                "type": "function",
                                "function": {
                                    "name": tc.name,
                                    "arguments": json.dumps(tc.arguments or {}),
                                },
                            }
                            for tc in msg.tool_calls
                        ],
                    }
                )
            else:
                out.append({"role": "assistant", "content": msg.content})
        elif msg.role == "tool":
            out.append(
                {
                    "role": "tool",
                    "tool_call_id": msg.tool_call_id or "",
                    "content": msg.content,
                }
            )
    return out

to_openai_tools

to_openai_tools(tools: list[ToolDef]) -> list[dict]

Convert tool definitions to OpenAI function-calling schemas.

Parameters:

Name	Type	Description	Default
tools	list[ToolDef]	Tools the agent may call this turn.	required

Returns:

Type	Description
list[dict]	Function schemas ready to send as the tools request field.

Source code in helioai/core/llm/openai_compat.py

def to_openai_tools(tools: list[ToolDef]) -> list[dict]:
    """Convert tool definitions to OpenAI function-calling schemas.

    Args:
        tools: Tools the agent may call this turn.

    Returns:
        Function schemas ready to send as the `tools` request field.
    """
    return [
        {
            "type": "function",
            "function": {
                "name": t.name,
                "description": t.description,
                "parameters": t.parameters or {"type": "object", "properties": {}},
            },
        }
        for t in tools
    ]

from_openai_response

from_openai_response(response: Any, provider: str = 'openai') -> Message

Convert an OpenAI chat-completions response to a neutral message.

Text content is preserved even when tool calls are present, and a tool call whose arguments are not valid JSON degrades to {} with a warning rather than raising — a malformed model output must not kill the agent loop. An inline <think>...</think> reasoning block, when a provider emits one, is stripped from the content before it reaches the agent loop or the user.

Parameters:

Name	Type	Description	Default
response	Any	The SDK response object.	required
provider	str	Provider name, used only to label log messages.	'openai'

Returns:

Type	Description
Message	The assistant's reply, with tool_calls set when the model requested any.

Source code in helioai/core/llm/openai_compat.py

def from_openai_response(response: Any, provider: str = "openai") -> Message:
    """Convert an OpenAI chat-completions response to a neutral message.

    Text content is preserved even when tool calls are present, and a tool call
    whose arguments are not valid JSON degrades to `{}` with a warning rather
    than raising — a malformed model output must not kill the agent loop. An
    inline `<think>...</think>` reasoning block, when a provider emits one, is
    stripped from the content before it reaches the agent loop or the user.

    Args:
        response: The SDK response object.
        provider: Provider name, used only to label log messages.

    Returns:
        The assistant's reply, with `tool_calls` set when the model requested any.
    """
    choice = response.choices[0]
    msg = choice.message
    raw_content = msg.content or ""
    content = _strip_reasoning(raw_content)
    tool_calls_raw = getattr(msg, "tool_calls", None) or []
    finish_reason = getattr(choice, "finish_reason", None)

    if not tool_calls_raw:
        if not content.strip():
            # A turn that produced nothing has exactly three causes and they need
            # different fixes: the budget ran out mid-generation (raise it), the model
            # spent the whole turn inside <think> and closed with nothing (shorten the
            # question), or it genuinely returned an empty completion (retry/provider).
            # The agent loop used to state the first as fact for all three. It is
            # visible here and nowhere else, so it is recorded here.
            log.warning(
                "%s empty completion: finish_reason=%s, %d raw chars, %d after stripping reasoning",
                provider,
                finish_reason,
                len(raw_content),
                len(content),
            )
        return Message(role="assistant", content=content)

    tool_calls: list[ToolCall] = []
    for tc in tool_calls_raw:
        try:
            args = json.loads(tc.function.arguments) if tc.function.arguments else {}
        except json.JSONDecodeError as e:
            # `finish_reason="length"` next to unparseable arguments is not "the model
            # emits bad JSON" — it is OUR output budget slicing a valid call mid-string.
            # The two need different fixes (raise max_output_tokens vs distrust the
            # model), and a log line that cannot tell them apart cost an hour of
            # diagnosis on a run where six 12k-char run_python calls all "lost" their code.
            log.warning(
                "%s tool_call %s args unparseable (finish_reason=%s, %d chars, %s): %r",
                provider,
                tc.function.name,
                finish_reason,
                len(tc.function.arguments or ""),
                "output budget truncated the call" if finish_reason == "length" else e,
                (tc.function.arguments or "")[:200],
            )
            args = {}
        tool_calls.append(ToolCall(id=tc.id, name=tc.function.name, arguments=args))
    return Message(role="assistant", content=content, tool_calls=tool_calls)

helioai.core.llm.azure_openai

Azure-hosted OpenAI.

Same wire format as every other OpenAI-compatible provider, so the conversion logic lives in openai_compat. Azure only differs in how the client is built — the endpoint embeds the deployment name and an api-version — plus two dialect details the base class already exposes as parameters: the system prompt is sent with the developer role, and temperature is omitted when unset because GPT-5 and the o-series reject it.

AzureOpenAIClient

Bases: OpenAICompatClient

Chat client for an Azure OpenAI deployment.

Parameters:

Name	Type	Description	Default
api_key	str	Azure OpenAI API key.	required
endpoint	str	Resource root, e.g. https://myresource.openai.azure.com.	required
api_version	str	Azure API version, e.g. 2024-12-01-preview.	required
deployment	str	Deployment name, sent as the request's model.	required
max_output_tokens	int	Cap on generated tokens.	4096
temperature	float | None	Sampling temperature; None omits the field.	None

Source code in helioai/core/llm/azure_openai.py

class AzureOpenAIClient(OpenAICompatClient):
    """Chat client for an Azure OpenAI deployment.

    Args:
        api_key: Azure OpenAI API key.
        endpoint: Resource root, e.g. `https://myresource.openai.azure.com`.
        api_version: Azure API version, e.g. `2024-12-01-preview`.
        deployment: Deployment name, sent as the request's `model`.
        max_output_tokens: Cap on generated tokens.
        temperature: Sampling temperature; `None` omits the field.
    """

    def __init__(
        self,
        api_key: str,
        endpoint: str,
        api_version: str,
        deployment: str,
        max_output_tokens: int = 4096,
        temperature: float | None = None,
    ):
        super().__init__(
            client=AsyncAzureOpenAI(
                api_key=api_key,
                azure_endpoint=endpoint,
                api_version=api_version,
                max_retries=0,
            ),
            model=deployment,
            system_role="developer",
            max_output_tokens=max_output_tokens,
            temperature=temperature,
            provider="azure",
        )

helioai.core.llm.gemini

GeminiClient: Google genai SDK → neutral Message model.

GeminiClient

Bases: LLMClient

Chat client for Google Gemini, using the native google-genai SDK.

Kept separate from OpenAICompatClient because the wire format genuinely differs: turns are Content objects with typed parts, the assistant role is called model, and tool results are matched by function name rather than by id. Since Gemini issues no call ids, this client synthesises name::hex and parses the name back out on the way in — a format that is persisted in existing sessions, so it must stay readable.

Parameters:

Name	Type	Description	Default
api_key	str	Gemini API key.	required
model	str	Model name, e.g. gemini-2.5-flash.	required
max_output_tokens	int	Cap on generated tokens.	4096
temperature	float	Sampling temperature.	0.2

Source code in helioai/core/llm/gemini.py

class GeminiClient(LLMClient):
    """Chat client for Google Gemini, using the native `google-genai` SDK.

    Kept separate from `OpenAICompatClient` because the wire format genuinely
    differs: turns are `Content` objects with typed parts, the assistant role is
    called `model`, and tool results are matched by function *name* rather than
    by id. Since Gemini issues no call ids, this client synthesises `name::hex`
    and parses the name back out on the way in — a format that is persisted in
    existing sessions, so it must stay readable.

    Args:
        api_key: Gemini API key.
        model: Model name, e.g. `gemini-2.5-flash`.
        max_output_tokens: Cap on generated tokens.
        temperature: Sampling temperature.
    """

    def __init__(
        self, api_key: str, model: str, max_output_tokens: int = 4096, temperature: float = 0.2
    ):
        self._client = genai.Client(api_key=api_key)
        self._model = model
        self._max_output_tokens = max_output_tokens
        self._temperature = temperature

    async def aclose(self) -> None:
        """Close the transport held by the genai client (its close() is sync)."""
        await close_sdk_client(self._client)

    async def chat(
        self,
        messages: list[Message],
        tools: list[ToolDef],
        system_prompt: str | None = None,
        tool_choice: str = "auto",
    ) -> Message:
        """Send one turn to Gemini and return the assistant's reply.

        When `system_prompt` is not given, it is recovered from any system message
        left in the history. `tool_choice="required"` maps to Gemini's ANY mode.
        """
        contents = self._to_gemini_contents(messages)
        gemini_tools = self._to_gemini_tools(tools) if tools else None

        if system_prompt is None:
            system_prompt = next((m.content for m in messages if m.role == "system"), None)

        tool_config = None
        if gemini_tools and tool_choice == "required":
            tool_config = gt.ToolConfig(
                function_calling_config=gt.FunctionCallingConfig(mode="ANY")
            )

        config = gt.GenerateContentConfig(
            max_output_tokens=self._max_output_tokens,
            temperature=self._temperature,
            tools=gemini_tools,
            tool_config=tool_config,
            system_instruction=system_prompt,
        )

        response = await call_with_retry(
            lambda: self._client.aio.models.generate_content(
                model=self._model,
                contents=contents,
                config=config,
            )
        )
        return self._from_gemini_response(response)

    @staticmethod
    def _to_gemini_contents(messages: list[Message]) -> list[gt.Content]:
        contents: list[gt.Content] = []
        for msg in messages:
            if msg.role == "system":
                continue
            if msg.role == "user":
                contents.append(gt.Content(role="user", parts=[gt.Part(text=msg.content)]))
            elif msg.role == "assistant":
                if msg.tool_calls:
                    parts = [
                        gt.Part(function_call=gt.FunctionCall(name=tc.name, args=tc.arguments))
                        for tc in msg.tool_calls
                    ]
                    contents.append(gt.Content(role="model", parts=parts))
                else:
                    contents.append(gt.Content(role="model", parts=[gt.Part(text=msg.content)]))
            elif msg.role == "tool":
                name, _, _ = (msg.tool_call_id or "::").partition("::")
                try:
                    response_payload = json.loads(msg.content)
                    if not isinstance(response_payload, dict):
                        response_payload = {"result": response_payload}
                except json.JSONDecodeError:
                    response_payload = {"result": msg.content}
                contents.append(
                    gt.Content(
                        role="user",
                        parts=[
                            gt.Part(
                                function_response=gt.FunctionResponse(
                                    name=name,
                                    response=response_payload,
                                )
                            )
                        ],
                    )
                )
        return contents

    @staticmethod
    def _to_gemini_tools(tools: list[ToolDef]) -> list[gt.Tool]:
        declarations = [
            gt.FunctionDeclaration(
                name=t.name,
                description=t.description,
                parameters=t.parameters or None,
            )
            for t in tools
        ]
        return [gt.Tool(function_declarations=declarations)]

    @staticmethod
    def _from_gemini_response(response) -> Message:
        tool_calls: list[ToolCall] = []
        text_chunks: list[str] = []
        candidate = response.candidates[0] if response.candidates else None
        parts = candidate.content.parts if candidate and candidate.content else []
        for part in parts or []:
            fc = getattr(part, "function_call", None)
            if fc and fc.name:
                args = dict(fc.args) if fc.args else {}
                call_id = f"{fc.name}::{uuid.uuid4().hex[:8]}"
                tool_calls.append(ToolCall(id=call_id, name=fc.name, arguments=args))
                continue
            text = getattr(part, "text", None)
            if text:
                text_chunks.append(text)
        if tool_calls:
            return Message(role="assistant", content="".join(text_chunks), tool_calls=tool_calls)
        return Message(role="assistant", content="".join(text_chunks))

aclose async

aclose() -> None

Close the transport held by the genai client (its close() is sync).

Source code in helioai/core/llm/gemini.py

async def aclose(self) -> None:
    """Close the transport held by the genai client (its close() is sync)."""
    await close_sdk_client(self._client)

chat async

chat(messages: list[Message], tools: list[ToolDef], system_prompt: str | None = None, tool_choice: str = 'auto') -> Message

Send one turn to Gemini and return the assistant's reply.

When system_prompt is not given, it is recovered from any system message left in the history. tool_choice="required" maps to Gemini's ANY mode.

Source code in helioai/core/llm/gemini.py

async def chat(
    self,
    messages: list[Message],
    tools: list[ToolDef],
    system_prompt: str | None = None,
    tool_choice: str = "auto",
) -> Message:
    """Send one turn to Gemini and return the assistant's reply.

    When `system_prompt` is not given, it is recovered from any system message
    left in the history. `tool_choice="required"` maps to Gemini's ANY mode.
    """
    contents = self._to_gemini_contents(messages)
    gemini_tools = self._to_gemini_tools(tools) if tools else None

    if system_prompt is None:
        system_prompt = next((m.content for m in messages if m.role == "system"), None)

    tool_config = None
    if gemini_tools and tool_choice == "required":
        tool_config = gt.ToolConfig(
            function_calling_config=gt.FunctionCallingConfig(mode="ANY")
        )

    config = gt.GenerateContentConfig(
        max_output_tokens=self._max_output_tokens,
        temperature=self._temperature,
        tools=gemini_tools,
        tool_config=tool_config,
        system_instruction=system_prompt,
    )

    response = await call_with_retry(
        lambda: self._client.aio.models.generate_content(
            model=self._model,
            contents=contents,
            config=config,
        )
    )
    return self._from_gemini_response(response)

helioai.core.llm.factory

Build the configured LLM client.

Providers that speak the OpenAI wire format are table entries, not classes — see OPENAI_COMPAT. Azure and Gemini need their own SDK client objects and stay explicit below.

build_llm_client

build_llm_client(provider: str | None = None) -> LLMClient

Return a client for the requested provider.

Parameters:

Name	Type	Description	Default
provider	str | None	Provider name. Defaults to HELIOAI_LLM_PROVIDER.	None

Returns:

Type	Description
LLMClient	A ready-to-use client.

Raises:

Type	Description
RuntimeError	If the provider is unknown or its API key is missing.

Example

llm = build_llm_client("groq") type(llm).name 'OpenAICompatClient'

Source code in helioai/core/llm/factory.py

def build_llm_client(provider: str | None = None) -> LLMClient:
    """Return a client for the requested provider.

    Args:
        provider: Provider name. Defaults to `HELIOAI_LLM_PROVIDER`.

    Returns:
        A ready-to-use client.

    Raises:
        RuntimeError: If the provider is unknown or its API key is missing.

    Example:
        >>> llm = build_llm_client("groq")
        >>> type(llm).__name__
        'OpenAICompatClient'
    """
    p = (provider or settings.llm.provider).lower()

    if p == "azure":
        from helioai.core.llm.azure_openai import AzureOpenAIClient

        cfg = settings.llm.azure
        if not cfg.api_key:
            raise RuntimeError("AZURE_OPENAI_API_KEY is not set")
        if not cfg.endpoint:
            raise RuntimeError("AZURE_OPENAI_ENDPOINT is not set")
        return AzureOpenAIClient(
            api_key=cfg.api_key,
            endpoint=cfg.endpoint,
            api_version=cfg.api_version,
            deployment=cfg.deployment,
            max_output_tokens=cfg.max_output_tokens,
            temperature=cfg.temperature,
        )

    if p == "gemini":
        from helioai.core.llm.gemini import GeminiClient

        cfg = settings.llm.gemini
        if not cfg.api_key:
            raise RuntimeError("GEMINI_API_KEY is not set")
        return GeminiClient(
            api_key=cfg.api_key,
            model=cfg.model,
            max_output_tokens=cfg.max_output_tokens,
            temperature=cfg.temperature,
        )

    if p in OPENAI_COMPAT:
        from helioai.core.llm.openai_compat import OpenAICompatClient

        spec = OPENAI_COMPAT[p]
        cfg = getattr(settings.llm, spec["config"])
        api_key = getattr(cfg, "api_key", "")
        if spec["key_env"] and not api_key:
            raise RuntimeError(f"{spec['key_env']} is not set")
        base_url = spec["base_url"] or f"{getattr(cfg, 'base_url', '').rstrip('/')}/v1"
        return OpenAICompatClient(
            provider=p,
            model=cfg.model,
            api_key=api_key,
            base_url=base_url,
            max_output_tokens=cfg.max_output_tokens,
            temperature=cfg.temperature,
        )

    known = "|".join(["azure", "gemini", *OPENAI_COMPAT])
    raise RuntimeError(f"Unknown LLM provider: {p!r}. Use {known}")

Storage

helioai.datastore

Session datastore — persist downloaded timeseries as .npz for reuse in run_python.

/data/.npz — compressed numpy archive

/data/manifest.json — index: name → metadata

Datasets are accessible in the sandbox via load_data("name"). The persisted data is the full-resolution download (before any downsampling). All I/O errors are silently swallowed — persistence must never break a tool call.

fill_mask

fill_mask(values, fillval: float | list | None = None)

Boolean mask of samples that carry no measurement.

Three conventions have to be caught at once, which is why every caller shares this one function instead of applying its own threshold:

• non-finite (NaN/inf) — already unusable;
• the ~1e31 magnitude convention, used by ACE among others;
• the value the dataset declares in its CDF FILLVAL, which is the only way to catch Wind/SWE's 99999.9. A blanket "reject >= 99999" rule is wrong: OMNI carries a real proton temperature of 99093 K in the 2003 Halloween window.

Parameters:

Name	Type	Description	Default
fillval	float | list | None	the declared FILLVAL, when the provider exposes it. A list as often as a scalar — Wind/SWE declares [99999.8984375], and a bare float(fillval) on it raises. Compared with rtol=1e-6 to survive a float32 round-trip while staying far tighter than the ~1% gap to real nearby values.	None

Example

fill_mask(np.array([1.2, -1e31, 3.4]), [-1e31]).tolist() [False, True, False]

Source code in helioai/datastore.py

def fill_mask(values, fillval: float | list | None = None):
    """Boolean mask of samples that carry no measurement.

    Three conventions have to be caught at once, which is why every caller shares
    this one function instead of applying its own threshold:

    - non-finite (NaN/inf) — already unusable;
    - the ~1e31 magnitude convention, used by ACE among others;
    - the value the dataset *declares* in its CDF FILLVAL, which is the only way
      to catch Wind/SWE's 99999.9. A blanket "reject >= 99999" rule is wrong: OMNI
      carries a real proton temperature of 99093 K in the 2003 Halloween window.

    Args:
        fillval: the declared FILLVAL, when the provider exposes it. A list as often
            as a scalar — Wind/SWE declares `[99999.8984375]`, and a bare
            `float(fillval)` on it raises. Compared with rtol=1e-6 to survive a
            float32 round-trip while staying far tighter than the ~1% gap to real
            nearby values.

    Example:
        >>> fill_mask(np.array([1.2, -1e31, 3.4]), [-1e31]).tolist()
        [False, True, False]
    """
    import numpy as np

    bad = ~np.isfinite(values) | (np.abs(values) >= 1e30)
    if fillval is not None:
        try:
            fv = float(np.asarray(fillval).ravel()[0])
            if np.isfinite(fv):
                bad = bad | np.isclose(values, fv, rtol=1e-6)
        except (TypeError, ValueError, IndexError):
            pass
    return bad

blank_fill

blank_fill(values, fillval=None)

Return (values with fill blanked to NaN, mask), or (values, None) if not numeric.

Applied at every point where downloaded data is persisted, so that anything reading it back — the sandbox, the exported notebook — sees NaN for "no measurement" rather than a sentinel that looks like a plausible reading. Leaving it to the reader meant remembering to call clean(), which cannot see FILLVAL, so one forgotten call put a 99999.9 "speed" into a plot and a mean.

Non-numeric parameters (string labels, epochs) have no fill convention and are passed through untouched.

Example

vals, mask = blank_fill(np.array([1.2, -1e31, 3.4]), [-1e31]) vals.tolist(), mask.tolist() ([1.2, nan, 3.4], [False, True, False])

Source code in helioai/datastore.py

def blank_fill(values, fillval=None):
    """Return (values with fill blanked to NaN, mask), or (values, None) if not numeric.

    Applied at every point where downloaded data is persisted, so that anything
    reading it back — the sandbox, the exported notebook — sees NaN for "no
    measurement" rather than a sentinel that looks like a plausible reading.
    Leaving it to the reader meant remembering to call clean(), which cannot see
    FILLVAL, so one forgotten call put a 99999.9 "speed" into a plot and a mean.

    Non-numeric parameters (string labels, epochs) have no fill convention and are
    passed through untouched.

    Example:
        >>> vals, mask = blank_fill(np.array([1.2, -1e31, 3.4]), [-1e31])
        >>> vals.tolist(), mask.tolist()
        ([1.2, nan, 3.4], [False, True, False])
    """
    import numpy as np

    try:
        numeric = np.array(values, dtype="float64")
    except (TypeError, ValueError):
        return values, None
    mask = fill_mask(numeric, fillval)
    numeric[mask] = np.nan
    return numeric, mask

read_manifest

read_manifest(session_dir: Path) -> dict

Return the manifest dict for a given session directory.

Parameters:

Name	Type	Description	Default
session_dir	Path	A session workspace directory (contains data/).	required

Returns:

Type	Description
dict	{"datasets": {name: {kind, param_id, start, stop, ...}}} — empty
dict	datasets dict when no manifest exists.

Source code in helioai/datastore.py

def read_manifest(session_dir: Path) -> dict:
    """Return the manifest dict for a given session directory.

    Args:
        session_dir: A session workspace directory (contains `data/`).

    Returns:
        {"datasets": {name: {kind, param_id, start, stop, ...}}} — empty
        datasets dict when no manifest exists.
    """
    return _read_manifest_file(session_dir / DATA_SUBDIR)

find_existing

find_existing(param_id: str, start: str, stop: str) -> str | None

Name of an already-persisted dataset for this exact param and window, or None.

The same matching _unique_name uses to reuse a slot — but consulted before the download rather than after, so a repeat request costs a dict lookup instead of a network round-trip. The prompt has always said "download each parameter ONCE"; a real run still re-fetched the same Wind field three times across three turns, with the dataset name sitting in plain sight in its own history. Discipline the model does not reliably apply belongs in the tool.

Source code in helioai/datastore.py

def find_existing(param_id: str, start: str, stop: str) -> str | None:
    """Name of an already-persisted dataset for this exact param and window, or None.

    The same matching `_unique_name` uses to reuse a slot — but consulted *before* the
    download rather than after, so a repeat request costs a dict lookup instead of a
    network round-trip. The prompt has always said "download each parameter ONCE"; a
    real run still re-fetched the same Wind field three times across three turns, with
    the dataset name sitting in plain sight in its own history. Discipline the model
    does not reliably apply belongs in the tool.
    """
    try:
        data_dir = _session_data_dir()
        if data_dir is None:
            return None
        for name, entry in _read_manifest_file(data_dir).get("datasets", {}).items():
            if (
                entry.get("param_id") == param_id
                and entry.get("start") == start
                and entry.get("stop") == stop
                and entry.get("kind") == "timeseries"
            ):
                return name
    except Exception as e:  # noqa: BLE001 — a cache miss must never block a download
        log.debug("find_existing failed: %s", e)
    return None

save_timeseries

save_timeseries(name_hint: str, *, time, values, param_id: str, units: str, start: str, stop: str, columns, source: str) -> dict | None

Persist a timeseries download as npz + a manifest entry.

Parameters:

Name	Type	Description	Default
name_hint	str	Basis for the dataset name (slugged, collision-suffixed).	required
time	array - like	Time axis as returned by speasy.	required
values	ndarray	Data array (fill already blanked).	required
param_id	str	Speasy id, recorded in the manifest for the export rewrite.	required
units	str	Physical units string.	required
start	str	ISO window start, recorded for the export rewrite.	required
stop	str	ISO window stop.	required
columns	list[str]	Component names.	required
source	str	Which tool produced the download.	required

Returns:

Type	Description
dict | None	{"dataset": } to reference in load_data(), or None when
dict | None	persisting failed (the download result is still usable in-memory).

Source code in helioai/datastore.py

def save_timeseries(
    name_hint: str,
    *,
    time,
    values,
    param_id: str,
    units: str,
    start: str,
    stop: str,
    columns,
    source: str,
) -> dict | None:
    """Persist a timeseries download as npz + a manifest entry.

    Args:
        name_hint: Basis for the dataset name (slugged, collision-suffixed).
        time (array-like): Time axis as returned by speasy.
        values (numpy.ndarray): Data array (fill already blanked).
        param_id: Speasy id, recorded in the manifest for the export rewrite.
        units: Physical units string.
        start: ISO window start, recorded for the export rewrite.
        stop: ISO window stop.
        columns (list[str]): Component names.
        source: Which tool produced the download.

    Returns:
        {"dataset": <final name>} to reference in `load_data()`, or None when
        persisting failed (the download result is still usable in-memory).
    """
    try:
        import time as _time

        import numpy as np

        data_dir = _session_data_dir()
        if data_dir is None:
            return None

        time_arr = np.asarray(time)
        values_arr = np.asarray(values, dtype=float)

        if values_arr.nbytes > _MAX_BYTES:
            log.warning(
                "datastore: skipping %r — %d MB exceeds cap",
                param_id,
                values_arr.nbytes // (1024 * 1024),
            )
            return None

        manifest = _read_manifest_file(data_dir)
        base = _slug(name_hint or param_id)
        name = _unique_name(manifest, base, param_id, start, stop)
        fname = f"{name}.npz"

        np.savez_compressed(data_dir / fname, time=time_arr, values=values_arr)

        cols = list(columns) if isinstance(columns, (list, tuple)) else []
        # Derived from what was actually written rather than passed in, so it can
        # never drift from the file: get_timeseries blanks fill values to NaN
        # before saving, so this is the fraction with no measurement.
        missing_pct = round(100 * float(np.isnan(values_arr).mean()), 1) if values_arr.size else 0.0
        manifest.setdefault("datasets", {})[name] = {
            "kind": "timeseries",
            "file": fname,
            "param_id": param_id,
            "units": units,
            "start": start,
            "stop": stop,
            "shape": list(values_arr.shape),
            "columns": cols,
            "missing_pct": missing_pct,
            "source": source,
            "created": str(int(_time.time())),
        }
        _write_manifest_file(data_dir, manifest)
        return {"dataset": name}
    except Exception as e:
        log.warning("datastore: save_timeseries failed for %r: %s", param_id, e)
        return None

save_event_collection

save_event_collection(name_hint: str, *, series, param_id: str, units: str, source: str) -> dict | None

Persist a batch of per-event timeseries. series = [(ev_start, ev_stop, ts|None), ...]. Returns {"dataset": name} or None on failure.

Source code in helioai/datastore.py

def save_event_collection(
    name_hint: str,
    *,
    series,
    param_id: str,
    units: str,
    source: str,
) -> dict | None:
    """Persist a batch of per-event timeseries. series = [(ev_start, ev_stop, ts|None), ...].
    Returns {"dataset": name} or None on failure.
    """
    try:
        import time as _time

        import numpy as np

        data_dir = _session_data_dir()
        if data_dir is None:
            return None

        arrays: dict[str, any] = {}
        events_meta: list[dict] = []
        total_bytes = 0
        capped = False

        for i, (ev_start, ev_stop, ts) in enumerate(series):
            if ts is None or not hasattr(ts, "time") or not hasattr(ts, "values"):
                events_meta.append(
                    {"idx": i, "start": ev_start, "stop": ev_stop, "status": "no_data"}
                )
                continue
            try:
                t_arr = np.asarray(ts.time)
                v_arr, _ = blank_fill(ts.values, (getattr(ts, "meta", {}) or {}).get("FILLVAL"))
                v_arr = np.asarray(v_arr, dtype=float)
                total_bytes += t_arr.nbytes + v_arr.nbytes
                if total_bytes > _MAX_BYTES:
                    if not capped:
                        log.warning(
                            "datastore: event collection for %r exceeds 100 MB cap, truncating",
                            param_id,
                        )
                        capped = True
                    events_meta.append(
                        {"idx": i, "start": ev_start, "stop": ev_stop, "status": "truncated"}
                    )
                    continue
                arrays[f"t{i}"] = t_arr
                arrays[f"v{i}"] = v_arr
                events_meta.append({"idx": i, "start": ev_start, "stop": ev_stop, "status": "ok"})
            except Exception:
                events_meta.append(
                    {"idx": i, "start": ev_start, "stop": ev_stop, "status": "no_data"}
                )

        if not arrays:
            return None

        manifest = _read_manifest_file(data_dir)
        base = _slug(name_hint or param_id) + "_events"
        name = base
        if name in manifest.get("datasets", {}):
            i2 = 2
            while f"{base}_{i2}" in manifest.get("datasets", {}):
                i2 += 1
            name = f"{base}_{i2}"
        fname = f"{name}.npz"

        np.savez_compressed(data_dir / fname, **arrays)

        manifest.setdefault("datasets", {})[name] = {
            "kind": "event_collection",
            "file": fname,
            "param_id": param_id,
            "units": units,
            "n_events": len(series),
            "events": events_meta,
            "source": source,
            "created": str(int(_time.time())),
        }
        _write_manifest_file(data_dir, manifest)
        return {"dataset": name}
    except Exception as e:
        log.warning("datastore: save_event_collection failed for %r: %s", param_id, e)
        return None

helioai.workspace

Workspace — stable output directory for sandbox figures and data.

Figures go to workspace//fig_N_M.png Code files go to workspace//code_N.py

The session label is a human-readable slug derived from the first user message, propagated via a contextvar set by stream_chat at the start of each request.

current_user

current_user() -> str

Return the user owning the current context, or the default user.

Source code in helioai/workspace.py

def current_user() -> str:
    """Return the user owning the current context, or the default user."""
    return _current_user.get() or DEFAULT_USER

set_user

set_user(user_id: str) -> object

Bind the user contextvar. Returns token for later reset.

Source code in helioai/workspace.py

def set_user(user_id: str) -> object:
    """Bind the user contextvar. Returns token for later reset."""
    return _current_user.set(user_id)

reset_user

reset_user(token: object) -> None

Restore the user contextvar from a token returned by set_user.

Source code in helioai/workspace.py

def reset_user(token: object) -> None:
    """Restore the user contextvar from a token returned by `set_user`."""
    _current_user.reset(token)  # type: ignore[arg-type]

user_home

user_home(user: str) -> Path

A user's private storage home: /users// (not created here).

Example

user_home("cli") PosixPath('.../data/users/cli')

Source code in helioai/workspace.py

def user_home(user: str) -> Path:
    """A user's private storage home: <data>/users/<user>/ (not created here).

    Example:
        >>> user_home("cli")
        PosixPath('.../data/users/cli')
    """
    return _users_root() / user

set_session

set_session(session_id: str) -> object

Bind the session contextvar. Returns token for later reset.

Source code in helioai/workspace.py

def set_session(session_id: str) -> object:
    """Bind the session contextvar. Returns token for later reset."""
    return _current_session.set(session_id)

reset_session

reset_session(token: object) -> None

Restore the session contextvar from a token returned by set_session.

Source code in helioai/workspace.py

def reset_session(token: object) -> None:
    """Restore the session contextvar from a token returned by `set_session`."""
    _current_session.reset(token)  # type: ignore[arg-type]

set_label

set_label(label: str) -> object

Bind the workspace label contextvar. Returns token for later reset.

Source code in helioai/workspace.py

def set_label(label: str) -> object:
    """Bind the workspace label contextvar. Returns token for later reset."""
    return _current_label.set(label)

reset_label

reset_label(token: object) -> None

Restore the label contextvar from a token returned by set_label.

Source code in helioai/workspace.py

def reset_label(token: object) -> None:
    """Restore the label contextvar from a token returned by `set_label`."""
    _current_label.reset(token)  # type: ignore[arg-type]

safe_id

safe_id(value: str, fallback: str = 'session') -> str

Reduce an identifier to something that cannot escape its parent directory.

Session ids are caller-supplied — a web request body, an MCP client, a CLI flag — and end up as path components here, in the export filename, and in the rmtree behind DELETE /api/sessions/{id}. Everything the project mints is a uuid4, so stripping to [A-Za-z0-9_-] is lossless in practice and turns ../.. into the fallback rather than a parent directory.

Example

safe_id("../../etc/passwd"), safe_id("sess-abc-123456") ('etcpasswd', 'sess-abc-123456')

Source code in helioai/workspace.py

def safe_id(value: str, fallback: str = "session") -> str:
    """Reduce an identifier to something that cannot escape its parent directory.

    Session ids are caller-supplied — a web request body, an MCP client, a CLI
    flag — and end up as path components here, in the export filename, and in the
    `rmtree` behind `DELETE /api/sessions/{id}`. Everything the project mints is a
    uuid4, so stripping to `[A-Za-z0-9_-]` is lossless in practice and turns
    `../..` into the fallback rather than a parent directory.

    Example:
        >>> safe_id("../../etc/passwd"), safe_id("sess-abc-123456")
        ('etcpasswd', 'sess-abc-123456')
    """
    cleaned = re.sub(r"[^A-Za-z0-9_-]", "", value)[:64]
    return cleaned or fallback

make_session_label

make_session_label(first_message: str, session_id: str) -> str

Build a human-readable slug for the session workspace folder.

Example: "Plot IMF Bz from ACE" + session abc123... → "plot-imf-bz-from_abc123"

Source code in helioai/workspace.py

def make_session_label(first_message: str, session_id: str) -> str:
    """Build a human-readable slug for the session workspace folder.

    Example: "Plot IMF Bz from ACE" + session abc123... → "plot-imf-bz-from_abc123"
    """
    words = re.sub(r"[^a-z0-9\s]", "", first_message.lower().strip()).split()
    slug = "-".join(words[:4]) if words else "session"
    return f"{slug[:25]}_{safe_id(session_id)[:6]}"

get_session_dir

get_session_dir() -> Path

Return the workspace directory for the current session.

Uses _current_label if set, falls back to _current_session UUID, then tmpdir. Creates the directory if it does not exist.

Source code in helioai/workspace.py

def get_session_dir() -> Path:
    """Return the workspace directory for the current session.

    Uses _current_label if set, falls back to _current_session UUID, then tmpdir.
    Creates the directory if it does not exist.
    """
    label = _current_label.get()
    if label:
        d = _root() / safe_id(label)
        d.mkdir(parents=True, exist_ok=True)
        return d
    session_id = _current_session.get()
    if session_id:
        d = _root() / safe_id(session_id)
        d.mkdir(parents=True, exist_ok=True)
        return d
    return _no_session_dir()

get_next_run_idx

get_next_run_idx(session_dir: Path) -> int

Return the next available run index for a session directory.

Scans for code_N.py files and returns max(N)+1, or 0 if none exist.

Source code in helioai/workspace.py

def get_next_run_idx(session_dir: Path) -> int:
    """Return the next available run index for a session directory.

    Scans for code_N.py files and returns max(N)+1, or 0 if none exist.
    """
    existing = list(session_dir.glob("code_*.py"))
    if not existing:
        return 0
    indices = []
    for p in existing:
        parts = p.stem.split("_")
        if len(parts) == 2 and parts[1].isdigit():
            indices.append(int(parts[1]))
    return max(indices) + 1 if indices else 0

get_run_dir_for_sandbox

get_run_dir_for_sandbox() -> str

Backward-compat: return session dir path as string (used by sandbox fallback).

Source code in helioai/workspace.py

def get_run_dir_for_sandbox() -> str:
    """Backward-compat: return session dir path as string (used by sandbox fallback)."""
    return str(get_session_dir())

is_under_workspace

is_under_workspace(path: str | Path) -> bool

True if path is safely under the per-user storage root (no traversal).

is_relative_to rather than a string prefix: comparing against str(root) + "/" hard-coded the POSIX separator, so on Windows the check never matched and /figure and /code returned 404 for every legitimate path. Fail-closed, so it was a dead web UI rather than a hole — but dead all the same.

Source code in helioai/workspace.py

def is_under_workspace(path: str | Path) -> bool:
    """True if path is safely under the per-user storage root (no traversal).

    `is_relative_to` rather than a string prefix: comparing against `str(root) + "/"`
    hard-coded the POSIX separator, so on Windows the check never matched and `/figure`
    and `/code` returned 404 for every legitimate path. Fail-closed, so it was a dead
    web UI rather than a hole — but dead all the same.
    """
    try:
        p = Path(path).resolve()
        return p.is_relative_to(_users_root().resolve())
    except (ValueError, OSError):
        return False

cleanup_old_runs

cleanup_old_runs(ttl_seconds: int | None = None) -> int

Purge session dirs older than ttl_seconds across all users. Returns count removed.

Source code in helioai/workspace.py

def cleanup_old_runs(ttl_seconds: int | None = None) -> int:
    """Purge session dirs older than ttl_seconds across all users. Returns count removed."""
    from helioai.config import settings

    if ttl_seconds is None:
        ttl_seconds = settings.workspace.ttl_seconds
    users_root = _users_root()
    if not users_root.exists():
        return 0
    cutoff = time.time() - ttl_seconds
    removed = 0
    for home in users_root.iterdir():
        ws = home / "workspace"
        if not ws.is_dir():
            continue
        for session_dir in ws.iterdir():
            if session_dir.is_dir() and session_dir.stat().st_mtime < cutoff:
                shutil.rmtree(session_dir, ignore_errors=True)
                removed += 1
    return removed

Export

helioai.export

Export a session as a reproducible Jupyter notebook.

A research result that cannot be re-run is worthless. The agent already saves every sandbox run as code_N.py in the session workspace; this module bundles those runs plus the conversation into a self-contained, re-executable .ipynb with a provenance header (parameter ids, time, library versions).

The saved runs are rewritten to standalone code (to_standalone): load_data() becomes a direct spz.get_data(...), the agent-only param_card()/ document_method() calls are dropped, and a minimal header supplies the imports plus real clean()/export() helpers — so each cell runs in a plain Jupyter kernel with no HelioAI sandbox around it.

export_session_notebook

export_session_notebook(user_id: str, session_id: str, out_path: Path | None = None) -> Path

Write the session as a .ipynb and return its path.

Default location: /.ipynb (or /.ipynb).

Example

export_session_notebook("cli", "8f3aa012-...") PosixPath('.../data/users/cli/workspace/plot-imf-bz_8f3aa0/plot-imf-bz_8f3aa0.ipynb')

The notebook opens with a provenance header (parameter ids, library versions), then one runnable cell per saved sandbox run, rewritten to standalone speasy calls.

Source code in helioai/export.py

def export_session_notebook(user_id: str, session_id: str, out_path: Path | None = None) -> Path:
    """Write the session as a .ipynb and return its path.

    Default location: <workspace>/<label>.ipynb (or <workspace>/<session>.ipynb).

    Example:
        >>> export_session_notebook("cli", "8f3aa012-...")
        PosixPath('.../data/users/cli/workspace/plot-imf-bz_8f3aa0/plot-imf-bz_8f3aa0.ipynb')

        The notebook opens with a provenance header (parameter ids, library versions),
        then one runnable cell per saved sandbox run, rewritten to standalone speasy calls.
    """
    import nbformat as nbf

    nb = build_notebook(user_id, session_id)

    if out_path is None:
        from helioai.workspace import safe_id, user_home

        label = store.get_workspace_dir(user_id, session_id) or session_id
        root = user_home(user_id) / "workspace"
        root.mkdir(parents=True, exist_ok=True)
        out_path = root / f"{safe_id(label)}.ipynb"
    out_path = Path(out_path)
    nbf.write(nb, str(out_path))
    return out_path

to_standalone

to_standalone(code_src: str, manifest: dict, *, with_header: bool = True) -> str

Turn a saved sandbox run into standalone, re-executable code.

Strips agent-only calls, rewrites load_data() → spz.get_data(), and (unless embedded in a notebook that already has a setup cell) prepends the imports plus real clean()/export() helpers it needs.

Parameters:

Name	Type	Description	Default
code_src	str	A code_N.py saved by the sandbox.	required
manifest	dict	The session manifest from read_manifest() — provides the param_id and window behind each load_data() name.	required
with_header	bool	Prepend the standalone imports/helpers header.	True

Example

manifest = {"datasets": {"imf_gsm": {"kind": "timeseries", ... "param_id": "amda/imf_gsm", "start": "2005-01-17", "stop": "2005-01-18"}}} to_standalone('d = load_data("imf_gsm")\nparam_card(d, "amda/imf_gsm")\n', ... manifest, with_header=False) 'd = spz.get_data("amda/imf_gsm", "2005-01-17", "2005-01-18")'

Source code in helioai/export.py

def to_standalone(code_src: str, manifest: dict, *, with_header: bool = True) -> str:
    """Turn a saved sandbox run into standalone, re-executable code.

    Strips agent-only calls, rewrites load_data() → spz.get_data(), and (unless
    embedded in a notebook that already has a setup cell) prepends the imports
    plus real clean()/export() helpers it needs.

    Args:
        code_src: A `code_N.py` saved by the sandbox.
        manifest: The session manifest from `read_manifest()` — provides the
            param_id and window behind each `load_data()` name.
        with_header: Prepend the standalone imports/helpers header.

    Example:
        >>> manifest = {"datasets": {"imf_gsm": {"kind": "timeseries",
        ...     "param_id": "amda/imf_gsm", "start": "2005-01-17", "stop": "2005-01-18"}}}
        >>> to_standalone('d = load_data("imf_gsm")\\nparam_card(d, "amda/imf_gsm")\\n',
        ...               manifest, with_header=False)
        'd = spz.get_data("amda/imf_gsm", "2005-01-17", "2005-01-18")'
    """
    body = _strip_known_imports(
        _rewrite_load_data_calls(_strip_agent_only_calls(code_src), manifest)
    )
    if not with_header:
        return body
    header = _standalone_header(body)
    return f"{header}\n\n\n{body}" if header else body

Configuration

helioai.config

Centralized configuration — loads .env once at startup.

settings is a module-level singleton imported everywhere. Fails fast at import time if the configured LLM provider is missing an API key.

AzureOpenAIConfig dataclass

Azure OpenAI deployment settings.

Azure routes by deployment name rather than model name, and reasoning models (GPT-5, o-series) reject an explicit temperature — hence temperature=None by default, which omits the field entirely.

Source code in helioai/config.py

@dataclass
class AzureOpenAIConfig:
    """Azure OpenAI deployment settings.

    Azure routes by deployment name rather than model name, and reasoning models
    (GPT-5, o-series) reject an explicit temperature — hence `temperature=None`
    by default, which omits the field entirely.
    """

    deployment: str = "models-gpt-53-chat"
    api_version: str = "2024-12-01-preview"
    # 8192, not the 2048 this used to be and not the 4096 the other providers use.
    # Azure draws reasoning tokens from this same allowance, so a reasoning
    # deployment can spend the whole budget thinking and return an empty message —
    # no text, no tool call. At 2048 that happened on any request that generates a
    # file: the standalone-script export in examples/02 produced nothing at all.
    max_output_tokens: int = 8192
    temperature: float | None = None
    api_key: str = ""
    endpoint: str = ""

GeminiConfig dataclass

Google Gemini settings, used with the native google-genai client.

Source code in helioai/config.py

@dataclass
class GeminiConfig:
    """Google Gemini settings, used with the native `google-genai` client."""

    model: str = "gemini-2.5-flash"
    max_output_tokens: int = 4096
    temperature: float = 0.2
    api_key: str = ""

GroqConfig dataclass

Groq settings. Reached through the shared OpenAI-compatible client.

Source code in helioai/config.py

@dataclass
class GroqConfig:
    """Groq settings. Reached through the shared OpenAI-compatible client."""

    model: str = "llama-3.3-70b-versatile"
    max_output_tokens: int = 4096
    temperature: float = 0.2
    api_key: str = ""

OpenCodeConfig dataclass

OpenCode's Zen gateway — OpenAI-compatible, whichever way you reach it: the flat-rate Go subscription, a BYOK-routed key, or any other model Zen hosts.

base_url defaults to the Go-plan endpoint, since that flat-rate tier is what most accounts actually have. It is a DIFFERENT catalogue from the general Zen endpoint (.../zen/v1, no /go/) — that one serves premium/BYOK-only models (Claude, ...) a Go subscription cannot reach, confirmed by querying both /models endpoints directly. Override HELIOAI_OPENCODE_URL to the plain Zen path if your access is not the Go plan.

No default model: what is reachable depends on your plan/BYOK setup and Zen's rotating catalogue (GLM, Kimi, DeepSeek, Qwen, MiniMax...). Set HELIOAI_OPENCODE_MODEL to the exact id from your dashboard — an empty string fails at the API with a clear "unknown model" rather than silently routing to a guessed default that may not exist on your plan.

16384 output tokens because everything this gateway serves is a reasoning model, and reasoning, prose AND tool-call arguments all draw on the same budget. At 4096, DeepSeek v4's run_python calls (~12k chars of JSON once a plot script is in them) were cut mid-string: the model saw "missing 1 required positional argument: 'code'", could not know why, and burned five turns re-sending the same truncated call. Same lesson as Azure's 2048→8192, one provider later.

Source code in helioai/config.py

@dataclass
class OpenCodeConfig:
    """OpenCode's Zen gateway — OpenAI-compatible, whichever way you reach it: the
    flat-rate Go subscription, a BYOK-routed key, or any other model Zen hosts.

    `base_url` defaults to the Go-plan endpoint, since that flat-rate tier is what
    most accounts actually have. It is a DIFFERENT catalogue from the general Zen
    endpoint (`.../zen/v1`, no `/go/`) — that one serves premium/BYOK-only models
    (Claude, ...) a Go subscription cannot reach, confirmed by querying both
    `/models` endpoints directly. Override `HELIOAI_OPENCODE_URL` to the plain Zen
    path if your access is not the Go plan.

    No default model: what is reachable depends on your plan/BYOK setup and Zen's
    rotating catalogue (GLM, Kimi, DeepSeek, Qwen, MiniMax...). Set
    `HELIOAI_OPENCODE_MODEL` to the exact id from your dashboard — an empty string
    fails at the API with a clear "unknown model" rather than silently routing to a
    guessed default that may not exist on your plan.

    16384 output tokens because everything this gateway serves is a reasoning model,
    and reasoning, prose AND tool-call arguments all draw on the same budget. At 4096,
    DeepSeek v4's run_python calls (~12k chars of JSON once a plot script is in them)
    were cut mid-string: the model saw "missing 1 required positional argument: 'code'",
    could not know why, and burned five turns re-sending the same truncated call.
    Same lesson as Azure's 2048→8192, one provider later.
    """

    base_url: str = "https://opencode.ai/zen/go"
    model: str = ""
    max_output_tokens: int = 16384
    temperature: float = 0.2
    api_key: str = ""

OllamaConfig dataclass

Local Ollama settings.

Ollama serves an OpenAI-compatible API on /v1, so it needs no client of its own and no API key. Point base_url elsewhere for any other local endpoint.

Source code in helioai/config.py

@dataclass
class OllamaConfig:
    """Local Ollama settings.

    Ollama serves an OpenAI-compatible API on `/v1`, so it needs no client of its
    own and no API key. Point `base_url` elsewhere for any other local endpoint.
    """

    base_url: str = "http://localhost:11434"
    model: str = "qwen2.5:14b-instruct"
    max_output_tokens: int = 4096
    temperature: float = 0.2
    api_key: str = ""

LLMConfig dataclass

Which provider to use, and the settings for each one.

Source code in helioai/config.py

@dataclass
class LLMConfig:
    """Which provider to use, and the settings for each one."""

    provider: str = "azure"
    azure: AzureOpenAIConfig = field(default_factory=AzureOpenAIConfig)
    gemini: GeminiConfig = field(default_factory=GeminiConfig)
    groq: GroqConfig = field(default_factory=GroqConfig)
    opencode: OpenCodeConfig = field(default_factory=OpenCodeConfig)
    ollama: OllamaConfig = field(default_factory=OllamaConfig)

AgentConfig dataclass

Agent loop limits.

max_iterations caps how many tool-calling rounds one question may take before the loop gives up, bounding both runtime and token spend.

Source code in helioai/config.py

@dataclass
class AgentConfig:
    """Agent loop limits.

    `max_iterations` caps how many tool-calling rounds one question may take
    before the loop gives up, bounding both runtime and token spend.
    """

    max_iterations: int = 10

RAGConfig dataclass

Parameter search settings.

Retrieval is hybrid: dense embeddings for descriptions, BM25 for exact tokens like BGSEc, fused by Reciprocal Rank Fusion with parameter rrf_k.

rerank_enabled stays False on purpose. A generic MS MARCO cross-encoder was measured to degrade results here: trained on web prose, it discards the dense+sparse consensus that makes exact-code matching work. Only a domain-tuned reranker would help.

Source code in helioai/config.py

@dataclass
class RAGConfig:
    """Parameter search settings.

    Retrieval is hybrid: dense embeddings for descriptions, BM25 for exact tokens
    like `BGSEc`, fused by Reciprocal Rank Fusion with parameter `rrf_k`.

    `rerank_enabled` stays False on purpose. A generic MS MARCO cross-encoder was
    measured to *degrade* results here: trained on web prose, it discards the
    dense+sparse consensus that makes exact-code matching work. Only a
    domain-tuned reranker would help.
    """

    chroma_dir: Path = field(default_factory=lambda: _DATA / "chroma")
    collection_name: str = "speasy_catalog"
    catalogs_collection_name: str = "speasy_catalogs"
    embed_model: str = "sentence-transformers/all-MiniLM-L6-v2"
    rerank_enabled: bool = False
    rerank_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"
    rerank_fetch_k: int = 20
    hybrid_enabled: bool = True
    hybrid_fetch_k: int = 50
    rrf_k: int = 60

WorkspaceConfig dataclass

Per-session working directories, cleaned up after ttl_seconds.

Source code in helioai/config.py

@dataclass
class WorkspaceConfig:
    """Per-session working directories, cleaned up after `ttl_seconds`."""

    workspace_dir: Path = field(default_factory=lambda: _DATA / "workspace")
    ttl_seconds: int = 86400 * 7  # 7 days

ProfileConfig dataclass

Location of the user profile injected into the system prompt.

Source code in helioai/config.py

@dataclass
class ProfileConfig:
    """Location of the user profile injected into the system prompt."""

    profile_path: Path = field(default_factory=lambda: _DATA / "profile.md")

RecipesConfig dataclass

Where scientific recipes are loaded from.

Defaults to the copy shipped inside the package so pip install works; override with HELIOAI_RECIPES_DIR to use your own set.

Source code in helioai/config.py

@dataclass
class RecipesConfig:
    """Where scientific recipes are loaded from.

    Defaults to the copy shipped inside the package so `pip install` works;
    override with `HELIOAI_RECIPES_DIR` to use your own set.
    """

    recipes_dir: Path = field(default_factory=lambda: _PKG_RECIPES)

CatalogsConfig dataclass

Where user-saved event catalogs are written, in speasy format.

Source code in helioai/config.py

@dataclass
class CatalogsConfig:
    """Where user-saved event catalogs are written, in speasy format."""

    catalogs_dir: Path = field(default_factory=lambda: _DATA / "catalogs")

LiteratureConfig dataclass

NASA ADS credentials for find_papers. Free token, no key means no tool.

Source code in helioai/config.py

@dataclass
class LiteratureConfig:
    """NASA ADS credentials for `find_papers`. Free token, no key means no tool."""

    ads_token: str = ""

MCPConfig dataclass

Remote MCP servers to mount, as a JSON object keyed by alias.

Source code in helioai/config.py

@dataclass
class MCPConfig:
    """Remote MCP servers to mount, as a JSON object keyed by alias."""

    servers_json: str = ""

VisionConfig dataclass

Multimodal review of generated figures.

A stateless side-call outside the agent loop: the image is downscaled, sent once, and only the text verdict enters the history — never the image, which would otherwise be resent on every subsequent turn. Off by default.

Source code in helioai/config.py

@dataclass
class VisionConfig:
    """Multimodal review of generated figures.

    A stateless side-call outside the agent loop: the image is downscaled, sent
    once, and only the text verdict enters the history — never the image, which
    would otherwise be resent on every subsequent turn. Off by default.
    """

    # Reviews sandbox figures with a multimodal side-call; only the text
    # verdict enters the history, never the image.
    enabled: bool = False
    provider: str = "azure"
    model: str = ""
    timeout_s: float = 20.0

DevConfig dataclass

Shared secret unlocking unrestricted mode past the heliophysics guardrail.

Empty by default, which means no token is valid and every request stays scoped. Compared in constant time.

Source code in helioai/config.py

@dataclass
class DevConfig:
    """Shared secret unlocking unrestricted mode past the heliophysics guardrail.

    Empty by default, which means no token is valid and every request stays
    scoped. Compared in constant time.
    """

    # Shared-secret that unlocks unrestricted LLM access (bypasses scope guardrail).
    # Empty (default) → no token is valid → all requests stay restricted.
    token: str = ""

WebAuthConfig dataclass

Nominative tokens for the web UI, parsed from HELIOAI_USERS.

Empty means no authentication and a single local user, which is the intended behaviour for local development only.

Source code in helioai/config.py

@dataclass
class WebAuthConfig:
    """Nominative tokens for the web UI, parsed from `HELIOAI_USERS`.

    Empty means no authentication and a single local user, which is the intended
    behaviour for local development only.
    """

    # Nominative tokens for the web UI: {token: user_id}. Parsed from
    # HELIOAI_USERS="tok1:vincent,tok2:alice". Empty → no auth, single local user.
    # ponytail: env-driven map, fine for a handful of researchers; move to a DB
    # table if tokens must be added/revoked at runtime.
    users: dict[str, str] = field(default_factory=dict)

Settings dataclass

Root settings object.

Imported as the module-level settings singleton and read everywhere; built once at import by _load(), which fails fast when the selected provider has no API key.

Source code in helioai/config.py

@dataclass
class Settings:
    """Root settings object.

    Imported as the module-level `settings` singleton and read everywhere; built
    once at import by `_load()`, which fails fast when the selected provider has
    no API key.
    """

    data_dir: Path = field(default_factory=lambda: _DATA)
    llm: LLMConfig = field(default_factory=LLMConfig)
    agent: AgentConfig = field(default_factory=AgentConfig)
    rag: RAGConfig = field(default_factory=RAGConfig)
    workspace: WorkspaceConfig = field(default_factory=WorkspaceConfig)
    profile: ProfileConfig = field(default_factory=ProfileConfig)
    recipes: RecipesConfig = field(default_factory=RecipesConfig)
    catalogs: CatalogsConfig = field(default_factory=CatalogsConfig)
    literature: LiteratureConfig = field(default_factory=LiteratureConfig)
    mcp: MCPConfig = field(default_factory=MCPConfig)
    vision: VisionConfig = field(default_factory=VisionConfig)
    dev: DevConfig = field(default_factory=DevConfig)
    web_auth: WebAuthConfig = field(default_factory=WebAuthConfig)

dev_unlock

dev_unlock(supplied: str | None) -> bool

True iff the supplied token matches the configured dev secret.

Returns False when the server-side token is empty (guards against accidentally unlocking an unconfigured instance).

Source code in helioai/config.py

def dev_unlock(supplied: str | None) -> bool:
    """True iff the supplied token matches the configured dev secret.

    Returns False when the server-side token is empty (guards against
    accidentally unlocking an unconfigured instance).
    """
    return (
        bool(settings.dev.token)
        and supplied is not None
        and hmac.compare_digest(supplied, settings.dev.token)
    )

Logging

helioai.logging_config

Structured logging via structlog.

Output format is selected by HELIOAI_LOG_FORMAT

• console (default): human-friendly, colourised.
• json: one JSON object per line.

setup_logging

setup_logging(level: str | int = 'INFO') -> None

Configure structlog and the root logger.

Output format follows HELIOAI_LOG_FORMAT: console (default) or json. Safe to call more than once — every entry point calls it, and repeated calls replace the handler rather than stacking duplicates.

Parameters:

Name	Type	Description	Default
level	str | int	Log level name or numeric value. Unknown names fall back to INFO.	'INFO'

Source code in helioai/logging_config.py

def setup_logging(level: str | int = "INFO") -> None:
    """Configure structlog and the root logger.

    Output format follows `HELIOAI_LOG_FORMAT`: `console` (default) or `json`.
    Safe to call more than once — every entry point calls it, and repeated calls
    replace the handler rather than stacking duplicates.

    Args:
        level: Log level name or numeric value. Unknown names fall back to INFO.
    """
    if isinstance(level, str):
        level = getattr(logging, level.upper(), logging.INFO)

    fmt = _format_from_env()

    shared_processors: list[Any] = [
        structlog.contextvars.merge_contextvars,
        structlog.processors.add_log_level,
        structlog.processors.TimeStamper(fmt="iso", utc=True),
        structlog.processors.StackInfoRenderer(),
        structlog.processors.format_exc_info,
    ]

    if fmt == "json":
        renderer: Any = structlog.processors.JSONRenderer()
    else:
        renderer = structlog.dev.ConsoleRenderer(colors=sys.stderr.isatty())

    structlog.configure(
        processors=[
            *shared_processors,
            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
        ],
        wrapper_class=structlog.make_filtering_bound_logger(level),
        logger_factory=structlog.stdlib.LoggerFactory(),
        cache_logger_on_first_use=True,
    )

    formatter = structlog.stdlib.ProcessorFormatter(
        foreign_pre_chain=shared_processors,
        processors=[
            structlog.stdlib.ProcessorFormatter.remove_processors_meta,
            renderer,
        ],
    )

    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(formatter)
    root = logging.getLogger()
    root.handlers = [handler]
    root.setLevel(level)
    _quiet_third_party_advisories()

get_logger

get_logger(name: str | None = None) -> Any

Return a structlog logger, optionally bound to a module name.

Source code in helioai/logging_config.py

def get_logger(name: str | None = None) -> Any:
    """Return a structlog logger, optionally bound to a module name."""
    return structlog.get_logger(name) if name else structlog.get_logger()