@deepagents/text2sql 0.3.1 → 0.6.0

package/dist/lib/adapters/groundings/{grounding.js → index.js}

@@ -1,3 +1,205 @@
+// packages/text2sql/src/lib/adapters/groundings/abstract.grounding.ts
+var AbstractGrounding = class {
+  tag;
+  constructor(tag) {
+    this.tag = tag;
+  }
+};
+
+// packages/text2sql/src/lib/adapters/groundings/column-stats.grounding.ts
+var ColumnStatsGrounding = class extends AbstractGrounding {
+  constructor(config = {}) {
+    super("column_stats");
+  }
+  /**
+   * Execute the grounding process.
+   * Annotates columns in ctx.tables and ctx.views with statistics.
+   */
+  async execute(ctx) {
+    const allContainers = [...ctx.tables, ...ctx.views];
+    for (const container of allContainers) {
+      for (const column of container.columns) {
+        try {
+          const stats = await this.collectStats(container.name, column);
+          if (stats) {
+            column.stats = stats;
+          }
+        } catch (error) {
+          console.warn(
+            "Error collecting stats for",
+            container.name,
+            column.name,
+            error
+          );
+        }
+      }
+    }
+    return () => this.#describe();
+  }
+  #describe() {
+    return null;
+  }
+};
+
+// packages/text2sql/src/lib/adapters/groundings/column-values.grounding.ts
+var ColumnValuesGrounding = class extends AbstractGrounding {
+  lowCardinalityLimit;
+  constructor(config = {}) {
+    super("column_values");
+    this.lowCardinalityLimit = config.lowCardinalityLimit ?? 20;
+  }
+  /**
+   * Get values for native ENUM type columns.
+   * Return undefined if column is not an ENUM type.
+   * Default implementation returns undefined (no native ENUM support).
+   */
+  async collectEnumValues(_tableName, _column) {
+    return void 0;
+  }
+  /**
+   * Parse CHECK constraint for enum-like IN clause.
+   * Extracts values from patterns like:
+   * - CHECK (status IN ('active', 'inactive'))
+   * - CHECK ((status)::text = ANY (ARRAY['a'::text, 'b'::text]))
+   * - CHECK (status = 'active' OR status = 'inactive')
+   */
+  parseCheckConstraint(constraint, columnName) {
+    if (constraint.type !== "CHECK" || !constraint.definition) {
+      return void 0;
+    }
+    if (constraint.columns && !constraint.columns.includes(columnName)) {
+      return void 0;
+    }
+    const def = constraint.definition;
+    const escapedCol = this.escapeRegex(columnName);
+    const colPattern = `(?:\\(?\\(?${escapedCol}\\)?(?:::(?:text|varchar|character varying))?\\)?)`;
+    const inMatch = def.match(
+      new RegExp(`${colPattern}\\s+IN\\s*\\(([^)]+)\\)`, "i")
+    );
+    if (inMatch) {
+      return this.extractStringValues(inMatch[1]);
+    }
+    const anyMatch = def.match(
+      new RegExp(
+        `${colPattern}\\s*=\\s*ANY\\s*\\(\\s*(?:ARRAY)?\\s*\\[([^\\]]+)\\]`,
+        "i"
+      )
+    );
+    if (anyMatch) {
+      return this.extractStringValues(anyMatch[1]);
+    }
+    const orPattern = new RegExp(
+      `\\b${this.escapeRegex(columnName)}\\b\\s*=\\s*'([^']*)'`,
+      "gi"
+    );
+    const orMatches = [...def.matchAll(orPattern)];
+    if (orMatches.length >= 2) {
+      return orMatches.map((m) => m[1]);
+    }
+    return void 0;
+  }
+  /**
+   * Extract string values from a comma-separated list.
+   */
+  extractStringValues(input) {
+    const values = [];
+    const matches = input.matchAll(/'([^']*)'/g);
+    for (const match of matches) {
+      values.push(match[1]);
+    }
+    return values.length > 0 ? values : void 0;
+  }
+  /**
+   * Escape special regex characters in a string.
+   */
+  escapeRegex(str) {
+    return str.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+  }
+  /**
+   * Get the table from context by name.
+   */
+  getTable(ctx, name) {
+    return ctx.tables.find((t) => t.name === name);
+  }
+  /**
+   * Execute the grounding process.
+   * Annotates columns in ctx.tables and ctx.views with values.
+   */
+  async execute(ctx) {
+    const allContainers = [...ctx.tables, ...ctx.views];
+    for (const container of allContainers) {
+      const table = this.getTable(ctx, container.name);
+      for (const column of container.columns) {
+        try {
+          const result = await this.resolveColumnValues(
+            container.name,
+            column,
+            table?.constraints
+          );
+          if (result) {
+            column.kind = result.kind;
+            column.values = result.values;
+          }
+        } catch (error) {
+          console.warn(
+            "Error collecting column values for",
+            container.name,
+            column.name,
+            error
+          );
+        }
+      }
+    }
+    return () => this.#describe();
+  }
+  /**
+   * Resolve column values from all sources in priority order.
+   */
+  async resolveColumnValues(tableName, column, constraints) {
+    const enumValues = await this.collectEnumValues(tableName, column);
+    if (enumValues?.length) {
+      return { kind: "Enum", values: enumValues };
+    }
+    if (constraints) {
+      for (const constraint of constraints) {
+        const checkValues = this.parseCheckConstraint(constraint, column.name);
+        if (checkValues?.length) {
+          return { kind: "Enum", values: checkValues };
+        }
+      }
+    }
+    const lowCardValues = await this.collectLowCardinality(tableName, column);
+    if (lowCardValues?.length) {
+      return { kind: "LowCardinality", values: lowCardValues };
+    }
+    return void 0;
+  }
+  #describe() {
+    return null;
+  }
+};
+
+// packages/text2sql/src/lib/adapters/groundings/constraint.grounding.ts
+var ConstraintGrounding = class extends AbstractGrounding {
+  constructor(config = {}) {
+    super("constraints");
+  }
+  /**
+   * Execute the grounding process.
+   * Annotates tables in ctx.tables with their constraints.
+   */
+  async execute(ctx) {
+    for (const table of ctx.tables) {
+      try {
+        table.constraints = await this.getConstraints(table.name);
+      } catch (error) {
+        console.warn("Error collecting constraints for", table.name, error);
+      }
+    }
+    return () => null;
+  }
+};
+
 // packages/text2sql/src/lib/adapters/groundings/context.ts
 function createGroundingContext() {
   return {
@@ -8,18 +210,159 @@ function createGroundingContext() {
   };
 }
 
-// packages/text2sql/src/lib/adapters/groundings/table.grounding.ts
-import pluralize from "pluralize";
+// packages/text2sql/src/lib/adapters/groundings/info.grounding.ts
+var InfoGrounding = class extends AbstractGrounding {
+  constructor(config = {}) {
+    super("dialect_info");
+  }
+  /**
+   * Execute the grounding process.
+   * Writes database info to ctx.info.
+   */
+  async execute(ctx) {
+    ctx.info = await this.collectInfo();
+    const lines = [`Dialect: ${ctx.info.dialect ?? "unknown"}`];
+    if (ctx.info.version) {
+      lines.push(`Version: ${ctx.info.version}`);
+    }
+    if (ctx.info.database) {
+      lines.push(`Database: ${ctx.info.database}`);
+    }
+    if (ctx.info.details && Object.keys(ctx.info.details).length) {
+      lines.push(`Details: ${JSON.stringify(ctx.info.details)}`);
+    }
+    return () => lines.join("\n");
+  }
+};
 
-// packages/text2sql/src/lib/adapters/grounding.ticket.ts
-var AbstractGrounding = class {
-  tag;
-  constructor(tag) {
-    this.tag = tag;
+// packages/text2sql/src/lib/adapters/groundings/report.grounding.ts
+import { groq } from "@ai-sdk/groq";
+import { tool } from "ai";
+import dedent from "dedent";
+import z from "zod";
+import {
+  agent,
+  generate,
+  toState,
+  user
+} from "@deepagents/agent";
+var ReportGrounding = class extends AbstractGrounding {
+  #adapter;
+  #model;
+  #cache;
+  #forceRefresh;
+  constructor(adapter, config = {}) {
+    super("business_context");
+    this.#adapter = adapter;
+    this.#model = config.model ?? groq("openai/gpt-oss-20b");
+    this.#cache = config.cache;
+    this.#forceRefresh = config.forceRefresh ?? false;
+  }
+  async execute(ctx) {
+    if (!this.#forceRefresh && this.#cache) {
+      const cached = await this.#cache.get();
+      if (cached) {
+        ctx.report = cached;
+        return () => cached;
+      }
+    }
+    const report = await this.#generateReport();
+    ctx.report = report;
+    if (this.#cache) {
+      await this.#cache.set(report);
+    }
+    return () => report;
+  }
+  async #generateReport() {
+    const reportAgent = agent({
+      name: "db-report-agent",
+      model: this.#model,
+      prompt: () => dedent`
+        <identity>
+          You are a database analyst expert. Your job is to understand what
+          a database represents and provide business context about it.
+          You have READ-ONLY access to the database.
+        </identity>
+
+        <instructions>
+          Write a business context that helps another agent answer questions accurately.
+
+          For EACH table, do queries ONE AT A TIME:
+          1. SELECT COUNT(*) to get row count
+          2. SELECT * LIMIT 3 to see sample data
+
+          Then write a report with:
+          - What business this database is for
+          - For each table: purpose, row count, and example of what the data looks like
+
+          Include concrete examples like "Track prices are $0.99",
+          "Customer names like 'Luís Gonçalves'", etc.
+
+          Keep it 400-600 words, conversational style.
+        </instructions>
+      `,
+      tools: {
+        query_database: tool({
+          description: "Execute a SELECT query to explore the database and gather insights.",
+          inputSchema: z.object({
+            sql: z.string().describe("The SELECT query to execute"),
+            purpose: z.string().describe(
+              "What insight you are trying to gather with this query"
+            )
+          }),
+          execute: ({ sql }, options) => {
+            const state = toState(options);
+            return state.adapter.execute(sql);
+          }
+        })
+      }
+    });
+    const { text } = await generate(
+      reportAgent,
+      [
+        user(
+          "Please analyze the database and write a contextual report about what this database represents."
+        )
+      ],
+      { adapter: this.#adapter }
+    );
+    return text;
+  }
+};
+
+// packages/text2sql/src/lib/adapters/groundings/row-count.grounding.ts
+var RowCountGrounding = class extends AbstractGrounding {
+  constructor(config = {}) {
+    super("row_counts");
+  }
+  /**
+   * Execute the grounding process.
+   * Annotates tables in ctx.tables with row counts and size hints.
+   */
+  async execute(ctx) {
+    for (const table of ctx.tables) {
+      const count = await this.getRowCount(table.name);
+      if (count != null) {
+        table.rowCount = count;
+        table.sizeHint = this.#classifyRowCount(count);
+      }
+    }
+    return () => null;
+  }
+  /**
+   * Classify row count into a size hint category.
+   */
+  #classifyRowCount(count) {
+    if (count < 100) return "tiny";
+    if (count < 1e3) return "small";
+    if (count < 1e4) return "medium";
+    if (count < 1e5) return "large";
+    return "huge";
   }
 };
 
 // packages/text2sql/src/lib/adapters/groundings/table.grounding.ts
+import pluralize from "pluralize";
 var TableGrounding = class extends AbstractGrounding {
   #filter;
   #forward;
@@ -135,18 +478,61 @@ var TableGrounding = class extends AbstractGrounding {
     }
     return tables.map((table) => {
       const rowCountInfo = table.rowCount != null ? ` [rows: ${table.rowCount}${table.sizeHint ? `, size: ${table.sizeHint}` : ""}]` : "";
-      const pkConstraint = table.constraints?.find((c) => c.type === "PRIMARY_KEY");
+      const pkConstraint = table.constraints?.find(
+        (c) => c.type === "PRIMARY_KEY"
+      );
       const pkColumns = new Set(pkConstraint?.columns ?? []);
+      const notNullColumns = new Set(
+        table.constraints?.filter((c) => c.type === "NOT_NULL").flatMap((c) => c.columns ?? []) ?? []
+      );
+      const defaultByColumn = /* @__PURE__ */ new Map();
+      for (const c of table.constraints?.filter(
+        (c2) => c2.type === "DEFAULT"
+      ) ?? []) {
+        for (const col of c.columns ?? []) {
+          if (c.defaultValue != null) {
+            defaultByColumn.set(col, c.defaultValue);
+          }
+        }
+      }
+      const uniqueColumns = new Set(
+        table.constraints?.filter((c) => c.type === "UNIQUE" && c.columns?.length === 1).flatMap((c) => c.columns ?? []) ?? []
+      );
+      const fkByColumn = /* @__PURE__ */ new Map();
+      for (const c of table.constraints?.filter(
+        (c2) => c2.type === "FOREIGN_KEY"
+      ) ?? []) {
+        const cols = c.columns ?? [];
+        const refCols = c.referencedColumns ?? [];
+        for (let i = 0; i < cols.length; i++) {
+          const refCol = refCols[i] ?? refCols[0] ?? cols[i];
+          fkByColumn.set(cols[i], `${c.referencedTable}.${refCol}`);
+        }
+      }
       const columns = table.columns.map((column) => {
         const annotations = [];
         const isPrimaryKey = pkColumns.has(column.name);
         if (isPrimaryKey) {
           annotations.push("PK");
         }
+        if (fkByColumn.has(column.name)) {
+          annotations.push(`FK -> ${fkByColumn.get(column.name)}`);
+        }
+        if (uniqueColumns.has(column.name)) {
+          annotations.push("UNIQUE");
+        }
+        if (notNullColumns.has(column.name)) {
+          annotations.push("NOT NULL");
+        }
+        if (defaultByColumn.has(column.name)) {
+          annotations.push(`DEFAULT: ${defaultByColumn.get(column.name)}`);
+        }
         if (column.isIndexed && !isPrimaryKey) {
           annotations.push("Indexed");
         }
-        if (column.kind === "LowCardinality" && column.values?.length) {
+        if (column.kind === "Enum" && column.values?.length) {
+          annotations.push(`Enum: ${column.values.join(", ")}`);
+        } else if (column.kind === "LowCardinality" && column.values?.length) {
           annotations.push(`LowCardinality: ${column.values.join(", ")}`);
         }
         if (column.stats) {
@@ -181,9 +567,19 @@ ${table.indexes.map((index) => {
         const columnsText = index.columns?.length ? index.columns.join(", ") : "expression";
         return `    - ${index.name}${propsText}: ${columnsText}`;
       }).join("\n")}` : "";
+      const multiColumnUniques = table.constraints?.filter(
+        (c) => c.type === "UNIQUE" && (c.columns?.length ?? 0) > 1
+      ) ?? [];
+      const uniqueConstraints = multiColumnUniques.length ? `
+  Unique Constraints:
+${multiColumnUniques.map((c) => `    - ${c.name}: (${c.columns?.join(", ")})`).join("\n")}` : "";
+      const checkConstraints = table.constraints?.filter((c) => c.type === "CHECK") ?? [];
+      const checks = checkConstraints.length ? `
+  Check Constraints:
+${checkConstraints.map((c) => `    - ${c.name}: ${c.definition}`).join("\n")}` : "";
       return `- Table: ${table.name}${rowCountInfo}
   Columns:
-${columns}${indexes}`;
+${columns}${indexes}${uniqueConstraints}${checks}`;
     }).join("\n\n");
   }
   #formatTableLabel = (tableName) => {
@@ -228,105 +624,6 @@ ${columns}${indexes}`;
   };
 };
 
-// packages/text2sql/src/lib/adapters/groundings/column-stats.grounding.ts
-var ColumnStatsGrounding = class extends AbstractGrounding {
-  constructor(config = {}) {
-    super("column_stats");
-  }
-  /**
-   * Execute the grounding process.
-   * Annotates columns in ctx.tables and ctx.views with statistics.
-   */
-  async execute(ctx) {
-    const allContainers = [...ctx.tables, ...ctx.views];
-    for (const container of allContainers) {
-      for (const column of container.columns) {
-        try {
-          const stats = await this.collectStats(container.name, column);
-          if (stats) {
-            column.stats = stats;
-          }
-        } catch (error) {
-          console.warn(
-            "Error collecting stats for",
-            container.name,
-            column.name,
-            error
-          );
-        }
-      }
-    }
-    return () => this.#describe();
-  }
-  #describe() {
-    return null;
-  }
-};
-
-// packages/text2sql/src/lib/adapters/groundings/low-cardinality.grounding.ts
-var LowCardinalityGrounding = class extends AbstractGrounding {
-  constructor(config = {}) {
-    super("low_cardinality");
-  }
-  /**
-   * Execute the grounding process.
-   * Annotates columns in ctx.tables and ctx.views with low cardinality values.
-   */
-  async execute(ctx) {
-    const allContainers = [...ctx.tables, ...ctx.views];
-    for (const container of allContainers) {
-      for (const column of container.columns) {
-        try {
-          const lowCard = await this.collectLowCardinality(
-            container.name,
-            column
-          );
-          if (lowCard) {
-            column.kind = lowCard.kind;
-            column.values = lowCard.values;
-          }
-        } catch (error) {
-          console.warn(
-            "Error collecting low cardinality values for",
-            container.name,
-            column.name,
-            error
-          );
-        }
-      }
-    }
-    return () => this.#describe();
-  }
-  #describe() {
-    return null;
-  }
-};
-
-// packages/text2sql/src/lib/adapters/groundings/info.grounding.ts
-var InfoGrounding = class extends AbstractGrounding {
-  constructor(config = {}) {
-    super("dialect_info");
-  }
-  /**
-   * Execute the grounding process.
-   * Writes database info to ctx.info.
-   */
-  async execute(ctx) {
-    ctx.info = await this.collectInfo();
-    const lines = [`Dialect: ${ctx.info.dialect ?? "unknown"}`];
-    if (ctx.info.version) {
-      lines.push(`Version: ${ctx.info.version}`);
-    }
-    if (ctx.info.database) {
-      lines.push(`Database: ${ctx.info.database}`);
-    }
-    if (ctx.info.details && Object.keys(ctx.info.details).length) {
-      lines.push(`Details: ${JSON.stringify(ctx.info.details)}`);
-    }
-    return () => lines.join("\n");
-  }
-};
-
 // packages/text2sql/src/lib/adapters/groundings/view.grounding.ts
 var ViewGrounding = class extends AbstractGrounding {
   #filter;
@@ -400,108 +697,16 @@ ${columns}`;
     return names.filter(filter);
   }
 };
-
-// packages/text2sql/src/lib/adapters/groundings/report.grounding.ts
-import { groq } from "@ai-sdk/groq";
-import { tool } from "ai";
-import dedent from "dedent";
-import z from "zod";
-import {
-  agent,
-  generate,
-  toState,
-  user
-} from "@deepagents/agent";
-var ReportGrounding = class extends AbstractGrounding {
-  #adapter;
-  #model;
-  #cache;
-  #forceRefresh;
-  constructor(adapter, config = {}) {
-    super("business_context");
-    this.#adapter = adapter;
-    this.#model = config.model ?? groq("openai/gpt-oss-20b");
-    this.#cache = config.cache;
-    this.#forceRefresh = config.forceRefresh ?? false;
-  }
-  async execute(ctx) {
-    if (!this.#forceRefresh && this.#cache) {
-      const cached = await this.#cache.get();
-      if (cached) {
-        ctx.report = cached;
-        return () => cached;
-      }
-    }
-    const report = await this.#generateReport();
-    ctx.report = report;
-    if (this.#cache) {
-      await this.#cache.set(report);
-    }
-    return () => report;
-  }
-  async #generateReport() {
-    const reportAgent = agent({
-      name: "db-report-agent",
-      model: this.#model,
-      prompt: () => dedent`
-        <identity>
-          You are a database analyst expert. Your job is to understand what
-          a database represents and provide business context about it.
-          You have READ-ONLY access to the database.
-        </identity>
-
-        <instructions>
-          Write a business context that helps another agent answer questions accurately.
-
-          For EACH table, do queries ONE AT A TIME:
-          1. SELECT COUNT(*) to get row count
-          2. SELECT * LIMIT 3 to see sample data
-
-          Then write a report with:
-          - What business this database is for
-          - For each table: purpose, row count, and example of what the data looks like
-
-          Include concrete examples like "Track prices are $0.99",
-          "Customer names like 'Luís Gonçalves'", etc.
-
-          Keep it 400-600 words, conversational style.
-        </instructions>
-      `,
-      tools: {
-        query_database: tool({
-          description: "Execute a SELECT query to explore the database and gather insights.",
-          inputSchema: z.object({
-            sql: z.string().describe("The SELECT query to execute"),
-            purpose: z.string().describe(
-              "What insight you are trying to gather with this query"
-            )
-          }),
-          execute: ({ sql }, options) => {
-            const state = toState(options);
-            return state.adapter.execute(sql);
-          }
-        })
-      }
-    });
-    const { text } = await generate(
-      reportAgent,
-      [
-        user(
-          "Please analyze the database and write a contextual report about what this database represents."
-        )
-      ],
-      { adapter: this.#adapter }
-    );
-    return text;
-  }
-};
 export {
+  AbstractGrounding,
   ColumnStatsGrounding,
+  ColumnValuesGrounding,
+  ConstraintGrounding,
   InfoGrounding,
-  LowCardinalityGrounding,
   ReportGrounding,
+  RowCountGrounding,
   TableGrounding,
   ViewGrounding,
   createGroundingContext
 };
-//# sourceMappingURL=grounding.js.map
+//# sourceMappingURL=index.js.map