@dbx-tools/search 0.6.107 → 0.6.111

package/src/lakebase-plugin.ts

@@ -0,0 +1,258 @@
+/**
+ * AppKit-compatible AI Search provider backed by Lakebase full-text search.
+ *
+ * @module
+ */
+
+import {
+  lakebase,
+  Plugin,
+  toPlugin,
+  ValidationError,
+  type BasePluginConfig,
+  type IAppRouter,
+  type PluginManifest,
+} from "@databricks/appkit";
+import type { IndexConfig, SearchRequest, SearchResponse } from "@databricks/appkit/beta";
+import { plugin as appkitPlugin } from "@dbx-tools/appkit";
+import { log, object, string } from "@dbx-tools/shared-core";
+import type { SearchDocument } from "@dbx-tools/shared-search";
+import type express from "express";
+import type { JSONSchema7 } from "json-schema";
+import { LakebaseSearchBackend } from "./lakebase.ts";
+
+const logger = log.logger("search/lakebase-plugin");
+
+/** Per-alias Lakebase search configuration. */
+export interface LakebaseAiSearchIndexConfig extends Omit<
+  IndexConfig,
+  "embeddingFn" | "endpointName" | "pagination" | "reranker"
+> {
+  /** Documents used to seed an empty full-text table during setup. */
+  documents?: SearchDocument[];
+  /** Document field indexed first. Defaults to `text`. */
+  textColumn?: string;
+}
+
+/** Configuration for the Lakebase implementation of AppKit AI Search. */
+export interface LakebaseAiSearchConfig extends BasePluginConfig {
+  indexes?: Record<string, LakebaseAiSearchIndexConfig>;
+  /** PostgreSQL schema that owns the generated full-text tables. */
+  schema?: string;
+  /** Enable the document upsert route and exported method. */
+  allowWrite?: boolean;
+}
+
+const CONFIG_SCHEMA: JSONSchema7 = {
+  type: "object",
+  required: ["indexes"],
+  properties: {
+    indexes: {
+      type: "object",
+      additionalProperties: {
+        type: "object",
+        properties: {
+          indexName: { type: "string" },
+          columns: { type: "array", items: { type: "string" } },
+          numResults: { type: "number" },
+          queryType: { enum: ["full_text"] },
+          textColumn: { type: "string" },
+          documents: { type: "array", items: { type: "object" } },
+        },
+      },
+    },
+    schema: { type: "string" },
+    allowWrite: { type: "boolean" },
+  },
+};
+
+interface ResolvedLakebaseIndex {
+  alias: string;
+  indexName: string;
+  config: LakebaseAiSearchIndexConfig;
+}
+
+function routeParam(value: string | string[]): string {
+  return Array.isArray(value) ? (value[0] ?? "") : value;
+}
+
+/** Lakebase implementation of the native AppKit `aiSearch` contract. */
+export class LakebaseAiSearchPlugin extends Plugin<LakebaseAiSearchConfig> {
+  static manifest = {
+    name: "aiSearch",
+    displayName: "Lakebase AI Search",
+    description: "AppKit AI Search contract backed by PostgreSQL full-text search",
+    stability: "beta",
+    resources: { required: [], optional: [] },
+    config: { schema: CONFIG_SCHEMA },
+  } satisfies PluginManifest<"aiSearch">;
+
+  declare protected config: LakebaseAiSearchConfig;
+  private backend: LakebaseSearchBackend | undefined;
+
+  override async setup(): Promise<void> {
+    const lake = appkitPlugin.require(this.context, lakebase, this);
+    this.backend = new LakebaseSearchBackend(
+      () => lake.exports().getPgConfig(),
+      string.trimToNull(this.config.schema) ?? "public",
+    );
+    for (const index of this.indexes()) {
+      await this.backend.provision(index.indexName, {
+        textColumn: index.config.textColumn ?? "text",
+        ...(index.config.documents ? { seed: index.config.documents } : {}),
+      });
+    }
+    logger.info("ready", {
+      indexes: this.indexes().map((index) => index.alias),
+      schema: string.trimToNull(this.config.schema) ?? "public",
+    });
+  }
+
+  override injectRoutes(router: IAppRouter): void {
+    this.route(router, {
+      name: "query",
+      method: "post",
+      path: "/:alias/query",
+      handler: async (req: express.Request, res: express.Response) => {
+        try {
+          res.json(await this.query(routeParam(req.params.alias), req.body as SearchRequest));
+        } catch (error) {
+          res.status(400).json({
+            error: error instanceof Error ? error.message : "Search failed",
+            plugin: this.name,
+          });
+        }
+      },
+    });
+    this.route(router, {
+      name: "getConfig",
+      method: "get",
+      path: "/:alias/config",
+      handler: async (req: express.Request, res: express.Response) => {
+        const index = this.resolveIndex(routeParam(req.params.alias));
+        res.json({
+          alias: index.alias,
+          columns: index.config.columns,
+          queryType: "full_text",
+          numResults: index.config.numResults ?? 20,
+          reranker: false,
+          pagination: false,
+        });
+      },
+    });
+    if (this.config.allowWrite) {
+      this.route(router, {
+        name: "addDocuments",
+        method: "post",
+        path: "/:alias/documents",
+        handler: async (req: express.Request, res: express.Response) => {
+          const documents = Array.isArray(req.body?.documents)
+            ? req.body.documents.filter(object.isRecord)
+            : [];
+          res.json(await this.addDocuments(routeParam(req.params.alias), documents));
+        },
+      });
+    }
+  }
+
+  /** Query one configured Lakebase full-text index using AppKit's response shape. */
+  async query(
+    alias: string,
+    request: SearchRequest,
+  ): Promise<SearchResponse<Record<string, unknown>>> {
+    if (!request.queryText) {
+      throw new ValidationError("Lakebase AI Search requires queryText");
+    }
+    if (request.queryVector) {
+      throw new ValidationError("Lakebase AI Search does not accept queryVector");
+    }
+    if (request.queryType && request.queryType !== "full_text") {
+      throw new ValidationError("Lakebase AI Search supports only full_text queries");
+    }
+    const index = this.resolveIndex(alias);
+    const startedAt = performance.now();
+    const result = await this.requireBackend().search(index.indexName, request.queryText, {
+      limit: request.numResults ?? index.config.numResults ?? 20,
+      ...(request.filters ? { filter: request.filters } : {}),
+    });
+    const columns = request.columns ?? index.config.columns;
+    return {
+      results: result.hits.map((hit) => ({
+        score: hit.score,
+        data: this.project({ id: hit.id, ...hit.fields }, columns),
+      })),
+      totalCount: result.count,
+      queryTimeMs: Math.max(0, performance.now() - startedAt),
+      queryType: "full_text",
+      nextPageToken: null,
+    };
+  }
+
+  /** Add or update documents in one configured Lakebase full-text index. */
+  async addDocuments(alias: string, documents: SearchDocument[]) {
+    if (!this.config.allowWrite) {
+      throw new ValidationError("Lakebase AI Search writes are disabled");
+    }
+    const index = this.resolveIndex(alias);
+    return this.requireBackend().addDocuments(
+      index.indexName,
+      documents.filter(object.isRecord),
+      index.config.textColumn ?? "text",
+    );
+  }
+
+  override clientConfig() {
+    return {
+      indexes: this.indexes().map((index) => ({
+        alias: index.alias,
+        queryType: "full_text" as const,
+        pagination: false,
+      })),
+    };
+  }
+
+  exports() {
+    return {
+      providerKind: "lakebase" as const,
+      query: this.query.bind(this),
+      addDocuments: this.addDocuments.bind(this),
+    };
+  }
+
+  async shutdown(): Promise<void> {
+    await this.backend?.close();
+    this.backend = undefined;
+  }
+
+  private indexes(): ResolvedLakebaseIndex[] {
+    return Object.entries(this.config.indexes ?? {}).map(([alias, config]) => ({
+      alias,
+      indexName: string.trimToNull(config.indexName) ?? alias,
+      config,
+    }));
+  }
+
+  private resolveIndex(alias: string): ResolvedLakebaseIndex {
+    const index = this.indexes().find((candidate) => candidate.alias === alias);
+    if (!index) throw new ValidationError(`Unknown AI Search index alias "${alias}"`);
+    return index;
+  }
+
+  private requireBackend(): LakebaseSearchBackend {
+    if (!this.backend) throw new ValidationError("Lakebase AI Search is not initialized");
+    return this.backend;
+  }
+
+  private project(
+    data: Record<string, unknown>,
+    columns: string[] | undefined,
+  ): Record<string, unknown> {
+    if (!columns?.length) return data;
+    return Object.fromEntries(
+      columns.filter((column) => column in data).map((column) => [column, data[column]]),
+    );
+  }
+}
+
+/** AppKit-compatible Lakebase AI Search provider. */
+export const lakebaseAiSearch = toPlugin(LakebaseAiSearchPlugin);

package/src/lakebase.ts

@@ -1,7 +1,6 @@
 /**
- * A Lakebase (Postgres) full-text search backend - the FALLBACK used when no
- * Databricks Vector Search endpoint/index is configured but a Lakebase pool is
- * available. It provisions a single table per index, indexes a generated
+ * Lakebase (Postgres) full-text runtime behind `lakebaseAiSearch`. It provisions
+ * a single table per index alias, indexes a generated
  * `tsvector`, and answers queries with a prefix `to_tsquery` + `ts_rank`.
  *
  * Queries are compiled rather than passed through `websearch_to_tsquery`,
@@ -30,6 +29,7 @@ import { log, object, string } from "@dbx-tools/shared-core";
 import type {
   SearchDocument,
   SearchHit,
+  SearchRequest,
   SearchResult,
   UpsertResult,
 } from "@dbx-tools/shared-search";
@@ -81,6 +81,7 @@ export function toTsQuery(terms: readonly string[], operator: "&" | "|" = "&"):
 /** Options for a single-index Lakebase search (mirrors the client's `SearchOptions`). */
 export interface LakebaseSearchOptions {
   limit?: number;
+  filter?: SearchRequest["filter"];
   scoreThreshold?: number;
   signal?: AbortSignal;
 }
@@ -156,24 +157,33 @@ export class LakebaseSearchBackend {
     // An empty (or punctuation-only) box returns rows rather than nothing, so
     // the UI shows content before the user types - as a keyword index would.
     if (terms.length === 0) {
+      const filter = this.filterClause(options.filter, 1);
+      const limitPosition = filter.params.length + 1;
       const { rows } = await this.query<SearchRow>(
         pool,
-        `SELECT id, document, 0::float4 AS score FROM ${table} ORDER BY id LIMIT $1`,
-        [limit],
+        `SELECT id, document, 0::float4 AS score
+           FROM ${table}
+          WHERE TRUE${filter.sql}
+          ORDER BY id
+          LIMIT $${limitPosition}`,
+        [...filter.params, limit],
         options.signal,
       );
       return this.toResult(text, index, rows, options);
     }
 
     // Precise pass: every term must match, each as a prefix.
+    const strictFilter = this.filterClause(options.filter, 2);
+    const strictLimitPosition = strictFilter.params.length + 2;
     const strict = await this.query<SearchRow>(
       pool,
       `SELECT id, document, ts_rank(search_vector, to_tsquery('english', $1)) AS score
          FROM ${table}
         WHERE search_vector @@ to_tsquery('english', $1)
+          ${strictFilter.sql}
         ORDER BY score DESC
-        LIMIT $2`,
-      [toTsQuery(terms), limit],
+        LIMIT $${strictLimitPosition}`,
+      [toTsQuery(terms), ...strictFilter.params, limit],
       options.signal,
     );
     if (strict.rows.length > 0) return this.toResult(text, index, strict.rows, options);
@@ -183,15 +193,18 @@ export class LakebaseSearchBackend {
     // fragment that is not a prefix (`telligence`) or a token the text-search
     // parser split differently than expected. Substring matching cannot use
     // the GIN index, which is why it only runs once the indexed pass fails.
+    const relaxedFilter = this.filterClause(options.filter, 3);
+    const relaxedLimitPosition = relaxedFilter.params.length + 3;
     const relaxed = await this.query<SearchRow>(
       pool,
       `SELECT id, document, ts_rank(search_vector, to_tsquery('english', $1)) AS score
          FROM ${table}
-        WHERE search_vector @@ to_tsquery('english', $1)
-           OR search_text ILIKE ANY($2::text[])
+        WHERE (search_vector @@ to_tsquery('english', $1)
+           OR search_text ILIKE ANY($2::text[]))
+          ${relaxedFilter.sql}
         ORDER BY score DESC
-        LIMIT $3`,
-      [toTsQuery(terms, "|"), terms.map((term) => `%${term}%`), limit],
+        LIMIT $${relaxedLimitPosition}`,
+      [toTsQuery(terms, "|"), terms.map((term) => `%${term}%`), ...relaxedFilter.params, limit],
       options.signal,
     );
     return this.toResult(text, index, relaxed.rows, options);
@@ -283,6 +296,14 @@ export class LakebaseSearchBackend {
       [],
       signal,
     );
+    await this.query(
+      pool,
+      `ALTER TABLE ${table}
+         ADD COLUMN IF NOT EXISTS search_vector tsvector
+         GENERATED ALWAYS AS (to_tsvector('english', search_text)) STORED`,
+      [],
+      signal,
+    );
     await this.query(
       pool,
       `CREATE INDEX IF NOT EXISTS ${this.ident(`${this.bareName(table)}_fts`)}
@@ -342,6 +363,29 @@ export class LakebaseSearchBackend {
     return fields;
   }
 
+  /** Compile AppKit scalar/array filters against the stored JSON document. */
+  private filterClause(
+    filter: SearchRequest["filter"],
+    startPosition: number,
+  ): { sql: string; params: unknown[] } {
+    if (!filter || Object.keys(filter).length === 0) return { sql: "", params: [] };
+    const clauses: string[] = [];
+    const params: unknown[] = [];
+    for (const [key, value] of Object.entries(filter)) {
+      const keyPosition = startPosition + params.length;
+      params.push(key);
+      const valuePosition = startPosition + params.length;
+      if (Array.isArray(value)) {
+        params.push(value.map(String));
+        clauses.push(`document ->> $${keyPosition} = ANY($${valuePosition}::text[])`);
+      } else {
+        params.push(String(value));
+        clauses.push(`document ->> $${keyPosition} = $${valuePosition}`);
+      }
+    }
+    return { sql: ` AND ${clauses.join(" AND ")}`, params };
+  }
+
   /** The fully-qualified table name for an index reference. */
   private tableFor(index: string): string {
     return `${this.ident(this.schema)}.${this.ident(this.bareName(index))}`;
@@ -350,10 +394,7 @@ export class LakebaseSearchBackend {
   /** A safe bare table name derived from an index reference. */
   private bareName(reference: string): string {
     const last = reference.split(".").filter(Boolean).pop() ?? reference;
-    const slug = last
-      .toLowerCase()
-      .replace(/[^a-z0-9_]+/g, "_")
-      .replace(/^_+|_+$/g, "");
+    const slug = string.toSlug(last).replace(/-/g, "_");
     return slug.length > 0 ? slug : "documents";
   }
 

package/src/native.ts

@@ -0,0 +1,117 @@
+/**
+ * Adapter from AppKit's native AI Search provider to the dbx-tools extension
+ * contract used by federated search and agent tools.
+ *
+ * @module
+ */
+
+import { ValidationError } from "@databricks/appkit";
+import type {
+  SearchFilters as AppKitSearchFilters,
+  SearchRequest as AppKitSearchRequest,
+  SearchResponse as AppKitSearchResponse,
+} from "@databricks/appkit/beta";
+import {
+  search as sharedSearch,
+  type SearchDocument,
+  type SearchHit,
+  type SearchRequest as ExtensionSearchRequest,
+  type SearchResult,
+  type UpsertResult,
+} from "@dbx-tools/shared-search";
+import { string } from "@dbx-tools/shared-core";
+import type { SearchOptions, SearchReadBackend } from "./client.ts";
+import { indexConfigFor, type ResolvedSearchConfig } from "./config.ts";
+
+/** Minimal query surface exposed by AppKit's native `aiSearch` plugin. */
+export interface AiSearchProvider {
+  providerKind?: "lakebase";
+  query(alias: string, request: AppKitSearchRequest): Promise<AppKitSearchResponse>;
+  addDocuments?(alias: string, documents: SearchDocument[]): Promise<UpsertResult>;
+}
+
+function filters(value: ExtensionSearchRequest["filter"]): AppKitSearchFilters | undefined {
+  if (!value) return undefined;
+  const result: AppKitSearchFilters = {};
+  for (const [key, item] of Object.entries(value)) {
+    if (
+      typeof item === "string" ||
+      typeof item === "number" ||
+      typeof item === "boolean" ||
+      (Array.isArray(item) &&
+        item.every((entry) => typeof entry === "string" || typeof entry === "number"))
+    ) {
+      result[key] = item;
+      continue;
+    }
+    throw new ValidationError(
+      `AI Search filter "${key}" must be a string, number, boolean, or string/number array`,
+    );
+  }
+  return result;
+}
+
+function hitId(
+  data: Record<string, unknown>,
+  primaryKey: string | undefined,
+  index: number,
+): string {
+  const value = data[primaryKey ?? "id"] ?? Object.values(data)[0] ?? index;
+  return String(value);
+}
+
+function aliasFor(index: string, config: ResolvedSearchConfig): string {
+  const alias = indexConfigFor(config, index)?.alias ?? string.trimToNull(index);
+  if (!alias) throw new ValidationError("AI Search requires a configured index alias");
+  return alias;
+}
+
+/** Build a read backend that delegates every Vector Search query to AppKit. */
+export function nativeAiSearchBackend(
+  provider: AiSearchProvider,
+  config: ResolvedSearchConfig,
+): SearchReadBackend {
+  const backend: SearchReadBackend = {
+    supportsLifecycle: provider.providerKind !== "lakebase",
+    async search(index: string, query: string, options: SearchOptions = {}): Promise<SearchResult> {
+      options.signal?.throwIfAborted();
+      const known = indexConfigFor(config, index);
+      const alias = aliasFor(index, config);
+      const resolvedQueryType = sharedSearch.toAiSearchQueryType(options.mode);
+      const resolvedFilters = filters(options.filter);
+      const response = await provider.query(alias, {
+        queryText: query,
+        numResults: options.limit ?? config.pageSize,
+        ...(options.columns ? { columns: [...options.columns] } : {}),
+        ...(resolvedQueryType ? { queryType: resolvedQueryType } : {}),
+        ...(resolvedFilters ? { filters: resolvedFilters } : {}),
+      });
+      options.signal?.throwIfAborted();
+      const hits: SearchHit[] = response.results
+        .map((result, resultIndex) => ({
+          id: hitId(result.data, known?.primaryKey, resultIndex),
+          score: result.score,
+          fields: result.data,
+        }))
+        .filter(
+          (hit) => options.scoreThreshold === undefined || hit.score >= options.scoreThreshold,
+        );
+      return { query, index, hits, count: hits.length };
+    },
+  };
+  const addDocuments = provider.addDocuments?.bind(provider);
+  if (!addDocuments) return backend;
+  return {
+    ...backend,
+    async addDocuments(
+      index: string,
+      documents: SearchDocument[],
+      signal?: AbortSignal,
+    ): Promise<UpsertResult> {
+      signal?.throwIfAborted();
+      const result = await addDocuments(aliasFor(index, config), documents);
+      signal?.throwIfAborted();
+      return result;
+    },
+  };
+}