Adding an assistant#

PyHydroGeophysX can host several domain AI assistants side by side. AQUAH works on hydrogeophysics; GeoSAGE, being ported now, works on geological modelling from gravity and magnetic data. Each is a plug-in: a small package that describes itself, and the studio builds its chat, its Workflow page and its runs from that description. The user picks one at the top of the assistant panel.

This page is for whoever adds one. It uses GeoSAGE as the running example; its scaffold is in PyHydroGeophysX/agents/assistants/geosage/.

What an assistant is made of#

PyHydroGeophysX/agents/assistants/
    __init__.py          the registry, the Assistant description, the contract
    aquah/               AQUAH
        __init__.py      ASSISTANT = Assistant(...)
        workflow.py      run(payload, progress, ...)  and  tools()
    geosage/             GeoSAGE (scaffold)
        __init__.py      ASSISTANT = Assistant(..., status="in development")
        tools.py         TOOLS: one Tool per stage
        workflow.py      configure(payload), CONTROLLER_PROMPT, run(...)

Three things are the assistant’s own:

  1. A description, PyHydroGeophysX.agents.assistants.Assistant: its name and domain, the chat persona, example requests, the kinds of input file the Workflow page offers, its glow colours, and where its workflow and tools are. Importing it must be cheap, so the workflow and tools are given as "module:attribute" strings and loaded when a run starts.

  2. Tools: a dictionary of Tool, one per capability.

  3. A workflow: the function “Auto to report” runs, in a process of its own.

Everything else is shared: the controller loop that picks each step, the live timeline, step-by-step approval, questions to the user, the chat, the report view and the Light and Dark appearances.

The description#

from PyHydroGeophysX.agents.assistants import Assistant

ASSISTANT = Assistant(
    key="geosage",
    name="GeoSAGE",
    domain="Geological modeling",
    title="Geological Semantic Analysis for Geophysical Exploration",
    summary="Turns joint gravity–magnetic inversion models into ...",
    workflow="PyHydroGeophysX.agents.assistants.geosage.workflow:run",
    tools="PyHydroGeophysX.agents.assistants.geosage.tools:TOOLS",
    persona=PERSONA,                # chat rules that are GeoSAGE's own
    examples=("find serpentinite targets for natural hydrogen ...",),
    input_roles=(("Gravity data", "gravity_file"),
                 ("Magnetic data (TMI)", "magnetic_file"), ...),
    studio_modules=("one_click", "gravmag", "joint_inversion", "mesh3d"),
    colors=("#ff9500", "#ffcc00", "#34c759", "#30b0c7"),
    requires_packages=("simpeg",),
    status="in development",        # "ready" once the tools run
)

input_roles are what the Workflow page offers under Or add files yourself; a run receives them as payload["inputs"][key]. A key ending in _dir takes a folder, and a key listed in ordered_roles takes several files in order. folder_classifier=True means the workflow can sort a chosen folder’s files into these roles (mode='classify'); AQUAH does, GeoSAGE does not yet. studio_modules limits which studio modules the chat may drive. While status is not "ready", or a package in requires_packages is missing, the assistant is listed in the panel but greyed out with the reason.

Tools#

A tool is what the controller chooses between. It declares what it needs and what it makes, and is offered only once its needs exist, so the order of a run follows from the data:

Tool("run_joint_inversion",
     "Jointly invert gravity and magnetic data for density and magnetic "
     "susceptibility on a common mesh (deterministic, SimPEG, cross-gradient "
     "coupling), and report the data fit.",
     handler, requires=("field_data",), produces=("property_models",),
     agent="InversionAgent", label="Run joint gravity-magnetic inversion",
     module="joint_inversion")

The description is read by the model when it chooses, so write it for someone who does not know the code. label is what the timeline shows, and module is the studio page that shows this kind of work.

A handler takes the run’s RunContext and returns (summary, outputs):

def run_joint_inversion(ctx):
    data = ctx.get("field_data")            # an earlier tool's product
    settings = ctx.config.get("inversion", {})
    # ... SimPEG joint inversion, files written under ctx.output_dir ...
    return (f"Converged in {n} iterations; gravity R² {r2g:.3f}, "
            f"magnetic R² {r2m:.3f}.",
            {"property_models": {"density": rho, "susceptibility": kappa,
                                 "mesh": mesh, "fit": fit}})

The summary is shown to the user the moment the step ends and is what the controller reads before choosing again, so make it the step’s conclusion with its key numbers. A handler may raise: the failure is recorded, shown, and the run carries on deciding. An LLM-based stage reaches the model through ctx.settings (api_key, model, llm_provider). A tool that needs the user to choose calls ctx.ask(question, options): the question appears in the Live tab, and with nobody to answer the first option is taken and noted.

The workflow#

The contract a workflow function honours is PyHydroGeophysX.agents.assistants.WORKFLOW_CONTRACT. In short: it receives the run’s payload (request, inputs, model access, output_dir, step_mode) and returns a result with status, interpretation, warnings and report_files. Build it on PyHydroGeophysX.agents.runtime.entry.drive(), which runs the controller over your tools and produces the step events the studio needs:

def run(payload, progress, *, approve=None, on_event=None, events=None, **_):
    config = configure(payload)                       # the ContextAgent
    ctx = RunContext(goal=payload["request"], config=config,
                     output_dir=payload["output_dir"], settings={...})
    on_step = step_by_step(approve) if payload.get("step_mode") and approve else None
    drive(ctx, tools=TOOLS, ask=make_ask(key, model, provider),
          progress_callback=progress, on_step=on_step, on_event=on_event,
          prompt=CONTROLLER_PROMPT, finish="report_files")
    return {"status": ..., "interpretation": ..., "report_files": ctx.get("report_files")}

prompt frames the controller’s decision for your domain (it must keep the {transcript} and {menu} fields), and finish names the product that has to exist before the run may end. GeoSAGE’s is the reviewed report, so it cannot stop at the draft.

Registering it#

Inside this repository, add the package beside aquah and geosage and list it in BUILTIN in PyHydroGeophysX/agents/assistants/__init__.py.

From a separate package, such as the GeoSAGE repository, declare an entry point; nothing in PyHydroGeophysX needs editing:

[project.entry-points."pyhydrogeophysx.assistants"]
geosage = "geosage.pyhydrogeophysx:ASSISTANT"

A plug-in that fails to import is left out and reported by PyHydroGeophysX.agents.assistants.load_errors(); it never takes the others down.

What the studio does with it#

  • The assistant panel lists every registered assistant. Choosing one starts a new conversation with its persona and examples, and is remembered.

  • The Workflow page offers its input roles and names it in everything it says; the Live tab shows its steps, each with the controller’s reason, the module, the time taken, what it found and the figures it wrote.

  • The glow round the central area and the orb take its colours.

  • With Approve each step before it runs, every step waits for the user in the Live tab; a tool’s own question appears there too.

  • While the controller decides, its reasoning streams into the Live tab: make_ask() streams the reply, and drive sends the "why" field on as phase="thought" events. Keep the reasoning first in your controller prompt’s JSON, as GeoSAGE’s does.

  • Pause and notes from the user reach your run through drive too (PyHydroGeophysX.agents.runtime.steering): a note is added to ctx.guidance and the transcript, and the next decision may change the settings in ADJUSTABLE.

  • drive(..., flexible=True) lets the controller leave the standard route when the evidence calls for it (PyHydroGeophysX.agents.runtime.toolbox). It can look at a product, add a step whose when gate is closed, run a finished step again with changed settings, leave route steps out, write a temporary script, or ask the user. Leaving a step out is the model’s call, refused only when finish or a requested product would then not be made; every other change of course is put to the user with the model’s reason. Mark a tool proposable=False when it must never be proposed (a loader that reads a configured file), and give it applies when it fits only some runs. AQUAH runs flexibly; flexibility needs the default decision prompt, so an assistant with a prompt of its own keeps its fixed route.

  • The route across the top of the Live tab comes from your tools’ requires and produces: route_ahead() works out, from what the run has, which tools still stand between it and finish. A when gate that looks inside an artifact should check ctx.projected(key) first and answer from the configuration, since a projected artifact has nothing inside it yet.

Porting GeoSAGE#

The paper’s seven agents map onto the scaffold like this:

ContextAgent

workflow.configure: request to configuration

DataAgent

prepare_data: bounds, fields, gridding, uncertainties

PetrologyAgent

compile_priors, and unit/group naming in build_quasi_geology

InversionAgent

run_joint_inversion: settings and the SimPEG backend

GeoAgent

build_quasi_geology: prior-guided or GMM/BIC grouping

ReportAgent

write_report: interpretation, target ranking, uncertainty

ReviewAgent

review_report: checks, revision, the final report

A workable order is to fill in the deterministic stages first (prepare_data, run_joint_inversion, build_quasi_geology) so a run produces models and figures without a model key, then the LLM stages, and then:

  1. set status="ready" in geosage/__init__.py;

  2. replace the stand-in handlers in test_an_assistant_runs_its_own_tools_to_its_own_final_product (tests/test_agents.py) with a small real case, or add one such test;

  3. run it in the studio with Approve each step before it runs switched on, in Light and Dark.

To try the chain before the stages exist, replace the handlers with stand-ins that return the declared products, as that test does; the studio’s timeline then shows the whole run.