@genkit-ai/google-cloud/lib/model-armor.mjs.map

Genkit AI framework plugin for Google Cloud Platform including Firestore trace/state store and deployment helpers for Cloud Functions for Firebase.

{"version":3,"sources":["../src/model-armor.ts"],"sourcesContent":["/**\n * Copyright 2025 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport { ModelArmorClient, protos } from '@google-cloud/modelarmor';\nimport { GenkitError } from 'genkit';\nimport {\n  GenerateRequest,\n  GenerateResponseData,\n  MessageData,\n  ModelMiddleware,\n  Part,\n} from 'genkit/model';\nimport { runInNewSpan } from 'genkit/tracing';\n\nexport interface ModelArmorOptions {\n  templateName: string;\n  client?: ModelArmorClient;\n  /**\n   * Options for the Model Armor client (e.g. apiEndpoint).\n   */\n  clientOptions?: ConstructorParameters<typeof ModelArmorClient>[0];\n  /**\n   * What to sanitize. Defaults to 'all'.\n   */\n  protectionTarget?: 'all' | 'userPrompt' | 'modelResponse';\n  /**\n   * Whether to block on SDP match even if the content was successfully de-identified.\n   * Defaults to false (lenient).\n   */\n  strictSdpEnforcement?: boolean;\n  /**\n   * List of filters to enforce. If not specified, all filters are enforced.\n   * Possible values: 'rai', 'pi_and_jailbreak', 'malicious_uris', 'csam', 'sdp'.\n   */\n  filters?: (\n    | 'rai'\n    | 'pi_and_jailbreak'\n    | 'malicious_uris'\n    | 'csam'\n    | 'sdp'\n    | (string & {})\n  )[];\n  /**\n   * Whether to apply the de-identification results to the content.\n   * - If true, the default logic (replace text, preserve structure) is used.\n   * - If false, no changes are applied.\n   * - If a function, it is called with the messages and SDP result, and should return the new messages.\n   *\n   * Defaults to false.\n   */\n  applyDeidentificationResults?:\n    | boolean\n    | ((data: {\n        messages: MessageData[];\n        sdpResult: protos.google.cloud.modelarmor.v1.ISdpFilterResult;\n      }) => MessageData[] | undefined);\n}\n\nfunction extractText(parts: Part[]): string {\n  return parts.map((p) => p.text || '').join('');\n}\n\n/**\n * If SDP (Sensitive Data Protection) filter returns sanitized data,\n * we swap out the data with sanitized data.\n */\nfunction applySdp(\n  messages: MessageData[],\n  targetIndex: number,\n  result: protos.google.cloud.modelarmor.v1.ISanitizationResult,\n  options: ModelArmorOptions\n): { sdpApplied: boolean; messages: MessageData[] } {\n  const sdpFilterResult = result.filterResults?.['sdp']?.sdpFilterResult;\n\n  if (!sdpFilterResult) {\n    return { sdpApplied: false, messages };\n  }\n\n  // If user provided applyDeidentificationResults, we use it to apply\n  // the deidentification results.\n  if (typeof options.applyDeidentificationResults === 'function') {\n    const newMessages = options.applyDeidentificationResults({\n      messages,\n      sdpResult: sdpFilterResult,\n    });\n    if (!newMessages) {\n      return { sdpApplied: false, messages };\n    }\n    const sdpApplied = !!sdpFilterResult.deidentifyResult?.data?.text;\n    return { sdpApplied, messages: newMessages };\n  }\n\n  // if applyDeidentificationResults is set to true, we use the default/basic\n  // approach to apply the results.\n  if (options.applyDeidentificationResults === true) {\n    const deidentifyResult = sdpFilterResult.deidentifyResult;\n    if (deidentifyResult && deidentifyResult.data?.text) {\n      const targetMessage = messages[targetIndex];\n      const nonTextParts = targetMessage.content.filter((p) => !p.text);\n      const newContent = [\n        ...nonTextParts,\n        { text: deidentifyResult.data.text },\n      ];\n      const newMessages = [...messages];\n      newMessages[targetIndex] = { ...targetMessage, content: newContent };\n      return {\n        sdpApplied: true,\n        messages: newMessages,\n      };\n    }\n  }\n\n  return { sdpApplied: false, messages };\n}\n\nfunction shouldBlock(\n  result: protos.google.cloud.modelarmor.v1.ISanitizationResult,\n  options: ModelArmorOptions,\n  sdpApplied: boolean\n): boolean {\n  if (result.filterMatchState !== 'MATCH_FOUND') {\n    return false;\n  }\n  // Check if we should block.\n  // If strict SDP enforcement is enabled and SDP was applied, we must block.\n  if (options.strictSdpEnforcement && sdpApplied) {\n    return true;\n  }\n  // Otherwise, check if any active filter matched.\n  if (result.filterResults) {\n    for (const [key, filterResult] of Object.entries(result.filterResults)) {\n      if (options.filters && !options.filters.includes(key)) continue;\n      if (key === 'sdp' && sdpApplied) continue;\n\n      // Look for matchState in the nested object\n      // e.g. filterResult.raiFilterResult.matchState\n      const nestedResult = Object.values(filterResult)[0];\n      if (nestedResult?.matchState === 'MATCH_FOUND') {\n        return true;\n      }\n    }\n  }\n  return false;\n}\n\nasync function sanitizeUserPrompt(\n  req: GenerateRequest,\n  client: ModelArmorClient,\n  options: ModelArmorOptions\n) {\n  let targetMessageIndex = -1;\n  // Find the last user message to sanitize\n  for (let i = req.messages.length - 1; i >= 0; i--) {\n    if (req.messages[i].role === 'user') {\n      targetMessageIndex = i;\n      break;\n    }\n  }\n\n  if (targetMessageIndex !== -1) {\n    const userMessage = req.messages[targetMessageIndex];\n    const promptText = extractText(userMessage.content);\n\n    if (promptText) {\n      await runInNewSpan(\n        { metadata: { name: 'sanitizeUserPrompt' } },\n        async (meta) => {\n          meta.input = {\n            name: options.templateName,\n            userPromptData: {\n              text: promptText,\n            },\n          };\n          const [response] = await client.sanitizeUserPrompt({\n            name: options.templateName,\n            userPromptData: {\n              text: promptText,\n            },\n          });\n          meta.output = response;\n\n          if (response.sanitizationResult) {\n            const result = response.sanitizationResult;\n            const { sdpApplied, messages: modifiedMessages } = applySdp(\n              req.messages,\n              targetMessageIndex,\n              result,\n              options\n            );\n\n            if (\n              sdpApplied ||\n              typeof options.applyDeidentificationResults === 'function'\n            ) {\n              req.messages = modifiedMessages;\n            }\n\n            if (shouldBlock(result, options, sdpApplied)) {\n              throw new GenkitError({\n                status: 'PERMISSION_DENIED',\n                message: 'Model Armor blocked user prompt.',\n                detail: result,\n              });\n            }\n          }\n        }\n      );\n    }\n  }\n}\n\nasync function sanitizeModelResponse(\n  response: GenerateResponseData,\n  client: ModelArmorClient,\n  options: ModelArmorOptions\n) {\n  const usingMessageProp = !!response.message;\n  const candidates = response.message\n    ? [{ index: 0, message: response.message, finishReason: 'stop' }]\n    : response.candidates || [];\n\n  for (const candidate of candidates) {\n    const modelText = extractText(candidate.message.content);\n\n    if (modelText) {\n      await runInNewSpan(\n        { metadata: { name: 'sanitizeModelResponse' } },\n        async (meta) => {\n          meta.input = {\n            name: options.templateName,\n            modelResponseData: {\n              text: modelText,\n            },\n          };\n          const [apiResponse] = await client.sanitizeModelResponse({\n            name: options.templateName,\n            modelResponseData: {\n              text: modelText,\n            },\n          });\n          meta.output = apiResponse;\n\n          if (apiResponse.sanitizationResult) {\n            const result = apiResponse.sanitizationResult;\n            const { sdpApplied, messages: modifiedMessages } = applySdp(\n              [candidate.message],\n              0,\n              result,\n              options\n            );\n\n            if (\n              sdpApplied ||\n              typeof options.applyDeidentificationResults === 'function'\n            ) {\n              candidate.message = modifiedMessages[0];\n            }\n\n            if (shouldBlock(result, options, sdpApplied)) {\n              throw new GenkitError({\n                status: 'PERMISSION_DENIED',\n                message: 'Model Armor blocked model response.',\n                detail: result,\n              });\n            }\n          }\n        }\n      );\n    }\n  }\n\n  if (usingMessageProp && candidates.length > 0) {\n    response.message = candidates[0].message;\n  }\n}\n\n/**\n * Model Middleware that uses Google Cloud Model Armor to sanitize user prompts and model responses.\n */\nexport function modelArmor(options: ModelArmorOptions): ModelMiddleware {\n  const client = options.client || new ModelArmorClient(options.clientOptions);\n  const protectionTarget = options.protectionTarget ?? 'all';\n  const protectUserPrompt =\n    protectionTarget === 'all' || protectionTarget === 'userPrompt';\n  const protectModelResponse =\n    protectionTarget === 'all' || protectionTarget === 'modelResponse';\n\n  return async (req, next) => {\n    // 1. Sanitize User Prompt\n    if (protectUserPrompt) {\n      await sanitizeUserPrompt(req, client, options);\n    }\n\n    // 2. Call Model\n    const response = await next(req);\n\n    // 3. Sanitize Model Response\n    if (protectModelResponse) {\n      await sanitizeModelResponse(response, client, options);\n    }\n\n    return response;\n  };\n}\n"],"mappings":"AAgBA,SAAS,wBAAgC;AACzC,SAAS,mBAAmB;AAQ5B,SAAS,oBAAoB;AA8C7B,SAAS,YAAY,OAAuB;AAC1C,SAAO,MAAM,IAAI,CAAC,MAAM,EAAE,QAAQ,EAAE,EAAE,KAAK,EAAE;AAC/C;AAMA,SAAS,SACP,UACA,aACA,QACA,SACkD;AAClD,QAAM,kBAAkB,OAAO,gBAAgB,KAAK,GAAG;AAEvD,MAAI,CAAC,iBAAiB;AACpB,WAAO,EAAE,YAAY,OAAO,SAAS;AAAA,EACvC;AAIA,MAAI,OAAO,QAAQ,iCAAiC,YAAY;AAC9D,UAAM,cAAc,QAAQ,6BAA6B;AAAA,MACvD;AAAA,MACA,WAAW;AAAA,IACb,CAAC;AACD,QAAI,CAAC,aAAa;AAChB,aAAO,EAAE,YAAY,OAAO,SAAS;AAAA,IACvC;AACA,UAAM,aAAa,CAAC,CAAC,gBAAgB,kBAAkB,MAAM;AAC7D,WAAO,EAAE,YAAY,UAAU,YAAY;AAAA,EAC7C;AAIA,MAAI,QAAQ,iCAAiC,MAAM;AACjD,UAAM,mBAAmB,gBAAgB;AACzC,QAAI,oBAAoB,iBAAiB,MAAM,MAAM;AACnD,YAAM,gBAAgB,SAAS,WAAW;AAC1C,YAAM,eAAe,cAAc,QAAQ,OAAO,CAAC,MAAM,CAAC,EAAE,IAAI;AAChE,YAAM,aAAa;AAAA,QACjB,GAAG;AAAA,QACH,EAAE,MAAM,iBAAiB,KAAK,KAAK;AAAA,MACrC;AACA,YAAM,cAAc,CAAC,GAAG,QAAQ;AAChC,kBAAY,WAAW,IAAI,EAAE,GAAG,eAAe,SAAS,WAAW;AACnE,aAAO;AAAA,QACL,YAAY;AAAA,QACZ,UAAU;AAAA,MACZ;AAAA,IACF;AAAA,EACF;AAEA,SAAO,EAAE,YAAY,OAAO,SAAS;AACvC;AAEA,SAAS,YACP,QACA,SACA,YACS;AACT,MAAI,OAAO,qBAAqB,eAAe;AAC7C,WAAO;AAAA,EACT;AAGA,MAAI,QAAQ,wBAAwB,YAAY;AAC9C,WAAO;AAAA,EACT;AAEA,MAAI,OAAO,eAAe;AACxB,eAAW,CAAC,KAAK,YAAY,KAAK,OAAO,QAAQ,OAAO,aAAa,GAAG;AACtE,UAAI,QAAQ,WAAW,CAAC,QAAQ,QAAQ,SAAS,GAAG,EAAG;AACvD,UAAI,QAAQ,SAAS,WAAY;AAIjC,YAAM,eAAe,OAAO,OAAO,YAAY,EAAE,CAAC;AAClD,UAAI,cAAc,eAAe,eAAe;AAC9C,eAAO;AAAA,MACT;AAAA,IACF;AAAA,EACF;AACA,SAAO;AACT;AAEA,eAAe,mBACb,KACA,QACA,SACA;AACA,MAAI,qBAAqB;AAEzB,WAAS,IAAI,IAAI,SAAS,SAAS,GAAG,KAAK,GAAG,KAAK;AACjD,QAAI,IAAI,SAAS,CAAC,EAAE,SAAS,QAAQ;AACnC,2BAAqB;AACrB;AAAA,IACF;AAAA,EACF;AAEA,MAAI,uBAAuB,IAAI;AAC7B,UAAM,cAAc,IAAI,SAAS,kBAAkB;AACnD,UAAM,aAAa,YAAY,YAAY,OAAO;AAElD,QAAI,YAAY;AACd,YAAM;AAAA,QACJ,EAAE,UAAU,EAAE,MAAM,qBAAqB,EAAE;AAAA,QAC3C,OAAO,SAAS;AACd,eAAK,QAAQ;AAAA,YACX,MAAM,QAAQ;AAAA,YACd,gBAAgB;AAAA,cACd,MAAM;AAAA,YACR;AAAA,UACF;AACA,gBAAM,CAAC,QAAQ,IAAI,MAAM,OAAO,mBAAmB;AAAA,YACjD,MAAM,QAAQ;AAAA,YACd,gBAAgB;AAAA,cACd,MAAM;AAAA,YACR;AAAA,UACF,CAAC;AACD,eAAK,SAAS;AAEd,cAAI,SAAS,oBAAoB;AAC/B,kBAAM,SAAS,SAAS;AACxB,kBAAM,EAAE,YAAY,UAAU,iBAAiB,IAAI;AAAA,cACjD,IAAI;AAAA,cACJ;AAAA,cACA;AAAA,cACA;AAAA,YACF;AAEA,gBACE,cACA,OAAO,QAAQ,iCAAiC,YAChD;AACA,kBAAI,WAAW;AAAA,YACjB;AAEA,gBAAI,YAAY,QAAQ,SAAS,UAAU,GAAG;AAC5C,oBAAM,IAAI,YAAY;AAAA,gBACpB,QAAQ;AAAA,gBACR,SAAS;AAAA,gBACT,QAAQ;AAAA,cACV,CAAC;AAAA,YACH;AAAA,UACF;AAAA,QACF;AAAA,MACF;AAAA,IACF;AAAA,EACF;AACF;AAEA,eAAe,sBACb,UACA,QACA,SACA;AACA,QAAM,mBAAmB,CAAC,CAAC,SAAS;AACpC,QAAM,aAAa,SAAS,UACxB,CAAC,EAAE,OAAO,GAAG,SAAS,SAAS,SAAS,cAAc,OAAO,CAAC,IAC9D,SAAS,cAAc,CAAC;AAE5B,aAAW,aAAa,YAAY;AAClC,UAAM,YAAY,YAAY,UAAU,QAAQ,OAAO;AAEvD,QAAI,WAAW;AACb,YAAM;AAAA,QACJ,EAAE,UAAU,EAAE,MAAM,wBAAwB,EAAE;AAAA,QAC9C,OAAO,SAAS;AACd,eAAK,QAAQ;AAAA,YACX,MAAM,QAAQ;AAAA,YACd,mBAAmB;AAAA,cACjB,MAAM;AAAA,YACR;AAAA,UACF;AACA,gBAAM,CAAC,WAAW,IAAI,MAAM,OAAO,sBAAsB;AAAA,YACvD,MAAM,QAAQ;AAAA,YACd,mBAAmB;AAAA,cACjB,MAAM;AAAA,YACR;AAAA,UACF,CAAC;AACD,eAAK,SAAS;AAEd,cAAI,YAAY,oBAAoB;AAClC,kBAAM,SAAS,YAAY;AAC3B,kBAAM,EAAE,YAAY,UAAU,iBAAiB,IAAI;AAAA,cACjD,CAAC,UAAU,OAAO;AAAA,cAClB;AAAA,cACA;AAAA,cACA;AAAA,YACF;AAEA,gBACE,cACA,OAAO,QAAQ,iCAAiC,YAChD;AACA,wBAAU,UAAU,iBAAiB,CAAC;AAAA,YACxC;AAEA,gBAAI,YAAY,QAAQ,SAAS,UAAU,GAAG;AAC5C,oBAAM,IAAI,YAAY;AAAA,gBACpB,QAAQ;AAAA,gBACR,SAAS;AAAA,gBACT,QAAQ;AAAA,cACV,CAAC;AAAA,YACH;AAAA,UACF;AAAA,QACF;AAAA,MACF;AAAA,IACF;AAAA,EACF;AAEA,MAAI,oBAAoB,WAAW,SAAS,GAAG;AAC7C,aAAS,UAAU,WAAW,CAAC,EAAE;AAAA,EACnC;AACF;AAKO,SAAS,WAAW,SAA6C;AACtE,QAAM,SAAS,QAAQ,UAAU,IAAI,iBAAiB,QAAQ,aAAa;AAC3E,QAAM,mBAAmB,QAAQ,oBAAoB;AACrD,QAAM,oBACJ,qBAAqB,SAAS,qBAAqB;AACrD,QAAM,uBACJ,qBAAqB,SAAS,qBAAqB;AAErD,SAAO,OAAO,KAAK,SAAS;AAE1B,QAAI,mBAAmB;AACrB,YAAM,mBAAmB,KAAK,QAAQ,OAAO;AAAA,IAC/C;AAGA,UAAM,WAAW,MAAM,KAAK,GAAG;AAG/B,QAAI,sBAAsB;AACxB,YAAM,sBAAsB,UAAU,QAAQ,OAAO;AAAA,IACvD;AAEA,WAAO;AAAA,EACT;AACF;","names":[]}