@gentrace/openai
Version:
Gentrace OpenAI v4 plugin for Node.JS
539 lines (535 loc) • 27.7 kB
JavaScript
'use strict';
var core = require('@gentrace/core');
var Mustache = require('mustache');
var OpenAI = require('openai');
var completions = require('openai/resources/beta/chat/completions');
var util = require('./util.js');
var parser = require('openai/lib/parser');
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;
};
var __asyncValues = (undefined && undefined.__asyncValues) || function (o) {
if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
var m = o[Symbol.asyncIterator], i;
return m ? m.call(o) : (o = typeof __values === "function" ? __values(o) : o[Symbol.iterator](), i = {}, verb("next"), verb("throw"), verb("return"), i[Symbol.asyncIterator] = function () { return this; }, i);
function verb(n) { i[n] = o[n] && function (v) { return new Promise(function (resolve, reject) { v = o[n](v), settle(resolve, reject, v.done, v.value); }); }; }
function settle(resolve, reject, d, v) { Promise.resolve(v).then(function(v) { resolve({ value: v, done: d }); }, reject); }
};
var __await = (undefined && undefined.__await) || function (v) { return this instanceof __await ? (this.v = v, this) : new __await(v); };
var __asyncGenerator = (undefined && undefined.__asyncGenerator) || function (thisArg, _arguments, generator) {
if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
var g = generator.apply(thisArg, _arguments || []), i, q = [];
return i = {}, verb("next"), verb("throw"), verb("return"), i[Symbol.asyncIterator] = function () { return this; }, i;
function verb(n) { if (g[n]) i[n] = function (v) { return new Promise(function (a, b) { q.push([n, v, a, b]) > 1 || resume(n, v); }); }; }
function resume(n, v) { try { step(g[n](v)); } catch (e) { settle(q[0][3], e); } }
function step(r) { r.value instanceof __await ? Promise.resolve(r.value.v).then(fulfill, reject) : settle(q[0][2], r); }
function fulfill(value) { resume("next", value); }
function reject(value) { resume("throw", value); }
function settle(f, v) { if (f(v), q.shift(), q.length) resume(q[0][0], q[0][1]); }
};
function createRenderedChatMessages(messages) {
let newMessages = [];
for (let message of messages) {
if (message.contentTemplate && message.contentInputs) {
const { contentTemplate, contentInputs } = message, rest = __rest(message, ["contentTemplate", "contentInputs"]);
newMessages.push(Object.assign(Object.assign({}, rest), { content: Mustache.render(contentTemplate, contentInputs) }));
}
else {
newMessages.push(Object.assign({}, message));
}
}
return newMessages;
}
class GentraceStream {
constructor(stream, pipelineRun, partialStepRun, isSelfContained, aggregator) {
this.stream = stream;
this.pipelineRun = pipelineRun;
this.partialStepRun = partialStepRun;
this.isSelfContained = isSelfContained;
this.aggregator = aggregator;
}
[Symbol.asyncIterator]() {
var _a;
return __asyncGenerator(this, arguments, function* _b() {
var _c, e_1, _d, _e;
const allItems = [];
try {
for (var _f = true, _g = __asyncValues(this.stream), _h; _h = yield __await(_g.next()), _c = _h.done, !_c; _f = true) {
_e = _h.value;
_f = false;
const item = _e;
// Yield each item from original stream
yield yield __await(item);
allItems.push(item);
}
}
catch (e_1_1) { e_1 = { error: e_1_1 }; }
finally {
try {
if (!_f && !_c && (_d = _g.return)) yield __await(_d.call(_g));
}
finally { if (e_1) throw e_1.error; }
}
const consolidatedResponse = this.aggregator(allItems);
const endTime = Date.now();
const elapsedTime = endTime - Date.parse(this.partialStepRun.startTime);
this.partialStepRun.elapsedTime = elapsedTime;
this.partialStepRun.endTime = new Date(endTime).toISOString();
this.partialStepRun.outputs = consolidatedResponse;
(_a = this.pipelineRun) === null || _a === void 0 ? void 0 : _a.addStepRunNode(this.partialStepRun);
if (this.isSelfContained) {
yield __await(this.pipelineRun.submit());
}
});
}
}
class GentraceEmbeddings extends OpenAI.Embeddings {
constructor({ client, pipelineRun, gentraceConfig, }) {
super(client);
this.pipelineRun = pipelineRun;
this.gentraceConfig = gentraceConfig;
}
createInner(body, options) {
var _a;
return __awaiter(this, void 0, void 0, function* () {
const { pipelineSlug, gentrace } = body, newPayload = __rest(body, ["pipelineSlug", "gentrace"]);
const { model } = newPayload, inputParams = __rest(newPayload, ["model"]);
let isSelfContainedPullRequest = !this.pipelineRun && pipelineSlug;
let pipelineRun = this.pipelineRun;
if (isSelfContainedPullRequest) {
const pipeline = new core.Pipeline({
id: pipelineSlug,
slug: pipelineSlug,
apiKey: this.gentraceConfig.apiKey,
basePath: this.gentraceConfig.basePath,
logger: this.gentraceConfig.logger,
});
pipelineRun = new core.PipelineRun({
pipeline,
});
}
const startTime = Date.now();
const completion = (yield this._client.post("/embeddings", Object.assign({ body: newPayload }, options)));
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new OpenAICreateEmbeddingStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), Object.assign({}, inputParams), { model }, completion, (_a = body === null || body === void 0 ? void 0 : body.gentrace) !== null && _a !== void 0 ? _a : {}));
if (isSelfContainedPullRequest) {
const { pipelineRunId } = yield pipelineRun.submit();
completion.pipelineRunId =
pipelineRunId;
return completion;
}
return completion;
});
}
}
class GentraceModerations extends OpenAI.Moderations {
constructor({ client, pipelineRun, gentraceConfig, }) {
super(client);
this.pipelineRun = pipelineRun;
this.gentraceConfig = gentraceConfig;
}
createInner(body, options) {
var _a;
return __awaiter(this, void 0, void 0, function* () {
const { pipelineSlug, gentrace } = body, newPayload = __rest(body, ["pipelineSlug", "gentrace"]);
const { model } = newPayload, inputParams = __rest(newPayload, ["model"]);
let isSelfContainedPullRequest = !this.pipelineRun && pipelineSlug;
let pipelineRun = this.pipelineRun;
if (isSelfContainedPullRequest) {
const pipeline = new core.Pipeline({
id: pipelineSlug,
slug: pipelineSlug,
apiKey: this.gentraceConfig.apiKey,
basePath: this.gentraceConfig.basePath,
logger: this.gentraceConfig.logger,
});
pipelineRun = new core.PipelineRun({
pipeline,
});
}
const startTime = Date.now();
const completion = (yield this._client.post("/moderations", Object.assign({ body: newPayload }, options)));
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new OpenAICreateModerationStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), Object.assign({}, inputParams), { model }, completion, (_a = body === null || body === void 0 ? void 0 : body.gentrace) !== null && _a !== void 0 ? _a : {}));
if (isSelfContainedPullRequest) {
const { pipelineRunId } = yield pipelineRun.submit();
completion.pipelineRunId =
pipelineRunId;
return completion;
}
return completion;
});
}
}
class GentraceChatCompletions extends OpenAI.Chat.Completions {
constructor({ client, pipelineRun, gentraceConfig, }) {
super(client);
this.pipelineRun = pipelineRun;
this.gentraceConfig = gentraceConfig;
}
// @ts-ignore
createInner(body, requestOptions) {
var _a, _b, _c;
return __awaiter(this, void 0, void 0, function* () {
const { pipelineSlug } = body;
let isSelfContainedPipelineRun = !this.pipelineRun && !!pipelineSlug;
let pipelineRun = this.pipelineRun;
if (isSelfContainedPipelineRun) {
const pipeline = new core.Pipeline({
id: pipelineSlug,
slug: pipelineSlug,
apiKey: this.gentraceConfig.apiKey,
basePath: this.gentraceConfig.basePath,
logger: this.gentraceConfig.logger,
});
pipelineRun = new core.PipelineRun({
pipeline,
});
}
const { messages, pipelineSlug: _pipelineSlug, gentrace } = body, baseCompletionOptions = __rest(body, ["messages", "pipelineSlug", "gentrace"]);
const renderedMessages = createRenderedChatMessages(messages);
const contentTemplatesArray = messages.map((message) => {
var _a;
return (_a = message.contentTemplate) !== null && _a !== void 0 ? _a : null;
});
const contentInputsArray = messages.map((message) => {
var _a;
return (_a = message.contentInputs) !== null && _a !== void 0 ? _a : null;
});
const startTime = Date.now();
const completion = this._client.post("/chat/completions", Object.assign(Object.assign({ body: Object.assign({ messages: renderedMessages }, baseCompletionOptions) }, requestOptions), { stream: (_a = body.stream) !== null && _a !== void 0 ? _a : false }));
const data = yield completion;
let finalData = data;
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
// user parameter is an input, not a model parameter
const { user, tools } = baseCompletionOptions, modelParams = __rest(baseCompletionOptions, ["user", "tools"]);
if (body.stream) {
finalData = new GentraceStream(data, pipelineRun, new OpenAICreateChatCompletionStepRun(0, new Date(startTime).toISOString(), "", {
messages: renderedMessages,
tools,
user,
contentInputs: contentInputsArray,
}, Object.assign(Object.assign({}, modelParams), { contentTemplates: contentTemplatesArray }), finalData, (_b = body === null || body === void 0 ? void 0 : body.gentrace) !== null && _b !== void 0 ? _b : {}), !!isSelfContainedPipelineRun, createChatCompletionStreamResponse);
}
else {
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new OpenAICreateChatCompletionStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), {
messages: renderedMessages,
user,
contentInputs: contentInputsArray,
}, Object.assign(Object.assign({}, modelParams), { contentTemplates: contentTemplatesArray }), finalData, (_c = body === null || body === void 0 ? void 0 : body.gentrace) !== null && _c !== void 0 ? _c : {}));
}
if (isSelfContainedPipelineRun) {
let pipelineRunId = "";
if (!body.stream) {
const submitInfo = yield pipelineRun.submit();
pipelineRunId = submitInfo.pipelineRunId;
}
else {
pipelineRunId = pipelineRun.getId();
}
finalData.pipelineRunId =
pipelineRunId;
return finalData;
}
return finalData;
});
}
}
class GentraceCompletions extends OpenAI.Completions {
constructor({ client, pipelineRun, gentraceConfig, }) {
super(client);
this.pipelineRun = pipelineRun;
this.gentraceConfig = gentraceConfig;
}
createInner(body, requestOptions) {
var _a, _b, _c, _d;
return __awaiter(this, void 0, void 0, function* () {
const { pipelineSlug } = body;
let isSelfContainedPipelineRun = !this.pipelineRun && !!pipelineSlug;
let pipelineRun = this.pipelineRun;
if (isSelfContainedPipelineRun) {
const pipeline = new core.Pipeline({
id: pipelineSlug,
slug: pipelineSlug,
apiKey: this.gentraceConfig.apiKey,
basePath: this.gentraceConfig.basePath,
logger: this.gentraceConfig.logger,
});
pipelineRun = new core.PipelineRun({
pipeline,
});
}
const { promptTemplate, promptInputs, prompt, pipelineSlug: _pipelineSlug, gentrace } = body, baseCompletionOptions = __rest(body, ["promptTemplate", "promptInputs", "prompt", "pipelineSlug", "gentrace"]);
let renderedPrompt = prompt;
if (promptTemplate && promptInputs) {
renderedPrompt = Mustache.render(promptTemplate, promptInputs);
}
const newCompletionOptions = Object.assign(Object.assign({}, baseCompletionOptions), { prompt: renderedPrompt, stream: (_a = baseCompletionOptions.stream) !== null && _a !== void 0 ? _a : false });
const startTime = Date.now();
const completion = this._client.post("/completions", Object.assign(Object.assign({ body: newCompletionOptions }, requestOptions), { stream: (_b = body.stream) !== null && _b !== void 0 ? _b : false }));
const data = yield completion;
let finalData = data;
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
// User and suffix parameters are inputs not model parameters
const { user, suffix } = baseCompletionOptions, partialModelParams = __rest(baseCompletionOptions, ["user", "suffix"]);
if (body.stream) {
finalData = new GentraceStream(data, pipelineRun, new OpenAICreateCompletionStepRun(0, new Date(startTime).toISOString(), "", {
prompt: promptTemplate && promptInputs ? promptInputs : prompt,
user,
suffix,
}, Object.assign(Object.assign({}, partialModelParams), { promptTemplate }), finalData, (_c = body === null || body === void 0 ? void 0 : body.gentrace) !== null && _c !== void 0 ? _c : {}), !!isSelfContainedPipelineRun, createCompletionStreamResponse);
}
else {
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new OpenAICreateCompletionStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), {
prompt: promptTemplate && promptInputs ? promptInputs : prompt,
user,
suffix,
}, Object.assign(Object.assign({}, partialModelParams), { promptTemplate }), finalData, (_d = body === null || body === void 0 ? void 0 : body.gentrace) !== null && _d !== void 0 ? _d : {}));
}
if (isSelfContainedPipelineRun) {
let pipelineRunId = "";
if (!body.stream) {
const submitInfo = yield pipelineRun.submit();
pipelineRunId = submitInfo.pipelineRunId;
}
else {
pipelineRunId = pipelineRun.getId();
}
finalData.pipelineRunId =
pipelineRunId;
return finalData;
}
return finalData;
});
}
}
class GentraceBetaChatCompletions extends completions.Completions {
constructor({ client, pipelineRun, gentraceConfig, }) {
super(client);
this.pipelineRun = pipelineRun;
this.gentraceConfig = gentraceConfig;
}
parseInner(body, options) {
var _a;
return __awaiter(this, void 0, void 0, function* () {
const { pipelineSlug } = body;
let isSelfContainedPipelineRun = !this.pipelineRun && !!pipelineSlug;
let pipelineRun = this.pipelineRun;
if (isSelfContainedPipelineRun) {
const pipeline = new core.Pipeline({
id: pipelineSlug,
slug: pipelineSlug,
apiKey: this.gentraceConfig.apiKey,
basePath: this.gentraceConfig.basePath,
logger: this.gentraceConfig.logger,
});
pipelineRun = new core.PipelineRun({
pipeline,
});
}
const { messages, pipelineSlug: _pipelineSlug, gentrace } = body, baseCompletionOptions = __rest(body, ["messages", "pipelineSlug", "gentrace"]);
const renderedMessages = createRenderedChatMessages(messages);
const contentTemplatesArray = messages.map((message) => {
var _a;
return (_a = message.contentTemplate) !== null && _a !== void 0 ? _a : null;
});
const contentInputsArray = messages.map((message) => {
var _a;
return (_a = message.contentInputs) !== null && _a !== void 0 ? _a : null;
});
const startTime = Date.now();
parser.validateInputTools(body.tools);
// TODO: this is not sufficient for structuredt outputs. There's additional parsing that needs to happen.
const completion = this._client.post("/chat/completions", Object.assign({ body: Object.assign({ messages: renderedMessages }, baseCompletionOptions) }, options));
const rawData = yield completion;
// @ts-ignore
const data = parser.parseChatCompletion(rawData, body);
const endTime = Date.now();
const elapsedTime = Math.floor(endTime - startTime);
// user parameter is an input, not a model parameter
const { user, tools } = baseCompletionOptions, modelParams = __rest(baseCompletionOptions, ["user", "tools"]);
pipelineRun === null || pipelineRun === void 0 ? void 0 : pipelineRun.addStepRunNode(new OpenAICreateChatCompletionStepRun(elapsedTime, new Date(startTime).toISOString(), new Date(endTime).toISOString(), {
messages: renderedMessages,
user,
contentInputs: contentInputsArray,
}, Object.assign(Object.assign({}, modelParams), { contentTemplates: contentTemplatesArray }), data, (_a = body === null || body === void 0 ? void 0 : body.gentrace) !== null && _a !== void 0 ? _a : {}));
if (isSelfContainedPipelineRun) {
const submitInfo = yield pipelineRun.submit();
const pipelineRunId = submitInfo.pipelineRunId;
return Object.assign(Object.assign({}, data), { pipelineRunId });
}
return data;
});
}
}
class OpenAIPipelineHandler extends OpenAI {
constructor(_a) {
var { pipelineRun, gentraceConfig } = _a, config = __rest(_a, ["pipelineRun", "gentraceConfig"]);
super(config);
this.config = config;
this.pipelineRun = pipelineRun;
this.gentraceConfig = gentraceConfig;
}
}
class OpenAICreateChatCompletionStepRun extends core.StepRun {
constructor(elapsedTime, startTime, endTime, inputs, modelParams, response, context) {
super("openai", "openai_createChatCompletion", elapsedTime, startTime, endTime, inputs, modelParams, response, context, undefined);
}
}
class OpenAICreateCompletionStepRun extends core.StepRun {
constructor(elapsedTime, startTime, endTime, inputs, modelParams, response, context) {
super("openai", "openai_createCompletion", elapsedTime, startTime, endTime, inputs, modelParams, response, context, undefined);
}
}
class OpenAICreateEmbeddingStepRun extends core.StepRun {
constructor(elapsedTime, startTime, endTime, inputs, modelParams, response, context) {
super("openai", "openai_createEmbedding", elapsedTime, startTime, endTime, inputs, modelParams, response, context, undefined);
}
}
class OpenAICreateModerationStepRun extends core.StepRun {
constructor(elapsedTime, startTime, endTime, inputs, modelParams, response, context) {
super("openai", "openai_createModeration", elapsedTime, startTime, endTime, inputs, modelParams, response, context, undefined);
}
}
function createChatCompletionStreamResponse(streamList) {
var _a, _b, _c, _d;
let finalResponseString = "";
const toolIdToInfoMap = {};
let model = "";
let id = "";
let created = 0;
for (const value of streamList) {
model = value.model;
id = value.id;
created = value.created;
if (value.choices && value.choices.length > 0) {
const firstChoice = value.choices[0];
if (firstChoice.delta && firstChoice.delta.tool_calls) {
for (const toolCall of firstChoice.delta.tool_calls) {
if (toolCall.id) {
const existingToolInfo = toolIdToInfoMap[toolCall.id];
if (!existingToolInfo) {
toolIdToInfoMap[toolCall.id] = toolCall;
}
else if (util.isDefined((_a = existingToolInfo.function) === null || _a === void 0 ? void 0 : _a.arguments) &&
util.isDefined((_b = toolCall.function) === null || _b === void 0 ? void 0 : _b.arguments)) {
// Assume that the ID is provided to differentiate between different tool calls.
existingToolInfo.function.arguments +=
toolCall.function.arguments;
}
}
else {
// Associate with the tool call in the map
const toolId = Object.keys(toolIdToInfoMap)[0];
if (toolId) {
const existingToolInfo = toolIdToInfoMap[toolId];
if (util.isDefined((_c = existingToolInfo.function) === null || _c === void 0 ? void 0 : _c.arguments) &&
util.isDefined((_d = toolCall.function) === null || _d === void 0 ? void 0 : _d.arguments)) {
existingToolInfo.function.arguments +=
toolCall.function.arguments;
}
}
}
}
}
if (firstChoice.delta && firstChoice.delta.content) {
finalResponseString += firstChoice.delta.content;
}
else if (firstChoice.finish_reason) {
break;
}
}
}
const finalResponse = {
id,
// Override this so it doesn't show chat.completion.chunk
object: "chat.completion",
created,
model,
choices: [
{
finish_reason: null,
index: 0,
message: {
content: finalResponseString,
role: "assistant",
tool_calls: Object.values(toolIdToInfoMap),
},
},
],
};
return finalResponse;
}
function createCompletionStreamResponse(streamList) {
let finalResponseString = "";
let model = "";
let id = "";
let created = 0;
for (const value of streamList) {
model = value.model;
id = value.id;
created = value.created;
if (value.choices && value.choices.length > 0) {
const firstChoice = value.choices[0];
if (firstChoice.text) {
finalResponseString += firstChoice.text;
}
else if (firstChoice.delta && firstChoice.delta.content) {
finalResponseString += firstChoice.delta.content;
}
else if (firstChoice.finish_reason &&
firstChoice.finish_reason === "stop") {
break;
}
}
}
const finalResponse = {
id,
// Override this so it doesn't show chat.completion.chunk.
object: "text_completion",
created,
model,
choices: [
{
text: finalResponseString,
index: 0,
logprobs: null,
finish_reason: "length",
},
],
};
return finalResponse;
}
exports.GentraceBetaChatCompletions = GentraceBetaChatCompletions;
exports.GentraceChatCompletions = GentraceChatCompletions;
exports.GentraceCompletions = GentraceCompletions;
exports.GentraceEmbeddings = GentraceEmbeddings;
exports.GentraceModerations = GentraceModerations;
exports.GentraceStream = GentraceStream;
exports.OpenAICreateChatCompletionStepRun = OpenAICreateChatCompletionStepRun;
exports.OpenAICreateCompletionStepRun = OpenAICreateCompletionStepRun;
exports.OpenAICreateEmbeddingStepRun = OpenAICreateEmbeddingStepRun;
exports.OpenAICreateModerationStepRun = OpenAICreateModerationStepRun;
exports.OpenAIPipelineHandler = OpenAIPipelineHandler;
//# sourceMappingURL=openai.js.map