150 lines
4.9 KiB
JavaScript
150 lines
4.9 KiB
JavaScript
/* eslint-disable no-param-reassign */
|
|
/* eslint-disable no-plusplus */
|
|
import { BaseChatModel, } from "@langchain/core/language_models/chat_models";
|
|
async function transformTextMessageAsync(message, transformer) {
|
|
if (typeof message.content === "string") {
|
|
message.content = await transformer(message.content);
|
|
return message;
|
|
}
|
|
const redactedContent = await Promise.all(message.content.map(async (m) => {
|
|
if (m.type === "text") {
|
|
return {
|
|
...m,
|
|
text: await transformer(m.text),
|
|
};
|
|
}
|
|
else {
|
|
return Promise.resolve(m);
|
|
}
|
|
}));
|
|
message.content = redactedContent;
|
|
return message;
|
|
}
|
|
function transformTextMessage(message, transformer) {
|
|
if (typeof message.content === "string") {
|
|
message.content = transformer(message.content);
|
|
return message;
|
|
}
|
|
const redactedContent = message.content.map((m) => {
|
|
if (m.type === "text") {
|
|
return {
|
|
...m,
|
|
text: transformer(m.text),
|
|
};
|
|
}
|
|
else {
|
|
return m;
|
|
}
|
|
});
|
|
message.content = redactedContent;
|
|
return message;
|
|
}
|
|
export class ArcjetRedact extends BaseChatModel {
|
|
static lc_name() {
|
|
return "ArcjetRedact";
|
|
}
|
|
constructor(options) {
|
|
super(options);
|
|
Object.defineProperty(this, "chatModel", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "entities", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "contextWindowSize", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "detect", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "replace", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "index", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
if (options.entities && options.entities.length === 0) {
|
|
throw new Error("no entities configured for redaction");
|
|
}
|
|
this.chatModel = options.chatModel;
|
|
this.entities = options.entities;
|
|
this.contextWindowSize = options.contextWindowSize;
|
|
this.detect = options.detect;
|
|
this.replace = options.replace;
|
|
this.index = 0;
|
|
}
|
|
_createUniqueReplacement(entity) {
|
|
const userReplacement = typeof this.replace !== "undefined" ? this.replace(entity) : undefined;
|
|
if (typeof userReplacement !== "undefined") {
|
|
return userReplacement;
|
|
}
|
|
this.index++;
|
|
if (entity === "email") {
|
|
return `<Redacted email #${this.index}>`;
|
|
}
|
|
if (entity === "phone-number") {
|
|
return `<Redacted phone number #${this.index}>`;
|
|
}
|
|
if (entity === "ip-address") {
|
|
return `<Redacted IP address #${this.index}>`;
|
|
}
|
|
if (entity === "credit-card-number") {
|
|
return `<Redacted credit card number #${this.index}>`;
|
|
}
|
|
return `<Redacted ${entity} #${this.index}>`;
|
|
}
|
|
_llmType() {
|
|
return "arcjet_redact";
|
|
}
|
|
async _generate(messages, options, runManager) {
|
|
const ajOptions = {
|
|
entities: this.entities,
|
|
contextWindowSize: this.contextWindowSize,
|
|
detect: this.detect,
|
|
replace: this._createUniqueReplacement.bind(this),
|
|
};
|
|
const unredactors = [];
|
|
// Support CommonJS
|
|
const { redact } = await import("@arcjet/redact");
|
|
const redacted = await Promise.all(messages.map(async (message) => {
|
|
return await transformTextMessageAsync(message, async (message) => {
|
|
const [redacted, unredact] = await redact(message, ajOptions);
|
|
unredactors.push(unredact);
|
|
return redacted;
|
|
});
|
|
}));
|
|
const response = await this.chatModel._generate(redacted, options, runManager);
|
|
return {
|
|
...response,
|
|
generations: response.generations.map((resp) => {
|
|
return {
|
|
...resp,
|
|
message: transformTextMessage(resp.message, (message) => {
|
|
for (const unredact of unredactors) {
|
|
message = unredact(message);
|
|
}
|
|
return message;
|
|
}),
|
|
};
|
|
}),
|
|
};
|
|
}
|
|
}
|