-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
layerup_security.ts
163 lines (133 loc) Β· 4.64 KB
/
layerup_security.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import { LLM, type BaseLLMParams } from "@langchain/core/language_models/llms";
import {
GuardrailResponse,
LayerupSecurity as LayerupSecuritySDK,
LLMMessage,
} from "@layerup/layerup-security";
export interface LayerupSecurityOptions extends BaseLLMParams {
llm: LLM;
layerupApiKey?: string;
layerupApiBaseUrl?: string;
promptGuardrails?: string[];
responseGuardrails?: string[];
mask?: boolean;
metadata?: Record<string, unknown>;
handlePromptGuardrailViolation?: (violation: GuardrailResponse) => LLMMessage;
handleResponseGuardrailViolation?: (
violation: GuardrailResponse
) => LLMMessage;
}
function defaultGuardrailViolationHandler(
violation: GuardrailResponse
): LLMMessage {
if (violation.canned_response) return violation.canned_response;
const guardrailName = violation.offending_guardrail
? `Guardrail ${violation.offending_guardrail}`
: "A guardrail";
throw new Error(
`${guardrailName} was violated without a proper guardrail violation handler.`
);
}
export class LayerupSecurity extends LLM {
static lc_name() {
return "LayerupSecurity";
}
lc_serializable = true;
llm: LLM;
layerupApiKey: string;
layerupApiBaseUrl = "https://api.uselayerup.com/v1";
promptGuardrails: string[] = [];
responseGuardrails: string[] = [];
mask = false;
metadata: Record<string, unknown> = {};
handlePromptGuardrailViolation: (violation: GuardrailResponse) => LLMMessage =
defaultGuardrailViolationHandler;
handleResponseGuardrailViolation: (
violation: GuardrailResponse
) => LLMMessage = defaultGuardrailViolationHandler;
private layerup: LayerupSecuritySDK;
constructor(options: LayerupSecurityOptions) {
super(options);
if (!options.llm) {
throw new Error("Layerup Security requires an LLM to be provided.");
} else if (!options.layerupApiKey) {
throw new Error("Layerup Security requires an API key to be provided.");
}
this.llm = options.llm;
this.layerupApiKey = options.layerupApiKey;
this.layerupApiBaseUrl =
options.layerupApiBaseUrl || this.layerupApiBaseUrl;
this.promptGuardrails = options.promptGuardrails || this.promptGuardrails;
this.responseGuardrails =
options.responseGuardrails || this.responseGuardrails;
this.mask = options.mask || this.mask;
this.metadata = options.metadata || this.metadata;
this.handlePromptGuardrailViolation =
options.handlePromptGuardrailViolation ||
this.handlePromptGuardrailViolation;
this.handleResponseGuardrailViolation =
options.handleResponseGuardrailViolation ||
this.handleResponseGuardrailViolation;
this.layerup = new LayerupSecuritySDK({
apiKey: this.layerupApiKey,
baseURL: this.layerupApiBaseUrl,
});
}
_llmType() {
return "layerup_security";
}
async _call(input: string, options?: BaseLLMParams): Promise<string> {
// Since LangChain LLMs only support string inputs, we will wrap each call to Layerup in a single-message
// array of messages, then extract the string element when we need to access it.
let messages: LLMMessage[] = [
{
role: "user",
content: input,
},
];
let unmaskResponse;
if (this.mask) {
[messages, unmaskResponse] = await this.layerup.maskPrompt(
messages,
this.metadata
);
}
if (this.promptGuardrails.length > 0) {
const securityResponse = await this.layerup.executeGuardrails(
this.promptGuardrails,
messages,
this.metadata
);
// If there is a guardrail violation, extract the canned response and reply with that instead
if (!securityResponse.all_safe) {
const replacedResponse =
this.handlePromptGuardrailViolation(securityResponse);
return replacedResponse.content;
}
}
// Invoke the underlying LLM with the prompt and options
let result = await this.llm.invoke(messages[0].content, options);
if (this.mask && unmaskResponse) {
result = unmaskResponse(result);
}
// Add to messages array for response guardrail handler
messages.push({
role: "assistant",
content: result,
});
if (this.responseGuardrails.length > 0) {
const securityResponse = await this.layerup.executeGuardrails(
this.responseGuardrails,
messages,
this.metadata
);
// If there is a guardrail violation, extract the canned response and reply with that instead
if (!securityResponse.all_safe) {
const replacedResponse =
this.handleResponseGuardrailViolation(securityResponse);
return replacedResponse.content;
}
}
return result;
}
}