agsamantha/node_modules/@langchain/community/dist/chat_models/arcjet.js

151 lines
4.9 KiB
JavaScript
Raw Normal View History

2024-10-02 15:15:21 -05:00
/* 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;
}),
};
}),
};
}
}