Fix Bandit + Only create API Gateway when needed + Add multimodal test (

#555)

* chore: Fix bandit script and review the flagged issues

* chore: Review/update code analysis feedbacks

* feat: Add llm handlers logs to the dashboard + format

* test: Add multi modal test

* feat: Remove API Gateway when sagemaker multi modal is not used.

.github/workflows/build.yaml

@@ -29,7 +29,7 @@ jobs:
         run: |
           pip install -r pytest_requirements.txt
           flake8 .
-          bandit -r .
+          bandit -c bandit.yaml -r .
           pip-audit -r pytest_requirements.txt
           pip-audit -r lib/shared/web-crawler-batch-job/requirements.txt
           pip-audit -r lib/shared/file-import-batch-job/requirements.txt

bandit.yaml

@@ -1 +1 @@
-exclude_dirs: ['tests', 'cdk.out', 'dist', 'node_modules', 'lib/user-interface/react-app']
+exclude_dirs: ['tests', 'cdk.out', 'dist', 'node_modules', 'lib/user-interface/react-app']

docs/guide/deploy.md

@@ -165,12 +165,16 @@ GenAIChatBotStack.ApiKeysSecretNameXXXX = ApiKeysSecretName-xxxxxx
 **Step 11.** Login with the user created in **Step 8** and follow the instructions.
 
 **Step 12.** (Optional) Run the integration tests
-The tests require to be authenticated against your AWS Account because it will create cognito users. In addition, the tests will use `anthropic.claude-instant-v1` (Claude Instant) and `amazon.titan-embed-text-v1` (Titan Embeddings G1 - Text) which need to be enabled in Bedrock.
+The tests require to be authenticated against your AWS Account because it will create cognito users. In addition, the tests will use `anthropic.claude-instant-v1` (Claude Instant), `anthropic.claude-3-haiku-20240307-v1:0` (Claude 3 Haiku) and `amazon.titan-embed-text-v1` (Titan Embeddings G1 - Text) which need to be enabled in Bedrock.
 
 To run the tests (Replace the url with the one you used in the steps above)
 ```bash
 REACT_APP_URL=https://dxxxxxxxxxxxxx.cloudfront.net pytest integtests/ --ignore integtests/user_interface -n 3 --dist=loadfile 
 ```
+To run the UI tests, you will fist need to download and run [geckodriver](https://github.com/mozilla/geckodriver)
+```bash
+REACT_APP_URL=https://dxxxxxxxxxxxxx.cloudfront.net pytest integtests/user_interface 
+```
 
 ## Monitoring
 

integtests/chatbot-api/kendra_workspace_test.py

@@ -109,6 +109,8 @@ def test_semantic_search(client: AppSyncClient):
 
 
 def test_query_llm(client, default_model, default_provider):
+    if pytest.skip_flag == True:
+        pytest.skip("Kendra is not enabled.")
     session_id = str(uuid.uuid4())
     request = {
         "action": "run",

integtests/chatbot-api/multi_modal_test.py

@@ -0,0 +1,61 @@
+import json
+import os
+import time
+import uuid
+import boto3
+from pathlib import Path
+
+import pytest
+
+
+def test_multi_modal(
+    client, config, cognito_credentials, default_multimodal_model, default_provider
+):
+    bucket = config.get("Storage").get("AWSS3").get("bucket")
+    s3 = boto3.resource(
+        "s3",
+        # Use identity pool credentials to verify it owrks
+        aws_access_key_id=cognito_credentials.aws_access_key,
+        aws_secret_access_key=cognito_credentials.aws_secret_key,
+        aws_session_token=cognito_credentials.aws_token,
+    )
+    key = "INTEG_TEST" + str(uuid.uuid4()) + ".jpeg"
+    object = s3.Object(bucket, "public/" + key)
+    wrong_object = s3.Object(bucket, "private/notallowed/1.jpg")
+    current_dir = os.path.dirname(os.path.realpath(__file__))
+    object.put(Body=Path(current_dir + "/resources/powered-by-aws.png").read_bytes())
+    with pytest.raises(Exception, match="AccessDenied"):
+        wrong_object.put(
+            Body=Path(current_dir + "/resources/powered-by-aws.png").read_bytes()
+        )
+
+    session_id = str(uuid.uuid4())
+
+    request = {
+        "action": "run",
+        "modelInterface": "multimodal",
+        "data": {
+            "mode": "chain",
+            "text": "What is this image?",
+            "files": [{"key": key, "provider": "s3"}],
+            "modelName": default_multimodal_model,
+            "provider": default_provider,
+            "sessionId": session_id,
+        },
+    }
+
+    client.send_query(json.dumps(request))
+
+    content = None
+    retries = 0
+    while retries < 30:
+        time.sleep(1)
+        retries += 1
+        session = client.get_session(session_id)
+        if session != None and len(session.get("history")) == 2:
+            content = session.get("history")[1].get("content").lower()
+            break
+
+    assert "powered by" in content
+    client.delete_session(session_id)
+    object.delete()

integtests/chatbot-api/opensearch_workspace_test.py

@@ -108,6 +108,8 @@ def test_search_document(client: AppSyncClient):
 
 
 def test_query_llm(client, default_model, default_provider):
+    if pytest.skip_flag == True:
+        pytest.skip("Open search is not enabled.")
     session_id = str(uuid.uuid4())
     request = {
         "action": "run",

integtests/clients/cognito_client.py

@@ -10,20 +10,37 @@ class Credentials(BaseModel):
     id_token: str
     email: str
     password: str
+    aws_access_key: str
+    aws_secret_key: str
+    aws_token: str
+
+    def __repr__(self):
+        return "Credentials(********)"
+
+    def __str___(self):
+        return "*******"
 
 
 class CognitoClient:
-    def __init__(self, region: str, user_pool_id: str, client_id: str) -> None:
+    def __init__(
+        self, region: str, user_pool_id: str, client_id: str, identity_pool_id: str
+    ) -> None:
         self.user_pool_id = user_pool_id
+        self.identity_pool_id = identity_pool_id
         self.client_id = client_id
+        self.region = region
         self.cognito_idp_client = boto3.client("cognito-idp", region_name=region)
+        self.cognito_identity_client = boto3.client(
+            "cognito-identity", region_name=region
+        )
 
     def get_credentials(self, email: str) -> Credentials:
         try:
             self.cognito_idp_client.admin_get_user(
                 UserPoolId=self.user_pool_id,
                 Username=email,
             )
+
         except self.cognito_idp_client.exceptions.UserNotFoundException:
             self.cognito_idp_client.admin_create_user(
                 UserPoolId=self.user_pool_id,
@@ -50,18 +67,37 @@ def get_credentials(self, email: str) -> Credentials:
             AuthParameters={"USERNAME": email, "PASSWORD": password},
         )
 
+        login_key = "cognito-idp." + self.region + ".amazonaws.com/" + self.user_pool_id
+        identity_response = self.cognito_identity_client.get_id(
+            IdentityPoolId=self.identity_pool_id,
+            Logins={login_key: response["AuthenticationResult"]["IdToken"]},
+        )
+
+        aws_credentials_respose = (
+            self.cognito_identity_client.get_credentials_for_identity(
+                IdentityId=identity_response["IdentityId"],
+                Logins={login_key: response["AuthenticationResult"]["IdToken"]},
+            )
+        )
+
         return Credentials(
             **{
                 "id_token": response["AuthenticationResult"]["IdToken"],
                 "email": email,
                 "password": password,
+                # Credential with limited permissions (upload images for multi modal)
+                "aws_access_key": aws_credentials_respose["Credentials"]["AccessKeyId"],
+                "aws_secret_key": aws_credentials_respose["Credentials"]["SecretKey"],
+                "aws_token": aws_credentials_respose["Credentials"]["SessionToken"],
             }
         )
 
     def get_password(self):
         return "".join(
-            random.choices(string.ascii_uppercase, k=10)
-            + random.choices(string.ascii_lowercase, k=10)
-            + random.choices(string.digits, k=5)
-            + random.choices(string.punctuation, k=3)
+            random.choices(
+                string.ascii_uppercase, k=10
+            )  # NOSONAR Only used for testing. Temporary password
+            + random.choices(string.ascii_lowercase, k=10)  # NOSONAR
+            + random.choices(string.digits, k=5)  # NOSONAR
+            + random.choices(string.punctuation, k=3)  # NOSONAR
         )

integtests/conftest.py

@@ -19,9 +19,13 @@ def cognito_credentials(config, worker_id) -> Credentials:
     user_pool_id = config.get("aws_user_pools_id")
     region = config.get("aws_cognito_region")
     user_pool_client_id = config.get("aws_user_pools_web_client_id")
+    identity_pool_id = config.get("aws_cognito_identity_pool_id")
 
     cognito = CognitoClient(
-        region=region, user_pool_id=user_pool_id, client_id=user_pool_client_id
+        region=region,
+        user_pool_id=user_pool_id,
+        client_id=user_pool_client_id,
+        identity_pool_id=identity_pool_id,
     )
     email = "[email protected]" + worker_id
 
@@ -44,6 +48,11 @@ def default_embed_model():
     return "amazon.titan-embed-text-v1"
 
 
+@pytest.fixture(scope="session")
+def default_multimodal_model():
+    return "anthropic.claude-3-haiku-20240307-v1:0"
+
+
 @pytest.fixture(scope="session")
 def default_provider():
     return "bedrock"
@@ -62,7 +71,7 @@ def react_url():
     return os.environ["REACT_APP_URL"]
 
 
-@pytest.fixture(scope="class")
+@pytest.fixture(scope="module")
 def selenium_driver(react_url):
     options = webdriver.FirefoxOptions()
     if os.getenv("HEADLESS"):

integtests/user_interface/react_app/test_login.py

@@ -14,6 +14,15 @@ def test_login(selenium_driver, cognito_credentials):
 def test_invalid_credentials(selenium_driver):
     page = LoginPage(selenium_driver)
     page.login(
-        Credentials(**{"id_token": "", "email": "invalid", "password": "invalid"})
+        Credentials(
+            **{
+                "id_token": "",
+                "email": "invalid",
+                "password": "invalid",
+                "aws_access_key": "",
+                "aws_secret_key": "",
+                "aws_token": "",
+            }
+        )  # NOSONAR
     )
     assert page.get_error() != None

lib/aws-genai-llm-chatbot-stack.ts

@@ -13,6 +13,7 @@ import { IdeficsInterface } from "./model-interfaces/idefics";
 import * as subscriptions from "aws-cdk-lib/aws-sns-subscriptions";
 import * as sns from "aws-cdk-lib/aws-sns";
 import * as iam from "aws-cdk-lib/aws-iam";
+import * as lambda from "aws-cdk-lib/aws-lambda";
 import * as cr from "aws-cdk-lib/custom-resources";
 import { NagSuppressions } from "cdk-nag";
 import { LogGroup } from "aws-cdk-lib/aws-logs";
@@ -107,19 +108,21 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
     // IDEFICS Interface Construct
     // This is the model interface receiving messages from the websocket interface via the message topic
     // and interacting with IDEFICS visual language models
-    models.models.filter(
-      (model) => model.interface === ModelInterface.MultiModal
+    const ideficsModels = models.models.filter(
+      (model) =>
+        model.interface === ModelInterface.MultiModal &&
+        model.name.toLocaleLowerCase().includes("idefics")
     );
 
     // check if any deployed model requires idefics interface
-
     const ideficsInterface = new IdeficsInterface(this, "IdeficsInterface", {
       shared,
       config: props.config,
       messagesTopic: chatBotApi.messagesTopic,
       sessionsTable: chatBotApi.sessionsTable,
       byUserIdIndex: chatBotApi.byUserIdIndex,
       chatbotFilesBucket: chatBotApi.filesBucket,
+      createPrivateGateway: ideficsModels.length > 0,
     });
 
     // Route all incoming messages targeted to idefics to the idefics model interface queue
@@ -228,6 +231,18 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
           "/aws/lambda/" + r.functionName
         );
       }),
+      llmRequestHandlersLogGroups: [
+        ideficsInterface.requestHandler,
+        langchainInterface?.requestHandler,
+      ]
+        .filter((i) => i)
+        .map((r) => {
+          return LogGroup.fromLogGroupName(
+            monitoringStack,
+            "Log" + (r as lambda.Function).node.id,
+            "/aws/lambda/" + (r as lambda.Function).functionName
+          );
+        }),
       cognito: {
         userPoolId: authentication.userPool.userPoolId,
         clientId: authentication.userPoolClient.userPoolClientId,
@@ -294,16 +309,20 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
         `/${this.stackName}/ChatBotApi/Realtime/Resolvers/lambda-resolver/ServiceRole/Resource`,
         `/${this.stackName}/ChatBotApi/Realtime/Resolvers/outgoing-message-handler/ServiceRole/Resource`,
         `/${this.stackName}/ChatBotApi/Realtime/Resolvers/outgoing-message-handler/ServiceRole/DefaultPolicy/Resource`,
-        `/${this.stackName}/IdeficsInterface/IdeficsInterfaceRequestHandler/ServiceRole/DefaultPolicy/Resource`,
-        `/${this.stackName}/IdeficsInterface/IdeficsInterfaceRequestHandler/ServiceRole/Resource`,
-        `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/CloudWatchRole/Resource`,
-        `/${this.stackName}/IdeficsInterface/S3IntegrationRole/DefaultPolicy/Resource`,
+        `/${this.stackName}/IdeficsInterface/MultiModalInterfaceRequestHandler/ServiceRole/DefaultPolicy/Resource`,
+        `/${this.stackName}/IdeficsInterface/MultiModalInterfaceRequestHandler/ServiceRole/Resource`,
         ...(langchainInterface
           ? [
               `/${this.stackName}/LangchainInterface/RequestHandler/ServiceRole/Resource`,
               `/${this.stackName}/LangchainInterface/RequestHandler/ServiceRole/DefaultPolicy/Resource`,
             ]
           : []),
+        ...(ideficsModels.length > 0
+          ? [
+              `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/CloudWatchRole/Resource`,
+              `/${this.stackName}/IdeficsInterface/S3IntegrationRole/DefaultPolicy/Resource`,
+            ]
+          : []),
       ],
       [
         {
@@ -316,27 +335,30 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
         },
       ]
     );
-    NagSuppressions.addResourceSuppressionsByPath(
-      this,
-      `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/DeploymentStage.prod/Resource`,
-      [
-        {
-          id: "AwsSolutions-APIG3",
-          reason: "WAF not required due to configured Cognito auth.",
-        },
-      ]
-    );
-    NagSuppressions.addResourceSuppressionsByPath(
-      this,
-      [
-        `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/Default/{object}/ANY/Resource`,
-        `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/Default/{object}/ANY/Resource`,
-      ],
-      [
-        { id: "AwsSolutions-APIG4", reason: "Private API within a VPC." },
-        { id: "AwsSolutions-COG4", reason: "Private API within a VPC." },
-      ]
-    );
+
+    if (ideficsModels.length > 0) {
+      NagSuppressions.addResourceSuppressionsByPath(
+        this,
+        `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/DeploymentStage.prod/Resource`,
+        [
+          {
+            id: "AwsSolutions-APIG3",
+            reason: "WAF not required due to configured Cognito auth.",
+          },
+        ]
+      );
+      NagSuppressions.addResourceSuppressionsByPath(
+        this,
+        [
+          `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/Default/{object}/ANY/Resource`,
+          `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/Default/{object}/ANY/Resource`,
+        ],
+        [
+          { id: "AwsSolutions-APIG4", reason: "Private API within a VPC." },
+          { id: "AwsSolutions-COG4", reason: "Private API within a VPC." },
+        ]
+      );
+    }
 
     // RAG configuration
     if (props.config.rag.enabled) {