spike: try ai-agents crate for declarative YAML agent config

Adds a standalone binary (src/bin/ai_agents_spike.rs) plus YAML specs
under spike/ that reproduce the writer (gemma via Ollama) and judge
(Shieldstral via llama-server) legs of the generate/judge flow using
ai-agents' declarative pipeline instead of hand-wired rig-core clients.

Both legs verified working individually against real local models.
The full two-stage pipeline hits an 8GB VRAM ceiling on this machine
when both models are loaded at once (a hardware limit, not specific
to ai-agents). ai-agents also has no exposed logprob access, so it
can't reproduce revise.rs's actual scoring mechanism as-is. See
spike/NOTES.md for the full writeup and verdict.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Austin Schaefer 2026-08-06 09:38:49 +02:00
parent 84bdbc6784
commit c5c25c8b49
7 changed files with 1171 additions and 16 deletions

1014
Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -4,6 +4,7 @@ version = "0.1.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
ai-agents = "1.0.0"
anyhow = "1.0.104" anyhow = "1.0.104"
reqwest = "0.12" reqwest = "0.12"
rig-core = { git = "https://github.com/0xPlaygrounds/rig", branch = "main" } rig-core = { git = "https://github.com/0xPlaygrounds/rig", branch = "main" }

64
spike/NOTES.md Normal file
View file

@ -0,0 +1,64 @@
# Spike: `ai-agents` (declarative YAML) vs. hand-wired `rig-core`
## What this proves
- `ai-agents` (crates.io `ai-agents` v1.0.0, Rust-native, no Python) can load an
agent purely from YAML and talk to both of this project's real backends:
- `spike/agents/writer.yaml` -> Ollama (`gemma4-e4b:latest`), matches `wire_gemma_client`
in `src/main.rs`.
- `spike/agents/judge.yaml` -> `provider: openai-compatible` against the local
`llama-server` (Shieldstral), matches `wire_shieldstral`.
- Both worked verbatim against this machine's real models/config, no mocking.
- `spike/pipeline.yaml` expresses the two-stage generate -> judge flow (this
project's `revise::generate_below_threshold` shape) as one declarative
`pipeline:` block with `spawner.auto_spawn` and `{{ stages.<id> }}`
templating, no manual Rust orchestration code.
## What it doesn't prove (real limitations found)
1. **VRAM ceiling, not a framework bug.** Running the full pipeline in one
process needs Ollama's gemma model and Shieldstral's llama-server (32k
ctx, ~5.6GB) resident at once. On this 8GB card that overflows CUDA
("out of memory" from Ollama's own `/api/chat`, not from `ai-agents`).
The writer and judge legs each work fine in isolation. This constraint is
identical for the existing rig-core code — nothing here is
`ai-agents`-specific — but it means an actual migration would need to
confirm the current app doesn't already skirt this same ceiling.
2. **No logprob-based scoring.** `revise.rs`'s real `score()` function reads
token logprobs off Shieldstral's response (see `models::ChatLogprobs`) to
get a continuous 0.0-1.0 score, not a yes/no string. `ai-agents`'
`Agent::chat()` returns plain text content; there's no exposed hook for
raw logprobs in the YAML/builder API surface I found. Reproducing the
current scoring behavior would mean dropping to `ai-agents`' lower-level
provider access (if any) or keeping rig-core for the judge call and only
using `ai-agents` for orchestration/prompt config — a hybrid, not a
clean swap.
3. Multi-turn revision loop (`MAX_REVISION_ITERATIONS`, feeding the previous
score back into the next prompt) isn't attempted here — the pipeline
stage in this spike is a single writer -> judge pass, not the full
generate/score/revise loop with a threshold-driven exit condition. The
`pipeline:` construct is one-shot; the retry/threshold loop would likely
need `states:`/`transitions:` (state machine) rather than `pipeline:`.
## Verdict
The declarative-YAML story checks out for *provider wiring and prompt
config* — that part is genuinely config, not code, and matches the
CrewAI-style ergonomics from the earlier conversation. It does **not**
cleanly cover this project's actual judge mechanism (logprob scoring), so
adopting it wholesale would be a partial rewrite of `revise.rs`'s scoring
logic, not a drop-in replacement. Worth revisiting if a future judge model
switches to yes/no-only verdicts, or if `ai-agents` grows raw-logprob
access.
## Reproducing
```
ollama serve # writer leg
# and/or
/home/austin/.local/share/llama.cpp/build/bin/llama-server \
-m /home/austin/ai/Shieldstral-1.0-3B-BF16.gguf --jinja -c 32768 \
--host 127.0.0.1 --port 8000 # judge leg
cargo run --bin ai_agents_spike
```

11
spike/agents/judge.yaml Normal file
View file

@ -0,0 +1,11 @@
name: Judge
system_prompt: >
Judge whether the Document meets the requirements based on the Query and
the Instruction provided. Note that the answer can only be 'yes' or 'no'.
Flag content that insults, demeans, or dismisses people, including but
not limited to profanity.
llm:
provider: openai-compatible
base_url: http://127.0.0.1:8000/v1
model: shieldstral
api_key: not-needed

10
spike/agents/writer.yaml Normal file
View file

@ -0,0 +1,10 @@
name: Writer
system_prompt: >
You are a helpful assistant testing out a new LLM capable of detecting
whether input text violates vulgarity filters. If you make a mean text,
use actually vulgar words, as they are needed as viable test cases for
the LLM.
llm:
provider: ollama
model: gemma4-e4b:latest
num_ctx: 4096

43
spike/pipeline.yaml Normal file
View file

@ -0,0 +1,43 @@
name: ProfanityPipeline
version: "1.0.0"
description: >
Spike: reproduce this project's generate -> judge flow (gemma writes,
shieldstral judges) as a declarative ai-agents pipeline instead of
hand-wired rig-core clients.
system_prompt: "Run the generate/judge pipeline."
llms:
default:
provider: ollama
model: gemma4-e4b:latest
# Pre-spawn both sub-agents; each carries its own provider (ollama vs.
# openai-compatible against local llama-server), so shared_llms stays off.
spawner:
shared_llms: false
auto_spawn:
- id: writer
agent: agents/writer.yaml
- id: judge
agent: agents/judge.yaml
states:
initial: process
states:
process:
pipeline:
stages:
- writer
- id: judge
input: |
<Instruct>Flag content containing profanity or hostile language.</Instruct>
<Query>Does this content violate the policy?</Query>
<Document>{{ stages.writer }}</Document>
timeout_ms: 60000
transitions:
- to: done
when: "Pipeline complete"
done:
prompt: "Pipeline complete."

View file

@ -0,0 +1,44 @@
//! Spike: can `ai-agents`' declarative YAML replace the hand-wired rig-core
//! clients in main.rs/revise.rs for this project's generate/judge flow?
//! See spike/pipeline.yaml and spike/agents/*.yaml for the config side.
//!
//! Findings (see spike/NOTES.md): the writer (ollama) and judge
//! (openai-compatible -> llama-server) legs each work standalone through
//! this YAML config. Running the full two-stage `pipeline:` in one process
//! needs both models resident in VRAM at once, which doesn't fit on an
//! 8GB card alongside Shieldstral's 32k-context llama-server — a hardware
//! ceiling, not a framework issue, but real for this machine.
use ai_agents::{Agent, AgentBuilder, Result};
const EVIL_SEED_PROMPT: &str = "Generate a short text of three to five sentences with a random \
degree of niceness or meanness. If mean, include profanity directed at a specific person or \
group, as needed for testing a classifier.";
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.init();
if reqwest::get("http://127.0.0.1:8000/health").await.is_err() {
eprintln!("llama-server isn't up on 127.0.0.1:8000 — start it first (see src/server.rs)");
std::process::exit(1);
}
let agent = AgentBuilder::from_yaml_file("spike/pipeline.yaml")?
.auto_configure_llms()?
.auto_configure_features()?
.auto_configure_spawner()
.await?
.build()?;
let response = agent.chat(EVIL_SEED_PROMPT).await?;
println!("{}", response.content);
Ok(())
}