@dbx-tools/search 0.6.9

package/src/query.ts

@@ -0,0 +1,127 @@
+/**
+ * Translation between the browser-safe search contract and the Databricks
+ * Vector Search query API, kept in one place so the client, the tools, and the
+ * routes never hand-roll a `filters_json` string or unpack a `data_array` by
+ * hand.
+ *
+ *   - {@link toQueryType} maps the friendly {@link SearchMode} onto the API's
+ *     `query_type` (`HYBRID` / `ANN`). Keyword-only search rides on `ANN` with
+ *     text but no vector, which Databricks answers with its BM25 path.
+ *   - {@link compileFilter} turns the `{ column: value }` filter object (with
+ *     optional operator maps like `{ ">=": 10 }`) into the `filters_json`
+ *     string the API expects, so a caller never learns Databricks' filter
+ *     spelling.
+ *   - {@link toHits} unpacks the columnar `{ manifest, result }` response into
+ *     `{ id, score, fields }` hits, pulling the score out of the reserved
+ *     `__db_score` / `score` column and the id out of the primary-key column.
+ *
+ * @module
+ */
+
+import { ValidationError } from "@databricks/appkit";
+import type { SearchHit, SearchMode } from "@dbx-tools/shared-search";
+import { json, object } from "@dbx-tools/shared-core";
+
+/** The column name Databricks Vector Search returns the relevance score under. */
+const SCORE_COLUMN = "__db_score";
+
+/** Map a {@link SearchMode} onto the serving API `query_type`. */
+export function toQueryType(mode: SearchMode): string {
+  return mode === "hybrid" ? "HYBRID" : "ANN";
+}
+
+/**
+ * Compile a `{ column: value }` filter object into the `filters_json` string
+ * the query API expects. A scalar becomes an equality; an array becomes an
+ * IN-style match; an operator map (`{ ">=": 10, "<": 20 }`) expands to the
+ * `column operator` keys Databricks uses. Returns `undefined` for an empty
+ * filter so the field is omitted rather than sent as `{}`.
+ */
+export function compileFilter(filter: Record<string, unknown> | undefined): string | undefined {
+  if (!filter || Object.keys(filter).length === 0) return undefined;
+  const compiled: Record<string, unknown> = {};
+  for (const [column, value] of Object.entries(filter)) {
+    if (object.isRecord(value)) {
+      for (const [op, operand] of Object.entries(value)) {
+        compiled[`${column} ${op}`] = operand;
+      }
+    } else {
+      compiled[column] = value;
+    }
+  }
+  return JSON.stringify(compiled);
+}
+
+/** A minimal structural view of the Vector Search query response. */
+export interface QueryResponseLike {
+  manifest?: { columns?: Array<{ name?: string }> };
+  result?: { data_array?: Array<Array<unknown>> };
+  next_page_token?: string;
+}
+
+/**
+ * Unpack a columnar query response into {@link SearchHit}s. The manifest names
+ * the columns in order; each row is a positional array. The score comes from
+ * the reserved score column and the id from `primaryKey` (falling back to the
+ * first column when the key is unknown). The score column is stripped from
+ * `fields` so a hit's fields are just the document.
+ */
+export function toHits(
+  response: QueryResponseLike,
+  primaryKey: string | undefined,
+  indexName?: string,
+): SearchHit[] {
+  const columns = (response.manifest?.columns ?? []).map((c) => c.name ?? "");
+  const rows = response.result?.data_array ?? [];
+  const scoreIdx = columns.indexOf(SCORE_COLUMN);
+  const keyIdx = primaryKey ? columns.indexOf(primaryKey) : -1;
+  return rows.map((row, rowIndex) => {
+    const fields: Record<string, unknown> = {};
+    columns.forEach((name, i) => {
+      if (i === scoreIdx || !name) return;
+      fields[name] = row[i];
+    });
+    const scoreRaw = scoreIdx >= 0 ? Number(row[scoreIdx]) : NaN;
+    const idRaw =
+      keyIdx >= 0 ? row[keyIdx] : primaryKey ? fields[primaryKey] : (row[0] ?? rowIndex);
+    return {
+      id: String(idRaw ?? rowIndex),
+      score: Number.isFinite(scoreRaw) ? scoreRaw : 0,
+      fields,
+      ...(indexName ? { index: indexName } : {}),
+    };
+  });
+}
+
+/**
+ * The columns to request from an index for a search. When neither the request
+ * nor the index config names columns, the score column alone is requested and
+ * the primary key is appended so a hit always has an id. Callers that want the
+ * whole document should pass the index's own column list.
+ */
+export function toRequestColumns(
+  requested: readonly string[] | undefined,
+  fallback: readonly string[] | undefined,
+  primaryKey: string | undefined,
+): string[] {
+  const base = requested && requested.length > 0 ? requested : (fallback ?? []);
+  const columns = new Set<string>(base);
+  if (primaryKey) columns.add(primaryKey);
+  if (columns.size === 0 && primaryKey) columns.add(primaryKey);
+  return [...columns];
+}
+
+/**
+ * Parse a JSON document payload the model / a route supplied for a write.
+ * Accepts an already-parsed array/object or a JSON string, and always returns
+ * an array so a single document and a batch are handled the same way. Throws a
+ * {@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;
+  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);
+}

package/src/runtime.ts

@@ -0,0 +1,64 @@
+/**
+ * The shared AI Search runtime: the resolved config plus a single
+ * {@link SearchClient} instance the plugin, the Mastra tools, and the routes
+ * all read. Mirrors the web-search / email runtime pattern - the plugin primes
+ * it from its config at setup, and everything else reads the same instance so
+ * a tool invoked outside the agent still sees the deployment's config.
+ *
+ * @module
+ */
+
+import { createSearchClient, SearchClient } from "./client.ts";
+import {
+  resolveSearchConfig,
+  type SearchPluginConfig,
+  type ResolvedSearchConfig,
+} from "./config.ts";
+import { LakebaseSearchBackend } from "./lakebase.ts";
+
+/**
+ * How the runtime is built. `lakebase` (when present) is the Postgres full-text
+ * FALLBACK backend the plugin wires up when no Vector Search endpoint is
+ * configured but the AppKit `lakebase` plugin is registered. Every read/write
+ * returns the same shape either way.
+ */
+export interface SearchRuntimeOptions {
+  config?: SearchPluginConfig;
+  lakebase?: LakebaseSearchBackend;
+}
+
+/** The shared resolved config plus the client reads run through. */
+export interface SearchRuntime {
+  config: ResolvedSearchConfig;
+  client: SearchClient;
+  /** The Lakebase fallback backend, when the runtime is backed by Postgres. */
+  lakebase?: LakebaseSearchBackend;
+}
+
+let runtime: SearchRuntime | undefined;
+
+/**
+ * Return the shared runtime, building it on first use from the supplied config
+ * layered over environment defaults. Overrides are only read when the runtime
+ * is first created, so prime it from the plugin's config at setup; subsequent
+ * calls pass nothing and get the same instance.
+ */
+export function getSearchRuntime(options?: SearchRuntimeOptions): SearchRuntime {
+  if (!runtime) {
+    const config = resolveSearchConfig(options?.config);
+    const lakebase = options?.lakebase;
+    runtime = {
+      config,
+      client: createSearchClient(config, undefined, lakebase),
+      ...(lakebase ? { lakebase } : {}),
+    };
+  }
+  return runtime;
+}
+
+/** Drop the memoized runtime so the next {@link getSearchRuntime} rebuilds it. */
+export function resetSearchRuntime(): void {
+  const backend = runtime?.lakebase;
+  runtime = undefined;
+  if (backend) void backend.close();
+}

package/src/schema.ts

@@ -0,0 +1,75 @@
+/**
+ * Tool-facing descriptions and the request schemas the Mastra tools + AppKit
+ * tool provider validate against. The wire shapes themselves live in
+ * `@dbx-tools/shared-search`; this module adds the model-readable
+ * descriptions so both hosts describe the tools identically.
+ *
+ * @module
+ */
+
+import { search } from "@dbx-tools/shared-search";
+import { string } from "@dbx-tools/shared-core";
+
+/** Description the model reads for the `search` tool. */
+export const SEARCH_TOOL_DESCRIPTION = string.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
+  (product names, error codes) and paraphrases both work. Optionally name an
+  index (defaults to the app's configured index), a result limit, and attribute
+  filters. Use it to look up docs, knowledge-base articles, products, or any
+  indexed content before answering.
+`);
+
+/** Description the model reads for the `universal_search` tool. */
+export const UNIVERSAL_SEARCH_TOOL_DESCRIPTION = string.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
+  collections (docs, tickets, products).
+`);
+
+/** Description the model reads for the `add_documents` tool. */
+export const ADD_DOCUMENTS_TOOL_DESCRIPTION = string.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
+  content the user provides.
+`);
+
+/** Description the model reads for the `create_index` tool. */
+export const CREATE_INDEX_TOOL_DESCRIPTION = string.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
+  index you write vectors to yourself, pass an embedding dimension instead of a
+  source table. The endpoint, embedding model, primary key, and text column are
+  inferred when omitted. Only available when the app enables the write surface.
+  Creating an index provisions infrastructure - do this only when the user
+  explicitly asks to set up a new index.
+`);
+
+/** Description the model reads for the `sync_index` tool. */
+export const SYNC_INDEX_TOOL_DESCRIPTION = string.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 = search.searchRequestSchema;
+
+/** Schema for the `universal_search` tool input. */
+export const universalSearchToolSchema = search.universalSearchRequestSchema;
+
+/** Schema for the `search` / `universal_search` tool output. */
+export const searchResultSchema = search.searchResultSchema;
+
+/** Schema for the `create_index` tool input. */
+export const createIndexToolSchema = search.createIndexRequestSchema;
+
+/** Schema for the `create_index` tool output (a resolved index definition). */
+export const indexInfoSchema = search.indexInfoSchema;
+
+/** Schema for the `sync_index` tool input. */
+export const syncIndexToolSchema = search.syncIndexRequestSchema;

package/src/tool.ts

@@ -0,0 +1,169 @@
+/**
+ * The `search`, `universal_search`, and (opt-in) `add_documents`,
+ * `create_index`, and `sync_index` Mastra tools.
+ *
+ * All three read the shared runtime primed by the plugin, so a tool spread
+ * into an agent uses the deployment's default index, columns, page size, and
+ * mode without any per-tool wiring. They run under the caller's OBO scope (the
+ * client resolves the execution context's workspace client), so search runs as
+ * the requesting user and Unity Catalog ACLs apply.
+ *
+ * The same tools are exposed to AppKit's own agents through the plugin's
+ * `ToolProvider` (see `plugin.ts`); this module is the Mastra half.
+ *
+ * @module
+ */
+
+import { search as searchContract, type UpsertResult } from "@dbx-tools/shared-search";
+import { createTool } from "@mastra/core/tools";
+import { z } from "zod";
+import { toCreateIndexOptions } from "./index-tools.ts";
+import { toDocumentArray } from "./query.ts";
+import { getSearchRuntime } from "./runtime.ts";
+import {
+  ADD_DOCUMENTS_TOOL_DESCRIPTION,
+  CREATE_INDEX_TOOL_DESCRIPTION,
+  SEARCH_TOOL_DESCRIPTION,
+  SYNC_INDEX_TOOL_DESCRIPTION,
+  UNIVERSAL_SEARCH_TOOL_DESCRIPTION,
+  createIndexToolSchema,
+  indexInfoSchema,
+  searchResultSchema,
+  searchToolSchema,
+  syncIndexToolSchema,
+  universalSearchToolSchema,
+} from "./schema.ts";
+
+/** Common option accepted by every tool factory: override the tool id. */
+export interface SearchToolOptions {
+  /** Override the tool id (defaults per tool). */
+  id?: string;
+}
+
+/**
+ * Build the `search` tool. Spread it into any agent that should be able to look
+ * things up in an index.
+ *
+ * @example
+ * ```ts
+ * import { searchTool } from "@dbx-tools/search";
+ * import { createAgent } from "@dbx-tools/appkit-mastra";
+ *
+ * const support = createAgent({
+ *   instructions: "Answer from the docs. Use `search` to find them.",
+ *   tools: () => ({ search: searchTool() }),
+ * });
+ * ```
+ */
+export function searchTool(options: SearchToolOptions = {}) {
+  return createTool({
+    id: options.id ?? "search",
+    description: SEARCH_TOOL_DESCRIPTION,
+    inputSchema: searchToolSchema,
+    outputSchema: searchResultSchema,
+    execute: async (input, context) => {
+      const request = searchToolSchema.parse(input);
+      const { client } = getSearchRuntime();
+      return client.search(request.query, {
+        ...(request.index ? { index: request.index } : {}),
+        ...(request.limit ? { limit: request.limit } : {}),
+        ...(request.mode ? { mode: request.mode } : {}),
+        ...(request.columns ? { columns: request.columns } : {}),
+        ...(request.filter ? { filter: request.filter } : {}),
+        ...(request.scoreThreshold !== undefined ? { scoreThreshold: request.scoreThreshold } : {}),
+        ...(context?.abortSignal ? { signal: context.abortSignal } : {}),
+      });
+    },
+  });
+}
+
+/** Build the `universal_search` tool (federated search across every index). */
+export function universalSearchTool(options: SearchToolOptions = {}) {
+  return createTool({
+    id: options.id ?? "universal_search",
+    description: UNIVERSAL_SEARCH_TOOL_DESCRIPTION,
+    inputSchema: universalSearchToolSchema,
+    outputSchema: searchResultSchema,
+    execute: async (input, context) => {
+      const request = universalSearchToolSchema.parse(input);
+      const { client } = getSearchRuntime();
+      return client.universalSearch(request.query, {
+        ...(request.indexes ? { indexes: request.indexes } : {}),
+        ...(request.limit ? { limit: request.limit } : {}),
+        ...(request.mode ? { mode: request.mode } : {}),
+        ...(context?.abortSignal ? { signal: context.abortSignal } : {}),
+      });
+    },
+  });
+}
+
+/**
+ * Build the opt-in `add_documents` tool (write into a direct-access index).
+ * Only install it when the plugin's write surface is enabled.
+ */
+export function addDocumentsTool(options: SearchToolOptions = {}) {
+  const inputSchema = searchContract.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: searchContract.searchRequestSchema
+      .pick({ index: true })
+      .extend({ documents: inputSchema }),
+    outputSchema: searchContract.upsertResultSchema,
+    execute: async (input, context): Promise<UpsertResult> => {
+      const { client, config } = getSearchRuntime();
+      const record = input as { index?: string; documents: unknown };
+      const documents = toDocumentArray(record.documents);
+      const index = record.index ?? config.defaultIndex ?? "";
+      return client.addDocuments(index, documents, context?.abortSignal);
+    },
+  });
+}
+
+/**
+ * Build the opt-in `create_index` tool (provision a Vector Search index).
+ * Only install it when the plugin's write surface is enabled. Delegates to
+ * {@link SearchClient.createIndex}, inferring the endpoint, embedding model,
+ * key, and columns from the request + plugin config.
+ */
+export function createIndexTool(options: SearchToolOptions = {}) {
+  return createTool({
+    id: options.id ?? "create_index",
+    description: CREATE_INDEX_TOOL_DESCRIPTION,
+    inputSchema: createIndexToolSchema,
+    outputSchema: indexInfoSchema,
+    execute: async (input, context) => {
+      const request = createIndexToolSchema.parse(input);
+      const { client } = getSearchRuntime();
+      return client.createIndex(request.name, toCreateIndexOptions(request, context?.abortSignal));
+    },
+  });
+}
+
+/** Output schema for the `sync_index` tool. */
+const syncIndexResultSchema = z.object({
+  index: z.string().describe("The index that was synced."),
+  synced: z.boolean().describe("True once the sync was triggered."),
+});
+
+/**
+ * Build the opt-in `sync_index` tool (refresh a Delta Sync index from its
+ * source table). Only install it when the plugin's write surface is enabled.
+ */
+export function syncIndexTool(options: SearchToolOptions = {}) {
+  return createTool({
+    id: options.id ?? "sync_index",
+    description: SYNC_INDEX_TOOL_DESCRIPTION,
+    inputSchema: syncIndexToolSchema,
+    outputSchema: syncIndexResultSchema,
+    execute: async (input, context) => {
+      const request = syncIndexToolSchema.parse(input);
+      const { client, config } = getSearchRuntime();
+      const index = request.index ?? config.defaultIndex ?? "";
+      await client.syncIndex(index, context?.abortSignal);
+      return { index, synced: true };
+    },
+  });
+}