labs/aicharts: plain prompt first, JSON mode only as fallback (#38)

Live, the json_schema decoder filled every step with the same key soup (order, step)
and dropped agg/by. The prompt's grammar + example steers better.


Claude-Session: https://claude.ai/code/session_01VDkN7ukK1kXkqov8FaZxva

Co-authored-by: Claude <noreply@anthropic.com>

.github/time/labs/aicharts/test_aicharts.mjs

@@ -112,26 +112,29 @@ const seeded = env.DB.writes.filter((w) => w.startsWith("INSERT OR IGNORE INTO o
 await call("/api/data");
 assert.equal(env.DB.writes.filter((w) => w.startsWith("INSERT OR IGNORE INTO orders")).length, seeded, "seeding happens once");
 
-// a spec: structured output first; the system prompt carries schema + sample, never all rows
+// a spec: plain prompt; the system prompt carries schema + sample, never all rows
 const good = { title: "Revenue by region", mark: "bar", x: "region", y: "value",
   transform: [{ op: "group", by: ["region"], agg: "sum", field: "amount", as: "value" }] };
 answers.push({ response: good });
 r = await post("/api/spec", { dataset: "orders", prompt: "revenue by region" });
 assert.equal(r.status, 200);
 assert.deepEqual((await r.json()).spec, good);
-assert.ok(calls[0].response_format, "asks for structured output");
+assert.equal(calls.length, 1);
+assert.equal(calls[0].response_format, undefined, "plain prompt first");
 const sys = calls[0].messages[0].content;
 assert.ok(sys.includes("region (string)") && sys.includes('"North"'));
 assert.ok(sys.length < 4000 && !sys.includes(`"id":${data.length}`), "the prompt has a sample, not the table");
 
-// a model without response_format, prose around the JSON, a bad first answer, then the fix
-answers.push(new Error("response_format unsupported"), { response: 'Here: {"mark": "bar", "x": "month", "y": "amount", "transform": [{"op": "group", "by": ["region"], "agg": "sum", "field": "amount"}]}' },
-             { choices: [{ message: { content: JSON.stringify({ ...good, mark: "line", x: "region" }) } }] });
+// no JSON at all -> JSON mode; prose around JSON is fine; a bad answer gets one repair round with the errors
+answers.push({ response: "I think a line chart would be nice." },
+             { response: '{"mark": "bar", "x": "month", "y": "amount", "transform": [{"op": "group", "by": ["region"], "agg": "sum", "field": "amount"}]}' },
+             { choices: [{ message: { content: "Fixed: " + JSON.stringify({ ...good, mark: "line", x: "region" }) } }] });
 calls.length = 0;
 r = await post("/api/spec", { dataset: "orders", prompt: "trend", prior: good });
 assert.equal(r.status, 200);
 assert.equal((await r.json()).spec.mark, "line");
-assert.ok(calls[1].messages.some((m) => m.content.includes("Current chart spec")), "the prior chart is context");
+assert.deepEqual(calls.map((c) => c.response_format?.type ?? null), [null, "json_object", null]);
+assert.ok(calls[0].messages.some((m) => m.content.includes("Current chart spec")), "the prior chart is context");
 assert.ok(calls.at(-1).messages.at(-1).content.includes("fails these checks"), "one repair round with the errors");
 
 // two bad answers: a 422 with the reasons, nothing drawn

.github/time/labs/aicharts/worker.js

@@ -10,7 +10,7 @@
 //
 // Read-only except saved_charts and quota. Bindings: DB (D1), AI (Workers AI). Vars: AI_MODEL, AI_DAILY, SAVE_DAILY, REV.
 // deploy.py uploads main.js, which hands this module the library source: export default make(src).
-import { validateSpec, MARKS, AGGS, OPS } from "./aicharts.js";
+import { validateSpec, MARKS, AGGS } from "./aicharts.js";
 
 const SCHEMA = [
   "CREATE TABLE IF NOT EXISTS orders (id INTEGER PRIMARY KEY, date TEXT, month TEXT, region TEXT, product TEXT, channel TEXT, units INTEGER, amount REAL)",
@@ -97,18 +97,6 @@ export function describe(data, columns) {
   return { columns, sample: data.slice(0, 5), aggregates, rows: data.length };
 }
 
-const SPEC_SCHEMA = {
-  type: "object", required: ["mark", "x", "y"],
-  properties: {
-    title: { type: "string" }, mark: { type: "string", enum: MARKS }, x: { type: "string" }, y: { type: "string" },
-    series: { type: "string" },
-    transform: { type: "array", maxItems: 8, items: { type: "object", required: ["op"], properties: {
-      op: { type: "string", enum: OPS }, field: { type: "string" }, by: { type: "array", items: { type: "string" } },
-      agg: { type: "string", enum: AGGS }, as: { type: "string" }, order: { type: "string", enum: ["asc", "desc"] },
-      eq: {}, in: { type: "array" }, gte: {}, lte: {}, step: { type: "number" }, n: { type: "integer" } } } },
-  },
-};
-
 export function systemPrompt(dataset, info) {
   return `You turn a question about a table into a chart spec. Answer with ONE JSON object and nothing else.
 Table "${dataset}" (${info.rows} rows). Columns: ${info.columns.map((c) => `${c.name} (${c.type})`).join(", ")}.
@@ -171,15 +159,15 @@ export function normalizeSpec(spec) {
   return out;
 }
 
+// Plain prompt first: the prompt's grammar and example steer small models better than a constrained decoder does (a
+// one-size schema for every step made them drop agg and by). JSON mode only when no object came back at all.
 async function model(env, messages) {
   const input = { messages, max_tokens: 700, temperature: 0.1 };
-  let out;
-  try {
-    out = await env.AI.run(env.AI_MODEL, { ...input, response_format: { type: "json_schema", json_schema: SPEC_SCHEMA } });
-  } catch {
-    out = await env.AI.run(env.AI_MODEL, input); // not every model takes response_format
+  const text = (out) => out?.response ?? out?.choices?.[0]?.message?.content ?? "";
+  let raw = text(await env.AI.run(env.AI_MODEL, input));
+  if (!firstObject(raw)) {
+    try { raw = text(await env.AI.run(env.AI_MODEL, { ...input, response_format: { type: "json_object" } })); } catch { /* model has no JSON mode */ }
   }
-  const raw = out?.response ?? out?.choices?.[0]?.message?.content ?? "";
   return { spec: normalizeSpec(firstObject(raw)), raw: typeof raw === "string" ? raw : JSON.stringify(raw) };
 }