Merge pull request #631 from AikidoSec/add-support-for-converse

AWS Bedrock: Add support for converse command

Author: hansott

library/sinks/AwsSDKVersion3.ts

@@ -4,30 +4,51 @@ import { wrapExport } from "../agent/hooks/wrapExport";
 import { Wrapper } from "../agent/Wrapper";
 import { isPlainObject } from "../helpers/isPlainObject";
 
-type Usage = {
+type InvokeUsage = {
   input_tokens?: number;
   output_tokens?: number;
 };
 
-function isUsage(usage: unknown): usage is Usage {
+function isUsage(usage: unknown): usage is InvokeUsage {
   return (
     isPlainObject(usage) &&
     typeof usage.input_tokens === "number" &&
     typeof usage.output_tokens === "number"
   );
 }
 
-type Response = {
+type InvokeResponse = {
   body?: Uint8Array;
 };
 
-function isResponse(response: unknown): response is Response {
+function isInvokeResponse(response: unknown): response is InvokeResponse {
+  return isPlainObject(response);
+}
+
+type ConverseUsage = {
+  inputTokens?: number;
+  outputTokens?: number;
+};
+
+function isConverseUsage(usage: unknown): usage is ConverseUsage {
+  return (
+    isPlainObject(usage) &&
+    typeof usage.inputTokens === "number" &&
+    typeof usage.outputTokens === "number"
+  );
+}
+
+type ConverseResponse = {
+  usage?: ConverseUsage;
+};
+
+function isConverseResponse(response: unknown): response is ConverseResponse {
   return isPlainObject(response);
 }
 
 export class AwsSDKVersion3 implements Wrapper {
-  private processResponse(response: unknown, agent: Agent) {
-    if (!isResponse(response)) {
+  private processInvokeModelResponse(response: unknown, agent: Agent) {
+    if (!isInvokeResponse(response)) {
       return;
     }
 
@@ -61,25 +82,66 @@ export class AwsSDKVersion3 implements Wrapper {
     }
   }
 
+  private processConverseResponse(
+    response: unknown,
+    command: unknown,
+    agent: Agent
+  ) {
+    // @ts-expect-error We don't know the type of command
+    if (!command || !command.input || !command.input.modelId) {
+      return;
+    }
+
+    if (!isConverseResponse(response)) {
+      return;
+    }
+
+    // @ts-expect-error We don't know the type of command
+    const modelId: string = command.input.modelId;
+
+    let inputTokens = 0;
+    let outputTokens = 0;
+
+    if (isConverseUsage(response.usage)) {
+      inputTokens = response.usage.inputTokens || 0;
+      outputTokens = response.usage.outputTokens || 0;
+    }
+
+    const aiStats = agent.getAIStatistics();
+    aiStats.onAICall({
+      provider: "bedrock",
+      model: modelId,
+      inputTokens: inputTokens,
+      outputTokens: outputTokens,
+    });
+  }
+
   wrap(hooks: Hooks) {
-    // Note: Converse command is not supported yet
     hooks
       .addPackage("@aws-sdk/client-bedrock-runtime")
       .withVersion("^3.0.0")
       .onRequire((exports, pkgInfo) => {
-        if (exports.BedrockRuntimeClient && exports.InvokeModelCommand) {
+        if (exports.BedrockRuntimeClient) {
           wrapExport(exports.BedrockRuntimeClient.prototype, "send", pkgInfo, {
             kind: "ai_op",
             modifyReturnValue: (args, returnValue, agent) => {
               if (args.length > 0) {
-                if (
-                  returnValue instanceof Promise &&
-                  args[0] instanceof exports.InvokeModelCommand
-                ) {
+                const command = args[0];
+                if (returnValue instanceof Promise) {
                   // Inspect the response after the promise resolves, it won't change the original promise
                   returnValue.then((response) => {
                     try {
-                      this.processResponse(response, agent);
+                      if (
+                        exports.InvokeModelCommand &&
+                        command instanceof exports.InvokeModelCommand
+                      ) {
+                        this.processInvokeModelResponse(response, agent);
+                      } else if (
+                        exports.ConverseCommand &&
+                        command instanceof exports.ConverseCommand
+                      ) {
+                        this.processConverseResponse(response, command, agent);
+                      }
                     } catch {
                       // If we don't catch these errors, it will result in an unhandled promise rejection!
                     }