@dbx-tools/search 0.9.154 → 0.9.158

package/src/lakebase.ts

@@ -25,7 +25,15 @@
 
 import { ExecutionError } from "@databricks/appkit";
 import { postgresConnectionOptions } from "@dbx-tools/postgres";
-import { asyncUtils, log, object, stringUtils } from "@dbx-tools/shared-core";
+import {
+  log,
+  sleep,
+  isRecord,
+  tokenizeWithOptions,
+  trimToEmpty,
+  trimToNull,
+  toSlug,
+} from "@dbx-tools/shared-core";
 import type {
   SearchDocument,
   SearchHit,
@@ -61,12 +69,7 @@ const RESERVED_COLUMNS = new Set(["id", "search_text", "document", "search_vecto
  * built.
  */
 export function toSearchTerms(query: string): string[] {
-  return [
-    ...stringUtils.tokenizeWithOptions(
-      { lowerCase: true, camelCase: false, distinct: true },
-      query,
-    ),
-  ];
+  return [...tokenizeWithOptions({ lowerCase: true, camelCase: false, distinct: true }, query)];
 }
 
 /**
@@ -139,7 +142,7 @@ export class LakebaseSearchBackend {
       for (let attempt = 0; attempt < 200; attempt += 1) {
         this.pool = this.poolSource.managedPool() ?? undefined;
         if (this.pool) return this.pool;
-        await asyncUtils.sleep(50);
+        await sleep(50);
       }
       throw new ExecutionError("search (lakebase): managed pool did not initialize", {
         context: { operation: "connect" },
@@ -179,7 +182,7 @@ export class LakebaseSearchBackend {
     query: string,
     options: LakebaseSearchOptions = {},
   ): Promise<SearchResult> {
-    const text = stringUtils.trimToEmpty(query);
+    const text = trimToEmpty(query);
     const table = this.tableFor(index);
     const limit = options.limit ?? 10;
     const pool = await this.getPool();
@@ -359,7 +362,7 @@ export class LakebaseSearchBackend {
     signal?: AbortSignal,
   ): Promise<void> {
     for (const doc of documents) {
-      const id = stringUtils.trimToNull(String(doc.id ?? doc.ID ?? "")) ?? undefined;
+      const id = trimToNull(String(doc.id ?? doc.ID ?? "")) ?? undefined;
       if (id === undefined) {
         throw new ExecutionError("search (lakebase): a document is missing an `id`", {
           context: { operation: "addDocuments" },
@@ -380,7 +383,7 @@ export class LakebaseSearchBackend {
 
   /** The text a row is indexed by: the text column first, then any other string field. */
   private searchText(doc: SearchDocument, textColumn: string): string {
-    const primary = stringUtils.trimToEmpty(String(doc[textColumn] ?? ""));
+    const primary = trimToEmpty(String(doc[textColumn] ?? ""));
     const rest = Object.entries(doc)
       .filter(([key, value]) => key !== textColumn && key !== "id" && typeof value === "string")
       .map(([, value]) => value as string);
@@ -389,7 +392,7 @@ export class LakebaseSearchBackend {
 
   /** A hit's `fields`: the stored document minus the reserved/internal keys. */
   private toFields(document: unknown): Record<string, unknown> {
-    if (!object.isRecord(document)) return {};
+    if (!isRecord(document)) return {};
     const fields: Record<string, unknown> = {};
     for (const [key, value] of Object.entries(document)) {
       if (RESERVED_COLUMNS.has(key)) continue;
@@ -429,7 +432,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 = stringUtils.toSlug(last).replace(/-/g, "_");
+    const slug = toSlug(last).replace(/-/g, "_");
     return slug.length > 0 ? slug : "documents";
   }
 

package/src/native.ts

@@ -11,14 +11,14 @@ import type {
   SearchRequest as AppKitSearchRequest,
   SearchResponse as AppKitSearchResponse,
 } from "@databricks/appkit/beta";
-import { stringUtils } from "@dbx-tools/shared-core";
+import { trimToNull } from "@dbx-tools/shared-core";
 import {
-  searchSchemas,
   type SearchDocument,
   type SearchHit,
-  type SearchRequest as ExtensionSearchRequest,
+  type SearchRequest,
   type SearchResult,
   type UpsertResult,
+  toAiSearchQueryType,
 } from "@dbx-tools/shared-search";
 import type { SearchOptions, SearchReadBackend } from "./client.ts";
 import { indexConfigFor, type ResolvedSearchConfig } from "./config.ts";
@@ -30,7 +30,7 @@ export interface AiSearchProvider {
   addDocuments?(alias: string, documents: SearchDocument[]): Promise<UpsertResult>;
 }
 
-function filters(value: ExtensionSearchRequest["filter"]): AppKitSearchFilters | undefined {
+function filters(value: SearchRequest["filter"]): AppKitSearchFilters | undefined {
   if (!value) return undefined;
   const result: AppKitSearchFilters = {};
   for (const [key, item] of Object.entries(value)) {
@@ -61,7 +61,7 @@ function hitId(
 }
 
 function aliasFor(index: string, config: ResolvedSearchConfig): string {
-  const alias = indexConfigFor(config, index)?.alias ?? stringUtils.trimToNull(index);
+  const alias = indexConfigFor(config, index)?.alias ?? trimToNull(index);
   if (!alias) throw new ValidationError("AI Search requires a configured index alias");
   return alias;
 }
@@ -77,7 +77,7 @@ export function nativeAiSearchBackend(
       options.signal?.throwIfAborted();
       const known = indexConfigFor(config, index);
       const alias = aliasFor(index, config);
-      const resolvedQueryType = searchSchemas.toAiSearchQueryType(options.mode);
+      const resolvedQueryType = toAiSearchQueryType(options.mode);
       const resolvedFilters = filters(options.filter);
       const response = await provider.query(alias, {
         queryText: query,

package/src/plugin.ts

@@ -22,14 +22,22 @@
 
 import { toPlugin, type IAppRouter, type PluginManifest } from "@databricks/appkit";
 import { aiSearch, defineTool, type ToolRegistry } from "@databricks/appkit/beta";
-import { pluginRegistry, ToolRegistryPlugin } from "@dbx-tools/appkit";
-import { errorUtils, log, stringUtils } from "@dbx-tools/shared-core";
+import { ToolRegistryPlugin, require } from "@dbx-tools/appkit";
+import { log, toError, errorMessage, trimToNull } from "@dbx-tools/shared-core";
 import {
   SEARCH_DEFAULT_BASE_PATH,
   SEARCH_PLUGIN_NAME,
   SEARCH_ROUTES,
-  searchSchemas,
   type SearchClientConfig,
+  searchRequestSchema,
+  universalSearchRequestSchema,
+  searchDocumentSchema,
+  createIndexRequestSchema,
+  syncIndexRequestSchema,
+  type SearchRequest,
+  type UniversalSearchRequest,
+  type CreateIndexRequest,
+  type SyncIndexRequest,
 } from "@dbx-tools/shared-search";
 import { toSearchOptions, toUniversalSearchOptions } from "./_search-options.ts";
 import { SEARCH_CONFIG_SCHEMA, resolveSearchConfig, type SearchPluginConfig } from "./config.ts";
@@ -112,14 +120,14 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
     const registry: ToolRegistry = {
       search: defineTool({
         description: SEARCH_TOOL_DESCRIPTION,
-        schema: searchSchemas.searchRequestSchema,
+        schema: searchRequestSchema,
         annotations: { effect: "read", requiresUserContext: true },
         autoInheritable: false,
         execute: async (args, signal) => this.runSearch(args, signal),
       }),
       universal_search: defineTool({
         description: UNIVERSAL_SEARCH_TOOL_DESCRIPTION,
-        schema: searchSchemas.universalSearchRequestSchema,
+        schema: universalSearchRequestSchema,
         annotations: { effect: "read", requiresUserContext: true },
         autoInheritable: false,
         execute: async (args, signal) => this.runUniversalSearch(args, signal),
@@ -128,9 +136,9 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
     if (config.allowWrite) {
       registry.add_documents = defineTool({
         description: ADD_DOCUMENTS_TOOL_DESCRIPTION,
-        schema: searchSchemas.searchRequestSchema
+        schema: searchRequestSchema
           .pick({ index: true })
-          .extend({ documents: searchSchemas.searchDocumentSchema.array() }),
+          .extend({ documents: searchDocumentSchema.array() }),
         annotations: { effect: "write", requiresUserContext: true },
         autoInheritable: false,
         execute: async (args, signal) => this.runAddDocuments(args, signal),
@@ -138,14 +146,14 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
       if (readBackend?.supportsLifecycle) {
         registry.create_index = defineTool({
           description: CREATE_INDEX_TOOL_DESCRIPTION,
-          schema: searchSchemas.createIndexRequestSchema,
+          schema: createIndexRequestSchema,
           annotations: { effect: "write", requiresUserContext: true },
           autoInheritable: false,
           execute: async (args, signal) => this.runCreateIndex(args, signal),
         });
         registry.sync_index = defineTool({
           description: SYNC_INDEX_TOOL_DESCRIPTION,
-          schema: searchSchemas.syncIndexRequestSchema,
+          schema: syncIndexRequestSchema,
           annotations: { effect: "write", requiresUserContext: true },
           autoInheritable: false,
           execute: async (args, signal) => this.runSyncIndex(args, signal),
@@ -181,7 +189,7 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
 
   /** Resolve the registered AppKit-compatible AI Search provider. */
   private resolveProviderBackend(): ReturnType<typeof nativeAiSearchBackend> {
-    const provider = pluginRegistry.require(this.context, aiSearch, this).exports();
+    const provider = require(this.context, aiSearch, this).exports();
     return nativeAiSearchBackend(provider, resolveSearchConfig(this.config));
   }
 
@@ -194,7 +202,7 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
   private async runEnsureOnSetup(config: SearchRuntime["config"]): Promise<void> {
     const spec = config.ensureOnSetup;
     if (!spec) return;
-    const index = stringUtils.trimToNull(spec.index) ?? config.defaultIndex;
+    const index = trimToNull(spec.index) ?? config.defaultIndex;
     if (!index) {
       logger.warn("ensure-skipped", {
         reason: "no index name (set `index` or `ensureOnSetup.index`)",
@@ -227,7 +235,7 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
         rowCount: info.rowCount ?? 0,
       });
     } catch (cause) {
-      logger.warn("ensure-failed", { index, error: errorUtils.toError(cause) });
+      logger.warn("ensure-failed", { index, error: toError(cause) });
     }
   }
 
@@ -272,7 +280,7 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
       path: SEARCH_ROUTES.universal,
       handler: async (req, res) => {
         await this.respond(res, "universalSearch", () => {
-          const request = searchSchemas.universalSearchRequestSchema.parse(req.body ?? {});
+          const request = universalSearchRequestSchema.parse(req.body ?? {});
           return this.asUser(req).runUniversalSearch(request);
         });
       },
@@ -347,10 +355,10 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
     try {
       res.json(await run());
     } catch (cause) {
-      const message = errorUtils.errorMessage(cause);
+      const message = errorMessage(cause);
       logger.warn("route-failed", {
         operation,
-        error: errorUtils.toError(cause),
+        error: toError(cause),
       });
       res.status(500).json({ error: message });
     }
@@ -377,31 +385,30 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
   override exports() {
     return {
       /** Search one index (default when omitted). Runs as the current context user. */
-      search: (request: searchSchemas.SearchRequest, signal?: AbortSignal) =>
-        this.runSearch(request, signal),
+      search: (request: SearchRequest, signal?: AbortSignal) => this.runSearch(request, signal),
       /** Search across every configured index and merge the hits. */
-      universalSearch: (request: searchSchemas.UniversalSearchRequest, signal?: AbortSignal) =>
+      universalSearch: (request: UniversalSearchRequest, signal?: AbortSignal) =>
         this.runUniversalSearch(request, signal),
       /** Add or update documents in a direct-access index (throws when writes are disabled). */
       addDocuments: (request: { index?: string; documents: unknown }, signal?: AbortSignal) =>
         this.runAddDocuments(request, signal),
       /** Create a Vector Search index (throws when writes are disabled). */
-      createIndex: (request: searchSchemas.CreateIndexRequest, signal?: AbortSignal) =>
+      createIndex: (request: CreateIndexRequest, signal?: AbortSignal) =>
         this.runCreateIndex(request, signal),
       /** Sync a Delta Sync index from its source table (throws when writes are disabled). */
-      syncIndex: (request: searchSchemas.SyncIndexRequest, signal?: AbortSignal) =>
+      syncIndex: (request: SyncIndexRequest, signal?: AbortSignal) =>
         this.runSyncIndex(request, signal),
     };
   }
 
   private async runSearch(args: unknown, signal?: AbortSignal) {
-    const request = searchSchemas.searchRequestSchema.parse(args);
+    const request = searchRequestSchema.parse(args);
     const { client } = this.runtime;
     return client.search(request.query, toSearchOptions(request, signal));
   }
 
   private async runUniversalSearch(args: unknown, signal?: AbortSignal) {
-    const request = searchSchemas.universalSearchRequestSchema.parse(args);
+    const request = universalSearchRequestSchema.parse(args);
     const { client } = this.runtime;
     return client.universalSearch(request.query, toUniversalSearchOptions(request, signal));
   }
@@ -415,13 +422,13 @@ export class SearchPlugin extends ToolRegistryPlugin<SearchPluginConfig> {
   }
 
   private async runCreateIndex(args: unknown, signal?: AbortSignal) {
-    const request = searchSchemas.createIndexRequestSchema.parse(args);
+    const request = createIndexRequestSchema.parse(args);
     const { client } = this.runtime;
     return client.createIndex(request.name, toCreateIndexOptions(request, signal));
   }
 
   private async runSyncIndex(args: unknown, signal?: AbortSignal) {
-    const request = searchSchemas.syncIndexRequestSchema.parse(args);
+    const request = syncIndexRequestSchema.parse(args);
     const { client, config } = this.runtime;
     const index = request.index ?? config.defaultIndex ?? "";
     await client.syncIndex(index, signal);

package/src/query.ts

@@ -1,7 +1,7 @@
 /** Request-shape helpers for search extension write routes and tools. */
 
 import { ValidationError } from "@databricks/appkit";
-import { json, object } from "@dbx-tools/shared-core";
+import { parse, isRecord } from "@dbx-tools/shared-core";
 
 /**
  * Parse a JSON document payload the model / a route supplied for a write.
@@ -10,10 +10,10 @@ import { json, object } from "@dbx-tools/shared-core";
  * {@link ValidationError} on unparseable input so the caller can surface it.
  */
 export function toDocumentArray(input: unknown): Array<Record<string, unknown>> {
-  const value = typeof input === "string" ? json.parse(input, undefined) : input;
+  const value = typeof input === "string" ? parse(input, undefined) : input;
   if (value === undefined) {
     throw new ValidationError("documents must be a JSON object, array, or string");
   }
   const list = Array.isArray(value) ? value : [value];
-  return list.filter(object.isRecord);
+  return list.filter(isRecord);
 }

package/src/schema.ts

@@ -7,11 +7,10 @@
  * @module
  */
 
-import { stringUtils } from "@dbx-tools/shared-core";
-import { searchSchemas } from "@dbx-tools/shared-search";
+import { toDescription } from "@dbx-tools/shared-core";
 
 /** Description the model reads for the `search` tool. */
-export const SEARCH_TOOL_DESCRIPTION = stringUtils.toDescription(`
+export const SEARCH_TOOL_DESCRIPTION = toDescription(`
   Search a Databricks AI Search (Vector Search) index for the documents most
   relevant to a query. Pass a natural-language question or keywords; hybrid
   matching combines semantic similarity with keyword ranking, so exact terms
@@ -22,7 +21,7 @@ export const SEARCH_TOOL_DESCRIPTION = stringUtils.toDescription(`
 `);
 
 /** Description the model reads for the `universal_search` tool. */
-export const UNIVERSAL_SEARCH_TOOL_DESCRIPTION = stringUtils.toDescription(`
+export const UNIVERSAL_SEARCH_TOOL_DESCRIPTION = toDescription(`
   Search across every configured AI Search index at once and return the best
   matches from all of them, merged and ranked. Use it when the right index
   isn't known in advance or when an answer may live in any of several
@@ -30,7 +29,7 @@ export const UNIVERSAL_SEARCH_TOOL_DESCRIPTION = stringUtils.toDescription(`
 `);
 
 /** Description the model reads for the `add_documents` tool. */
-export const ADD_DOCUMENTS_TOOL_DESCRIPTION = stringUtils.toDescription(`
+export const ADD_DOCUMENTS_TOOL_DESCRIPTION = toDescription(`
   Add or update documents in a direct-access AI Search index. Pass an array of
   documents as JSON objects; each MUST include the index's primary-key column.
   Only available when the app enables the write surface. Use it to index new
@@ -38,7 +37,7 @@ export const ADD_DOCUMENTS_TOOL_DESCRIPTION = stringUtils.toDescription(`
 `);
 
 /** Description the model reads for the `create_index` tool. */
-export const CREATE_INDEX_TOOL_DESCRIPTION = stringUtils.toDescription(`
+export const CREATE_INDEX_TOOL_DESCRIPTION = toDescription(`
   Create a Databricks AI Search (Vector Search) index. For the common case pass
   a Delta source table (catalog.schema.table): Databricks computes embeddings
   from its text column and keeps the index synced. To create a direct-access
@@ -50,26 +49,8 @@ export const CREATE_INDEX_TOOL_DESCRIPTION = stringUtils.toDescription(`
 `);
 
 /** Description the model reads for the `sync_index` tool. */
-export const SYNC_INDEX_TOOL_DESCRIPTION = stringUtils.toDescription(`
+export const SYNC_INDEX_TOOL_DESCRIPTION = toDescription(`
   Refresh a Delta Sync AI Search index from its source table so newly added or
   changed rows become searchable. Optionally name the index (defaults to the
   app's default index). Only available when the app enables the write surface.
 `);
-
-/** Schema for the `search` tool input (the shared request schema). */
-export const searchToolSchema = searchSchemas.searchRequestSchema;
-
-/** Schema for the `universal_search` tool input. */
-export const universalSearchToolSchema = searchSchemas.universalSearchRequestSchema;
-
-/** Schema for the `search` / `universal_search` tool output. */
-export const searchResultSchema = searchSchemas.searchResultSchema;
-
-/** Schema for the `create_index` tool input. */
-export const createIndexToolSchema = searchSchemas.createIndexRequestSchema;
-
-/** Schema for the `create_index` tool output (a resolved index definition). */
-export const indexInfoSchema = searchSchemas.indexInfoSchema;
-
-/** Schema for the `sync_index` tool input. */
-export const syncIndexToolSchema = searchSchemas.syncIndexRequestSchema;

package/src/tool.ts

@@ -15,7 +15,17 @@
  * @module
  */
 
-import { searchSchemas, type UpsertResult } from "@dbx-tools/shared-search";
+import {
+  createIndexRequestSchema,
+  indexInfoSchema,
+  searchDocumentSchema,
+  searchRequestSchema,
+  searchResultSchema,
+  syncIndexRequestSchema,
+  universalSearchRequestSchema,
+  upsertResultSchema,
+  type UpsertResult,
+} from "@dbx-tools/shared-search";
 import { createTool } from "@mastra/core/tools";
 import { z } from "zod";
 import { toSearchOptions, toUniversalSearchOptions } from "./_search-options.ts";
@@ -30,12 +40,6 @@ import {
   SEARCH_TOOL_DESCRIPTION,
   SYNC_INDEX_TOOL_DESCRIPTION,
   UNIVERSAL_SEARCH_TOOL_DESCRIPTION,
-  createIndexToolSchema,
-  indexInfoSchema,
-  searchResultSchema,
-  searchToolSchema,
-  syncIndexToolSchema,
-  universalSearchToolSchema,
 } from "./schema.ts";
 
 /** Options accepted by every standalone tool factory. */
@@ -84,10 +88,10 @@ export function searchTool(options: SearchToolOptions = {}) {
   return createTool({
     id: options.id ?? "search",
     description: SEARCH_TOOL_DESCRIPTION,
-    inputSchema: searchToolSchema,
+    inputSchema: searchRequestSchema,
     outputSchema: searchResultSchema,
     execute: async (input, context) => {
-      const request = searchToolSchema.parse(input);
+      const request = searchRequestSchema.parse(input);
       return runtime.client.search(request.query, toSearchOptions(request, context?.abortSignal));
     },
   });
@@ -99,10 +103,10 @@ export function universalSearchTool(options: SearchToolOptions = {}) {
   return createTool({
     id: options.id ?? "universal_search",
     description: UNIVERSAL_SEARCH_TOOL_DESCRIPTION,
-    inputSchema: universalSearchToolSchema,
+    inputSchema: universalSearchRequestSchema,
     outputSchema: searchResultSchema,
     execute: async (input, context) => {
-      const request = universalSearchToolSchema.parse(input);
+      const request = universalSearchRequestSchema.parse(input);
       return runtime.client.universalSearch(
         request.query,
         toUniversalSearchOptions(request, context?.abortSignal),
@@ -117,16 +121,14 @@ export function universalSearchTool(options: SearchToolOptions = {}) {
  */
 export function addDocumentsTool(options: SearchToolOptions = {}) {
   const runtime = toolRuntime(options);
-  const inputSchema = searchSchemas.searchDocumentSchema
+  const inputSchema = searchDocumentSchema
     .array()
     .describe("Documents to add or update. Each must include the index primary key.");
   return createTool({
     id: options.id ?? "add_documents",
     description: ADD_DOCUMENTS_TOOL_DESCRIPTION,
-    inputSchema: searchSchemas.searchRequestSchema
-      .pick({ index: true })
-      .extend({ documents: inputSchema }),
-    outputSchema: searchSchemas.upsertResultSchema,
+    inputSchema: searchRequestSchema.pick({ index: true }).extend({ documents: inputSchema }),
+    outputSchema: upsertResultSchema,
     execute: async (input, context): Promise<UpsertResult> => {
       const { client, config } = runtime;
       const record = input as { index?: string; documents: unknown };
@@ -148,10 +150,10 @@ export function createIndexTool(options: SearchToolOptions = {}) {
   return createTool({
     id: options.id ?? "create_index",
     description: CREATE_INDEX_TOOL_DESCRIPTION,
-    inputSchema: createIndexToolSchema,
+    inputSchema: createIndexRequestSchema,
     outputSchema: indexInfoSchema,
     execute: async (input, context) => {
-      const request = createIndexToolSchema.parse(input);
+      const request = createIndexRequestSchema.parse(input);
       return runtime.client.createIndex(
         request.name,
         toCreateIndexOptions(request, context?.abortSignal),
@@ -175,10 +177,10 @@ export function syncIndexTool(options: SearchToolOptions = {}) {
   return createTool({
     id: options.id ?? "sync_index",
     description: SYNC_INDEX_TOOL_DESCRIPTION,
-    inputSchema: syncIndexToolSchema,
+    inputSchema: syncIndexRequestSchema,
     outputSchema: syncIndexResultSchema,
     execute: async (input, context) => {
-      const request = syncIndexToolSchema.parse(input);
+      const request = syncIndexRequestSchema.parse(input);
       const { client, config } = runtime;
       const index = request.index ?? config.defaultIndex ?? "";
       await client.syncIndex(index, context?.abortSignal);