@gentrace/pinecone
Version:
Gentrace Pinecone v1 plugin for Node.JS
214 lines (211 loc) • 13.4 kB
JavaScript
import { StepRun, Pipeline, PipelineRun } from '@gentrace/core';
import { Pinecone, Index } from '@pinecone-database/pinecone';
var __awaiter = (undefined && undefined.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
var __rest = (undefined && undefined.__rest) || function (s, e) {
var t = {};
for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
t[p] = s[p];
if (s != null && typeof Object.getOwnPropertySymbols === "function")
for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
t[p[i]] = s[p[i]];
}
return t;
};
class PineconePipelineHandler extends Pinecone {
constructor({ pipelineRun, config, gentraceConfig, }) {
super(config);
this.configProtected = config;
this.pipelineRun = pipelineRun;
this.gentraceConfig = gentraceConfig;
}
setPipelineRun(pipelineRun) {
this.pipelineRun = pipelineRun;
}
setupSelfContainedPipelineRun(pipelineId, coreLogic) {
return __awaiter(this, void 0, void 0, function* () {
let isSelfContainedPipelineRun = !this.pipelineRun && pipelineId;
let pipelineRun = this.pipelineRun;
if (isSelfContainedPipelineRun) {
const pipeline = new Pipeline({
slug: pipelineId,
apiKey: this.gentraceConfig.apiKey,
basePath: this.gentraceConfig.basePath,
logger: this.gentraceConfig.logger,
});
pipelineRun = new PipelineRun({
pipeline,
});
}
let returnValue = yield coreLogic(pipelineRun);
if (isSelfContainedPipelineRun) {
const { pipelineRunId } = yield pipelineRun.submit();
// Return value could be void (e.g. in upsert)
if (!returnValue) {
// @ts-ignore
returnValue = {};
}
returnValue.pipelineRunId =
pipelineRunId;
return returnValue;
}
return returnValue;
});
}
// @ts-ignore: hack to avoid base class inheritance issues
indexInner(index, namespace = "") {
const apiHandler = new Index(index, this.configProtected, namespace);
const boundFetch = apiHandler.fetch.bind(apiHandler);
const fetch = (options, gentraceParams) => __awaiter(this, void 0, void 0, function* () {
return this.setupSelfContainedPipelineRun(gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.pipelineSlug, (pipelineRun) => __awaiter(this, void 0, void 0, function* () {
var _a;
const startTime = Date.now();
const response = yield boundFetch(options);
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new PineconeFetchStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), { ids: options }, response, (_a = gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.gentrace) !== null && _a !== void 0 ? _a : {}));
return response;
}));
});
// @ts-ignore
apiHandler.fetch = fetch;
const boundUpdate = apiHandler.update.bind(apiHandler);
const update = (options) => __awaiter(this, void 0, void 0, function* () {
return this.setupSelfContainedPipelineRun(options.pipelineSlug, (pipelineRun) => __awaiter(this, void 0, void 0, function* () {
var _b;
const originalOptions = __rest(options, ["gentrace", "pipelineSlug"]);
const startTime = Date.now();
const response = yield boundUpdate(options);
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new PineconeUpdateStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), Object.assign({}, originalOptions), response, (_b = options.gentrace) !== null && _b !== void 0 ? _b : {}));
return response;
}));
});
apiHandler.update = update;
const boundQuery = apiHandler.query.bind(apiHandler);
const query = (options) => __awaiter(this, void 0, void 0, function* () {
return this.setupSelfContainedPipelineRun(options.pipelineSlug, (pipelineRun) => __awaiter(this, void 0, void 0, function* () {
var _c;
const originalOptions = __rest(options, ["gentrace", "pipelineSlug"]);
const startTime = Date.now();
const response = yield boundQuery(originalOptions);
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
const { topK, filter } = originalOptions, inputs = __rest(originalOptions, ["topK", "filter"]);
const modelParams = { topK, filter };
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new PineconeQueryStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), Object.assign({}, inputs), Object.assign({}, modelParams), response, (_c = options.gentrace) !== null && _c !== void 0 ? _c : {}));
return response;
}));
});
apiHandler.query = query;
const boundUpsert = apiHandler.upsert.bind(apiHandler);
const upsert = (records, gentraceParams) => __awaiter(this, void 0, void 0, function* () {
return this.setupSelfContainedPipelineRun(gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.pipelineSlug, (pipelineRun) => __awaiter(this, void 0, void 0, function* () {
var _d;
const startTime = Date.now();
const returnValue = yield boundUpsert(records);
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new PineconeUpsertStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), { records }, (_d = gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.gentrace) !== null && _d !== void 0 ? _d : {}));
return returnValue;
}));
});
// @ts-ignore
apiHandler.upsert = upsert;
const boundOneDelete = apiHandler.deleteOne.bind(apiHandler);
const deleteOne = (recordId, gentraceParams) => __awaiter(this, void 0, void 0, function* () {
return this.setupSelfContainedPipelineRun(gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.pipelineSlug, (pipelineRun) => __awaiter(this, void 0, void 0, function* () {
var _e;
const startTime = Date.now();
const returnValue = yield boundOneDelete(recordId);
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new PineconeDeleteOneStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), { recordId }, (_e = gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.gentrace) !== null && _e !== void 0 ? _e : {}));
return returnValue;
}));
});
// @ts-ignore
apiHandler.deleteOne = deleteOne;
const boundManyDelete = apiHandler.deleteMany.bind(apiHandler);
const deleteMany = (options, gentraceParams) => __awaiter(this, void 0, void 0, function* () {
return this.setupSelfContainedPipelineRun(gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.pipelineSlug, (pipelineRun) => __awaiter(this, void 0, void 0, function* () {
var _f;
const startTime = Date.now();
const returnValue = yield boundManyDelete(options);
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new PineconeDeleteManyStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), { options }, (_f = gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.gentrace) !== null && _f !== void 0 ? _f : {}));
return returnValue;
}));
});
// @ts-ignore
apiHandler.deleteMany = deleteMany;
const boundAllDelete = apiHandler.deleteAll.bind(apiHandler);
const deleteAll = (gentraceParams) => __awaiter(this, void 0, void 0, function* () {
return this.setupSelfContainedPipelineRun(gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.pipelineSlug, (pipelineRun) => __awaiter(this, void 0, void 0, function* () {
var _g;
const startTime = Date.now();
const returnValue = yield boundAllDelete();
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new PineconeDeleteAllStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), (_g = gentraceParams === null || gentraceParams === void 0 ? void 0 : gentraceParams.gentrace) !== null && _g !== void 0 ? _g : {}));
return returnValue;
}));
});
// @ts-ignore
apiHandler.deleteAll = deleteAll;
const namespaceFn = (namespace) => {
return this.indexInner(index, namespace);
};
// @ts-ignore
apiHandler.namespace = namespaceFn;
// @ts-ignore
return apiHandler;
}
}
class PineconeFetchStepRun extends StepRun {
constructor(elapsedTime, startTime, endTime, inputs, response, context) {
super("pinecone", "pinecone_indexFetch", elapsedTime, startTime, endTime, inputs, {}, response, context !== null && context !== void 0 ? context : {}, undefined);
}
}
class PineconeQueryStepRun extends StepRun {
constructor(elapsedTime, startTime, endTime, inputs, modelParams, response, context) {
super("pinecone", "pinecone_indexQuery", elapsedTime, startTime, endTime, inputs, modelParams, response, context !== null && context !== void 0 ? context : {}, undefined);
}
}
class PineconeUpdateStepRun extends StepRun {
constructor(elapsedTime, startTime, endTime, inputs, response, context) {
super("pinecone", "pinecone_indexUpdate", elapsedTime, startTime, endTime, inputs, {}, response, context !== null && context !== void 0 ? context : {}, undefined);
}
}
class PineconeUpsertStepRun extends StepRun {
constructor(elapsedTime, startTime, endTime, inputs, context) {
super("pinecone", "pinecone_indexUpsert", elapsedTime, startTime, endTime, inputs, {}, {}, context !== null && context !== void 0 ? context : {}, undefined);
}
}
class PineconeDeleteOneStepRun extends StepRun {
constructor(elapsedTime, startTime, endTime, inputs, context) {
super("pinecone", "pinecone_indexDelete1", elapsedTime, startTime, endTime, inputs, {}, {}, context !== null && context !== void 0 ? context : {}, undefined);
}
}
class PineconeDeleteManyStepRun extends StepRun {
constructor(elapsedTime, startTime, endTime, inputs, context) {
super("pinecone", "pinecone_indexDeleteMany", elapsedTime, startTime, endTime, inputs, {}, {}, context !== null && context !== void 0 ? context : {}, undefined);
}
}
class PineconeDeleteAllStepRun extends StepRun {
constructor(elapsedTime, startTime, endTime, context) {
super("pinecone", "pinecone_indexDeleteAll", elapsedTime, startTime, endTime, {}, {}, {}, context !== null && context !== void 0 ? context : {}, undefined);
}
}
export { PineconeDeleteAllStepRun, PineconeDeleteManyStepRun, PineconeDeleteOneStepRun, PineconeFetchStepRun, PineconePipelineHandler, PineconeQueryStepRun, PineconeUpdateStepRun, PineconeUpsertStepRun };
//# sourceMappingURL=pinecone.mjs.map