Merge branch 'main' into feat/models-endpoint

Author: VisargD

plugins/azure/azure.test.ts

@@ -45,6 +45,29 @@ describe('Azure Plugins', () => {
         expect(result.error).toBeNull();
         expect(result.verdict).toBe(true);
         expect(result.transformed).toBe(true);
+      }, 10000);
+
+      it('should not redact anything if text has no PII', async () => {
+        const context = structuredClone(mockContext);
+        context.request.text = "hello, I'm a harmless string";
+        context.request.json = {
+          messages: [{ role: 'user', content: "hello, I'm a harmless string" }],
+        };
+        const result = await piiHandler(context, params, 'beforeRequestHook');
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(true);
+        expect(result.transformed).toBe(false);
+      });
+
+      it('should not redact anything if redact is false', async () => {
+        const result = await piiHandler(
+          mockContext,
+          { ...params, redact: false },
+          'beforeRequestHook'
+        );
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(false);
+        expect(result.transformed).toBe(false);
       });
 
       it('should handle API errors gracefully', async () => {

plugins/azure/contentSafety.ts

@@ -17,7 +17,7 @@ export const handler: PluginHandler<{
   context: PluginContext,
   parameters: PluginParameters<{ contentSafety: AzureCredentials }>,
   eventType: HookEventType,
-  options
+  pluginOptions?: Record<string, any>
 ) => {
   let error = null;
   let verdict = true;
@@ -69,8 +69,8 @@ export const handler: PluginHandler<{
   const { token, error: tokenError } = await getAccessToken(
     credentials as any,
     'contentSafety',
-    options,
-    options?.env
+    pluginOptions,
+    pluginOptions?.env
   );
 
   if (tokenError) {
@@ -110,14 +110,13 @@ export const handler: PluginHandler<{
   };
 
   const timeout = parameters.timeout || 5000;
+  const requestOptions: Record<string, any> = { headers };
+  if (agent) {
+    requestOptions.dispatcher = agent;
+  }
   let response;
   try {
-    response = await post(
-      url,
-      request,
-      { headers, dispatcher: agent },
-      timeout
-    );
+    response = await post(url, request, requestOptions, timeout);
   } catch (e) {
     return { error: e, verdict: true, data };
   }

plugins/azure/pii.ts

@@ -12,7 +12,7 @@ import { getAccessToken } from './utils';
 const redact = async (
   documents: any[],
   parameters: PluginParameters<{ pii: AzureCredentials }>,
-  options?: Record<string, any>
+  pluginOptions?: Record<string, any>
 ) => {
   const body = {
     kind: 'PiiEntityRecognition',
@@ -35,8 +35,8 @@ const redact = async (
   const { token, error: tokenError } = await getAccessToken(
     credentials as any,
     'pii',
-    options,
-    options?.env
+    pluginOptions,
+    pluginOptions?.env
   );
 
   const headers: Record<string, string> = {
@@ -66,20 +66,19 @@ const redact = async (
   }
 
   const timeout = parameters.timeout || 5000;
-  const response = await post(
-    url,
-    body,
-    { headers, dispatcher: agent },
-    timeout
-  );
+  const requestOptions: Record<string, any> = { headers };
+  if (agent) {
+    requestOptions.dispatcher = agent;
+  }
+  const response = await post(url, body, requestOptions, timeout);
   return response;
 };
 
 export const handler: PluginHandler<{ pii: AzureCredentials }> = async (
   context: PluginContext,
   parameters: PluginParameters<{ pii: AzureCredentials }>,
   eventType: HookEventType,
-  options?: Record<string, any>
+  pluginOptions?: Record<string, any>
 ) => {
   let error = null;
   let verdict = true;
@@ -134,9 +133,18 @@ export const handler: PluginHandler<{ pii: AzureCredentials }> = async (
   }));
 
   try {
-    const response = await redact(documents, parameters, options);
+    const response = await redact(documents, parameters, pluginOptions);
+    if (!response?.results?.documents) {
+      throw new Error('Invalid response from Azure PII API');
+    }
     data = response.results.documents;
-    if (parameters.redact) {
+    const containsPII =
+      data.length > 0 && data.some((doc: any) => doc.entities.length > 0);
+    if (containsPII) {
+      verdict = false;
+    }
+    if (parameters.redact && containsPII) {
+      verdict = true;
       const redactedData = (response.results.documents ?? []).map(
         (doc: any) => doc.redactedText
       );

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

@@ -49,6 +49,7 @@ import type {
 import {
   getMimeType,
   recursivelyDeleteUnsupportedParameters,
+  transformGeminiToolParameters,
   transformVertexLogprobs,
 } from './utils';
 
@@ -301,6 +302,11 @@ export const VertexGoogleChatCompleteConfig: ProviderConfig = {
           ) {
             tools.push(buildGoogleSearchRetrievalTool(tool));
           } else {
+            if (tool.function?.parameters) {
+              tool.function.parameters = transformGeminiToolParameters(
+                tool.function.parameters
+              );
+            }
             functionDeclarations.push(tool.function);
           }
         }

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

@@ -1,5 +1,8 @@
 import { Params } from '../../types/requestBody';
-import { derefer, recursivelyDeleteUnsupportedParameters } from './utils';
+import {
+  recursivelyDeleteUnsupportedParameters,
+  transformGeminiToolParameters,
+} from './utils';
 import { GoogleEmbedParams } from './embed';
 import { EmbedInstancesData } from './types';
 /**
@@ -39,20 +42,11 @@ export function transformGenerationConfig(params: Params) {
   }
   if (params?.response_format?.type === 'json_schema') {
     generationConfig['responseMimeType'] = 'application/json';
-    recursivelyDeleteUnsupportedParameters(
-      params?.response_format?.json_schema?.schema
-    );
     let schema =
       params?.response_format?.json_schema?.schema ??
       params?.response_format?.json_schema;
-    if (Object.keys(schema).includes('$defs')) {
-      schema = derefer(schema);
-      delete schema['$defs'];
-    }
-    if (Object.hasOwn(schema, '$schema')) {
-      delete schema['$schema'];
-    }
-    generationConfig['responseSchema'] = schema;
+    recursivelyDeleteUnsupportedParameters(schema);
+    generationConfig['responseSchema'] = transformGeminiToolParameters(schema);
   }
 
   if (params?.thinking) {