@pgsage/core 0.1.0

package/dist/orchestrator/index.js

@@ -0,0 +1,126 @@
+/**
+ * Orchestrator
+ *
+ * Ties retriever → planner → validator → estimator → executor → corrector
+ * into a single agent loop and exposes the SDK's primary entrypoint:
+ * `agent.query(question)`.
+ *
+ * Each stage is composable — callers can import individual stage factories
+ * directly if they need finer control. The orchestrator just wires them up
+ * with sensible defaults.
+ */
+import { retrieve } from "../retriever/index.js";
+import { createPlanner } from "../planner/index.js";
+import { validate } from "../validator/index.js";
+import { createEstimator } from "../estimator/index.js";
+import { createExecutor } from "../executor/index.js";
+import { createCorrector } from "../corrector/index.js";
+import { createExplainer } from "../explainer/index.js";
+// ---------------------------------------------------------------------------
+// Factory
+// ---------------------------------------------------------------------------
+/**
+ * Create the pgsage agent. All stage instances are created once and reused
+ * across query() calls.
+ *
+ * @example
+ * ```ts
+ * const agent = createAgent({
+ *   pool,
+ *   readonlyPool,
+ *   embedder: new VoyageEmbeddingProvider({ apiKey: process.env.VOYAGE_API_KEY }),
+ * });
+ * const response = await agent.query("What is the median household income in RI?");
+ * ```
+ */
+export function createAgent(config) {
+    const topK = config.topK ?? 10;
+    const planner = createPlanner({ model: config.model });
+    const estimator = createEstimator({
+        pool: config.pool,
+        costThreshold: config.costThreshold,
+        rowThreshold: config.rowThreshold,
+    });
+    const executor = createExecutor({
+        pool: config.readonlyPool,
+        rowLimit: config.rowLimit,
+        statementTimeout: config.statementTimeout,
+    });
+    const corrector = createCorrector({
+        planner,
+        maxRetries: config.maxRetries,
+    });
+    const explainer = createExplainer({
+        model: config.explainerModel ?? config.model,
+    });
+    return {
+        async query(question, options = {}) {
+            const progress = options.onProgress ?? (() => undefined);
+            // Step 0: Retrieve relevant schema chunks via pgvector similarity search
+            progress(0, "Searching schema\u2026");
+            const context = await retrieve(config.pool, question, config.embedder, { topK });
+            // Step 1: Generate the initial SQL plan
+            progress(1, "Generating SQL\u2026");
+            const initialPlan = await planner.plan(question, context);
+            // Step 2: Emit before the first execute attempt
+            progress(2, "Validating & executing\u2026");
+            // Build the execute pipeline: validate → estimate → execute
+            // (Used by both the corrector and the final execution)
+            const executePipeline = async (sql) => {
+                // 2a. AST safety validation
+                const validation = validate(sql);
+                if (!validation.valid) {
+                    throw new Error(`Validation failed: ${validation.reason ?? "unknown reason"}`);
+                }
+                // 2b. Cost estimation (reject if too expensive)
+                const estimate = await estimator.estimate(sql);
+                if (estimate.exceedsThreshold) {
+                    throw new Error(`Cost check failed: ${estimate.thresholdReason ?? "exceeds threshold"}`);
+                }
+                // 2c. Execute on the read-only pool
+                return executor.execute(sql);
+            };
+            // Step 3 (conditional): Run through the corrector
+            // The corrector emits step 3 ("Correcting query…") only when it retries.
+            let correctorAttempt = 0;
+            const correction = await corrector.correct(question, context, initialPlan, async (sql) => {
+                // Emit "Correcting" on second+ attempts (first attempt is covered by step 2 above)
+                if (correctorAttempt > 0) {
+                    progress(3, "Correcting query\u2026");
+                }
+                correctorAttempt++;
+                return executePipeline(sql);
+            });
+            if (correction.exhausted) {
+                throw new Error(`pgsage could not produce a valid query after ${config.maxRetries ?? 3} retries. ` +
+                    `Last error: ${correction.attempts.at(-1)?.error ?? "unknown"}`);
+            }
+            // Final execution: re-run the winning plan to get the result rows
+            const finalResult = await executePipeline(correction.plan.sql);
+            // Step 4: Generate a natural-language answer via the explainer LLM
+            progress(4, "Summarizing results\u2026");
+            const answer = await explainer.explain({
+                question,
+                sql: correction.plan.sql,
+                assumptions: correction.plan.assumptions,
+                displayHint: correction.plan.displayHint,
+                fields: finalResult.fields,
+                rows: finalResult.rows,
+                rowCount: finalResult.rowCount,
+                truncated: finalResult.truncated,
+            });
+            return {
+                answer,
+                sql: correction.plan.sql,
+                assumptions: correction.plan.assumptions,
+                displayHint: correction.plan.displayHint,
+                rows: finalResult.rows,
+                rowCount: finalResult.rowCount,
+                fields: finalResult.fields,
+                truncated: finalResult.truncated,
+                attempts: correction.attempts,
+            };
+        },
+    };
+}
+//# sourceMappingURL=index.js.map

package/dist/orchestrator/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/orchestrator/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;GAUG;AAIH,OAAO,EAAE,QAAQ,EAAE,MAAM,uBAAuB,CAAC;AACjD,OAAO,EAAE,aAAa,EAAE,MAAM,qBAAqB,CAAC;AACpD,OAAO,EAAE,QAAQ,EAAE,MAAM,uBAAuB,CAAC;AACjD,OAAO,EAAE,eAAe,EAAE,MAAM,uBAAuB,CAAC;AACxD,OAAO,EAAE,cAAc,EAAE,MAAM,sBAAsB,CAAC;AACtD,OAAO,EAAE,eAAe,EAAE,MAAM,uBAAuB,CAAC;AACxD,OAAO,EAAE,eAAe,EAAE,MAAM,uBAAuB,CAAC;AA4FxD,8EAA8E;AAC9E,UAAU;AACV,8EAA8E;AAE9E;;;;;;;;;;;;;GAaG;AACH,MAAM,UAAU,WAAW,CAAC,MAA0B;IACpD,MAAM,IAAI,GAAG,MAAM,CAAC,IAAI,IAAI,EAAE,CAAC;IAE/B,MAAM,OAAO,GAAG,aAAa,CAAC,EAAE,KAAK,EAAE,MAAM,CAAC,KAAK,EAAE,CAAC,CAAC;IAEvD,MAAM,SAAS,GAAG,eAAe,CAAC;QAChC,IAAI,EAAE,MAAM,CAAC,IAAI;QACjB,aAAa,EAAE,MAAM,CAAC,aAAa;QACnC,YAAY,EAAE,MAAM,CAAC,YAAY;KAClC,CAAC,CAAC;IAEH,MAAM,QAAQ,GAAG,cAAc,CAAC;QAC9B,IAAI,EAAE,MAAM,CAAC,YAAY;QACzB,QAAQ,EAAE,MAAM,CAAC,QAAQ;QACzB,gBAAgB,EAAE,MAAM,CAAC,gBAAgB;KAC1C,CAAC,CAAC;IAEH,MAAM,SAAS,GAAG,eAAe,CAAC;QAChC,OAAO;QACP,UAAU,EAAE,MAAM,CAAC,UAAU;KAC9B,CAAC,CAAC;IAEH,MAAM,SAAS,GAAG,eAAe,CAAC;QAChC,KAAK,EAAE,MAAM,CAAC,cAAc,IAAI,MAAM,CAAC,KAAK;KAC7C,CAAC,CAAC;IAEH,OAAO;QACL,KAAK,CAAC,KAAK,CAAC,QAAgB,EAAE,UAAwB,EAAE;YACtD,MAAM,QAAQ,GAAG,OAAO,CAAC,UAAU,IAAI,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,CAAC;YAEzD,yEAAyE;YACzE,QAAQ,CAAC,CAAC,EAAE,wBAAwB,CAAC,CAAC;YACtC,MAAM,OAAO,GAAG,MAAM,QAAQ,CAAC,MAAM,CAAC,IAAI,EAAE,QAAQ,EAAE,MAAM,CAAC,QAAQ,EAAE,EAAE,IAAI,EAAE,CAAC,CAAC;YAEjF,wCAAwC;YACxC,QAAQ,CAAC,CAAC,EAAE,sBAAsB,CAAC,CAAC;YACpC,MAAM,WAAW,GAAG,MAAM,OAAO,CAAC,IAAI,CAAC,QAAQ,EAAE,OAAO,CAAC,CAAC;YAE1D,gDAAgD;YAChD,QAAQ,CAAC,CAAC,EAAE,8BAA8B,CAAC,CAAC;YAE5C,4DAA4D;YAC5D,uDAAuD;YACvD,MAAM,eAAe,GAAG,KAAK,EAAE,GAAW,EAA4B,EAAE;gBACtE,4BAA4B;gBAC5B,MAAM,UAAU,GAAG,QAAQ,CAAC,GAAG,CAAC,CAAC;gBACjC,IAAI,CAAC,UAAU,CAAC,KAAK,EAAE,CAAC;oBACtB,MAAM,IAAI,KAAK,CAAC,sBAAsB,UAAU,CAAC,MAAM,IAAI,gBAAgB,EAAE,CAAC,CAAC;gBACjF,CAAC;gBAED,gDAAgD;gBAChD,MAAM,QAAQ,GAAG,MAAM,SAAS,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC;gBAC/C,IAAI,QAAQ,CAAC,gBAAgB,EAAE,CAAC;oBAC9B,MAAM,IAAI,KAAK,CAAC,sBAAsB,QAAQ,CAAC,eAAe,IAAI,mBAAmB,EAAE,CAAC,CAAC;gBAC3F,CAAC;gBAED,oCAAoC;gBACpC,OAAO,QAAQ,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC;YAC/B,CAAC,CAAC;YAEF,kDAAkD;YAClD,yEAAyE;YACzE,IAAI,gBAAgB,GAAG,CAAC,CAAC;YACzB,MAAM,UAAU,GAAG,MAAM,SAAS,CAAC,OAAO,CACxC,QAAQ,EACR,OAAO,EACP,WAAW,EACX,KAAK,EAAE,GAAW,EAAE,EAAE;gBACpB,mFAAmF;gBACnF,IAAI,gBAAgB,GAAG,CAAC,EAAE,CAAC;oBACzB,QAAQ,CAAC,CAAC,EAAE,wBAAwB,CAAC,CAAC;gBACxC,CAAC;gBACD,gBAAgB,EAAE,CAAC;gBACnB,OAAO,eAAe,CAAC,GAAG,CAAC,CAAC;YAC9B,CAAC,CACF,CAAC;YAEF,IAAI,UAAU,CAAC,SAAS,EAAE,CAAC;gBACzB,MAAM,IAAI,KAAK,CACb,gDAAgD,MAAM,CAAC,UAAU,IAAI,CAAC,YAAY;oBAChF,eAAe,UAAU,CAAC,QAAQ,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,EAAE,KAAK,IAAI,SAAS,EAAE,CAClE,CAAC;YACJ,CAAC;YAED,kEAAkE;YAClE,MAAM,WAAW,GAAG,MAAM,eAAe,CAAC,UAAU,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;YAE/D,mEAAmE;YACnE,QAAQ,CAAC,CAAC,EAAE,2BAA2B,CAAC,CAAC;YACzC,MAAM,MAAM,GAAG,MAAM,SAAS,CAAC,OAAO,CAAC;gBACrC,QAAQ;gBACR,GAAG,EAAE,UAAU,CAAC,IAAI,CAAC,GAAG;gBACxB,WAAW,EAAE,UAAU,CAAC,IAAI,CAAC,WAAW;gBACxC,WAAW,EAAE,UAAU,CAAC,IAAI,CAAC,WAAW;gBACxC,MAAM,EAAE,WAAW,CAAC,MAAM;gBAC1B,IAAI,EAAE,WAAW,CAAC,IAAI;gBACtB,QAAQ,EAAE,WAAW,CAAC,QAAQ;gBAC9B,SAAS,EAAE,WAAW,CAAC,SAAS;aACjC,CAAC,CAAC;YAEH,OAAO;gBACL,MAAM;gBACN,GAAG,EAAE,UAAU,CAAC,IAAI,CAAC,GAAG;gBACxB,WAAW,EAAE,UAAU,CAAC,IAAI,CAAC,WAAW;gBACxC,WAAW,EAAE,UAAU,CAAC,IAAI,CAAC,WAAW;gBACxC,IAAI,EAAE,WAAW,CAAC,IAAI;gBACtB,QAAQ,EAAE,WAAW,CAAC,QAAQ;gBAC9B,MAAM,EAAE,WAAW,CAAC,MAAM;gBAC1B,SAAS,EAAE,WAAW,CAAC,SAAS;gBAChC,QAAQ,EAAE,UAAU,CAAC,QAAQ;aAC9B,CAAC;QACJ,CAAC;KACF,CAAC;AACJ,CAAC"}

package/dist/planner/index.d.ts

@@ -0,0 +1,83 @@
+/**
+ * Planner
+ *
+ * Calls Claude (via Mastra Agent) with the retrieved schema context and
+ * few-shot examples to generate a SQL query plus stated assumptions about
+ * ambiguous terms.
+ *
+ * Usage:
+ *   const planner = createPlanner({ model: "anthropic/claude-sonnet-4-6" });
+ *   const result = await planner.plan(question, retrievedChunks);
+ *   // result.sql — the generated SELECT query
+ *   // result.assumptions — list of interpretive assumptions
+ */
+import type { RetrievedSchemaChunk } from "../retriever/index.js";
+import { type FewShotExample } from "./prompt.js";
+/** Display hint values — matches the four visualization types in the web UI. */
+export type DisplayHint = "big" | "statrow" | "bar" | "grouped" | "table";
+export interface PlanResult {
+    /** The generated PostgreSQL SELECT query. Empty string if no query could be produced. */
+    sql: string;
+    /** Plain-English assumptions made about ambiguous terms or requirements. */
+    assumptions: string[];
+    /**
+     * Suggested UI visualization layout for the result.
+     * - `"big"` — single scalar answer (one row, one numeric value)
+     * - `"statrow"` — 2–4 related metrics about the same subject
+     * - `"bar"` — ranked list of one metric across many geographies
+     * - `"grouped"` — two-series comparison across geographies
+     * - `"table"` — default fallback for general tabular data
+     */
+    displayHint: DisplayHint;
+}
+export interface PlannerConfig {
+    /**
+     * Mastra model router string.
+     * @default "anthropic/claude-sonnet-4-6"
+     */
+    model?: string;
+    /**
+     * Override the default few-shot examples sent with every plan call.
+     * Pass an empty array to disable examples.
+     */
+    examples?: FewShotExample[];
+}
+/** One prior attempt that failed, used by the corrector to re-plan. */
+export interface PriorAttempt {
+    /** The SQL that was generated and failed (or returned suspicious results). */
+    sql: string;
+    /** The error message or zero-row explanation from the executor/corrector. */
+    error: string;
+}
+export interface PlanOptions {
+    /**
+     * Per-call few-shot example override. Falls back to the config-level
+     * examples, then the built-in Census examples.
+     */
+    examples?: FewShotExample[];
+    /**
+     * Prior failed attempts to include in the prompt so the model can learn
+     * from its mistakes. Populated by the corrector on retry calls.
+     */
+    errorContext?: PriorAttempt[];
+}
+export interface Planner {
+    /**
+     * Generate a SQL query and assumptions for a natural-language question.
+     *
+     * @param question  The user's natural-language question.
+     * @param context   Schema chunks retrieved from pgvector similarity search.
+     * @param options   Optional per-call overrides.
+     */
+    plan(question: string, context: RetrievedSchemaChunk[], options?: PlanOptions): Promise<PlanResult>;
+}
+/**
+ * Creates a reusable Planner that wraps a Mastra Agent.
+ * The Agent is instantiated once and reused across all plan() calls.
+ *
+ * Requires the ANTHROPIC_API_KEY environment variable to be set.
+ */
+export declare function createPlanner(config?: PlannerConfig): Planner;
+export { SYSTEM_PROMPT, DEFAULT_FEW_SHOT_EXAMPLES, formatSchemaContext, buildUserMessage } from "./prompt.js";
+export type { FewShotExample } from "./prompt.js";
+//# sourceMappingURL=index.d.ts.map

package/dist/planner/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/planner/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;GAYG;AAIH,OAAO,KAAK,EAAE,oBAAoB,EAAE,MAAM,uBAAuB,CAAC;AAClE,OAAO,EAIL,KAAK,cAAc,EACpB,MAAM,aAAa,CAAC;AAMrB,gFAAgF;AAChF,MAAM,MAAM,WAAW,GAAG,KAAK,GAAG,SAAS,GAAG,KAAK,GAAG,SAAS,GAAG,OAAO,CAAC;AAE1E,MAAM,WAAW,UAAU;IACzB,yFAAyF;IACzF,GAAG,EAAE,MAAM,CAAC;IACZ,4EAA4E;IAC5E,WAAW,EAAE,MAAM,EAAE,CAAC;IACtB;;;;;;;OAOG;IACH,WAAW,EAAE,WAAW,CAAC;CAC1B;AAED,MAAM,WAAW,aAAa;IAC5B;;;OAGG;IACH,KAAK,CAAC,EAAE,MAAM,CAAC;IACf;;;OAGG;IACH,QAAQ,CAAC,EAAE,cAAc,EAAE,CAAC;CAC7B;AAED,uEAAuE;AACvE,MAAM,WAAW,YAAY;IAC3B,8EAA8E;IAC9E,GAAG,EAAE,MAAM,CAAC;IACZ,6EAA6E;IAC7E,KAAK,EAAE,MAAM,CAAC;CACf;AAED,MAAM,WAAW,WAAW;IAC1B;;;OAGG;IACH,QAAQ,CAAC,EAAE,cAAc,EAAE,CAAC;IAC5B;;;OAGG;IACH,YAAY,CAAC,EAAE,YAAY,EAAE,CAAC;CAC/B;AAED,MAAM,WAAW,OAAO;IACtB;;;;;;OAMG;IACH,IAAI,CACF,QAAQ,EAAE,MAAM,EAChB,OAAO,EAAE,oBAAoB,EAAE,EAC/B,OAAO,CAAC,EAAE,WAAW,GACpB,OAAO,CAAC,UAAU,CAAC,CAAC;CACxB;AA2BD;;;;;GAKG;AACH,wBAAgB,aAAa,CAAC,MAAM,GAAE,aAAkB,GAAG,OAAO,CA2BjE;AAMD,OAAO,EAAE,aAAa,EAAE,yBAAyB,EAAE,mBAAmB,EAAE,gBAAgB,EAAE,MAAM,aAAa,CAAC;AAC9G,YAAY,EAAE,cAAc,EAAE,MAAM,aAAa,CAAC"}

package/dist/planner/index.js

@@ -0,0 +1,67 @@
+/**
+ * Planner
+ *
+ * Calls Claude (via Mastra Agent) with the retrieved schema context and
+ * few-shot examples to generate a SQL query plus stated assumptions about
+ * ambiguous terms.
+ *
+ * Usage:
+ *   const planner = createPlanner({ model: "anthropic/claude-sonnet-4-6" });
+ *   const result = await planner.plan(question, retrievedChunks);
+ *   // result.sql — the generated SELECT query
+ *   // result.assumptions — list of interpretive assumptions
+ */
+import { Agent } from "@mastra/core/agent";
+import { z } from "zod";
+import { SYSTEM_PROMPT, DEFAULT_FEW_SHOT_EXAMPLES, buildUserMessage, } from "./prompt.js";
+// ---------------------------------------------------------------------------
+// Zod schema for structured output
+// ---------------------------------------------------------------------------
+const planResultSchema = z.object({
+    sql: z.string().describe("A valid PostgreSQL SELECT query, or empty string if none can be produced"),
+    assumptions: z
+        .array(z.string())
+        .describe("Plain-English assumptions about ambiguous terms or interpretations"),
+    displayHint: z
+        .enum(["big", "statrow", "bar", "grouped", "table"])
+        .describe("Suggested UI visualization layout. " +
+        "Use 'big' for a single scalar answer (one row, one primary numeric value, e.g. a county's median income). " +
+        "Use 'statrow' for 2-4 related metrics about the same subject (e.g. a housing snapshot with home value, rent, and homeownership rate). " +
+        "Use 'bar' for a ranked list of one metric across many geographies (e.g. top 10 counties by income). " +
+        "Use 'grouped' for two-series comparisons across geographies (e.g. rent in 2020 vs 2024 by state). " +
+        "Use 'table' as the default fallback for general tabular data that does not fit the above."),
+});
+// ---------------------------------------------------------------------------
+// Factory
+// ---------------------------------------------------------------------------
+/**
+ * Creates a reusable Planner that wraps a Mastra Agent.
+ * The Agent is instantiated once and reused across all plan() calls.
+ *
+ * Requires the ANTHROPIC_API_KEY environment variable to be set.
+ */
+export function createPlanner(config = {}) {
+    const modelId = config.model ?? "anthropic/claude-sonnet-4-6";
+    const configExamples = config.examples ?? DEFAULT_FEW_SHOT_EXAMPLES;
+    const agent = new Agent({
+        id: "pgsage-planner",
+        name: "pgsage-planner",
+        instructions: SYSTEM_PROMPT,
+        model: modelId,
+    });
+    return {
+        async plan(question, context, options = {}) {
+            const examples = options.examples ?? configExamples;
+            const userMessage = buildUserMessage(question, context, examples, options.errorContext);
+            const result = await agent.generate(userMessage, {
+                structuredOutput: { schema: planResultSchema },
+            });
+            return result.object;
+        },
+    };
+}
+// ---------------------------------------------------------------------------
+// Convenience re-exports
+// ---------------------------------------------------------------------------
+export { SYSTEM_PROMPT, DEFAULT_FEW_SHOT_EXAMPLES, formatSchemaContext, buildUserMessage } from "./prompt.js";
+//# sourceMappingURL=index.js.map

package/dist/planner/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/planner/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;GAYG;AAEH,OAAO,EAAE,KAAK,EAAE,MAAM,oBAAoB,CAAC;AAC3C,OAAO,EAAE,CAAC,EAAE,MAAM,KAAK,CAAC;AAExB,OAAO,EACL,aAAa,EACb,yBAAyB,EACzB,gBAAgB,GAEjB,MAAM,aAAa,CAAC;AA0ErB,8EAA8E;AAC9E,mCAAmC;AACnC,8EAA8E;AAE9E,MAAM,gBAAgB,GAAG,CAAC,CAAC,MAAM,CAAC;IAChC,GAAG,EAAE,CAAC,CAAC,MAAM,EAAE,CAAC,QAAQ,CAAC,0EAA0E,CAAC;IACpG,WAAW,EAAE,CAAC;SACX,KAAK,CAAC,CAAC,CAAC,MAAM,EAAE,CAAC;SACjB,QAAQ,CAAC,oEAAoE,CAAC;IACjF,WAAW,EAAE,CAAC;SACX,IAAI,CAAC,CAAC,KAAK,EAAE,SAAS,EAAE,KAAK,EAAE,SAAS,EAAE,OAAO,CAAC,CAAC;SACnD,QAAQ,CACP,qCAAqC;QACnC,4GAA4G;QAC5G,wIAAwI;QACxI,sGAAsG;QACtG,oGAAoG;QACpG,2FAA2F,CAC9F;CACJ,CAAC,CAAC;AAEH,8EAA8E;AAC9E,UAAU;AACV,8EAA8E;AAE9E;;;;;GAKG;AACH,MAAM,UAAU,aAAa,CAAC,SAAwB,EAAE;IACtD,MAAM,OAAO,GAAG,MAAM,CAAC,KAAK,IAAI,6BAA6B,CAAC;IAC9D,MAAM,cAAc,GAAG,MAAM,CAAC,QAAQ,IAAI,yBAAyB,CAAC;IAEpE,MAAM,KAAK,GAAG,IAAI,KAAK,CAAC;QACtB,EAAE,EAAE,gBAAgB;QACpB,IAAI,EAAE,gBAAgB;QACtB,YAAY,EAAE,aAAa;QAC3B,KAAK,EAAE,OAAO;KACf,CAAC,CAAC;IAEH,OAAO;QACL,KAAK,CAAC,IAAI,CACR,QAAgB,EAChB,OAA+B,EAC/B,UAAuB,EAAE;YAEzB,MAAM,QAAQ,GAAG,OAAO,CAAC,QAAQ,IAAI,cAAc,CAAC;YACpD,MAAM,WAAW,GAAG,gBAAgB,CAAC,QAAQ,EAAE,OAAO,EAAE,QAAQ,EAAE,OAAO,CAAC,YAAY,CAAC,CAAC;YAExF,MAAM,MAAM,GAAG,MAAM,KAAK,CAAC,QAAQ,CAAC,WAAW,EAAE;gBAC/C,gBAAgB,EAAE,EAAE,MAAM,EAAE,gBAAgB,EAAE;aAC/C,CAAC,CAAC;YAEH,OAAO,MAAM,CAAC,MAAM,CAAC;QACvB,CAAC;KACF,CAAC;AACJ,CAAC;AAED,8EAA8E;AAC9E,yBAAyB;AACzB,8EAA8E;AAE9E,OAAO,EAAE,aAAa,EAAE,yBAAyB,EAAE,mBAAmB,EAAE,gBAAgB,EAAE,MAAM,aAAa,CAAC"}

package/dist/planner/prompt.d.ts

@@ -0,0 +1,37 @@
+/**
+ * Planner prompt utilities
+ *
+ * Contains:
+ *  - SYSTEM_PROMPT: the base system prompt for the SQL planner
+ *  - FEW_SHOT_EXAMPLES: representative Q→SQL pairs for the Census schema
+ *  - formatSchemaContext(): converts RetrievedSchemaChunk[] into a compact,
+ *    LLM-readable schema block
+ */
+import type { RetrievedSchemaChunk } from "../retriever/index.js";
+export declare const SYSTEM_PROMPT = "You are an expert PostgreSQL query writer for a schema-aware SQL agent.\n\n## Your task\nGiven a natural-language question and the relevant schema context, produce:\n1. A single, valid PostgreSQL SELECT query that answers the question.\n2. A list of assumptions you made about any ambiguous terms or requirements.\n\n## Rules\n- Output ONLY a SELECT statement. Never produce INSERT, UPDATE, DELETE, DROP, ALTER, TRUNCATE, CREATE, GRANT, REVOKE, COPY, or any other non-SELECT statement.\n- Never produce multiple statements separated by semicolons.\n- Never use EXECUTE, PERFORM, or dynamic SQL.\n- Always use schema-qualified table names (e.g. public.geography, not just geography).\n- Prefer explicit column names over SELECT *.\n- Always write the estimate column as `e.estimate` \u2014 never alias it away from that name (e.g. do NOT write `e.estimate AS median_income`). Other columns may use aliases for readability.\n- Add a LIMIT clause if the question asks for top-N results; default to LIMIT 100 if a row cap is not otherwise specified and the query could return unbounded rows.\n- When sorting on a nullable numeric column, always add NULLS LAST (e.g. ORDER BY e.estimate DESC NULLS LAST) so that NULL values do not appear at the top of the result.\n- Use standard PostgreSQL functions only. Do not use vendor extensions.\n- If you cannot produce a valid SELECT query from the given context, return an empty string for sql and explain the issue in assumptions.\n\n## Variable categories\nThe variables table uses exactly these 8 category values. Use the correct one \u2014 do not guess:\n\n  income          \u2014 household income, per-capita income, Gini index, income brackets\n  poverty         \u2014 population below the poverty line, poverty ratios\n  housing         \u2014 housing units, home values, gross rent, occupancy, rent burden\n  employment      \u2014 labor force participation, unemployment counts, industry, occupation\n  demographics    \u2014 total population, age, sex, race, Hispanic/Latino origin\n  education       \u2014 educational attainment (high school, bachelor's, graduate degrees)\n  transportation  \u2014 commute mode (drove alone, transit, worked from home), travel time\n  health_insurance \u2014 health insurance coverage by age and sex\n\nCommon mistakes to avoid:\n  - Bachelor's degree \u2192 category = 'education' (NOT 'demographics')\n  - Poverty / below poverty line \u2192 category = 'poverty' (NOT 'income')\n  - Commute mode, work from home, travel time \u2192 category = 'transportation' (NOT 'employment')\n  - Health insurance coverage \u2192 category = 'health_insurance'\n\n## Estimates are raw counts, not percentages\nMost estimate values are raw COUNTS (number of people, number of housing units), NOT percentages or rates.\nTo express a rate or percentage, you must compute it by dividing a count by its relevant total:\n\n  homeownership rate = 100.0 * owner_occupied.estimate / NULLIF(total_occupied.estimate, 0)\n  poverty rate       = 100.0 * below_poverty.estimate  / NULLIF(total_for_poverty.estimate, 0)\n  drove-alone rate   = 100.0 * drove_alone.estimate    / NULLIF(total_workers.estimate, 0)\n\nNever compare a raw count estimate directly to a percentage threshold. For example:\n  WRONG: owner_e.estimate < 60        \u2014 this compares a unit count to 60, not 60%\n  RIGHT: 100.0 * owner_e.estimate / NULLIF(total_e.estimate, 0) < 60\n\nWhen a question asks for a rate or percentage, always join in the denominator total variable\nand compute the ratio explicitly.\n\n## ILIKE precision\nILIKE label patterns can match multiple variables. When a single specific metric is needed:\n- Use a tighter pattern (e.g. `v.label ILIKE '%median household income in the past 12 months%'` rather than `'%median household income%'` which also matches age-group breakdowns).\n- Or add `AND v.concept ILIKE '...'` to narrow further.\n- Or filter to a known variable_code directly (e.g. `e.variable_code = 'B19013_001E'`) when the code is clear from the question or schema context.\nAmbiguous ILIKE matches combined with ORDER BY will surface whichever matched variable has the highest/lowest value \u2014 which may not be the intended one.\n\n## Assumption surfacing\nFor every ambiguous term, geographic scope decision, join assumption, or data interpretation, add a plain-English sentence to the assumptions list. Examples:\n- \"Interpreting 'large counties' as counties with an estimated population > 100,000.\"\n- \"Using the most recent ACS 5-year estimate (2024 vintage).\"\n- \"Joining estimates to geography on geo_id to resolve county names.\"\n\n## Display hint\nChoose the best UI visualization layout for the result:\n- \"big\" \u2014 a single scalar answer: one row, one primary numeric value (e.g. a county's median income or total population).\n- \"statrow\" \u2014 2\u20134 related metrics about the same subject in one row (e.g. a housing snapshot with home value, rent, and homeownership rate).\n- \"bar\" \u2014 a ranked list of one metric across many geographies (e.g. top 10 counties by income).\n- \"grouped\" \u2014 a two-series comparison across geographies (e.g. rent in 2020 vs 2024 by state).\n- \"table\" \u2014 default fallback for anything that does not fit the above.\n\n## Output format\nRespond with a JSON object matching this exact shape:\n{\n  \"sql\": \"<your SELECT query or empty string>\",\n  \"assumptions\": [\"<assumption 1>\", \"<assumption 2>\"],\n  \"displayHint\": \"<big | statrow | bar | grouped | table>\"\n}\nDo not include any text outside the JSON object.";
+export interface FewShotExample {
+    question: string;
+    schemaContext: string;
+    result: {
+        sql: string;
+        assumptions: string[];
+        displayHint: "big" | "statrow" | "bar" | "grouped" | "table";
+    };
+}
+export declare const DEFAULT_FEW_SHOT_EXAMPLES: FewShotExample[];
+/**
+ * Converts an array of RetrievedSchemaChunk into a compact, LLM-readable
+ * schema context block.  Table-level chunks are listed first, followed by
+ * column chunks grouped under their parent table.
+ */
+export declare function formatSchemaContext(chunks: RetrievedSchemaChunk[]): string;
+/**
+ * Builds the full user message to send to the planning agent.
+ * Includes the schema context, optional few-shot examples, optional prior
+ * failed attempts (errorContext), and the question.
+ */
+export declare function buildUserMessage(question: string, context: RetrievedSchemaChunk[], examples: FewShotExample[], errorContext?: {
+    sql: string;
+    error: string;
+}[]): string;
+//# sourceMappingURL=prompt.d.ts.map

package/dist/planner/prompt.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"prompt.d.ts","sourceRoot":"","sources":["../../src/planner/prompt.ts"],"names":[],"mappings":"AAAA;;;;;;;;GAQG;AAEH,OAAO,KAAK,EAAE,oBAAoB,EAAE,MAAM,uBAAuB,CAAC;AAMlE,eAAO,MAAM,aAAa,68KAiFuB,CAAC;AAMlD,MAAM,WAAW,cAAc;IAC7B,QAAQ,EAAE,MAAM,CAAC;IACjB,aAAa,EAAE,MAAM,CAAC;IACtB,MAAM,EAAE;QAAE,GAAG,EAAE,MAAM,CAAC;QAAC,WAAW,EAAE,MAAM,EAAE,CAAC;QAAC,WAAW,EAAE,KAAK,GAAG,SAAS,GAAG,KAAK,GAAG,SAAS,GAAG,OAAO,CAAA;KAAE,CAAC;CAC9G;AAED,eAAO,MAAM,yBAAyB,EAAE,cAAc,EAwOrD,CAAC;AAMF;;;;GAIG;AACH,wBAAgB,mBAAmB,CAAC,MAAM,EAAE,oBAAoB,EAAE,GAAG,MAAM,CA6D1E;AAMD;;;;GAIG;AACH,wBAAgB,gBAAgB,CAC9B,QAAQ,EAAE,MAAM,EAChB,OAAO,EAAE,oBAAoB,EAAE,EAC/B,QAAQ,EAAE,cAAc,EAAE,EAC1B,YAAY,CAAC,EAAE;IAAE,GAAG,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAA;CAAE,EAAE,GAC9C,MAAM,CAkDR"}