galaxyproject / galaxyproject/brc-analytics

Assistant: emit machine state via typed tool calls, not parsed from the reply text

Open
#1,429 1 comment 0 reactions 1 assignee Claimed by @dannon View on GitHub
Dominant language
TypeScript
Stars
7
Forks
11
Avg merge
2d 12h
Merged PRs (30d)
16

Description

Parent epic: #1288

## Problem

The assistant produces two different things in one blob of text: a human-readable
message, and machine state (the Analysis Setup tracker fields, suggestion chips,
and the handoff signal). Today the machine state rides *inside the model's prose* as
trailing `SCHEMA_UPDATE:` / `SUGGESTIONS:` lines that we regex out afterward
(`_parse_structured_output`, `assistant_agent.py:832`). The prompt instructs the
model to append these lines (`assistant_agent.py:240-296`).

This coupling is the root cause of a recurring class of bugs: when the model phrases
the state line even slightly off — inline instead of on its own line, split across
lines, wrapped in markdown, mangled brackets, or dropped entirely — the parser
either misses the update or leaks raw JSON into the chat as bare content.

On current `main` the parser only matches a marker at the **start of a line** and
requires valid JSON; on any miss it appends the offending line straight back into
the visible reply (`assistant_agent.py:863-864, 876-877`). That fallback *is* the
leak. And because the parse happens *after* the model turn is done, a malformed
payload can't be corrected — there's nothing to hand back to the model. It's either
salvaged by a regex heuristic or lost/leaked.

## Why hardening the parser isn't the long-term fix

PR #1408 is the latest round of hardening the same channel, and its findings make a
good case for evolving the approach rather than continuing to patch it:

- **Small / self-hosted models drop the line.** Against the TACC MiniMax / gpt-oss
endpoints the model intermittently skips `SCHEMA_UPDATE` on the very turn a
decision is made, leaving the tracker behind. #1408 addressed this well by asking
the model to re-emit every committed field each turn, which is self-healing. It
does mean re-stating state on every reply and still relies on the model formatting
it correctly — both of which a typed channel would remove.
- **Raw JSON could still leak.** #1408 improved the parser to scan the whole reply,
decode JSON wherever it sits, and excise an uppercase marker's payload even when it
can't parse it, while leaving lowercase prose like "a few suggestions: [...]"
alone. Each of these refinements is sound, but together they're a growing set of
heuristics guarding a channel that was never designed to carry structured data.
- **Some fields shouldn't come from the model at all.** #1408 also showed that
`data_characteristics` and `gene_annotation` are properties of the chosen workflow
and assembly, not user decisions — the model correctly never emits them, so they
sat Pending and blocked handoff until we derived them from the catalog. The same
principle applies to the rest of the state: the system can own what it already
knows, rather than asking the model to format it.

Each of these is a symptom of state traveling through the prose channel, and each fix
makes the parser more forgiving rather than removing the need to parse.

## Proposed change

Move machine state onto **dedicated, schema-validated tool calls that the harness
intercepts** — typed, gate-able, and decoupled from how the model phrases its prose.

Concretely:
- Replace the `SCHEMA_UPDATE:` line convention with a `set_analysis_state` (or
similarly named) tool whose arguments are a typed Pydantic model of the schema
fields. The tool handler applies the same logic `_apply_schema_updates` does
today, including the dependent-field clearing chain.
- Do the same for suggestion chips (a `propose_suggestions` tool), keeping the
existing catalog-grounding validation in `_build_suggestion_chips` /
`_chip_entities_in_catalog`.
- The visible reply becomes *only* the message. Nothing structured can leak into it
because nothing structured travels through it.
- Keep the catalog-derived fills from #1408 (`data_characteristics`,
`gene_annotation`) — those stay system-owned regardless of channel.

## The key win: pydantic-ai validates and retries in-loop

The agent already runs on pydantic-ai. When state is a **typed tool argument**, a
malformed value isn't a post-hoc parsing failure we have to salvage or swallow — it's
a Pydantic `ValidationError` raised *inside the agent loop*, before the turn ends.
pydantic-ai catches it and hands the error back to the model as a tool-retry, so the
model gets a chance to correct the call and produce a valid one *on the same turn*.

That changes the failure model:
- **Today:** parse runs after the turn. Malformed → regex-salvage or lose/leak.
No feedback to the model, no correction.
- **With tool calls:** malformed → validation error → automatic retry with the
error message as context → the model fixes it. Bad state is rejected at the
boundary and repaired, never silently shown and never silently dropped.

This gives us the same self-healing that #1408's re-emit-every-turn approach was
reaching for, enforced by the framework rather than coaxed from the prompt, and
without spending tokens re-stating state on every reply.

## Cost / tradeoff

This can cost an extra model turn (tool call → tool result → final message) on turns
that both change state and answer the user. That's an accepted cost:
- State stops leaking into the chat structurally, not heuristically.
- Malformed state is caught and retried in-loop (above), never shown, never dropped.
- "Answer the question directly" and "set the right state" become independently
checkable, which also makes them independently testable in evals.
- We can stop asking the model to re-emit every field every turn.

## Scope

- **In:** schema/tracker updates, suggestion chips, and the handoff signal moved to
tool calls; prompt sections at `assistant_agent.py:240-296` rewritten; parser
fallback removed; tests migrated from the string-parsing cases to tool-call cases.
- **Out (for now):** the catalog query IR and other existing tools — this is only
about the state that currently rides in the reply text.

## Acceptance criteria

- No `SCHEMA_UPDATE:` / `SUGGESTIONS:` marker parsing remains in the reply path.
- Raw structured payloads can no longer appear in a visible reply under any
malformed-output scenario the #1408 tests exercise (inline / multi-line /
mangled-JSON / dropped-line).
- Tracker fields update via the tool with schema validation; dependent-field
clearing preserved.
- A malformed state tool call triggers a pydantic-ai validation retry rather than a
silent drop or a leak (covered by a test).
- Suggestion chips remain catalog-grounded (#1297 behavior).
- Handoff still fires only when the schema is complete.
- Behavior verified against a smaller/self-hosted model (the class of model that
motivated #1408), not just a strong one.

## References

- Current parser: `backend/api/app/services/assistant_agent.py:832`
(`_parse_structured_output`), prompt at `:240-296`, apply at `:959`
(`_apply_schema_updates`).
- Motivating hardening PR: #1408. Related chip grounding: #1297. Related tracker
auto-fill: #1324 / #1331.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.