UNPKG

@datastax/astra-db-ts

Version:
152 lines (151 loc) 5.93 kB
"use strict"; // Copyright Datastax, Inc // SPDX-License-Identifier: Apache-2.0 Object.defineProperty(exports, "__esModule", { value: true }); exports.FindAndRerankCursor = exports.RerankedResult = void 0; const index_js_1 = require("../../documents/index.js"); const abstract_cursor_js_1 = require("../../documents/cursors/abstract-cursor.js"); const constants_js_1 = require("../../lib/constants.js"); const ctx_js_1 = require("../../lib/api/ser-des/ctx.js"); const common_js_1 = require("../../documents/cursors/common.js"); const utils_js_1 = require("../../lib/utils.js"); class RerankedResult { constructor(document, scores) { Object.defineProperty(this, "document", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "scores", { enumerable: true, configurable: true, writable: true, value: void 0 }); this.document = document; this.scores = scores; } } exports.RerankedResult = RerankedResult; class FindAndRerankCursor extends abstract_cursor_js_1.AbstractCursor { constructor(parent, serdes, filter, options, mapping) { super(options ?? {}, mapping); Object.defineProperty(this, "_httpClient", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "_serdes", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "_parent", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "_filter", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "_sortVector", { enumerable: true, configurable: true, writable: true, value: new utils_js_1.QueryState() }); this._parent = parent; this._httpClient = parent._httpClient; this._serdes = serdes; this._filter = filter; } [constants_js_1.$CustomInspect]() { return `${this.constructor.name}(source="${this._parent.keyspace}.${this._parent.name}",state="${this._state}",consumed=${this.consumed()},buffered=${this.buffered()})`; } filter(filter) { return (0, common_js_1.buildFLCFilter)(this, filter); } sort(sort) { return (0, common_js_1.buildFLCSort)(this, sort); } limit(limit) { return (0, common_js_1.buildFLCOption)(this, 'limit', limit || undefined); } hybridLimits(hybridLimits) { return (0, common_js_1.buildFLCOption)(this, 'hybridLimits', hybridLimits); } rerankOn(rerankOn) { return (0, common_js_1.buildFLCOption)(this, 'rerankOn', rerankOn); } rerankQuery(rerankQuery) { return (0, common_js_1.buildFLCOption)(this, 'rerankQuery', rerankQuery); } includeScores(includeScores) { return (0, common_js_1.buildFLCOption)(this, 'includeScores', includeScores ?? true); } includeSortVector(includeSortVector) { return (0, common_js_1.buildFLCOption)(this, 'includeSortVector', includeSortVector ?? true); } project(projection) { return (0, common_js_1.buildFLCPreMapOption)(this, 'projection', projection); } map(map) { return (0, common_js_1.buildFLCMap)(this, map); } async getSortVector() { if (this._sortVector.state === utils_js_1.QueryState.Unattempted && this._options.includeSortVector) { const reset2idle = this._state === 'idle'; await this._next(true, '.getSortVector'); if (reset2idle) { this._state = 'idle'; } } return this._sortVector.unwrap(); } clone() { return (0, common_js_1.cloneFLC)(this, this._filter, this._options, this._mapping); } async _nextPage(extra, tm) { const command = { findAndRerank: { filter: this._filter[0], projection: this._options.projection, sort: this._options.sort, options: { limit: this._options.limit, hybridLimits: this._options.hybridLimits, rerankOn: this._options.rerankOn, rerankQuery: this._options.rerankQuery, includeScores: this._options.includeScores, includeSortVector: this._options.includeSortVector, }, }, }; const raw = await this._httpClient.executeCommand(command, { timeoutManager: tm ?? this._httpClient.tm.single('generalMethodTimeoutMs', this._options), bigNumsPresent: this._filter[1], extraLogInfo: extra, }); this._nextPageState.swap(raw.data?.nextPageState); const buffer = raw.data?.documents ?? []; for (let i = 0, n = buffer.length; i < n; i++) { const deserialized = this._serdes.deserialize(buffer[i], raw, ctx_js_1.SerDesTarget.Record); buffer[i] = new RerankedResult(deserialized, raw.status?.documentResponses?.[i]?.scores ?? {}); } const sortVector = raw.status?.sortVector; this._sortVector.swap(sortVector ? (0, index_js_1.vector)(sortVector) : sortVector); this._options.includeSortVector = false; return buffer; } _tm() { return this._httpClient.tm; } } exports.FindAndRerankCursor = FindAndRerankCursor;