UNPKG

@langchain/community

Version:
1 lines 7.32 kB
{"version":3,"file":"arcjet.cjs","names":["BaseChatModel"],"sources":["../../src/chat_models/arcjet.ts"],"sourcesContent":["/* eslint-disable no-param-reassign */\nimport {\n BaseChatModel,\n type BaseChatModelParams,\n} from \"@langchain/core/language_models/chat_models\";\nimport type { ArcjetSensitiveInfoType, RedactOptions } from \"@arcjet/redact\";\nimport { CallbackManagerForLLMRun } from \"@langchain/core/callbacks/manager\";\nimport { BaseMessage } from \"@langchain/core/messages\";\nimport { ChatResult } from \"@langchain/core/outputs\";\n\ntype DetectSensitiveInfoEntities<T> = (\n tokens: string[]\n) => Array<ArcjetSensitiveInfoType | T | undefined>;\ntype ValidEntities<Detect> = Array<\n undefined extends Detect\n ? ArcjetSensitiveInfoType\n : Detect extends DetectSensitiveInfoEntities<infer CustomEntities>\n ? ArcjetSensitiveInfoType | CustomEntities\n : never\n>;\n\nexport interface ArcjetRedactOptions<Detect> extends BaseChatModelParams {\n chatModel: BaseChatModel;\n entities?: ValidEntities<Detect>;\n contextWindowSize?: number;\n detect?: Detect;\n replace?: (entity: ValidEntities<Detect>[number]) => string | undefined;\n}\n\nexport type { ArcjetSensitiveInfoType, RedactOptions };\n\nasync function transformTextMessageAsync(\n message: BaseMessage,\n transformer: (text: string) => Promise<string>\n): Promise<BaseMessage> {\n if (typeof message.content === \"string\") {\n message.content = await transformer(message.content);\n return message;\n }\n\n const redactedContent = await Promise.all(\n message.content.map(async (m) => {\n if (m.type === \"text\") {\n return {\n ...m,\n text: await transformer(m.text),\n };\n } else {\n return Promise.resolve(m);\n }\n })\n );\n message.content = redactedContent;\n return message;\n}\n\nfunction transformTextMessage(\n message: BaseMessage,\n transformer: (text: string) => string\n): BaseMessage {\n if (typeof message.content === \"string\") {\n message.content = transformer(message.content);\n return message;\n }\n\n const redactedContent = message.content.map((m) => {\n if (m.type === \"text\") {\n return {\n ...m,\n text: transformer(m.text),\n };\n } else {\n return m;\n }\n });\n message.content = redactedContent;\n return message;\n}\n\nexport class ArcjetRedact<\n Detect extends DetectSensitiveInfoEntities<CustomEntities> | undefined,\n CustomEntities extends string,\n> extends BaseChatModel {\n static lc_name() {\n return \"ArcjetRedact\";\n }\n\n chatModel: BaseChatModel;\n\n entities?: ValidEntities<Detect>;\n\n contextWindowSize?: number;\n\n detect?: Detect;\n\n replace?: (entity: ValidEntities<Detect>[number]) => string | undefined;\n\n index: number;\n\n constructor(options: ArcjetRedactOptions<Detect>) {\n super(options);\n\n if (options.entities && options.entities.length === 0) {\n throw new Error(\"no entities configured for redaction\");\n }\n\n this.chatModel = options.chatModel;\n this.entities = options.entities;\n this.contextWindowSize = options.contextWindowSize;\n this.detect = options.detect;\n this.replace = options.replace;\n this.index = 0;\n }\n\n _createUniqueReplacement(entity: ValidEntities<Detect>[number]): string {\n const userReplacement =\n typeof this.replace !== \"undefined\" ? this.replace(entity) : undefined;\n if (typeof userReplacement !== \"undefined\") {\n return userReplacement;\n }\n\n this.index++;\n\n if (entity === \"email\") {\n return `<Redacted email #${this.index}>`;\n }\n\n if (entity === \"phone-number\") {\n return `<Redacted phone number #${this.index}>`;\n }\n\n if (entity === \"ip-address\") {\n return `<Redacted IP address #${this.index}>`;\n }\n\n if (entity === \"credit-card-number\") {\n return `<Redacted credit card number #${this.index}>`;\n }\n\n return `<Redacted ${entity} #${this.index}>`;\n }\n\n _llmType() {\n return \"arcjet_redact\";\n }\n\n async _generate(\n messages: BaseMessage[],\n options: this[\"ParsedCallOptions\"],\n runManager?: CallbackManagerForLLMRun | undefined\n ): Promise<ChatResult> {\n const ajOptions: RedactOptions<Detect> = {\n entities: this.entities,\n contextWindowSize: this.contextWindowSize,\n detect: this.detect,\n replace: this._createUniqueReplacement.bind(this),\n };\n\n const unredactors: Array<(message: string) => string> = [];\n // Support CommonJS\n const { redact } = await import(\"@arcjet/redact\");\n const redacted = await Promise.all(\n messages.map(async (message) => {\n return await transformTextMessageAsync(message, async (message) => {\n const [redacted, unredact] = await redact(message, ajOptions);\n unredactors.push(unredact);\n return redacted;\n });\n })\n );\n\n const response = await this.chatModel._generate(\n redacted,\n options,\n runManager\n );\n\n return {\n ...response,\n generations: response.generations.map((resp) => {\n return {\n ...resp,\n message: transformTextMessage(resp.message, (message: string) => {\n for (const unredact of unredactors) {\n message = unredact(message);\n }\n return message;\n }),\n };\n }),\n };\n }\n}\n"],"mappings":";;;;;AA+BA,eAAe,0BACb,SACA,aACsB;AACtB,KAAI,OAAO,QAAQ,YAAY,UAAU;AACvC,UAAQ,UAAU,MAAM,YAAY,QAAQ,QAAQ;AACpD,SAAO;;AAeT,SAAQ,UAZgB,MAAM,QAAQ,IACpC,QAAQ,QAAQ,IAAI,OAAO,MAAM;AAC/B,MAAI,EAAE,SAAS,OACb,QAAO;GACL,GAAG;GACH,MAAM,MAAM,YAAY,EAAE,KAAK;GAChC;MAED,QAAO,QAAQ,QAAQ,EAAE;GAE3B,CACH;AAED,QAAO;;AAGT,SAAS,qBACP,SACA,aACa;AACb,KAAI,OAAO,QAAQ,YAAY,UAAU;AACvC,UAAQ,UAAU,YAAY,QAAQ,QAAQ;AAC9C,SAAO;;AAaT,SAAQ,UAVgB,QAAQ,QAAQ,KAAK,MAAM;AACjD,MAAI,EAAE,SAAS,OACb,QAAO;GACL,GAAG;GACH,MAAM,YAAY,EAAE,KAAK;GAC1B;MAED,QAAO;GAET;AAEF,QAAO;;AAGT,IAAa,eAAb,cAGUA,4CAAAA,cAAc;CACtB,OAAO,UAAU;AACf,SAAO;;CAGT;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA,YAAY,SAAsC;AAChD,QAAM,QAAQ;AAEd,MAAI,QAAQ,YAAY,QAAQ,SAAS,WAAW,EAClD,OAAM,IAAI,MAAM,uCAAuC;AAGzD,OAAK,YAAY,QAAQ;AACzB,OAAK,WAAW,QAAQ;AACxB,OAAK,oBAAoB,QAAQ;AACjC,OAAK,SAAS,QAAQ;AACtB,OAAK,UAAU,QAAQ;AACvB,OAAK,QAAQ;;CAGf,yBAAyB,QAA+C;EACtE,MAAM,kBACJ,OAAO,KAAK,YAAY,cAAc,KAAK,QAAQ,OAAO,GAAG,KAAA;AAC/D,MAAI,OAAO,oBAAoB,YAC7B,QAAO;AAGT,OAAK;AAEL,MAAI,WAAW,QACb,QAAO,oBAAoB,KAAK,MAAM;AAGxC,MAAI,WAAW,eACb,QAAO,2BAA2B,KAAK,MAAM;AAG/C,MAAI,WAAW,aACb,QAAO,yBAAyB,KAAK,MAAM;AAG7C,MAAI,WAAW,qBACb,QAAO,iCAAiC,KAAK,MAAM;AAGrD,SAAO,aAAa,OAAO,IAAI,KAAK,MAAM;;CAG5C,WAAW;AACT,SAAO;;CAGT,MAAM,UACJ,UACA,SACA,YACqB;EACrB,MAAM,YAAmC;GACvC,UAAU,KAAK;GACf,mBAAmB,KAAK;GACxB,QAAQ,KAAK;GACb,SAAS,KAAK,yBAAyB,KAAK,KAAK;GAClD;EAED,MAAM,cAAkD,EAAE;EAE1D,MAAM,EAAE,WAAW,MAAM,OAAO;EAChC,MAAM,WAAW,MAAM,QAAQ,IAC7B,SAAS,IAAI,OAAO,YAAY;AAC9B,UAAO,MAAM,0BAA0B,SAAS,OAAO,YAAY;IACjE,MAAM,CAAC,UAAU,YAAY,MAAM,OAAO,SAAS,UAAU;AAC7D,gBAAY,KAAK,SAAS;AAC1B,WAAO;KACP;IACF,CACH;EAED,MAAM,WAAW,MAAM,KAAK,UAAU,UACpC,UACA,SACA,WACD;AAED,SAAO;GACL,GAAG;GACH,aAAa,SAAS,YAAY,KAAK,SAAS;AAC9C,WAAO;KACL,GAAG;KACH,SAAS,qBAAqB,KAAK,UAAU,YAAoB;AAC/D,WAAK,MAAM,YAAY,YACrB,WAAU,SAAS,QAAQ;AAE7B,aAAO;OACP;KACH;KACD;GACH"}