@pgsage/core 0.1.0

package/dist/planner/prompt.js

@@ -0,0 +1,436 @@
+/**
+ * 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
+ */
+// ---------------------------------------------------------------------------
+// System prompt
+// ---------------------------------------------------------------------------
+export const SYSTEM_PROMPT = `\
+You are an expert PostgreSQL query writer for a schema-aware SQL agent.
+
+## Your task
+Given a natural-language question and the relevant schema context, produce:
+1. A single, valid PostgreSQL SELECT query that answers the question.
+2. A list of assumptions you made about any ambiguous terms or requirements.
+
+## Rules
+- Output ONLY a SELECT statement. Never produce INSERT, UPDATE, DELETE, DROP, ALTER, TRUNCATE, CREATE, GRANT, REVOKE, COPY, or any other non-SELECT statement.
+- Never produce multiple statements separated by semicolons.
+- Never use EXECUTE, PERFORM, or dynamic SQL.
+- Always use schema-qualified table names (e.g. public.geography, not just geography).
+- Prefer explicit column names over SELECT *.
+- Always write the estimate column as \`e.estimate\` — never alias it away from that name (e.g. do NOT write \`e.estimate AS median_income\`). Other columns may use aliases for readability.
+- 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.
+- 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.
+- Use standard PostgreSQL functions only. Do not use vendor extensions.
+- If you cannot produce a valid SELECT query from the given context, return an empty string for sql and explain the issue in assumptions.
+
+## Variable categories
+The variables table uses exactly these 8 category values. Use the correct one — do not guess:
+
+  income          — household income, per-capita income, Gini index, income brackets
+  poverty         — population below the poverty line, poverty ratios
+  housing         — housing units, home values, gross rent, occupancy, rent burden
+  employment      — labor force participation, unemployment counts, industry, occupation
+  demographics    — total population, age, sex, race, Hispanic/Latino origin
+  education       — educational attainment (high school, bachelor's, graduate degrees)
+  transportation  — commute mode (drove alone, transit, worked from home), travel time
+  health_insurance — health insurance coverage by age and sex
+
+Common mistakes to avoid:
+  - Bachelor's degree → category = 'education' (NOT 'demographics')
+  - Poverty / below poverty line → category = 'poverty' (NOT 'income')
+  - Commute mode, work from home, travel time → category = 'transportation' (NOT 'employment')
+  - Health insurance coverage → category = 'health_insurance'
+
+## Estimates are raw counts, not percentages
+Most estimate values are raw COUNTS (number of people, number of housing units), NOT percentages or rates.
+To express a rate or percentage, you must compute it by dividing a count by its relevant total:
+
+  homeownership rate = 100.0 * owner_occupied.estimate / NULLIF(total_occupied.estimate, 0)
+  poverty rate       = 100.0 * below_poverty.estimate  / NULLIF(total_for_poverty.estimate, 0)
+  drove-alone rate   = 100.0 * drove_alone.estimate    / NULLIF(total_workers.estimate, 0)
+
+Never compare a raw count estimate directly to a percentage threshold. For example:
+  WRONG: owner_e.estimate < 60        — this compares a unit count to 60, not 60%
+  RIGHT: 100.0 * owner_e.estimate / NULLIF(total_e.estimate, 0) < 60
+
+When a question asks for a rate or percentage, always join in the denominator total variable
+and compute the ratio explicitly.
+
+## ILIKE precision
+ILIKE label patterns can match multiple variables. When a single specific metric is needed:
+- 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).
+- Or add \`AND v.concept ILIKE '...'\` to narrow further.
+- 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.
+Ambiguous ILIKE matches combined with ORDER BY will surface whichever matched variable has the highest/lowest value — which may not be the intended one.
+
+## Assumption surfacing
+For every ambiguous term, geographic scope decision, join assumption, or data interpretation, add a plain-English sentence to the assumptions list. Examples:
+- "Interpreting 'large counties' as counties with an estimated population > 100,000."
+- "Using the most recent ACS 5-year estimate (2024 vintage)."
+- "Joining estimates to geography on geo_id to resolve county names."
+
+## Display hint
+Choose the best UI visualization layout for the result:
+- "big" — a single scalar answer: one row, one primary numeric value (e.g. a county's median income or total population).
+- "statrow" — 2–4 related metrics about the same subject in one row (e.g. a housing snapshot with home value, rent, and homeownership rate).
+- "bar" — a ranked list of one metric across many geographies (e.g. top 10 counties by income).
+- "grouped" — a two-series comparison across geographies (e.g. rent in 2020 vs 2024 by state).
+- "table" — default fallback for anything that does not fit the above.
+
+## Output format
+Respond with a JSON object matching this exact shape:
+{
+  "sql": "<your SELECT query or empty string>",
+  "assumptions": ["<assumption 1>", "<assumption 2>"],
+  "displayHint": "<big | statrow | bar | grouped | table>"
+}
+Do not include any text outside the JSON object.`;
+export const DEFAULT_FEW_SHOT_EXAMPLES = [
+    // ---------------------------------------------------------------------------
+    // Example 1: Simple lookup — single variable, named geography
+    // Demonstrates: correct category, e.estimate not aliased, NULLS LAST
+    // ---------------------------------------------------------------------------
+    {
+        question: "What is the median household income in Providence County, Rhode Island?",
+        schemaContext: `TABLE public.geography
+  geo_id text PK — FIPS-based composite id
+  state_name text NOT NULL — e.g. "Rhode Island"
+  county_name text — e.g. "Providence County"
+  level text NOT NULL — "state" or "county"
+
+TABLE public.variables
+  variable_code text PK — ACS code e.g. "B19013_001E"
+  label text NOT NULL — e.g. "Median household income in the past 12 months (in 2024 inflation-adjusted dollars)"
+  concept text
+  category text — income | poverty | housing | employment | demographics | education | transportation | health_insurance
+
+TABLE public.estimates
+  geo_id text FK→geography.geo_id
+  variable_code text FK→variables.variable_code
+  estimate numeric — the estimated value (raw count or median dollar amount)
+  margin_error numeric
+  PK (geo_id, variable_code)`,
+        result: {
+            sql: `SELECT
+  g.county_name,
+  g.state_name,
+  v.label,
+  e.estimate,
+  e.margin_error
+FROM public.estimates e
+JOIN public.geography g ON g.geo_id = e.geo_id
+JOIN public.variables v ON v.variable_code = e.variable_code
+WHERE g.county_name = 'Providence County'
+  AND g.state_name = 'Rhode Island'
+  AND g.level = 'county'
+  AND v.category = 'income'
+  AND v.variable_code = 'B19013_001E'
+ORDER BY e.estimate DESC NULLS LAST
+LIMIT 10;`,
+            assumptions: [
+                "Filtering to Providence County, Rhode Island by matching county_name and state_name exactly.",
+                "Using variable_code = 'B19013_001E' (Median household income in the past 12 months) to avoid matching age-group income breakdowns.",
+                "Keeping the column named e.estimate (not aliased) as required.",
+            ],
+            displayHint: "big",
+        },
+    },
+    // ---------------------------------------------------------------------------
+    // Example 2: Ranking — education category (common misclassification target)
+    // Demonstrates: education category, NULLS LAST, tight variable_code filter
+    // ---------------------------------------------------------------------------
+    {
+        question: "Which 5 states have the highest percentage of adults with a bachelor's degree?",
+        schemaContext: `TABLE public.geography
+  geo_id text PK
+  state_name text NOT NULL
+  level text NOT NULL — "state" or "county"
+
+TABLE public.variables
+  variable_code text PK
+  label text NOT NULL
+  category text — income | poverty | housing | employment | demographics | education | transportation | health_insurance
+
+TABLE public.estimates
+  geo_id text FK→geography.geo_id
+  variable_code text FK→variables.variable_code
+  estimate numeric`,
+        result: {
+            sql: `SELECT
+  g.state_name,
+  ROUND(100.0 * bach.estimate / NULLIF(total.estimate, 0), 1) AS pct_bachelors,
+  bach.estimate,
+  total.estimate AS total_adults_25_plus
+FROM public.geography g
+JOIN public.estimates bach ON bach.geo_id = g.geo_id
+JOIN public.estimates total ON total.geo_id = g.geo_id
+WHERE g.level = 'state'
+  AND bach.variable_code = 'B15003_022E'
+  AND total.variable_code = 'B15003_001E'
+ORDER BY pct_bachelors DESC NULLS LAST
+LIMIT 5;`,
+            assumptions: [
+                "Bachelor's degree variables are in the 'education' category, not 'demographics'.",
+                "Using B15003_022E (bachelor's degree count) divided by B15003_001E (total population 25+) to compute the percentage.",
+                "Filtering to state-level geographies (level = 'state').",
+                "Ordering by percentage descending with NULLS LAST; returning top 5.",
+            ],
+            displayHint: "bar",
+        },
+    },
+    // ---------------------------------------------------------------------------
+    // Example 3: Rate computation — homeownership rate (raw count ÷ total)
+    // Demonstrates: estimates are counts not percentages, must divide to get rate
+    // ---------------------------------------------------------------------------
+    {
+        question: "Which 5 states have the lowest homeownership rate?",
+        schemaContext: `TABLE public.geography
+  geo_id text PK
+  state_name text NOT NULL
+  level text NOT NULL
+
+TABLE public.variables
+  variable_code text PK
+  label text NOT NULL
+  category text
+
+TABLE public.estimates
+  geo_id text FK→geography.geo_id
+  variable_code text FK→variables.variable_code
+  estimate numeric — raw COUNT of housing units, not a percentage`,
+        result: {
+            sql: `SELECT
+  g.state_name,
+  ROUND(100.0 * owner.estimate / NULLIF(total.estimate, 0), 1) AS homeownership_rate,
+  owner.estimate AS owner_occupied_units,
+  total.estimate AS total_occupied_units
+FROM public.geography g
+JOIN public.estimates owner ON owner.geo_id = g.geo_id
+JOIN public.estimates total ON total.geo_id = g.geo_id
+WHERE g.level = 'state'
+  AND owner.variable_code = 'B25003_002E'
+  AND total.variable_code = 'B25003_001E'
+ORDER BY homeownership_rate ASC NULLS LAST
+LIMIT 5;`,
+            assumptions: [
+                "Homeownership estimates (B25003_002E, B25003_001E) are raw unit COUNTS, not percentages.",
+                "Computing the rate as 100 * owner_occupied / total_occupied using NULLIF to avoid division by zero.",
+                "B25003_002E = owner-occupied units; B25003_001E = total occupied housing units.",
+                "Filtering to state-level geographies; ordering ASC NULLS LAST for lowest rates first.",
+            ],
+            displayHint: "bar",
+        },
+    },
+    // ---------------------------------------------------------------------------
+    // Example 4: Transportation category + worked from home
+    // Demonstrates: transportation (not employment), exact label wording
+    // ---------------------------------------------------------------------------
+    {
+        question: "Which 10 counties have the most people with commutes of 90 minutes or more?",
+        schemaContext: `TABLE public.geography
+  geo_id text PK
+  county_name text
+  state_name text NOT NULL
+  level text NOT NULL
+
+TABLE public.variables
+  variable_code text PK
+  label text NOT NULL
+  category text — transportation category covers commute mode and travel time
+
+TABLE public.estimates
+  geo_id text FK→geography.geo_id
+  variable_code text FK→variables.variable_code
+  estimate numeric`,
+        result: {
+            sql: `SELECT
+  g.county_name,
+  g.state_name,
+  e.estimate,
+  e.margin_error
+FROM public.estimates e
+JOIN public.geography g ON g.geo_id = e.geo_id
+WHERE g.level = 'county'
+  AND e.variable_code = 'B08303_013E'
+ORDER BY e.estimate DESC NULLS LAST
+LIMIT 10;`,
+            assumptions: [
+                "Commute time variables are in the 'transportation' category, not 'employment'.",
+                "Using variable_code B08303_013E (Workers with commute 90 or more minutes) directly — the label wording does not contain '90 minutes or more', so ILIKE would miss it.",
+                "Filtering to county-level geographies; ordering by estimate DESC NULLS LAST.",
+            ],
+            displayHint: "bar",
+        },
+    },
+    // ---------------------------------------------------------------------------
+    // Example 5: Cross-category with computed poverty rate
+    // Demonstrates: poverty category, rate-from-counts, multi-variable join
+    // ---------------------------------------------------------------------------
+    {
+        question: "Show counties with population over 500,000 where more than 15% of residents live below the poverty line.",
+        schemaContext: `TABLE public.geography
+  geo_id text PK
+  county_name text
+  state_name text NOT NULL
+  level text NOT NULL
+
+TABLE public.variables
+  variable_code text PK
+  label text NOT NULL
+  category text
+
+TABLE public.estimates
+  geo_id text FK→geography.geo_id
+  variable_code text FK→variables.variable_code
+  estimate numeric — raw count`,
+        result: {
+            sql: `SELECT
+  g.county_name,
+  g.state_name,
+  pop.estimate AS total_population,
+  pov.estimate AS below_poverty_count,
+  ROUND(100.0 * pov.estimate / NULLIF(pov_total.estimate, 0), 1) AS poverty_rate_pct
+FROM public.geography g
+JOIN public.estimates pop ON pop.geo_id = g.geo_id
+JOIN public.estimates pov ON pov.geo_id = g.geo_id
+JOIN public.estimates pov_total ON pov_total.geo_id = g.geo_id
+WHERE g.level = 'county'
+  AND pop.variable_code = 'B01003_001E'
+  AND pov.variable_code = 'B17001_002E'
+  AND pov_total.variable_code = 'B17001_001E'
+  AND pop.estimate > 500000
+  AND 100.0 * pov.estimate / NULLIF(pov_total.estimate, 0) > 15
+ORDER BY poverty_rate_pct DESC NULLS LAST
+LIMIT 100;`,
+            assumptions: [
+                "Poverty variables are in the 'poverty' category, not 'income'.",
+                "B17001_002E = population with income below poverty level (count); B17001_001E = total population for poverty determination (denominator).",
+                "Poverty rate is computed as 100 * below_poverty / total_for_poverty — estimates are raw counts, not percentages.",
+                "Population filter uses B01003_001E (total population); poverty rate threshold is > 15 percentage points.",
+                "Filtering to county-level geographies; ordering by poverty rate descending.",
+            ],
+            displayHint: "table",
+        },
+    },
+];
+// ---------------------------------------------------------------------------
+// Schema context formatter
+// ---------------------------------------------------------------------------
+/**
+ * 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 function formatSchemaContext(chunks) {
+    if (chunks.length === 0) {
+        return "(no schema context retrieved)";
+    }
+    const tableChunks = chunks.filter((c) => c.objectType === "table");
+    const columnChunks = chunks.filter((c) => c.objectType === "column");
+    const lines = [];
+    // Group column chunks by their parent table reference
+    const columnsByTable = new Map();
+    for (const col of columnChunks) {
+        // objectRef for columns: "{schema}.{table}.{column}"
+        const parts = col.objectRef.split(".");
+        const tableRef = parts.slice(0, 2).join(".");
+        const existing = columnsByTable.get(tableRef) ?? [];
+        existing.push(col);
+        columnsByTable.set(tableRef, existing);
+    }
+    // Build output: tables with their columns interleaved
+    const renderedTables = new Set();
+    for (const tableChunk of tableChunks) {
+        lines.push(tableChunk.content);
+        renderedTables.add(tableChunk.objectRef);
+        const cols = columnsByTable.get(tableChunk.objectRef) ?? [];
+        for (const col of cols) {
+            // Indent column details under the table
+            const indented = col.content
+                .split("\n")
+                .map((l) => `  ${l}`)
+                .join("\n");
+            lines.push(indented);
+        }
+        lines.push(""); // blank line between tables
+    }
+    // Any column chunks whose parent table wasn't retrieved as a table chunk
+    for (const [tableRef, cols] of columnsByTable.entries()) {
+        if (renderedTables.has(tableRef))
+            continue;
+        lines.push(`TABLE ${tableRef} (partial — columns only):`);
+        for (const col of cols) {
+            const indented = col.content
+                .split("\n")
+                .map((l) => `  ${l}`)
+                .join("\n");
+            lines.push(indented);
+        }
+        lines.push("");
+    }
+    // If only column chunks exist and no table chunks, handle gracefully above
+    if (tableChunks.length === 0 && columnChunks.length === 0) {
+        return chunks.map((c) => c.content).join("\n\n");
+    }
+    return lines.join("\n").trimEnd();
+}
+// ---------------------------------------------------------------------------
+// User message builder
+// ---------------------------------------------------------------------------
+/**
+ * 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 function buildUserMessage(question, context, examples, errorContext) {
+    const parts = [];
+    // Schema context section
+    parts.push("## Relevant Schema Context");
+    parts.push(formatSchemaContext(context));
+    parts.push("");
+    // Few-shot examples section
+    if (examples.length > 0) {
+        parts.push("## Examples");
+        for (const ex of examples) {
+            parts.push(`### Question: ${ex.question}`);
+            parts.push("Schema context:");
+            parts.push("```");
+            parts.push(ex.schemaContext);
+            parts.push("```");
+            parts.push("Response:");
+            parts.push("```json");
+            parts.push(JSON.stringify(ex.result, null, 2));
+            parts.push("```");
+            parts.push("");
+        }
+    }
+    // Prior failed attempts (correction context)
+    if (errorContext && errorContext.length > 0) {
+        parts.push("## Prior Attempts (Failed)");
+        parts.push("The following SQL queries were generated previously but failed. " +
+            "Study the errors and produce a corrected query that avoids these mistakes.");
+        parts.push("");
+        for (let i = 0; i < errorContext.length; i++) {
+            const attempt = errorContext[i];
+            parts.push(`### Attempt ${i + 1}`);
+            parts.push("SQL:");
+            parts.push("```sql");
+            parts.push(attempt.sql);
+            parts.push("```");
+            parts.push(`Error: ${attempt.error}`);
+            parts.push("");
+        }
+    }
+    // The actual question
+    parts.push("## Question");
+    parts.push(question);
+    return parts.join("\n");
+}
+//# sourceMappingURL=prompt.js.map

package/dist/planner/prompt.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"prompt.js","sourceRoot":"","sources":["../../src/planner/prompt.ts"],"names":[],"mappings":"AAAA;;;;;;;;GAQG;AAIH,8EAA8E;AAC9E,gBAAgB;AAChB,8EAA8E;AAE9E,MAAM,CAAC,MAAM,aAAa,GAAG;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;iDAiFoB,CAAC;AAYlD,MAAM,CAAC,MAAM,yBAAyB,GAAqB;IACzD,8EAA8E;IAC9E,8DAA8D;IAC9D,qEAAqE;IACrE,8EAA8E;IAC9E;QACE,QAAQ,EAAE,yEAAyE;QACnF,aAAa,EAAE;;;;;;;;;;;;;;;;;6BAiBU;QACzB,MAAM,EAAE;YACN,GAAG,EAAE;;;;;;;;;;;;;;;UAeD;YACJ,WAAW,EAAE;gBACX,8FAA8F;gBAC9F,oIAAoI;gBACpI,gEAAgE;aACjE;YACD,WAAW,EAAE,KAAK;SACnB;KACF;IAED,8EAA8E;IAC9E,4EAA4E;IAC5E,2EAA2E;IAC3E,8EAA8E;IAC9E;QACE,QAAQ,EAAE,gFAAgF;QAC1F,aAAa,EAAE;;;;;;;;;;;;;mBAaA;QACf,MAAM,EAAE;YACN,GAAG,EAAE;;;;;;;;;;;;SAYF;YACH,WAAW,EAAE;gBACX,kFAAkF;gBAClF,sHAAsH;gBACtH,yDAAyD;gBACzD,qEAAqE;aACtE;YACD,WAAW,EAAE,KAAK;SACnB;KACF;IAED,8EAA8E;IAC9E,uEAAuE;IACvE,8EAA8E;IAC9E,8EAA8E;IAC9E;QACE,QAAQ,EAAE,oDAAoD;QAC9D,aAAa,EAAE;;;;;;;;;;;;;kEAa+C;QAC9D,MAAM,EAAE;YACN,GAAG,EAAE;;;;;;;;;;;;SAYF;YACH,WAAW,EAAE;gBACX,0FAA0F;gBAC1F,qGAAqG;gBACrG,iFAAiF;gBACjF,uFAAuF;aACxF;YACD,WAAW,EAAE,KAAK;SACnB;KACF;IAED,8EAA8E;IAC9E,wDAAwD;IACxD,qEAAqE;IACrE,8EAA8E;IAC9E;QACE,QAAQ,EAAE,6EAA6E;QACvF,aAAa,EAAE;;;;;;;;;;;;;;mBAcA;QACf,MAAM,EAAE;YACN,GAAG,EAAE;;;;;;;;;;UAUD;YACJ,WAAW,EAAE;gBACX,gFAAgF;gBAChF,uKAAuK;gBACvK,8EAA8E;aAC/E;YACD,WAAW,EAAE,KAAK;SACnB;KACF;IAED,8EAA8E;IAC9E,uDAAuD;IACvD,wEAAwE;IACxE,8EAA8E;IAC9E;QACE,QAAQ,EACN,0GAA0G;QAC5G,aAAa,EAAE;;;;;;;;;;;;;;+BAcY;QAC3B,MAAM,EAAE;YACN,GAAG,EAAE;;;;;;;;;;;;;;;;;WAiBA;YACL,WAAW,EAAE;gBACX,gEAAgE;gBAChE,2IAA2I;gBAC3I,kHAAkH;gBAClH,0GAA0G;gBAC1G,6EAA6E;aAC9E;YACD,WAAW,EAAE,OAAO;SACrB;KACF;CACF,CAAC;AAEF,8EAA8E;AAC9E,2BAA2B;AAC3B,8EAA8E;AAE9E;;;;GAIG;AACH,MAAM,UAAU,mBAAmB,CAAC,MAA8B;IAChE,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;QACxB,OAAO,+BAA+B,CAAC;IACzC,CAAC;IAED,MAAM,WAAW,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,UAAU,KAAK,OAAO,CAAC,CAAC;IACnE,MAAM,YAAY,GAAG,MAAM,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,UAAU,KAAK,QAAQ,CAAC,CAAC;IAErE,MAAM,KAAK,GAAa,EAAE,CAAC;IAE3B,sDAAsD;IACtD,MAAM,cAAc,GAAG,IAAI,GAAG,EAAkC,CAAC;IACjE,KAAK,MAAM,GAAG,IAAI,YAAY,EAAE,CAAC;QAC/B,qDAAqD;QACrD,MAAM,KAAK,GAAG,GAAG,CAAC,SAAS,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;QACvC,MAAM,QAAQ,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;QAC7C,MAAM,QAAQ,GAAG,cAAc,CAAC,GAAG,CAAC,QAAQ,CAAC,IAAI,EAAE,CAAC;QACpD,QAAQ,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;QACnB,cAAc,CAAC,GAAG,CAAC,QAAQ,EAAE,QAAQ,CAAC,CAAC;IACzC,CAAC;IAED,sDAAsD;IACtD,MAAM,cAAc,GAAG,IAAI,GAAG,EAAU,CAAC;IAEzC,KAAK,MAAM,UAAU,IAAI,WAAW,EAAE,CAAC;QACrC,KAAK,CAAC,IAAI,CAAC,UAAU,CAAC,OAAO,CAAC,CAAC;QAC/B,cAAc,CAAC,GAAG,CAAC,UAAU,CAAC,SAAS,CAAC,CAAC;QAEzC,MAAM,IAAI,GAAG,cAAc,CAAC,GAAG,CAAC,UAAU,CAAC,SAAS,CAAC,IAAI,EAAE,CAAC;QAC5D,KAAK,MAAM,GAAG,IAAI,IAAI,EAAE,CAAC;YACvB,wCAAwC;YACxC,MAAM,QAAQ,GAAG,GAAG,CAAC,OAAO;iBACzB,KAAK,CAAC,IAAI,CAAC;iBACX,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,KAAK,CAAC,EAAE,CAAC;iBACpB,IAAI,CAAC,IAAI,CAAC,CAAC;YACd,KAAK,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;QACvB,CAAC;QAED,KAAK,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC,CAAC,4BAA4B;IAC9C,CAAC;IAED,yEAAyE;IACzE,KAAK,MAAM,CAAC,QAAQ,EAAE,IAAI,CAAC,IAAI,cAAc,CAAC,OAAO,EAAE,EAAE,CAAC;QACxD,IAAI,cAAc,CAAC,GAAG,CAAC,QAAQ,CAAC;YAAE,SAAS;QAC3C,KAAK,CAAC,IAAI,CAAC,SAAS,QAAQ,4BAA4B,CAAC,CAAC;QAC1D,KAAK,MAAM,GAAG,IAAI,IAAI,EAAE,CAAC;YACvB,MAAM,QAAQ,GAAG,GAAG,CAAC,OAAO;iBACzB,KAAK,CAAC,IAAI,CAAC;iBACX,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,KAAK,CAAC,EAAE,CAAC;iBACpB,IAAI,CAAC,IAAI,CAAC,CAAC;YACd,KAAK,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;QACvB,CAAC;QACD,KAAK,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;IACjB,CAAC;IAED,2EAA2E;IAC3E,IAAI,WAAW,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;QAC1D,OAAO,MAAM,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;IACnD,CAAC;IAED,OAAO,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,OAAO,EAAE,CAAC;AACpC,CAAC;AAED,8EAA8E;AAC9E,uBAAuB;AACvB,8EAA8E;AAE9E;;;;GAIG;AACH,MAAM,UAAU,gBAAgB,CAC9B,QAAgB,EAChB,OAA+B,EAC/B,QAA0B,EAC1B,YAA+C;IAE/C,MAAM,KAAK,GAAa,EAAE,CAAC;IAE3B,yBAAyB;IACzB,KAAK,CAAC,IAAI,CAAC,4BAA4B,CAAC,CAAC;IACzC,KAAK,CAAC,IAAI,CAAC,mBAAmB,CAAC,OAAO,CAAC,CAAC,CAAC;IACzC,KAAK,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;IAEf,4BAA4B;IAC5B,IAAI,QAAQ,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;QACxB,KAAK,CAAC,IAAI,CAAC,aAAa,CAAC,CAAC;QAC1B,KAAK,MAAM,EAAE,IAAI,QAAQ,EAAE,CAAC;YAC1B,KAAK,CAAC,IAAI,CAAC,iBAAiB,EAAE,CAAC,QAAQ,EAAE,CAAC,CAAC;YAC3C,KAAK,CAAC,IAAI,CAAC,iBAAiB,CAAC,CAAC;YAC9B,KAAK,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;YAClB,KAAK,CAAC,IAAI,CAAC,EAAE,CAAC,aAAa,CAAC,CAAC;YAC7B,KAAK,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;YAClB,KAAK,CAAC,IAAI,CAAC,WAAW,CAAC,CAAC;YACxB,KAAK,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;YACtB,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,SAAS,CAAC,EAAE,CAAC,MAAM,EAAE,IAAI,EAAE,CAAC,CAAC,CAAC,CAAC;YAC/C,KAAK,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;YAClB,KAAK,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;QACjB,CAAC;IACH,CAAC;IAED,6CAA6C;IAC7C,IAAI,YAAY,IAAI,YAAY,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;QAC5C,KAAK,CAAC,IAAI,CAAC,4BAA4B,CAAC,CAAC;QACzC,KAAK,CAAC,IAAI,CACR,kEAAkE;YAChE,4EAA4E,CAC/E,CAAC;QACF,KAAK,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;QACf,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,YAAY,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YAC7C,MAAM,OAAO,GAAG,YAAY,CAAC,CAAC,CAAE,CAAC;YACjC,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;YACnC,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;YACnB,KAAK,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;YACrB,KAAK,CAAC,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC;YACxB,KAAK,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;YAClB,KAAK,CAAC,IAAI,CAAC,UAAU,OAAO,CAAC,KAAK,EAAE,CAAC,CAAC;YACtC,KAAK,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;QACjB,CAAC;IACH,CAAC;IAED,sBAAsB;IACtB,KAAK,CAAC,IAAI,CAAC,aAAa,CAAC,CAAC;IAC1B,KAAK,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC;IAErB,OAAO,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;AAC1B,CAAC"}

package/dist/retriever/index.d.ts

@@ -0,0 +1,90 @@
+/**
+ * Retriever
+ *
+ * Two responsibilities:
+ *
+ *  1. indexSchema() — converts an IntrospectedSchema into text descriptions,
+ *     embeds them via an EmbeddingProvider, and upserts the results into the
+ *     schema_embeddings table.  Safe to run repeatedly (idempotent via ON
+ *     CONFLICT ... DO UPDATE).
+ *
+ *  2. retrieve() — embeds a natural-language question and performs a pgvector
+ *     cosine similarity search against schema_embeddings, returning the
+ *     top-K most relevant schema chunks.
+ */
+import type { Pool } from "pg";
+import type { EmbeddingProvider } from "../embeddings/types.js";
+import type { IntrospectedSchema, IntrospectedTable } from "../introspector/index.js";
+export type SchemaObjectType = "table" | "column";
+export interface RetrievedSchemaChunk {
+    /** Row id from schema_embeddings. */
+    id: string;
+    /** The type of schema object this chunk describes. */
+    objectType: SchemaObjectType;
+    /**
+     * Stable reference for the object.
+     * Tables:  "{schema}.{table}"
+     * Columns: "{schema}.{table}.{column}"
+     */
+    objectRef: string;
+    /** The text that was embedded (useful for debugging / prompt building). */
+    content: string;
+    /** Cosine similarity score in [0, 1]. Higher is more relevant. */
+    score: number;
+}
+export interface IndexSchemaResult {
+    /** Total rows upserted into schema_embeddings. */
+    rowsUpserted: number;
+    /** Total tokens consumed by embedding calls. */
+    totalTokens: number;
+}
+export interface IndexSchemaOptions {
+    /**
+     * If true, re-embed and overwrite entries even if they already exist.
+     * Default: false (upsert only updates the embedding + content).
+     */
+    force?: boolean;
+}
+export interface RetrieveOptions {
+    /**
+     * Maximum number of results to return.
+     * Default: 10.
+     */
+    topK?: number;
+    /**
+     * Minimum cosine similarity score (0–1).
+     * Results below this threshold are filtered out.
+     * Default: 0 (no threshold).
+     */
+    minScore?: number;
+    /**
+     * Limit results to specific object types.
+     * Default: all types.
+     */
+    objectTypes?: SchemaObjectType[];
+}
+/**
+ * Generate a human-readable description of a table suitable for embedding.
+ * Includes: schema/table name, type, row estimate, comment, and a column
+ * summary with types, nullability, PK/FK annotations.
+ */
+export declare function describeTable(table: IntrospectedTable): string;
+/**
+ * Generate a human-readable description of a single column suitable for
+ * embedding.  Includes table context, type info, PK/FK annotations, and
+ * comment.
+ */
+export declare function describeColumn(table: IntrospectedTable, columnIndex: number): string;
+/**
+ * Embed an IntrospectedSchema and upsert the results into schema_embeddings.
+ *
+ * Generates one chunk per table and one chunk per column, embeds them all
+ * in batches via the provider, then bulk-upserts into the DB.
+ */
+export declare function indexSchema(pool: Pool, schema: IntrospectedSchema, provider: EmbeddingProvider, _options?: IndexSchemaOptions): Promise<IndexSchemaResult>;
+/**
+ * Embed a natural-language question and return the most relevant schema
+ * chunks from schema_embeddings, ranked by cosine similarity.
+ */
+export declare function retrieve(pool: Pool, question: string, provider: EmbeddingProvider, options?: RetrieveOptions): Promise<RetrievedSchemaChunk[]>;
+//# sourceMappingURL=index.d.ts.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/retriever/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;GAaG;AAEH,OAAO,KAAK,EAAE,IAAI,EAAE,MAAM,IAAI,CAAC;AAC/B,OAAO,KAAK,EAAE,iBAAiB,EAAE,MAAM,wBAAwB,CAAC;AAChE,OAAO,KAAK,EAAE,kBAAkB,EAAE,iBAAiB,EAAE,MAAM,0BAA0B,CAAC;AAMtF,MAAM,MAAM,gBAAgB,GAAG,OAAO,GAAG,QAAQ,CAAC;AAElD,MAAM,WAAW,oBAAoB;IACnC,qCAAqC;IACrC,EAAE,EAAE,MAAM,CAAC;IACX,sDAAsD;IACtD,UAAU,EAAE,gBAAgB,CAAC;IAC7B;;;;OAIG;IACH,SAAS,EAAE,MAAM,CAAC;IAClB,2EAA2E;IAC3E,OAAO,EAAE,MAAM,CAAC;IAChB,kEAAkE;IAClE,KAAK,EAAE,MAAM,CAAC;CACf;AAED,MAAM,WAAW,iBAAiB;IAChC,kDAAkD;IAClD,YAAY,EAAE,MAAM,CAAC;IACrB,gDAAgD;IAChD,WAAW,EAAE,MAAM,CAAC;CACrB;AAMD,MAAM,WAAW,kBAAkB;IACjC;;;OAGG;IACH,KAAK,CAAC,EAAE,OAAO,CAAC;CACjB;AAED,MAAM,WAAW,eAAe;IAC9B;;;OAGG;IACH,IAAI,CAAC,EAAE,MAAM,CAAC;IACd;;;;OAIG;IACH,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB;;;OAGG;IACH,WAAW,CAAC,EAAE,gBAAgB,EAAE,CAAC;CAClC;AAMD;;;;GAIG;AACH,wBAAgB,aAAa,CAAC,KAAK,EAAE,iBAAiB,GAAG,MAAM,CAsB9D;AAED;;;;GAIG;AACH,wBAAgB,cAAc,CAC5B,KAAK,EAAE,iBAAiB,EACxB,WAAW,EAAE,MAAM,GAClB,MAAM,CAoBR;AAMD;;;;;GAKG;AACH,wBAAsB,WAAW,CAC/B,IAAI,EAAE,IAAI,EACV,MAAM,EAAE,kBAAkB,EAC1B,QAAQ,EAAE,iBAAiB,EAC3B,QAAQ,GAAE,kBAAuB,GAChC,OAAO,CAAC,iBAAiB,CAAC,CA6D5B;AAcD;;;GAGG;AACH,wBAAsB,QAAQ,CAC5B,IAAI,EAAE,IAAI,EACV,QAAQ,EAAE,MAAM,EAChB,QAAQ,EAAE,iBAAiB,EAC3B,OAAO,GAAE,eAAoB,GAC5B,OAAO,CAAC,oBAAoB,EAAE,CAAC,CAqCjC"}