UNPKG

@upstash/search

Version:

An HTTP/REST based AI Search client built on top of Upstash REST API.

359 lines (351 loc) 11.4 kB
// src/client/error.ts var UpstashError = class extends Error { constructor(message) { super(message); this.name = "UpstashError"; } }; // src/search.ts import { Index as VectorIndex } from "@upstash/vector"; // src/client/metadata.ts var operationMap = { equals: "=", notEquals: "!=", lessThan: "<", lessThanOrEquals: "<=", greaterThan: ">", greaterThanOrEquals: ">=", glob: "GLOB", notGlob: "NOT GLOB", in: "IN", notIn: "NOT IN", contains: "CONTAINS", notContains: "NOT CONTAINS" }; var valueFormatter = (value) => { return Array.isArray(value) ? `(${value.map((v) => typeof v === "string" ? `'${v}'` : v).join(", ")})` : typeof value === "string" ? `'${value}'` : value; }; function constructFilterString(filterTree) { if ("OR" in filterTree) { return `(${filterTree.OR.map((node) => constructFilterString(node)).join(" OR ")})`; } if ("AND" in filterTree) { return `(${filterTree.AND.map((node) => constructFilterString(node)).join(" AND ")})`; } const field = Object.keys(filterTree)[0]; const operationObj = filterTree[field]; const operation = Object.keys(operationObj)[0]; const value = operationObj[operation]; if (!operation || value === void 0) { throw new Error( `Invalid filter operation for field ${String(field)}: ${JSON.stringify(operationObj)}` ); } const mappedOperation = operationMap[operation]; if (!mappedOperation) { throw new Error(`Invalid filter operation for field ${String(field)}: ${operation}`); } const formattedValue = valueFormatter(value); return `${String(field)} ${mappedOperation} ${formattedValue}`; } // src/search-index.ts var SearchIndex = class { /** * Initializes a new SearchIndex instance for the specified index. * * @param vectorIndex - The underlying vector index used for search operations. * @param indexName - The name to use for this index. Must be a non-empty string. * @throws Will throw an error if the indexn name is not provided. */ constructor(httpClient, vectorIndex, indexName) { this.httpClient = httpClient; this.vectorIndex = vectorIndex; this.indexName = indexName; if (!indexName) { throw new Error("indexName is required when defining a SearchIndex"); } } /** * Inserts or updates documents in the index. * * Documents are identified by their unique IDs. If a document with the same ID exists, it will be updated. * * @param params - A document or array of documents to upsert, including `id`, `content`, and optional `metadata`. * @returns A promise resolving to the result of the upsert operation. */ upsert = async (params) => { const upsertParams = Array.isArray(params) ? params : [params]; const path = ["upsert-data", this.indexName]; const { result } = await this.httpClient.request({ path, body: upsertParams }); return result; }; /** * Searches for documents matching a query string. * * Returns documents that best match the provided query, optionally filtered and limited in number. * * @param query - Text string used to find matching documents within the index. * @param limit - Maximum number of results to retrieve (defaults to 5 documents). * @param filter - Optional search constraint using either a string expression or structured filter object. * @param reranking - Optional boolean to enhance search result ordering. * @param semanticWeight - Optional relevance balance between semantic and keyword search (0-1 range, defaults to 0.75). * For instance, 0.2 applies 20% semantic matching with 80% full-text matching. * You can learn more about how Upstash Search works from [our docs](https://upstash.com/docs/search/features/algorithm). * @param inputEnrichment - Optional boolean to enhance queries before searching (enabled by default). * @returns Promise that resolves to an array of documents matching the */ search = async (params) => { const { query, limit = 5, filter, reranking, semanticWeight, inputEnrichment } = params; if (semanticWeight && (semanticWeight < 0 || semanticWeight > 1)) { throw new UpstashError("semanticWeight must be between 0 and 1"); } const path = ["search", this.indexName]; const { result } = await this.httpClient.request({ path, body: { query, topK: limit, includeData: true, includeMetadata: true, filter: typeof filter === "string" || filter === void 0 ? filter : constructFilterString(filter), reranking, semanticWeight, inputEnrichment } }); return result.map(({ id, content, metadata, score }) => ({ id, content, metadata, score })); }; /** * Fetches documents by their IDs from the index. * * @param params - An array of document IDs to retrieve. * @returns A promise resolving to an array of documents or `null` if a document is not found. */ fetch = async (params) => { const result = await this.vectorIndex.fetch(params, { namespace: this.indexName, includeData: true, includeMetadata: true }); return result.map((fetchResult) => { if (!fetchResult) return fetchResult; return { id: fetchResult.id, content: fetchResult.content, metadata: fetchResult.metadata }; }); }; /** * Deletes documents by their IDs from the index. * * @param params - An array of document IDs to delete. * @returns A promise resolving to the result of the deletion operation. */ delete = async (params) => { return await this.vectorIndex.delete(params, { namespace: this.indexName }); }; /** * Retrieves documents within a specific range, with pagination support. * * Useful for paginating through large result sets by providing a `cursor`. * * @param params - Range parameters including `cursor`, `limit`, and ID `prefix`. * @returns A promise resolving to the next cursor and documents in the range. */ range = async (params) => { const { nextCursor, vectors } = await this.vectorIndex.range( { ...params, includeData: true, includeMetadata: true }, { namespace: this.indexName } ); return { nextCursor, documents: vectors.map( ({ id, content, metadata }) => ({ id, content, metadata }) ) }; }; /** * Clears all documents in the current index. * * Useful for resetting the index before or after tests, or when a clean state is needed. * * @returns A promise resolving to the result of the reset operation. */ reset = async () => { return await this.vectorIndex.reset({ namespace: this.indexName }); }; /** * Deletes the entire index and all its documents. * * Use with caution, as this operation is irreversible. * * @returns A promise resolving to the result of the delete operation. */ deleteIndex = async () => { return await this.vectorIndex.deleteNamespace(this.indexName); }; /** * Retrieves information about the current index. * * Provides document count and pending document count, indicating documents that are awaiting indexing. * * @returns A promise resolving to index information with document counts. */ info = async () => { const info = await this.vectorIndex.info(); const { pendingVectorCount, vectorCount } = info.namespaces[this.indexName] ?? { pendingVectorCount: 0, vectorCount: 0 }; return { pendingDocumentCount: pendingVectorCount, documentCount: vectorCount }; }; }; // src/search.ts var Search = class { /** * Creates a new Search instance. * * @param vectorIndex - The underlying index used for search operations. */ constructor(client) { this.client = client; this.vectorIndex = new VectorIndex(client); } vectorIndex; /** * Returns a SearchIndex instance for a given index. * * Each index is an isolated collection where documents can be added, * retrieved, searched, and deleted. * * @param indexName - The name to use as an index. * @returns A SearchIndex instance for managing documents within the index. */ index = (indexName) => { return new SearchIndex(this.client, this.vectorIndex, indexName); }; /** * Retrieves a list of all available indexes. * * @returns An array of strings representing the names of available indexes. */ listIndexes = async () => { return await this.vectorIndex.listNamespaces(); }; /** * Retrieves overall search index statistics. * * This includes disk usage, total document count, pending document count, * and details about each available index. * * @returns An object containing search system metrics and index details. */ info = async () => { const { indexSize, namespaces, pendingVectorCount, vectorCount } = await this.vectorIndex.info(); const indexes = Object.fromEntries( Object.entries(namespaces).map((namespace) => [ namespace[0], { pendingDocumentCount: namespace[1].pendingVectorCount, documentCount: namespace[1].vectorCount } ]) ); return { diskSize: indexSize, pendingDocumentCount: pendingVectorCount, documentCount: vectorCount, indexes }; }; }; // src/client/search-client.ts var HttpClient = class { baseUrl; headers; options; retry; constructor(config) { this.options = { cache: config.cache }; this.baseUrl = config.baseUrl.replace(/\/$/, ""); this.headers = { "Content-Type": "application/json", ...config.headers }; this.retry = typeof config?.retry === "boolean" && config?.retry === false ? { attempts: 1, backoff: () => 0 } : { attempts: config?.retry?.retries ?? 5, backoff: config?.retry?.backoff ?? ((retryCount) => Math.exp(retryCount) * 50) }; } async request(req) { const requestOptions = { cache: this.options.cache, method: "POST", headers: this.headers, body: JSON.stringify(req.body), keepalive: true }; let res = null; let error = null; for (let i = 0; i <= this.retry.attempts; i++) { try { res = await fetch([this.baseUrl, ...req.path ?? []].join("/"), requestOptions); break; } catch (error_) { error = error_; if (i < this.retry.attempts) { await new Promise((r) => setTimeout(r, this.retry.backoff(i))); } } } if (!res) { throw error ?? new Error("Exhausted all retries"); } const body = await res.json(); if (!res.ok) { throw new UpstashError(`${body.error}`); } return { result: body.result, error: body.error }; } }; // src/client/telemetry.ts var VERSION = "0.1.0"; function getRuntime() { if (typeof process === "object" && typeof process.versions == "object" && process.versions.bun) return `bun@${process.versions.bun}`; if (typeof process === "object" && typeof process.version === "string") { return `node@${process.version}`; } if (typeof EdgeRuntime === "string") { return "edge-light"; } return "undetermined"; } export { UpstashError, Search, HttpClient, VERSION, getRuntime };