import { ModelArmorClient, protos } from '@google-cloud/modelarmor';
import { MessageData, ModelMiddleware } from 'genkit/model';

/**
 * @license
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
/**
 * Model Armor middleware for Genkit.
 *
 * @module model-armor
 */

interface ModelArmorOptions {
    templateName: string;
    client?: ModelArmorClient;
    /**
     * Options for the Model Armor client (e.g. apiEndpoint).
     */
    clientOptions?: ConstructorParameters<typeof ModelArmorClient>[0];
    /**
     * What to sanitize. Defaults to 'all'.
     */
    protectionTarget?: 'all' | 'userPrompt' | 'modelResponse';
    /**
     * Whether to block on SDP match even if the content was successfully de-identified.
     * Defaults to false (lenient).
     */
    strictSdpEnforcement?: boolean;
    /**
     * List of filters to enforce. If not specified, all filters are enforced.
     * Possible values: 'rai', 'pi_and_jailbreak', 'malicious_uris', 'csam', 'sdp'.
     */
    filters?: ('rai' | 'pi_and_jailbreak' | 'malicious_uris' | 'csam' | 'sdp' | (string & {}))[];
    /**
     * Whether to apply the de-identification results to the content.
     * - If true, the default logic (replace text, preserve structure) is used.
     * - If false, no changes are applied.
     * - If a function, it is called with the messages and SDP result, and should return the new messages.
     *
     * Defaults to false.
     */
    applyDeidentificationResults?: boolean | ((data: {
        messages: MessageData[];
        sdpResult: protos.google.cloud.modelarmor.v1.ISdpFilterResult;
    }) => MessageData[] | undefined);
}
/**
 * Model Middleware that uses Google Cloud Model Armor to sanitize user prompts and model responses.
 */
declare function modelArmor(options: ModelArmorOptions): ModelMiddleware;

export { type ModelArmorOptions, modelArmor };