@gentrace/core
Version:
Core Gentrace Node.JS library
362 lines (359 loc) • 17.7 kB
JavaScript
import _ from 'lodash';
import { v4 } from 'uuid';
import { V1Api } from '../api/v1-api.mjs';
import { RunRequestCollectionMethodEnum } from '../models/run-request.mjs';
import { globalGentraceApiV2, globalGentraceConfig, globalRequestBuffer } from './init.mjs';
import { StepRun } from './step-run.mjs';
import { getParamNames, zip, safeJsonParse, getTestCounter } from './utils.mjs';
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;
};
const getRun = (id) => __awaiter(void 0, void 0, void 0, function* () {
if (!globalGentraceApiV2) {
throw new Error("Gentrace API key not initialized. Call init() first.");
}
const response = yield globalGentraceApiV2.v2RunsIdGet(id);
const run = response.data;
return run;
});
class PipelineRun {
constructor({ pipeline, context, }) {
this.evaluations = [];
this.path = [];
this.id = v4();
this.instantiationTime = new Date().toISOString();
this.pipeline = pipeline;
this.stepRuns = [];
this.context = context;
}
getPipeline() {
return this.pipeline;
}
getId() {
return this.id;
}
getContext() {
return this.context;
}
getError() {
return this.error;
}
setError(error) {
this.error = error;
}
updateContext(updatedContext) {
this.context = Object.assign(Object.assign({}, this.context), updatedContext);
return this.context;
}
addStepRunNode(stepRun) {
return __awaiter(this, void 0, void 0, function* () {
this.stepRuns.push(stepRun);
});
}
/**
* Creates a checkpoint by recording a `StepRun` instance with execution metadata and pushes it to `this.stepRuns`.
* If no prior `StepRun` instances exist, the elapsed time is set to 0 and the start and end times are set to the
* current timestamp. If it is empty, elapsed time is set to 0 and start time and end time are set to the current
* timestamp.
*
* @param {PRStepRunType & { inputs: any; outputs: any; }} step The information about the step to checkpoint.
* This includes the inputs and outputs of the step, as well as optional provider, invocation and modelParams metadata.
*
* @example
* const stepInfo = {
* providerName: 'MyProvider',
* invocation: 'doSomething',
* inputs: { x: 10, y: 20 },
* outputs: { result: 30 }
* };
* checkpoint(stepInfo);
*
* @returns {void} The function does not return anything.
*
* @throws {Error} If the `StepRun` constructor or any other operations throw an error, it will be propagated.
*/
checkpoint(step) {
var _a, _b, _c, _d, _e, _f, _g, _h;
const lastElement = this.stepRuns[this.stepRuns.length - 1];
if (lastElement) {
const { endTime: stepStartTime } = lastElement;
const elapsedTime = new Date().getTime() - new Date(stepStartTime).getTime();
const endTimeNew = new Date().toISOString();
this.stepRuns.push(new StepRun((_a = step.providerName) !== null && _a !== void 0 ? _a : "undeclared", (_b = step.invocation) !== null && _b !== void 0 ? _b : "undeclared", elapsedTime, stepStartTime, endTimeNew, step.inputs, (_c = step.modelParams) !== null && _c !== void 0 ? _c : {}, step.outputs, (_d = step.context) !== null && _d !== void 0 ? _d : {}, step.error));
}
else {
const endTime = new Date().toISOString();
const elapsedTime = new Date(endTime).getTime() -
new Date(this.instantiationTime).getTime();
this.stepRuns.push(new StepRun((_e = step.providerName) !== null && _e !== void 0 ? _e : "undeclared", (_f = step.invocation) !== null && _f !== void 0 ? _f : "undeclared", elapsedTime, this.instantiationTime, endTime, step.inputs, (_g = step.modelParams) !== null && _g !== void 0 ? _g : {}, step.outputs, (_h = step.context) !== null && _h !== void 0 ? _h : {}, step.error));
}
}
/**
* Asynchronously measures the execution time of a function.
*
* @template F Function type that extends (...args: any[]) => any
* @param {F} func The function to be measured.
* @param {Parameters<F>} inputs The parameters to be passed to the function.
* @param {Omit<PRStepRunType, "inputs" | "outputs">} [stepInfo] Optional metadata for the function execution.
* @returns {Promise<ReturnType<F>>} Returns a promise that resolves to the return type of the function.
*
* @example
* async function foo(n: number) {
* return n * 2;
* }
* const result = await measure(foo, [2]); // result will be 4
*
* The function also records a `StepRun` instance with execution metadata and pushes it to `this.stepRuns`.
* The recorded `StepRun` includes information such as the elapsed time, start and end time,
* resolved inputs, and model parameters if provided.
*/
measure(func, inputs, stepInfo) {
var _a, _b, _c, _d, _e, _f, _g, _h;
return __awaiter(this, void 0, void 0, function* () {
const startTime = Date.now();
const paramNames = getParamNames(func);
const resolvedInputs = zip(paramNames, inputs).reduce((acc, current) => {
const [key, value] = current;
acc[key] = value;
return acc;
}, {});
try {
const returnValue = yield func(...inputs);
// Our server only accepts outputs as an object.
let modifiedOutput = returnValue;
if (typeof returnValue !== "object" || Array.isArray(returnValue)) {
modifiedOutput = { value: returnValue };
}
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
this.stepRuns.push(new StepRun((_a = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.providerName) !== null && _a !== void 0 ? _a : "undeclared", (_b = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.invocation) !== null && _b !== void 0 ? _b : "undeclared", elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), resolvedInputs, (_c = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.modelParams) !== null && _c !== void 0 ? _c : {}, modifiedOutput, (_d = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.context) !== null && _d !== void 0 ? _d : {}, stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.error));
return returnValue;
}
catch (e) {
try {
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
const errorString = e.toString();
this.stepRuns.push(new StepRun((_e = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.providerName) !== null && _e !== void 0 ? _e : "undeclared", (_f = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.invocation) !== null && _f !== void 0 ? _f : "undeclared", elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), resolvedInputs, (_g = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.modelParams) !== null && _g !== void 0 ? _g : {}, {}, (_h = stepInfo === null || stepInfo === void 0 ? void 0 : stepInfo.context) !== null && _h !== void 0 ? _h : {}, errorString));
}
finally {
throw e;
}
}
});
}
/**
* Adds an evaluation to the pipeline run.
* @param evaluation The evaluation to add.
*/
addEval(evaluation) {
this.evaluations.push(evaluation);
}
toObject() {
var _a;
let mergedMetadata = {};
const updatedStepRuns = this.stepRuns.map(({ providerName: providerName, elapsedTime, startTime, endTime, invocation, modelParams, inputs, outputs, context: stepRunContext, error, }) => {
var _a;
let _b = (_a = this.context) !== null && _a !== void 0 ? _a : {}, { metadata: thisContextMetadata, previousRunId: _prPreviousRunId } = _b, restThisContext = __rest(_b, ["metadata", "previousRunId"]);
let _c = stepRunContext !== null && stepRunContext !== void 0 ? stepRunContext : {}, { metadata: stepRunContextMetadata, previousRunId: _srPreviousRunId } = _c, restStepRunContext = __rest(_c, ["metadata", "previousRunId"]);
// Merge metadata
mergedMetadata = Object.assign(Object.assign(Object.assign({}, mergedMetadata), thisContextMetadata), stepRunContextMetadata);
return {
providerName,
elapsedTime,
startTime,
endTime,
invocation,
modelParams,
inputs,
outputs,
context: Object.assign(Object.assign({}, restThisContext), restStepRunContext),
error,
};
});
return {
id: this.id,
slug: this.pipeline.slug,
metadata: mergedMetadata,
previousRunId: (_a = this.context) === null || _a === void 0 ? void 0 : _a.previousRunId,
collectionMethod: RunRequestCollectionMethodEnum.Runner,
stepRuns: updatedStepRuns,
evaluations: this.evaluations,
error: this.error,
};
}
toJson() {
return JSON.stringify(this.toObject(), null, 2);
}
static getRedactedRunFromJson(json, options) {
const selectFields = options === null || options === void 0 ? void 0 : options.selectFields;
const pipeline = options === null || options === void 0 ? void 0 : options.pipeline;
const pipelineRunObject = (typeof json === "string" ? safeJsonParse(json) : json);
if (!pipelineRunObject) {
return null;
}
function isSingleStepRunWhitelist(value) {
return typeof value !== "function";
}
if (selectFields) {
const stepRuns = pipelineRunObject.stepRuns;
let selectedFields;
if (!isSingleStepRunWhitelist(selectFields)) {
selectedFields = selectFields(stepRuns);
if (selectedFields.length !== stepRuns.length && pipeline) {
pipeline.logWarn("The selectFields function did not return the correct number of fields.");
}
if (selectedFields.length < stepRuns.length && pipeline) {
// Fill rest with the last element
const last = selectedFields[selectedFields.length - 1];
selectedFields = selectedFields.concat(Array.from({ length: stepRuns.length - selectedFields.length }).map(() => last));
}
}
else {
selectedFields = Array.from({ length: stepRuns.length }).map(() => selectFields);
}
pipelineRunObject.stepRuns = stepRuns.map((stepRun, index) => {
const selectedFieldsForStep = selectedFields[index];
const updatedStepRun = Object.assign({}, stepRun);
if (selectedFieldsForStep) {
for (const [fieldKey, selector] of Object.entries(selectedFieldsForStep)) {
const fieldKeyTyped = fieldKey;
const stepRunValue = updatedStepRun[fieldKeyTyped];
if (!stepRunValue) {
continue;
}
const fieldType = typeof stepRunValue;
const isSelectorArray = Array.isArray(selector);
const isSelectorString = typeof selector === "string";
const isSelectorBoolean = typeof selector === "boolean";
if (fieldType === "object" &&
(isSelectorArray || isSelectorString)) {
const newValue = {};
if (isSelectorString) {
const whiteListValue = _.get(stepRunValue, selector);
_.set(newValue, selector, whiteListValue);
}
else if (isSelectorArray) {
for (const key of selector) {
const whiteListValue = _.get(stepRunValue, key);
_.set(newValue, key, whiteListValue);
}
}
// @ts-ignore
updatedStepRun[fieldKeyTyped] = newValue;
}
else if (fieldType === "object" && isSelectorBoolean) {
if (!selector) {
// @ts-ignore
updatedStepRun[fieldKeyTyped] = {};
}
}
}
}
return updatedStepRun;
});
}
if (pipeline) {
pipeline.logInfo("Submitting PipelineRun to Gentrace");
}
return pipelineRunObject;
}
static submitFromJson(json, options) {
var _a;
return __awaiter(this, void 0, void 0, function* () {
const pipeline = options === null || options === void 0 ? void 0 : options.pipeline;
const waitForServer = (_a = options === null || options === void 0 ? void 0 : options.waitForServer) !== null && _a !== void 0 ? _a : false;
const api = new V1Api(pipeline ? pipeline.config : globalGentraceConfig);
const pipelineRunObject = PipelineRun.getRedactedRunFromJson(json, {
waitForServer: false,
pipeline,
selectFields: options === null || options === void 0 ? void 0 : options.selectFields,
});
if (!pipelineRunObject) {
if (pipeline) {
pipeline.logWarn("Invalid JSON passed to submitFromJson");
}
return {};
}
const submission = api.v1RunPost(pipelineRunObject);
if (!waitForServer) {
globalRequestBuffer[pipelineRunObject.id] = submission;
submission
.catch((e) => {
if (pipeline) {
pipeline.logWarn(e);
}
})
.then(() => {
if (pipeline) {
pipeline.logInfo("Successfully submitted PipelineRun to Gentrace");
}
})
.finally(() => {
delete globalRequestBuffer[pipelineRunObject.id];
});
const data = {
pipelineRunId: pipelineRunObject.id,
};
return data;
}
try {
const pipelinePostResponse = yield submission;
if (pipeline) {
pipeline.logInfo("Successfully submitted PipelineRun to Gentrace");
}
return pipelinePostResponse.data;
}
catch (e) {
if (pipeline) {
pipeline.logWarn(e);
}
throw e;
}
});
}
submit({ waitForServer, selectFields, } = { waitForServer: false }) {
return __awaiter(this, void 0, void 0, function* () {
const testCounter = getTestCounter();
if (testCounter > 0) {
const data = {
pipelineRunId: this.id,
};
return data;
}
const pipelineRunObject = this.toObject();
const response = yield PipelineRun.submitFromJson(pipelineRunObject, {
waitForServer,
pipeline: this.pipeline,
selectFields,
});
return response;
});
}
/**
* Returns the current evaluations stored in the PipelineRun instance.
* @returns {LocalEvaluation[]} An array of LocalEvaluation objects.
*/
getLocalEvaluations() {
return this.evaluations;
}
}
export { PipelineRun, getRun };
//# sourceMappingURL=pipeline-run.mjs.map