Skip to content

Commit 45a56ed

Browse files
authored
Merge branch 'main' into feature/add-qualifire-integration
2 parents b0e1227 + 7a1792c commit 45a56ed

File tree

12 files changed

+263
-12
lines changed

12 files changed

+263
-12
lines changed

plugins/default/regexReplace.ts

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import {
2+
HookEventType,
3+
PluginContext,
4+
PluginHandler,
5+
PluginParameters,
6+
} from '../types';
7+
import { getCurrentContentPart, setCurrentContentPart } from '../utils';
8+
9+
export const handler: PluginHandler = async (
10+
context: PluginContext,
11+
parameters: PluginParameters,
12+
eventType: HookEventType
13+
) => {
14+
let error = null;
15+
let verdict = true;
16+
let data: any = null;
17+
const transformedData: Record<string, any> = {
18+
request: {
19+
json: null,
20+
},
21+
response: {
22+
json: null,
23+
},
24+
};
25+
let transformed = false;
26+
27+
try {
28+
const regexPattern = parameters.rule;
29+
const redactText = parameters.redactText || '[REDACTED]';
30+
const failOnDetection = parameters.failOnDetection || false;
31+
32+
const { content, textArray } = getCurrentContentPart(context, eventType);
33+
34+
if (!regexPattern) {
35+
throw new Error('Missing regex pattern');
36+
}
37+
if (!content) {
38+
throw new Error('Missing text to match');
39+
}
40+
41+
const regex = new RegExp(regexPattern, 'g');
42+
43+
// Process all text items in the array
44+
let hasMatches = false;
45+
const mappedTextArray: Array<string | null> = [];
46+
textArray.forEach((text) => {
47+
if (!text) {
48+
mappedTextArray.push(null);
49+
return;
50+
}
51+
52+
// Reset regex for each text when using global flag
53+
regex.lastIndex = 0;
54+
55+
const matches = text.match(regex);
56+
if (matches && matches.length > 0) {
57+
hasMatches = true;
58+
}
59+
const replacedText = text.replace(regex, redactText);
60+
mappedTextArray.push(replacedText);
61+
});
62+
63+
// Handle transformation
64+
if (hasMatches) {
65+
setCurrentContentPart(
66+
context,
67+
eventType,
68+
transformedData,
69+
mappedTextArray
70+
);
71+
transformed = true;
72+
}
73+
if (failOnDetection && hasMatches) {
74+
verdict = false;
75+
}
76+
data = {
77+
regexPattern,
78+
verdict,
79+
explanation: transformed
80+
? `Pattern '${regexPattern}' matched and was replaced with '${redactText}'`
81+
: `The regex pattern '${regexPattern}' did not match any text.`,
82+
};
83+
} catch (e: any) {
84+
error = e;
85+
data = {
86+
explanation: `An error occurred while processing the regex: ${e.message}`,
87+
regexPattern: parameters.rule,
88+
};
89+
}
90+
91+
return { error, verdict, data, transformedData, transformed };
92+
};

plugins/index.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ import { handler as panwPrismaAirsintercept } from './panw-prisma-airs/intercept
6161
import { handler as defaultjwt } from './default/jwt';
6262
import { handler as defaultrequiredMetadataKeys } from './default/requiredMetadataKeys';
6363
import { handler as walledaiguardrails } from './walledai/guardrails';
64+
import { handler as defaultregexReplace } from './default/regexReplace';
6465

6566
export const plugins = {
6667
default: {
@@ -81,6 +82,7 @@ export const plugins = {
8182
modelWhitelist: defaultmodelWhitelist,
8283
jwt: defaultjwt,
8384
requiredMetadataKeys: defaultrequiredMetadataKeys,
85+
regexReplace: defaultregexReplace,
8486
},
8587
qualifire: {
8688
dangerousContent: qualifireDangerousContent,

src/globals.ts

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ export const FEATHERLESS_AI: string = 'featherless-ai';
102102
export const KRUTRIM: string = 'krutrim';
103103
export const QDRANT: string = 'qdrant';
104104
export const THREE_ZERO_TWO_AI: string = '302ai';
105+
export const TRIPO3D: string = 'tripo3d';
105106

106107
export const VALID_PROVIDERS = [
107108
ANTHROPIC,
@@ -167,6 +168,7 @@ export const VALID_PROVIDERS = [
167168
KRUTRIM,
168169
QDRANT,
169170
THREE_ZERO_TWO_AI,
171+
TRIPO3D,
170172
];
171173

172174
export const CONTENT_TYPES = {
@@ -215,6 +217,8 @@ export const fileExtensionMimeTypeMap = {
215217
mpegps: 'video/mpegps',
216218
flv: 'video/flv',
217219
webm: 'video/webm',
220+
mkv: 'video/mkv',
221+
threegpp: 'video/three_gpp',
218222
};
219223

220224
export const imagesMimeTypes = [
@@ -238,6 +242,19 @@ export const documentMimeTypes = [
238242
fileExtensionMimeTypeMap.txt,
239243
];
240244

245+
export const videoMimeTypes = [
246+
fileExtensionMimeTypeMap.mkv,
247+
fileExtensionMimeTypeMap.mov,
248+
fileExtensionMimeTypeMap.mp4,
249+
fileExtensionMimeTypeMap.webm,
250+
fileExtensionMimeTypeMap.flv,
251+
fileExtensionMimeTypeMap.mpeg,
252+
fileExtensionMimeTypeMap.mpg,
253+
fileExtensionMimeTypeMap.wmv,
254+
fileExtensionMimeTypeMap.threegpp,
255+
fileExtensionMimeTypeMap.avi,
256+
];
257+
241258
export enum BatchEndpoints {
242259
CHAT_COMPLETIONS = '/v1/chat/completions',
243260
COMPLETIONS = '/v1/completions',

src/providers/bedrock/chatComplete.ts

Lines changed: 41 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import {
22
BEDROCK,
3-
documentMimeTypes,
43
fileExtensionMimeTypeMap,
54
imagesMimeTypes,
5+
videoMimeTypes,
66
} from '../../globals';
77
import {
88
Message,
@@ -190,7 +190,16 @@ const getMessageContent = (message: Message) => {
190190
format: fileFormat,
191191
},
192192
});
193-
} else if (documentMimeTypes.includes(mimeType)) {
193+
} else if (videoMimeTypes.includes(mimeType)) {
194+
out.push({
195+
video: {
196+
format: fileFormat,
197+
source: {
198+
bytes,
199+
},
200+
},
201+
});
202+
} else {
194203
out.push({
195204
document: {
196205
format: fileFormat,
@@ -204,25 +213,46 @@ const getMessageContent = (message: Message) => {
204213
} else if (item.type === 'file') {
205214
const mimeType = item.file?.mime_type || fileExtensionMimeTypeMap.pdf;
206215
const fileFormat = mimeType.split('/')[1];
207-
if (item.file?.file_url) {
216+
if (imagesMimeTypes.includes(mimeType)) {
208217
out.push({
209-
document: {
218+
image: {
219+
source: {
220+
...(item.file?.file_data && { bytes: item.file.file_data }),
221+
...(item.file?.file_url && {
222+
s3Location: {
223+
uri: item.file.file_url,
224+
},
225+
}),
226+
},
227+
format: fileFormat,
228+
},
229+
});
230+
} else if (videoMimeTypes.includes(mimeType)) {
231+
out.push({
232+
video: {
210233
format: fileFormat,
211-
name: item.file.file_name || crypto.randomUUID(),
212234
source: {
213-
s3Location: {
214-
uri: item.file.file_url,
215-
},
235+
...(item.file?.file_data && { bytes: item.file.file_data }),
236+
...(item.file?.file_url && {
237+
s3Location: {
238+
uri: item.file.file_url,
239+
},
240+
}),
216241
},
217242
},
218243
});
219-
} else if (item.file?.file_data) {
244+
} else {
220245
out.push({
221246
document: {
222247
format: fileFormat,
223-
name: item.file.file_name || crypto.randomUUID(),
248+
name: item.file?.file_name || crypto.randomUUID(),
224249
source: {
225-
bytes: item.file.file_data,
250+
...(item.file?.file_data && { bytes: item.file.file_data }),
251+
...(item.file?.file_url && {
252+
s3Location: {
253+
uri: item.file.file_url,
254+
},
255+
}),
226256
},
227257
},
228258
});

src/providers/bedrock/types.ts

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,11 @@ export type BedrockContentItem = {
136136
};
137137
image?: {
138138
source: {
139-
bytes: string;
139+
bytes?: string;
140+
s3Location?: {
141+
uri: string;
142+
bucketOwner?: string;
143+
};
140144
};
141145
format: string;
142146
};
@@ -147,6 +151,17 @@ export type BedrockContentItem = {
147151
bytes?: string;
148152
s3Location?: {
149153
uri: string;
154+
bucketOwner?: string;
155+
};
156+
};
157+
};
158+
video?: {
159+
format: string;
160+
source: {
161+
bytes?: string;
162+
s3Location?: {
163+
uri: string;
164+
bucketOwner?: string;
150165
};
151166
};
152167
};

src/providers/google-vertex-ai/chatComplete.ts

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,10 @@ export const VertexGoogleChatCompleteConfig: ProviderConfig = {
354354
param: 'generationConfig',
355355
transform: (params: Params) => transformGenerationConfig(params),
356356
},
357+
modalities: {
358+
param: 'generationConfig',
359+
transform: (params: Params) => transformGenerationConfig(params),
360+
},
357361
};
358362

359363
interface AnthorpicTextContentItem {
@@ -479,6 +483,13 @@ export const GoogleChatCompleteResponseTransform: (
479483
content = part.text;
480484
contentBlocks.push({ type: 'text', text: part.text });
481485
}
486+
} else if (part.inlineData) {
487+
contentBlocks.push({
488+
type: 'image_url',
489+
image_url: {
490+
url: `data:${part.inlineData.mimeType};base64,${part.inlineData.data}`,
491+
},
492+
});
482493
}
483494
}
484495

@@ -674,6 +685,23 @@ export const GoogleChatCompleteStreamChunkTransform: (
674685
}
675686
}),
676687
};
688+
} else if (generation.content?.parts[0]?.inlineData) {
689+
const part = generation.content.parts[0];
690+
const contentBlocks = [
691+
{
692+
index: streamState.containsChainOfThoughtMessage ? 1 : 0,
693+
delta: {
694+
type: 'image_url',
695+
image_url: {
696+
url: `data:${part.inlineData?.mimeType};base64,${part.inlineData?.data}`,
697+
},
698+
},
699+
},
700+
];
701+
message = {
702+
role: 'assistant',
703+
content_blocks: contentBlocks,
704+
};
677705
}
678706
return {
679707
delta: message,

src/providers/google-vertex-ai/transformGenerationConfig.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,11 @@ export function transformGenerationConfig(params: Params) {
5757
thinkingConfig['thinking_budget'] = budget_tokens;
5858
generationConfig['thinking_config'] = thinkingConfig;
5959
}
60+
if (params.modalities) {
61+
generationConfig['responseModalities'] = params.modalities.map((modality) =>
62+
modality.toUpperCase()
63+
);
64+
}
6065

6166
return generationConfig;
6267
}

src/providers/google-vertex-ai/types.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ export interface GoogleResponseCandidate {
2020
text?: string;
2121
thought?: string; // for models like gemini-2.0-flash-thinking-exp refer: https://ai.google.dev/gemini-api/docs/thinking-mode#streaming_model_thinking
2222
functionCall?: GoogleGenerateFunctionCall;
23+
inlineData?: {
24+
mimeType: string;
25+
data: string;
26+
};
2327
}[];
2428
};
2529
logprobsResult?: {

0 commit comments

Comments
 (0)