@dbx-tools/search 0.6.9

package/lib/src/client.d.ts

@@ -0,0 +1,304 @@
+/**
+ * A small, Meilisearch-shaped client over Databricks AI Search (Vector
+ * Search). The whole point of this module is ergonomics: the Databricks SDK's
+ * `vectorSearchIndexes.queryIndex({ index_name, columns, query_text,
+ * query_type, num_results, filters_json })` is powerful but verbose, and the
+ * response is columnar. This client hides all of that behind two objects:
+ *
+ * ```ts
+ * const client = createSearchClient();
+ * const index = client.index("main.support.docs");
+ * const { hits } = await index.search("reset my password", { limit: 5 });
+ * await index.addDocuments([{ id: "42", title: "Reset", body: "..." }]);
+ * ```
+ *
+ * `client.search(query, opts)` searches the default index; `client.index(name)`
+ * returns a handle bound to one index; `client.universalSearch(query)` fans a
+ * query across several indexes and merges the results (Meilisearch's federated
+ * / multi-search, the "universal search" the caller asked for). Everything is
+ * async and cancellable, resolves the OBO workspace client from the active
+ * AppKit execution context (falling back to a service-principal client outside
+ * a request), and returns the browser-safe shapes from
+ * `@dbx-tools/shared-search`.
+ *
+ * Autocomplete is just a search with a small `limit` and the raw query text -
+ * hybrid mode already prefix-matches, so no separate endpoint is needed; the
+ * {@link SearchIndex.autocomplete} helper is a thin, self-documenting alias.
+ *
+ * @module
+ */
+import { appkit } from "@dbx-tools/appkit";
+import type { SearchDocument, SearchMode, SearchResult, UpsertResult } from "@dbx-tools/shared-search";
+import { type ResolvedSearchConfig } from "./config.ts";
+import type { LakebaseSearchBackend } from "./lakebase.ts";
+type WorkspaceClientLike = appkit.WorkspaceClientLike;
+/** Options accepted by a single-index search. */
+export interface SearchOptions {
+    /** Maximum hits to return. Defaults to the configured page size. */
+    limit?: number;
+    /** Match mode. Defaults to the configured mode (hybrid). */
+    mode?: SearchMode;
+    /** Columns to return per hit. Defaults to the index's configured columns. */
+    columns?: readonly string[];
+    /** Attribute filters as `{ column: value }` or `{ column: { ">=": n } }`. */
+    filter?: Record<string, unknown>;
+    /** Drop hits below this score. */
+    scoreThreshold?: number;
+    /** External cancellation. */
+    signal?: AbortSignal;
+}
+/** Options accepted by a universal (federated) search. */
+export interface UniversalSearchOptions {
+    /** Indexes to search. Defaults to every known index. */
+    indexes?: readonly string[];
+    /** Maximum hits per index before merging. */
+    limit?: number;
+    /** Match mode. */
+    mode?: SearchMode;
+    /** External cancellation. */
+    signal?: AbortSignal;
+}
+/**
+ * A handle bound to one index. Modeled on Meilisearch's `client.index(uid)`:
+ * `search` / `autocomplete` read, `addDocuments` / `deleteDocuments` write (to
+ * a direct-access index), and `info` fetches the live definition.
+ */
+export declare class SearchIndex {
+    readonly name: string;
+    private readonly client;
+    constructor(name: string, client: SearchClient);
+    /** Search this index. See {@link SearchClient.search}. */
+    search(query: string, options?: SearchOptions): Promise<SearchResult>;
+    /**
+     * Autocomplete against this index: a search with a small default `limit` and
+     * the raw prefix as the query. Hybrid mode already handles prefixes, so this
+     * is a self-documenting alias rather than a separate code path.
+     */
+    autocomplete(prefix: string, options?: SearchOptions): Promise<SearchResult>;
+    /** Add or update documents in this (direct-access) index. */
+    addDocuments(documents: SearchDocument[], signal?: AbortSignal): Promise<UpsertResult>;
+    /** Delete documents from this (direct-access) index by primary key. */
+    deleteDocuments(ids: Array<string | number>, signal?: AbortSignal): Promise<UpsertResult>;
+    /** Fetch this index's live definition (primary key, columns, readiness). */
+    info(signal?: AbortSignal): Promise<IndexInfo>;
+    /** Trigger a sync of this (Delta Sync) index from its source table. */
+    sync(signal?: AbortSignal): Promise<void>;
+    /** Delete this index. */
+    delete(signal?: AbortSignal): Promise<void>;
+    /** Create this index if it does not exist, otherwise return the existing one. */
+    ensure(options?: CreateIndexOptions): Promise<IndexInfo>;
+}
+/** A resolved live index definition. */
+export interface IndexInfo {
+    name: string;
+    endpoint?: string;
+    primaryKey?: string;
+    columns: string[];
+    ready: boolean;
+    rowCount?: number;
+    /** True for a DIRECT_ACCESS index (you supply vectors; queries embed client-side). */
+    directAccess?: boolean;
+}
+/**
+ * Options for {@link SearchClient.createIndex} / {@link SearchClient.ensureIndex}.
+ * The goal is the same as the rest of the client: name a source table and a
+ * text column, and everything else - the endpoint, the embedding model, the
+ * index type, the sync mode - has a sensible default that infers from the
+ * workspace, overridable when a deployment needs to go deeper.
+ */
+export interface CreateIndexOptions {
+    /**
+     * Source Delta table (catalog.schema.table) for a Delta Sync index. Provide
+     * this for the common case: Databricks computes and syncs embeddings from it.
+     * Omit it (and pass {@link embeddingDimension}) to create a direct-access
+     * index you write vectors to yourself.
+     */
+    sourceTable?: string;
+    /** Primary-key column. Defaults to `id`. */
+    primaryKey?: string;
+    /**
+     * The text column embeddings are computed from (Delta Sync). Defaults to the
+     * first of `text` / `content` / `body` present, else the caller must set it.
+     */
+    embeddingSourceColumn?: string;
+    /**
+     * Embedding model endpoint. A loose name is fuzzy-matched; when omitted the
+     * best embedding endpoint in the workspace is chosen ({@link resolveEmbeddingModel}).
+     */
+    embeddingModel?: string;
+    /** Vector Search endpoint to host the index on. Defaults to the plugin's `endpoint`. */
+    endpoint?: string;
+    /** For a direct-access index: the embedding vector dimension (self-managed vectors, no source table). */
+    embeddingDimension?: number;
+    /** For a direct-access index: the column the vector is stored in. Defaults to `embedding`. */
+    embeddingVectorColumn?: string;
+    /**
+     * The column-name -> type map used to build a direct-access index's
+     * `schema_json` (types: `string`, `int`, `long`, `float`, `double`,
+     * `boolean`, `date`, `timestamp`). Must include the primary key and any
+     * columns you upsert. Defaults to `{ id: "string", text: "string" }`. The
+     * embedding vector column is added automatically.
+     */
+    schema?: Record<string, string>;
+    /** Sync mode for a Delta Sync index. `TRIGGERED` (default) syncs on demand; `CONTINUOUS` keeps fresh. */
+    pipelineType?: "TRIGGERED" | "CONTINUOUS";
+    /** Extra columns to sync alongside the embedding source (Delta Sync). */
+    columnsToSync?: string[];
+    /** External cancellation. */
+    signal?: AbortSignal;
+}
+/** Options for {@link SearchClient.ensureEndpoint}. */
+export interface EnsureEndpointOptions {
+    /** Wait for the endpoint to come online before returning. Defaults to false. */
+    wait?: boolean;
+    /** External cancellation. */
+    signal?: AbortSignal;
+}
+/**
+ * Options for {@link SearchClient.provision} - a one-call "make this index real
+ * and searchable" used at boot or in a seed script. It ensures the endpoint,
+ * ensures the index, and (optionally) seeds documents when the index is empty.
+ */
+export interface ProvisionOptions extends CreateIndexOptions {
+    /** Documents to seed when the index has no rows yet. Skipped if it already has data. */
+    seed?: SearchDocument[];
+    /**
+     * Wait for the endpoint AND index to come online before returning (needed
+     * before seeding). Defaults to true. Endpoint creation can take many minutes.
+     */
+    wait?: boolean;
+    /** How long to wait for readiness before giving up. Defaults to 20 minutes. */
+    timeoutMs?: number;
+}
+/**
+ * The AI Search client. Construct it with {@link createSearchClient} (which
+ * reads a resolved config) or directly for one-off use. All reads resolve the
+ * OBO workspace client from the active execution context.
+ */
+export declare class SearchClient {
+    private readonly config;
+    private readonly workspaceClientFactory;
+    /**
+     * Optional Lakebase full-text FALLBACK backend. Present only when no Vector
+     * Search endpoint is configured but a Lakebase pool is available; when set,
+     * search / provision / write operations delegate to it and return the exact
+     * same shapes, so nothing downstream can tell which backend answered.
+     */
+    private readonly lakebase?;
+    constructor(config?: ResolvedSearchConfig, workspaceClientFactory?: () => WorkspaceClientLike, 
+    /**
+     * Optional Lakebase full-text FALLBACK backend. Present only when no Vector
+     * Search endpoint is configured but a Lakebase pool is available; when set,
+     * search / provision / write operations delegate to it and return the exact
+     * same shapes, so nothing downstream can tell which backend answered.
+     */
+    lakebase?: LakebaseSearchBackend | undefined);
+    /** True when this client is answering out of the Lakebase fallback backend. */
+    get usesLakebase(): boolean;
+    /** A handle bound to one index (by full UC name or configured alias). */
+    index(reference: string): SearchIndex;
+    /**
+     * Search one index. `index` may be a full UC name, a configured alias, or
+     * omitted to use the default index. Returns hits sorted most-relevant-first.
+     */
+    search(query: string, options?: SearchOptions & {
+        index?: string;
+    }): Promise<SearchResult>;
+    /** Cache of index-name -> is-DIRECT_ACCESS, so a search embeds its query only when needed. */
+    private readonly directAccessCache;
+    /** Whether an index is DIRECT_ACCESS (memoized); a lookup failure assumes Delta Sync. */
+    private isDirectAccess;
+    /**
+     * Fan a query across several indexes and merge the hits, sorted by score -
+     * the "universal search" a single box over many collections needs. Each
+     * index is searched concurrently; an index that errors is logged and skipped
+     * so one bad index does not sink the whole search.
+     */
+    universalSearch(query: string, options?: UniversalSearchOptions): Promise<SearchResult>;
+    /**
+     * Resolve an index reference (full UC name or configured alias) to the name
+     * the API expects, failing fast when it resolves to nothing. Without this a
+     * blank reference - what an omitted `index` becomes when no default is
+     * configured - reaches the SDK as an empty `index_name`, which builds a
+     * URL with the name segment missing and comes back as a confusing
+     * `ENDPOINT_NOT_FOUND` instead of naming the real problem.
+     */
+    private requireIndexName;
+    /** Fetch an index's live definition. */
+    getIndex(reference: string, signal?: AbortSignal): Promise<IndexInfo>;
+    /** Add or update documents in a direct-access index. */
+    addDocuments(reference: string, documents: SearchDocument[], signal?: AbortSignal): Promise<UpsertResult>;
+    /** Delete documents from a direct-access index by primary key. */
+    deleteDocuments(reference: string, ids: Array<string | number>, signal?: AbortSignal): Promise<UpsertResult>;
+    /**
+     * Resolve an embedding endpoint id for creating a Delta Sync index. Reuses
+     * the model resolver: a configured / passed name is fuzzy-matched against the
+     * live catalogue, otherwise the highest-ranked embedding endpoint is chosen.
+     */
+    resolveEmbeddingModel(requested?: string, signal?: AbortSignal): Promise<string | null>;
+    /**
+     * Embed text via a Databricks embedding serving endpoint, returning one
+     * vector per input. Used to seed a direct-access index and to turn a search
+     * query into a query vector - Databricks only manages embeddings for Delta
+     * Sync indexes, so a direct-access index (no Delta table, no warehouse) needs
+     * the client to embed on write and on query. The endpoint is resolved the
+     * same way as for index creation when not named.
+     */
+    embed(texts: string[], model?: string, signal?: AbortSignal): Promise<number[][]>;
+    /** The vector dimension a model produces (embeds a probe string once). */
+    private embeddingDimension;
+    /**
+     * Create an AI Search index with as little ceremony as possible. Two shapes:
+     *
+     *   - **Delta Sync** (the default): pass `sourceTable`; Databricks computes
+     *     embeddings from the text column and keeps the index synced. The
+     *     embedding model is resolved automatically when not named.
+     *   - **Direct Access**: omit `sourceTable` and pass `embeddingDimension`;
+     *     you write vectors yourself via {@link addDocuments}.
+     *
+     * Everything else infers: the endpoint from the plugin config, the primary
+     * key (`id`), the text column (`text` / `content` / `body`), and the vector
+     * column (`embedding`). Returns the created index's {@link IndexInfo}.
+     */
+    createIndex(name: string, options?: CreateIndexOptions): Promise<IndexInfo>;
+    /**
+     * Create the index if it does not already exist, otherwise return the
+     * existing one. Idempotent - safe to call on every boot to guarantee an
+     * index is present.
+     */
+    ensureIndex(name: string, options?: CreateIndexOptions): Promise<IndexInfo>;
+    /**
+     * Ensure an index exists, is online, and (optionally) holds seed data - the
+     * "wire up a real index on boot" path. Idempotent and cheap to call every
+     * boot: it creates the endpoint and index only if missing, waits for them to
+     * come online, and seeds documents ONLY when the index is still empty.
+     *
+     * For the demo/dummy-data case this needs no Delta table and no warehouse:
+     * the default is a MANAGED direct-access index (Databricks embeds the `text`
+     * column on write and query), so seeding is just an `addDocuments` of plain
+     * rows and search-by-text works immediately.
+     */
+    provision(name: string, options?: ProvisionOptions): Promise<IndexInfo>;
+    /** Poll an index until it reports ready, or throw after {@link timeoutMs}. */
+    private waitForIndexReady;
+    /** Trigger a sync of a Delta Sync index from its source table. */
+    syncIndex(reference: string, signal?: AbortSignal): Promise<void>;
+    /** Delete an index. */
+    deleteIndex(reference: string, signal?: AbortSignal): Promise<void>;
+    /** List the indexes hosted on a Vector Search endpoint (name + type only). */
+    listIndexes(endpoint?: string, signal?: AbortSignal): Promise<string[]>;
+    /**
+     * Ensure a Vector Search endpoint exists, creating a `STANDARD` one when it
+     * does not. Optionally wait for it to come online. Idempotent.
+     */
+    ensureEndpoint(name?: string, options?: EnsureEndpointOptions): Promise<void>;
+    /**
+     * Resolve the workspace client and run one call under a bounded timeout. The
+     * caller's signal (if any) and a timeout are merged into one SDK `Context`
+     * via {@link appkit.databricks.toContext}, so either unwinds the request.
+     */
+    private withClient;
+}
+/** Construct a {@link SearchClient} from a resolved config. */
+export declare function createSearchClient(config?: ResolvedSearchConfig, workspaceClientFactory?: () => WorkspaceClientLike, lakebase?: LakebaseSearchBackend): SearchClient;
+export {};