From 08f899beb1c2422c670fe3993d7c5b560e3d44da Mon Sep 17 00:00:00 2001 From: Charles Marion Date: Thu, 19 Sep 2024 10:00:31 -0500 Subject: [PATCH] feat: Replace the use of identity pool by s3 signed urls. (#568) --- integtests/chatbot-api/multi_modal_test.py | 46 ++-- integtests/clients/appsync_client.py | 4 + integtests/clients/cognito_client.py | 20 -- integtests/security/unauthorized_test.py | 2 + lib/authentication/index.ts | 22 -- lib/aws-genai-llm-chatbot-stack.ts | 5 +- .../functions/api-handler/routes/documents.py | 35 ++- .../functions/api-handler/routes/sessions.py | 20 ++ lib/chatbot-api/index.ts | 1 + lib/chatbot-api/rest-api.ts | 3 + lib/chatbot-api/schema/schema.graphql | 4 +- .../request-handler/adapters/claude.py | 21 +- .../request-handler/adapters/idefics.py | 9 +- .../functions/request-handler/index.py | 1 + lib/model-interfaces/idefics/index.ts | 49 ++-- .../python-sdk/python/genai_core/presign.py | 95 ++++++++ .../python-sdk/python/genai_core/upload.py | 35 --- lib/user-interface/index.ts | 44 ---- lib/user-interface/private-website.ts | 2 - lib/user-interface/public-website.ts | 30 --- .../src/common/api-client/sessions-client.ts | 35 ++- .../react-app/src/common/types.ts | 7 - .../components/chatbot/chat-input-panel.tsx | 29 ++- .../src/components/chatbot/chat-message.tsx | 30 ++- .../react-app/src/components/chatbot/chat.tsx | 8 +- .../src/components/chatbot/image-dialog.tsx | 52 ++-- .../react-app/src/components/chatbot/utils.ts | 6 - .../autentication-construct.test.ts.snap | 162 ------------- .../chatbot-api-construct.test.ts.snap | 225 +++++------------- .../api-handler/routes/documents_test.py | 4 +- .../api-handler/routes/sessions_test.py | 7 + 31 files changed, 400 insertions(+), 613 deletions(-) create mode 100644 lib/shared/layers/python-sdk/python/genai_core/presign.py delete mode 100644 lib/shared/layers/python-sdk/python/genai_core/upload.py diff --git a/integtests/chatbot-api/multi_modal_test.py b/integtests/chatbot-api/multi_modal_test.py index e2c30ba73..99464d5f3 100644 --- a/integtests/chatbot-api/multi_modal_test.py +++ b/integtests/chatbot-api/multi_modal_test.py @@ -2,32 +2,30 @@ import os import time import uuid -import boto3 from pathlib import Path import pytest +import requests +from gql.transport.exceptions import TransportQueryError -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, - ) +def test_multi_modal(client, default_multimodal_model, default_provider): + key = "INTEG_TEST" + str(uuid.uuid4()) + ".jpeg" - object = s3.Object(bucket, "public/" + key) - wrong_object = s3.Object(bucket, "private/notallowed/1.jpg") + result = client.add_file( + input={ + "fileName": key, + } + ) + + fields = result.get("fields") + cleaned_fields = fields.replace("{", "").replace("}", "") + pairs = [pair.strip() for pair in cleaned_fields.split(",")] + fields_dict = dict(pair.split("=", 1) for pair in pairs) 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() - ) + files = {"file": Path(current_dir + "/resources/powered-by-aws.png").read_bytes()} + response = requests.post(result.get("url"), data=fields_dict, files=files) + assert response.status_code == 204 session_id = str(uuid.uuid4()) @@ -58,4 +56,12 @@ def test_multi_modal( assert "powered by" in content client.delete_session(session_id) - object.delete() + + # Verify it can get the file + url = client.get_file_url(key) + assert url.startswith("https://") + + +def test_unknown_file(client): + with pytest.raises(TransportQueryError, match="File does not exist"): + client.get_file_url("file") diff --git a/integtests/clients/appsync_client.py b/integtests/clients/appsync_client.py index da27603fb..7d320471d 100644 --- a/integtests/clients/appsync_client.py +++ b/integtests/clients/appsync_client.py @@ -205,6 +205,10 @@ def add_file(self, input): ) return self.client.execute(query).get("getUploadFileURL") + def get_file_url(self, fileName): + query = dsl_gql(DSLQuery(self.schema.Query.getFileURL.args(fileName=fileName))) + return self.client.execute(query).get("getFileURL") + def get_document(self, input): query = dsl_gql( DSLQuery( diff --git a/integtests/clients/cognito_client.py b/integtests/clients/cognito_client.py index ae8430961..912105ca4 100644 --- a/integtests/clients/cognito_client.py +++ b/integtests/clients/cognito_client.py @@ -10,9 +10,6 @@ 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(********)" @@ -67,28 +64,11 @@ 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"], } ) diff --git a/integtests/security/unauthorized_test.py b/integtests/security/unauthorized_test.py index 3f8e436f1..a48208bff 100644 --- a/integtests/security/unauthorized_test.py +++ b/integtests/security/unauthorized_test.py @@ -126,6 +126,8 @@ def test_unauthenticated(unauthenticated_client: AppSyncClient): unauthenticated_client.start_kendra_data_sync("id") with pytest.raises(TransportQueryError, match=match): unauthenticated_client.is_kendra_data_synching("id") + with pytest.raises(TransportQueryError, match=match): + unauthenticated_client.get_file_url("file") with pytest.raises(TransportQueryError, match=match): unauthenticated_client.list_kendra_indexes() with pytest.raises(TransportQueryError, match=match): diff --git a/lib/authentication/index.ts b/lib/authentication/index.ts index 5a9528fbc..816aff41d 100644 --- a/lib/authentication/index.ts +++ b/lib/authentication/index.ts @@ -1,4 +1,3 @@ -import * as cognitoIdentityPool from "@aws-cdk/aws-cognito-identitypool-alpha"; import * as cdk from "aws-cdk-lib"; import { SystemConfig } from "../shared/types"; import * as cognito from "aws-cdk-lib/aws-cognito"; @@ -12,7 +11,6 @@ import * as logs from "aws-cdk-lib/aws-logs"; export class Authentication extends Construct { public readonly userPool: cognito.UserPool; public readonly userPoolClient: cognito.UserPoolClient; - public readonly identityPool: cognitoIdentityPool.IdentityPool; public readonly cognitoDomain: cognito.UserPoolDomain; public readonly updateUserPoolClient: lambda.Function; public readonly customOidcProvider: cognito.UserPoolIdentityProviderOidc; @@ -53,21 +51,6 @@ export class Authentication extends Construct { this.cognitoDomain = userPooldomain; } - const identityPool = new cognitoIdentityPool.IdentityPool( - this, - "IdentityPool", - { - authenticationProviders: { - userPools: [ - new cognitoIdentityPool.UserPoolAuthenticationProvider({ - userPool, - userPoolClient, - }), - ], - }, - } - ); - if (config.cognitoFederation?.enabled) { // Create an IAM Role for the Lambda function const lambdaRoleUpdateClient = new iam.Role( @@ -243,16 +226,11 @@ export class Authentication extends Construct { this.userPool = userPool; this.userPoolClient = userPoolClient; - this.identityPool = identityPool; new cdk.CfnOutput(this, "UserPoolId", { value: userPool.userPoolId, }); - new cdk.CfnOutput(this, "IdentityPoolId", { - value: identityPool.identityPoolId, - }); - new cdk.CfnOutput(this, "UserPoolWebClientId", { value: userPoolClient.userPoolClientId, }); diff --git a/lib/aws-genai-llm-chatbot-stack.ts b/lib/aws-genai-llm-chatbot-stack.ts index 5fd0ed056..12691fdeb 100644 --- a/lib/aws-genai-llm-chatbot-stack.ts +++ b/lib/aws-genai-llm-chatbot-stack.ts @@ -156,7 +156,6 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack { userPoolId: authentication.userPool.userPoolId, userPoolClient: authentication.userPoolClient, userPoolClientId: authentication.userPoolClient.userPoolClientId, - identityPool: authentication.identityPool, api: chatBotApi, chatbotFilesBucket: chatBotApi.filesBucket, crossEncodersEnabled: @@ -296,7 +295,6 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack { NagSuppressions.addResourceSuppressionsByPath( this, [ - `/${this.stackName}/Authentication/IdentityPool/AuthenticatedRole/DefaultPolicy/Resource`, `/${this.stackName}/Authentication/UserPool/smsRole/Resource`, `/${this.stackName}/Custom::CDKBucketDeployment8693BB64968944B69AAFB0CC9EB8756C/ServiceRole/DefaultPolicy/Resource`, `/${this.stackName}/LogRetentionaae0aa3c5b4d4f87b02d85b201efdd8a/ServiceRole/Resource`, @@ -351,8 +349,7 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack { NagSuppressions.addResourceSuppressionsByPath( this, [ - `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/Default/{object}/ANY/Resource`, - `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/Default/{object}/ANY/Resource`, + `/${this.stackName}/IdeficsInterface/ChatbotFilesPrivateApi/Default/{folder}/{key}/GET/Resource`, ], [ { id: "AwsSolutions-APIG4", reason: "Private API within a VPC." }, diff --git a/lib/chatbot-api/functions/api-handler/routes/documents.py b/lib/chatbot-api/functions/api-handler/routes/documents.py index 9ccb6916c..36fb3d4ce 100644 --- a/lib/chatbot-api/functions/api-handler/routes/documents.py +++ b/lib/chatbot-api/functions/api-handler/routes/documents.py @@ -7,8 +7,9 @@ SAFE_SHORT_STR_VALIDATION, ) import genai_core.types -import genai_core.upload +import genai_core.presign import genai_core.documents +import genai_core.auth from pydantic import BaseModel, Field from aws_lambda_powertools import Logger, Tracer from aws_lambda_powertools.event_handler.appsync import Router @@ -20,7 +21,7 @@ class FileUploadRequest(BaseModel): - workspaceId: str = ID_FIELD_VALIDATION + workspaceId: Optional[str] = ID_FIELD_VALIDATION fileName: str = Field(min_length=1, max_length=500, pattern=SAFE_STR_REGEX) @@ -104,7 +105,7 @@ class DocumentSubscriptionStatusRequest(BaseModel): ) -allowed_extensions = set( +allowed_workspace_extensions = set( [ ".csv", ".doc", @@ -128,18 +129,36 @@ class DocumentSubscriptionStatusRequest(BaseModel): ] ) +allowed_session_extensions = set( + [ + ".jpg", + ".jpeg", + ".png", + ] +) + @router.resolver(field_name="getUploadFileURL") @tracer.capture_method def file_upload(input: dict): request = FileUploadRequest(**input) _, extension = os.path.splitext(request.fileName) - if extension not in allowed_extensions: - raise genai_core.types.CommonError("Invalid file extension") - result = genai_core.upload.generate_presigned_post( - request.workspaceId, request.fileName - ) + if "workspaceId" in input: + if extension not in allowed_workspace_extensions: + raise genai_core.types.CommonError("Invalid file extension") + + result = genai_core.presign.generate_workspace_presigned_post( + request.workspaceId, request.fileName + ) + else: + if extension not in allowed_session_extensions: + raise genai_core.types.CommonError("Invalid file extension") + + user_id = genai_core.auth.get_user_id(router) + result = genai_core.presign.generate_user_presigned_post( + user_id, request.fileName + ) logger.info("Generated pre-signed for " + request.fileName) return result diff --git a/lib/chatbot-api/functions/api-handler/routes/sessions.py b/lib/chatbot-api/functions/api-handler/routes/sessions.py index 959d4e33a..c98bd60fc 100644 --- a/lib/chatbot-api/functions/api-handler/routes/sessions.py +++ b/lib/chatbot-api/functions/api-handler/routes/sessions.py @@ -1,4 +1,7 @@ +from pydantic import BaseModel, Field +from common.constant import SAFE_STR_REGEX from common.validation import WorkspaceIdValidation +import genai_core.presign import genai_core.sessions import genai_core.types import genai_core.auth @@ -12,6 +15,23 @@ logger = Logger() +class FileURequestValidation(BaseModel): + fileName: str = Field(min_length=1, max_length=500, pattern=SAFE_STR_REGEX) + + +@router.resolver(field_name="getFileURL") +@tracer.capture_method +def get_file(fileName: str): + FileURequestValidation(**{"fileName": fileName}) + user_id = genai_core.auth.get_user_id(router) + result = genai_core.presign.generate_user_presigned_get( + user_id, fileName, expiration=600 + ) + + logger.info("Generated pre-signed for " + fileName) + return result + + @router.resolver(field_name="listSessions") @tracer.capture_method def get_sessions(): diff --git a/lib/chatbot-api/index.ts b/lib/chatbot-api/index.ts index ed3526043..5e890906c 100644 --- a/lib/chatbot-api/index.ts +++ b/lib/chatbot-api/index.ts @@ -95,6 +95,7 @@ export class ChatBotApi extends Construct { byUserIdIndex: chatTables.byUserIdIndex, api, userFeedbackBucket: chatBuckets.userFeedbackBucket, + filesBucket: chatBuckets.filesBucket, }); this.resolvers.push(apiResolvers.appSyncLambdaResolver); diff --git a/lib/chatbot-api/rest-api.ts b/lib/chatbot-api/rest-api.ts index 39607a75b..901212b56 100644 --- a/lib/chatbot-api/rest-api.ts +++ b/lib/chatbot-api/rest-api.ts @@ -23,6 +23,7 @@ export interface ApiResolversProps { readonly userPool: cognito.UserPool; readonly sessionsTable: dynamodb.Table; readonly byUserIdIndex: string; + readonly filesBucket: s3.Bucket; readonly userFeedbackBucket: s3.Bucket; readonly modelsParameter: ssm.StringParameter; readonly models: SageMakerModelEndpoint[]; @@ -69,6 +70,7 @@ export class ApiResolvers extends Construct { SESSIONS_BY_USER_ID_INDEX_NAME: props.byUserIdIndex, USER_FEEDBACK_BUCKET_NAME: props.userFeedbackBucket?.bucketName ?? "", UPLOAD_BUCKET_NAME: props.ragEngines?.uploadBucket?.bucketName ?? "", + CHATBOT_FILES_BUCKET_NAME: props.filesBucket.bucketName, PROCESSING_BUCKET_NAME: props.ragEngines?.processingBucket?.bucketName ?? "", AURORA_DB_SECRET_ID: props.ragEngines?.auroraPgVector?.database @@ -296,6 +298,7 @@ export class ApiResolvers extends Construct { props.modelsParameter.grantRead(apiHandler); props.sessionsTable.grantReadWriteData(apiHandler); props.userFeedbackBucket.grantReadWrite(apiHandler); + props.filesBucket.grantReadWrite(apiHandler); props.ragEngines?.uploadBucket.grantReadWrite(apiHandler); props.ragEngines?.processingBucket.grantReadWrite(apiHandler); diff --git a/lib/chatbot-api/schema/schema.graphql b/lib/chatbot-api/schema/schema.graphql index 4601bfdb3..69f1f4e06 100644 --- a/lib/chatbot-api/schema/schema.graphql +++ b/lib/chatbot-api/schema/schema.graphql @@ -131,7 +131,7 @@ type EmbeddingModel @aws_cognito_user_pools { } input FileUploadInput { - workspaceId: String! + workspaceId: String fileName: String! } @@ -363,6 +363,8 @@ type Query { checkHealth: Boolean @aws_cognito_user_pools getUploadFileURL(input: FileUploadInput!): FileUploadResult @aws_cognito_user_pools + getFileURL(fileName: String!): String + @aws_cognito_user_pools listModels: [Model!]! @aws_cognito_user_pools listWorkspaces: [Workspace!]! @aws_cognito_user_pools getWorkspace(workspaceId: String!): Workspace @aws_cognito_user_pools diff --git a/lib/model-interfaces/idefics/functions/request-handler/adapters/claude.py b/lib/model-interfaces/idefics/functions/request-handler/adapters/claude.py index d86d40137..cc525a36a 100644 --- a/lib/model-interfaces/idefics/functions/request-handler/adapters/claude.py +++ b/lib/model-interfaces/idefics/functions/request-handler/adapters/claude.py @@ -12,21 +12,16 @@ s3 = boto3.resource("s3") -def get_image_message( - file: dict, -): +def get_image_message(file: dict, user_id: str): if file["key"] is None: raise Exception("Invalid S3 Key " + file["key"]) + key = "private/" + user_id + "/" + file["key"] logger.info( - "Fetching image", - bucket=os.environ["CHATBOT_FILES_BUCKET_NAME"], - key="public/" + file["key"], + "Fetching image", bucket=os.environ["CHATBOT_FILES_BUCKET_NAME"], key=key ) - response = s3.Object( - os.environ["CHATBOT_FILES_BUCKET_NAME"], "public/" + file["key"] - ) + response = s3.Object(os.environ["CHATBOT_FILES_BUCKET_NAME"], key) img = str(b64encode(response.get()["Body"].read()), "ascii") return { "type": "image", @@ -46,7 +41,9 @@ def __init__(self, model_id: str): self.model_id = model_id self.client = get_bedrock_client() - def format_prompt(self, prompt: str, messages: list, files: list) -> str: + def format_prompt( + self, prompt: str, messages: list, files: list, user_id: str + ) -> str: prompts = [] # Chat history @@ -59,7 +56,7 @@ def format_prompt(self, prompt: str, messages: list, files: list) -> str: prompts.append(user_msg) message_files = message.additional_kwargs.get("files", []) for message_file in message_files: - user_msg["content"].append(get_image_message(message_file)) + user_msg["content"].append(get_image_message(message_file, user_id)) if message.type.lower() == ChatbotMessageType.AI.value.lower(): prompts.append({"role": "assistant", "content": message.content}) @@ -70,7 +67,7 @@ def format_prompt(self, prompt: str, messages: list, files: list) -> str: } prompts.append(user_msg) for file in files: - user_msg["content"].append(get_image_message(file)) + user_msg["content"].append(get_image_message(file, user_id)) return json.dumps( { diff --git a/lib/model-interfaces/idefics/functions/request-handler/adapters/idefics.py b/lib/model-interfaces/idefics/functions/request-handler/adapters/idefics.py index ce6331e59..7b110069e 100644 --- a/lib/model-interfaces/idefics/functions/request-handler/adapters/idefics.py +++ b/lib/model-interfaces/idefics/functions/request-handler/adapters/idefics.py @@ -16,7 +16,9 @@ class Idefics(MultiModalModelBase): def __init__(self, model_id: str): self.model_id = model_id - def format_prompt(self, prompt: str, messages: list, files: list) -> str: + def format_prompt( + self, prompt: str, messages: list, files: list, user_id: str + ) -> str: human_prompt_template = "User:{prompt}" human_prompt_with_image = "User:{prompt}![]({image})" @@ -30,7 +32,8 @@ def format_prompt(self, prompt: str, messages: list, files: list) -> str: prompts.append(human_prompt_template.format(prompt=message.content)) for message_file in message_files: image = urljoin( - os.environ["CHATBOT_FILES_PRIVATE_API"], message_file["key"] + os.environ["CHATBOT_FILES_PRIVATE_API"], + user_id + "/" + message_file["key"], ) prompts.append( human_prompt_with_image.format( @@ -45,7 +48,7 @@ def format_prompt(self, prompt: str, messages: list, files: list) -> str: prompts.append(human_prompt_template.format(prompt=prompt)) for file in files: - key = file["key"] + key = user_id + "/" + file["key"] prompts.append( human_prompt_with_image.format( prompt=prompt, diff --git a/lib/model-interfaces/idefics/functions/request-handler/index.py b/lib/model-interfaces/idefics/functions/request-handler/index.py index c350fe3b2..3527dbbe4 100644 --- a/lib/model-interfaces/idefics/functions/request-handler/index.py +++ b/lib/model-interfaces/idefics/functions/request-handler/index.py @@ -53,6 +53,7 @@ def handle_run(record): prompt=prompt, messages=messages, files=files, + user_id=user_id, ) mlm_response = model.handle_run(prompt=prompt_template, model_kwargs=model_kwargs) diff --git a/lib/model-interfaces/idefics/index.ts b/lib/model-interfaces/idefics/index.ts index d46e682df..5714f58e2 100644 --- a/lib/model-interfaces/idefics/index.ts +++ b/lib/model-interfaces/idefics/index.ts @@ -234,19 +234,9 @@ export class IdeficsInterface extends Construct { }); integrationRole.addToPolicy( new iam.PolicyStatement({ - actions: ["s3:Get*", "s3:List*"], + actions: ["s3:GetObject*"], effect: iam.Effect.ALLOW, - resources: [ - `${this.props.chatbotFilesBucket.bucketArn}/*`, - `${this.props.chatbotFilesBucket.bucketArn}/*/*`, - ], - }) - ); - integrationRole.addToPolicy( - new iam.PolicyStatement({ - actions: ["kms:Decrypt", "kms:ReEncryptFrom"], - effect: iam.Effect.ALLOW, - resources: ["arn:aws:kms:*"], + resources: [`${this.props.chatbotFilesBucket.bucketArn}/private/*`], }) ); @@ -254,11 +244,12 @@ export class IdeficsInterface extends Construct { service: "s3", integrationHttpMethod: "GET", region: cdk.Aws.REGION, - path: `${this.props.chatbotFilesBucket.bucketName}/public/{object}`, + path: `${this.props.chatbotFilesBucket.bucketName}/private/{folder}/{key}`, options: { credentialsRole: integrationRole, requestParameters: { - "integration.request.path.object": "method.request.path.object", + "integration.request.path.folder": "method.request.path.folder", + "integration.request.path.key": "method.request.path.key", }, integrationResponses: [ { @@ -272,21 +263,25 @@ export class IdeficsInterface extends Construct { }, }); - const fileResource = api.root.addResource("{object}"); - fileResource.addMethod("ANY", s3Integration, { - methodResponses: [ - { - statusCode: "200", - responseParameters: { - "method.response.header.Content-Type": true, + api.root + .addResource("{folder}") + .addResource("{key}") + .addMethod("GET", s3Integration, { + methodResponses: [ + { + statusCode: "200", + responseParameters: { + "method.response.header.Content-Type": true, + }, }, + ], + requestParameters: { + "method.request.path.folder": true, + "method.request.path.key": true, + "method.request.header.Content-Type": true, }, - ], - requestParameters: { - "method.request.path.object": true, - "method.request.header.Content-Type": true, - }, - }); + }); + /** * CDK NAG suppression */ diff --git a/lib/shared/layers/python-sdk/python/genai_core/presign.py b/lib/shared/layers/python-sdk/python/genai_core/presign.py new file mode 100644 index 000000000..ad06d3e16 --- /dev/null +++ b/lib/shared/layers/python-sdk/python/genai_core/presign.py @@ -0,0 +1,95 @@ +import os +import boto3 +import botocore +import genai_core.workspaces +import genai_core.types +import unicodedata + +UPLOAD_BUCKET_NAME = os.environ.get("UPLOAD_BUCKET_NAME") +CHATBOT_FILES_BUCKET_NAME = os.environ.get("CHATBOT_FILES_BUCKET_NAME") +MAX_FILE_SIZE = 100 * 1000 * 1000 # 100Mb + + +def generate_workspace_presigned_post( + workspace_id: str, file_name: str, expiration=3600 +): + s3_client = boto3.client("s3") + + file_name = unicodedata.normalize("NFC", file_name) + workspace = genai_core.workspaces.get_workspace(workspace_id) + if not workspace: + raise genai_core.types.CommonError("Workspace not found") + + file_name = os.path.basename(file_name) + object_name = f"{workspace_id}/{file_name}" + + conditions = [ + ["content-length-range", 0, MAX_FILE_SIZE], + ] + + response = s3_client.generate_presigned_post( + UPLOAD_BUCKET_NAME, object_name, Conditions=conditions, ExpiresIn=expiration + ) + + if not response: + return None + + response["url"] = f"https://{UPLOAD_BUCKET_NAME}.s3-accelerate.amazonaws.com" + + return response + + +def generate_user_presigned_post(user_id: str, file_name: str, expiration=3600): + s3_client = boto3.client("s3") + + file_name = unicodedata.normalize("NFC", file_name) + if not user_id or len(user_id) < 10: + raise genai_core.types.CommonError("User not set") + + file_name = os.path.basename(file_name) + object_name = f"private/{user_id}/{file_name}" + + conditions = [ + ["content-length-range", 0, MAX_FILE_SIZE], + ] + + response = s3_client.generate_presigned_post( + CHATBOT_FILES_BUCKET_NAME, + object_name, + Conditions=conditions, + ExpiresIn=expiration, + ) + + if not response: + return None + + response["url"] = f"https://{CHATBOT_FILES_BUCKET_NAME}.s3-accelerate.amazonaws.com" + + return response + + +def generate_user_presigned_get(user_id: str, file_name: str, expiration=3600): + s3_client = boto3.client("s3") + + file_name = unicodedata.normalize("NFC", file_name) + if not user_id or len(user_id) < 10: + raise genai_core.types.CommonError("User not set") + + file_name = os.path.basename(file_name) + object_name = f"private/{user_id}/{file_name}" + try: + s3_client.head_object(Bucket=CHATBOT_FILES_BUCKET_NAME, Key=object_name) + except botocore.exceptions.ClientError as e: + if e.response["Error"]["Code"] == "404": + raise genai_core.types.CommonError("File does not exist") + else: + raise e + response = s3_client.generate_presigned_url( + "get_object", + Params={"Bucket": CHATBOT_FILES_BUCKET_NAME, "Key": object_name}, + ExpiresIn=expiration, + ) + + if not response: + return None + return response diff --git a/lib/shared/layers/python-sdk/python/genai_core/upload.py b/lib/shared/layers/python-sdk/python/genai_core/upload.py deleted file mode 100644 index 3d3f63884..000000000 --- a/lib/shared/layers/python-sdk/python/genai_core/upload.py +++ /dev/null @@ -1,35 +0,0 @@ -import os -import boto3 -import genai_core.workspaces -import genai_core.types -import unicodedata - -UPLOAD_BUCKET_NAME = os.environ.get("UPLOAD_BUCKET_NAME") -MAX_FILE_SIZE = 100 * 1000 * 1000 # 100Mb - - -def generate_presigned_post(workspace_id: str, file_name: str, expiration=3600): - s3_client = boto3.client("s3") - - file_name = unicodedata.normalize("NFC", file_name) - workspace = genai_core.workspaces.get_workspace(workspace_id) - if not workspace: - raise genai_core.types.CommonError("Workspace not found") - - file_name = os.path.basename(file_name) - object_name = f"{workspace_id}/{file_name}" - - conditions = [ - ["content-length-range", 0, MAX_FILE_SIZE], - ] - - response = s3_client.generate_presigned_post( - UPLOAD_BUCKET_NAME, object_name, Conditions=conditions, ExpiresIn=expiration - ) - - if not response: - return None - - response["url"] = f"https://{UPLOAD_BUCKET_NAME}.s3-accelerate.amazonaws.com" - - return response diff --git a/lib/user-interface/index.ts b/lib/user-interface/index.ts index 4276923e5..17b1e4b9f 100644 --- a/lib/user-interface/index.ts +++ b/lib/user-interface/index.ts @@ -1,6 +1,4 @@ -import * as cognitoIdentityPool from "@aws-cdk/aws-cognito-identitypool-alpha"; import * as cdk from "aws-cdk-lib"; -import * as iam from "aws-cdk-lib/aws-iam"; import * as s3 from "aws-cdk-lib/aws-s3"; import * as cognito from "aws-cdk-lib/aws-cognito"; import * as s3deploy from "aws-cdk-lib/aws-s3-deployment"; @@ -24,7 +22,6 @@ export interface UserInterfaceProps { readonly userPoolId: string; readonly userPoolClientId: string; readonly userPoolClient: cognito.UserPoolClient; - readonly identityPool: cognitoIdentityPool.IdentityPool; readonly api: ChatBotApi; readonly chatbotFilesBucket: s3.Bucket; readonly crossEncodersEnabled: boolean; @@ -86,12 +83,10 @@ export class UserInterface extends Construct { aws_cognito_region: cdk.Aws.REGION, aws_user_pools_id: props.userPoolId, aws_user_pools_web_client_id: props.userPoolClientId, - aws_cognito_identity_pool_id: props.identityPool.identityPoolId, Auth: { region: cdk.Aws.REGION, userPoolId: props.userPoolId, userPoolWebClientId: props.userPoolClientId, - identityPoolId: props.identityPool.identityPoolId, }, oauth: props.config.cognitoFederation?.enabled ? { @@ -105,12 +100,6 @@ export class UserInterface extends Construct { aws_appsync_graphqlEndpoint: props.api.graphqlApi.graphqlUrl, aws_appsync_region: cdk.Aws.REGION, aws_appsync_authenticationType: "AMAZON_COGNITO_USER_POOLS", - Storage: { - AWSS3: { - bucket: props.chatbotFilesBucket.bucketName, - region: cdk.Aws.REGION, - }, - }, config: { auth_federated_provider: props.config.cognitoFederation?.enabled ? { @@ -130,39 +119,6 @@ export class UserInterface extends Construct { }, }); - // Allow authenticated web users to read upload data to the attachments bucket for their chat files - // ref: https://docs.amplify.aws/lib/storage/getting-started/q/platform/js/#using-amazon-s3 - props.identityPool.authenticatedRole.addToPrincipalPolicy( - new iam.PolicyStatement({ - effect: iam.Effect.ALLOW, - actions: ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"], - resources: [ - `${props.chatbotFilesBucket.bucketArn}/public/*`, - `${props.chatbotFilesBucket.bucketArn}/protected/\${cognito-identity.amazonaws.com:sub}/*`, - `${props.chatbotFilesBucket.bucketArn}/private/\${cognito-identity.amazonaws.com:sub}/*`, - ], - }) - ); - props.identityPool.authenticatedRole.addToPrincipalPolicy( - new iam.PolicyStatement({ - effect: iam.Effect.ALLOW, - actions: ["s3:ListBucket"], - resources: [`${props.chatbotFilesBucket.bucketArn}`], - conditions: { - StringLike: { - "s3:prefix": [ - "public/", - "public/*", - "protected/", - "protected/*", - "private/${cognito-identity.amazonaws.com:sub}/", - "private/${cognito-identity.amazonaws.com:sub}/*", - ], - }, - }, - }) - ); - // Enable CORS for the attachments bucket to allow uploads from the user interface // ref: https://docs.amplify.aws/lib/storage/getting-started/q/platform/js/#amazon-s3-bucket-cors-policy-setup props.chatbotFilesBucket.addCorsRule({ diff --git a/lib/user-interface/private-website.ts b/lib/user-interface/private-website.ts index 450485f3b..fa566b7e4 100644 --- a/lib/user-interface/private-website.ts +++ b/lib/user-interface/private-website.ts @@ -1,4 +1,3 @@ -import * as cognitoIdentityPool from "@aws-cdk/aws-cognito-identitypool-alpha"; import * as cdk from "aws-cdk-lib"; import * as iam from "aws-cdk-lib/aws-iam"; import * as s3 from "aws-cdk-lib/aws-s3"; @@ -17,7 +16,6 @@ export interface PrivateWebsiteProps { readonly shared: Shared; readonly userPoolId: string; readonly userPoolClientId: string; - readonly identityPool: cognitoIdentityPool.IdentityPool; readonly api: ChatBotApi; readonly chatbotFilesBucket: s3.Bucket; readonly crossEncodersEnabled: boolean; diff --git a/lib/user-interface/public-website.ts b/lib/user-interface/public-website.ts index fa2f37b13..aa3e449b8 100644 --- a/lib/user-interface/public-website.ts +++ b/lib/user-interface/public-website.ts @@ -1,4 +1,3 @@ -import * as cognitoIdentityPool from "@aws-cdk/aws-cognito-identitypool-alpha"; import * as cdk from "aws-cdk-lib"; import * as cf from "aws-cdk-lib/aws-cloudfront"; import * as s3 from "aws-cdk-lib/aws-s3"; @@ -14,9 +13,7 @@ export interface PublicWebsiteProps { readonly shared: Shared; readonly userPoolId: string; readonly userPoolClientId: string; - readonly identityPool: cognitoIdentityPool.IdentityPool; readonly api: ChatBotApi; - readonly chatbotFilesBucket: s3.Bucket; readonly crossEncodersEnabled: boolean; readonly sagemakerEmbeddingsEnabled: boolean; readonly websiteBucket: s3.Bucket; @@ -34,7 +31,6 @@ export class PublicWebsite extends Construct { const originAccessIdentity = new cf.OriginAccessIdentity(this, "S3OAI"); props.websiteBucket.grantRead(originAccessIdentity); - props.chatbotFilesBucket.grantRead(originAccessIdentity); const cfGeoRestrictEnable = props.config.cfGeoRestrictEnable; const cfGeoRestrictList = props.config.cfGeoRestrictList; @@ -88,32 +84,6 @@ export class PublicWebsite extends Construct { originAccessIdentity, }, }, - { - behaviors: [ - { - pathPattern: "/chabot/files/*", - allowedMethods: cf.CloudFrontAllowedMethods.ALL, - viewerProtocolPolicy: cf.ViewerProtocolPolicy.REDIRECT_TO_HTTPS, - defaultTtl: cdk.Duration.seconds(0), - forwardedValues: { - queryString: true, - headers: [ - "Referer", - "Origin", - "Authorization", - "Content-Type", - "x-forwarded-user", - "Access-Control-Request-Headers", - "Access-Control-Request-Method", - ], - }, - }, - ], - s3OriginSource: { - s3BucketSource: props.chatbotFilesBucket, - originAccessIdentity, - }, - }, ], geoRestriction: cfGeoRestrictEnable ? cf.GeoRestriction.allowlist(...cfGeoRestrictList) diff --git a/lib/user-interface/react-app/src/common/api-client/sessions-client.ts b/lib/user-interface/react-app/src/common/api-client/sessions-client.ts index 727c2dbee..e35c2625f 100644 --- a/lib/user-interface/react-app/src/common/api-client/sessions-client.ts +++ b/lib/user-interface/react-app/src/common/api-client/sessions-client.ts @@ -1,15 +1,48 @@ import { API } from "aws-amplify"; import { GraphQLQuery, GraphQLResult } from "@aws-amplify/api"; -import { listSessions, getSession } from "../../graphql/queries"; +import { + listSessions, + getSession, + getUploadFileURL, + getFileURL, +} from "../../graphql/queries"; import { deleteSession, deleteUserSessions } from "../../graphql/mutations"; import { ListSessionsQuery, GetSessionQuery, DeleteSessionMutation, DeleteUserSessionsMutation, + GetUploadFileURLQuery, + GetFileURLQuery, } from "../../API"; export class SessionsClient { + async getFileUploadSignedUrl( + fileName: string + ): Promise>> { + const result = API.graphql>({ + query: getUploadFileURL, + variables: { + input: { + fileName, + }, + }, + }); + return result; + } + + async getFileSignedUrl( + fileName: string + ): Promise>> { + const result = API.graphql>({ + query: getFileURL, + variables: { + fileName, + }, + }); + return result; + } + async getSessions(): Promise>> { const result = await API.graphql>({ query: listSessions, diff --git a/lib/user-interface/react-app/src/common/types.ts b/lib/user-interface/react-app/src/common/types.ts index 5396ce5f7..6de940cfb 100644 --- a/lib/user-interface/react-app/src/common/types.ts +++ b/lib/user-interface/react-app/src/common/types.ts @@ -3,7 +3,6 @@ import { CognitoHostedUIIdentityProvider } from "@aws-amplify/auth"; export interface AppConfig { aws_project_region: string; - aws_cognito_identity_pool_id: string; aws_user_pools_id: string; aws_user_pools_web_client_id: string; config: { @@ -30,12 +29,6 @@ export interface AppConfig { default_cross_encoder_model: string; privateWebsite: boolean; }; - Storage: { - AWSS3: { - bucket: string; - region: string; - }; - }; } export interface NavigationPanelState { diff --git a/lib/user-interface/react-app/src/components/chatbot/chat-input-panel.tsx b/lib/user-interface/react-app/src/components/chatbot/chat-input-panel.tsx index b125a8410..8888b004a 100644 --- a/lib/user-interface/react-app/src/components/chatbot/chat-input-panel.tsx +++ b/lib/user-interface/react-app/src/components/chatbot/chat-input-panel.tsx @@ -50,11 +50,7 @@ import { ChatBotToken, } from "./types"; import { sendQuery } from "../../graphql/mutations"; -import { - getSelectedModelMetadata, - getSignedUrl, - updateMessageHistoryRef, -} from "./utils"; +import { getSelectedModelMetadata, updateMessageHistoryRef } from "./utils"; import { receiveMessages } from "../../graphql/subscriptions"; import { Utils } from "../../common/utils"; @@ -309,15 +305,22 @@ export default function ChatInputPanel(props: ChatInputPanelProps) { }, [props.messageHistory]); useEffect(() => { + if (!appContext) return; + + const apiClient = new ApiClient(appContext); const getSignedUrls = async () => { if (props.configuration?.files as ImageFile[]) { const files: ImageFile[] = []; for await (const file of props.configuration?.files ?? []) { - const signedUrl = await getSignedUrl(file.key); - files.push({ - ...file, - url: signedUrl, - }); + const signedUrl = ( + await apiClient.sessions.getFileSignedUrl(file.key) + ).data?.getFileURL; + if (signedUrl) { + files.push({ + ...file, + url: signedUrl, + }); + } } setFiles(files); @@ -325,9 +328,11 @@ export default function ChatInputPanel(props: ChatInputPanelProps) { }; if (props.configuration.files?.length) { - getSignedUrls(); + getSignedUrls().catch((e) => { + console.log("Unable to get signed URL", e); + }); } - }, [props.configuration]); + }, [appContext, props.configuration]); const hasImagesInChatHistory = function (): boolean { return ( diff --git a/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx b/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx index 6c3d40a7e..c66033e88 100644 --- a/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx +++ b/lib/user-interface/react-app/src/components/chatbot/chat-message.tsx @@ -10,7 +10,7 @@ import { TextContent, Textarea, } from "@cloudscape-design/components"; -import { useEffect, useState } from "react"; +import { useContext, useEffect, useState } from "react"; import { JsonView, darkStyles } from "react-json-view-lite"; import ReactMarkdown from "react-markdown"; import remarkGfm from "remark-gfm"; @@ -23,10 +23,10 @@ import { RagDocument, } from "./types"; -import { getSignedUrl } from "./utils"; - import "react-json-view-lite/dist/index.css"; import "../../styles/app.scss"; +import { AppContext } from "../../common/app-context"; +import { ApiClient } from "../../common/api-client/api-client"; export interface ChatMessageProps { message: ChatBotHistoryItem; @@ -37,6 +37,7 @@ export interface ChatMessageProps { } export default function ChatMessage(props: ChatMessageProps) { + const appContext = useContext(AppContext); const [loading, setLoading] = useState(false); const [message] = useState(props.message); const [files, setFiles] = useState([] as ImageFile[]); @@ -45,16 +46,23 @@ export default function ChatMessage(props: ChatMessageProps) { const [selectedIcon, setSelectedIcon] = useState<1 | 0 | null>(null); useEffect(() => { + if (!appContext) return; + + const apiClient = new ApiClient(appContext); const getSignedUrls = async () => { setLoading(true); if (message.metadata?.files as ImageFile[]) { const files: ImageFile[] = []; for await (const file of message.metadata?.files as ImageFile[]) { - const signedUrl = await getSignedUrl(file.key); - files.push({ - ...file, - url: signedUrl as string, - }); + const signedUrl = ( + await apiClient.sessions.getFileSignedUrl(file.key) + ).data?.getFileURL; + if (signedUrl) { + files.push({ + ...file, + url: signedUrl, + }); + } } setLoading(false); @@ -63,9 +71,11 @@ export default function ChatMessage(props: ChatMessageProps) { }; if (message.metadata?.files as ImageFile[]) { - getSignedUrls(); + getSignedUrls().catch((e) => { + console.log("Unable to get signed URL", e); + }); } - }, [message]); + }, [appContext, message]); let content = ""; if (props.message.content && props.message.content.length > 0) { diff --git a/lib/user-interface/react-app/src/components/chatbot/chat.tsx b/lib/user-interface/react-app/src/components/chatbot/chat.tsx index cb315a90e..2d18bdb82 100644 --- a/lib/user-interface/react-app/src/components/chatbot/chat.tsx +++ b/lib/user-interface/react-app/src/components/chatbot/chat.tsx @@ -21,7 +21,9 @@ import { CHATBOT_NAME } from "../../common/constants"; export default function Chat(props: { sessionId?: string }) { const appContext = useContext(AppContext); const [running, setRunning] = useState(false); - const [session, setSession] = useState<{ id: string; loading: boolean } | undefined>(); + const [session, setSession] = useState< + { id: string; loading: boolean } | undefined + >(); const [initError, setInitError] = useState(undefined); const [configuration, setConfiguration] = useState( () => ({ @@ -149,7 +151,7 @@ export default function Chat(props: { sessionId?: string }) { )}
- {session && + {session && ( - } + )}
); diff --git a/lib/user-interface/react-app/src/components/chatbot/image-dialog.tsx b/lib/user-interface/react-app/src/components/chatbot/image-dialog.tsx index 3141d0ebd..c531a2dd6 100644 --- a/lib/user-interface/react-app/src/components/chatbot/image-dialog.tsx +++ b/lib/user-interface/react-app/src/components/chatbot/image-dialog.tsx @@ -10,10 +10,13 @@ import { } from "@cloudscape-design/components"; import { useForm } from "../../common/hooks/use-form"; -import { Storage } from "aws-amplify"; -import { Dispatch, useState } from "react"; +import { Dispatch, useContext, useState } from "react"; import { v4 as uuidv4 } from "uuid"; import { ChatBotConfiguration, FileStorageProvider, ImageFile } from "./types"; +import { AppContext } from "../../common/app-context"; +import { ApiClient } from "../../common/api-client/api-client"; +import { FileUploader } from "../../common/file-uploader"; +import { Utils } from "../../common/utils"; export interface ImageDialogProps { sessionId: string; @@ -26,6 +29,7 @@ export interface ImageDialogProps { const ALLOWED_MIME_TYPES = ["image/png", "image/jpg", "image/jpeg"]; export default function ImageDialog(props: ImageDialogProps) { + const appContext = useContext(AppContext); const [error, setError] = useState(null); const [loading, setLoading] = useState(false); const [files, setFiles] = useState([] as File[]); @@ -55,10 +59,14 @@ export default function ImageDialog(props: ImageDialogProps) { }); const saveConfig = async () => { - if (!validate()) return; + if (!validate() || !appContext) return; setLoading(true); + const apiClient = new ApiClient(appContext); - const files: ImageFile[] = (await uploadFiles(data.files)) as ImageFile[]; + const files: ImageFile[] = (await uploadFiles( + data.files, + apiClient + )) as ImageFile[]; props.setConfiguration({ ...props.configuration, @@ -102,18 +110,20 @@ export default function ImageDialog(props: ImageDialogProps) { return true; }; - const uploadFiles = async (files: File[]) => { + const uploadFiles = async (files: File[], client: ApiClient) => { const s3Files = []; + const uploader = new FileUploader(); for await (const file of files) { try { - const response = await uploadFile(file); + const response = await uploadFile(file, client, uploader); s3Files.push({ - key: `${response.key}`, + key: `${response}`, provider: FileStorageProvider.S3, }); } catch (error) { - const errorMessage = "Error uploading file: " + error; - console.log(errorMessage); + const errorMessage = + "Error uploading file: " + Utils.getErrorMessage(error); + console.log(errorMessage, error); setError(errorMessage); } } @@ -125,20 +135,22 @@ export default function ImageDialog(props: ImageDialogProps) { return s3Files; }; - const uploadFile = async (file: File) => { - console.log(file); + const uploadFile = async ( + file: File, + client: ApiClient, + uploader: FileUploader + ) => { const id = uuidv4(); - const shortId = id.split("-")[0]; // get the extension of the file and content type const extension = file.name.split(".").pop(); - const contentType = file.type; - - const response = await Storage.put(`${shortId}.${extension}`, file, { - contentType, - }); - return { - ...response, - }; + const url = ( + await client.sessions.getFileUploadSignedUrl(`${id}.${extension}`) + ).data?.getUploadFileURL; + if (!url) { + throw new Error("Unable to get the upload url."); + } + await uploader.upload(file, url, () => {}); + return `${id}.${extension}`; }; return ( diff --git a/lib/user-interface/react-app/src/components/chatbot/utils.ts b/lib/user-interface/react-app/src/components/chatbot/utils.ts index 1d589003a..1c3783820 100644 --- a/lib/user-interface/react-app/src/components/chatbot/utils.ts +++ b/lib/user-interface/react-app/src/components/chatbot/utils.ts @@ -1,4 +1,3 @@ -import { Storage } from "aws-amplify"; import { Dispatch, SetStateAction } from "react"; import { ChatBotAction, @@ -309,11 +308,6 @@ export function updateChatSessions( } } -export async function getSignedUrl(key: string) { - const signedUrl = await Storage.get(key as string); - return signedUrl; -} - export function getSelectedModelMetadata( models: Model[] | undefined, selectedModelOption: SelectProps.Option | null diff --git a/tests/authentication/__snapshots__/autentication-construct.test.ts.snap b/tests/authentication/__snapshots__/autentication-construct.test.ts.snap index 6d758f2af..d66f32aa9 100644 --- a/tests/authentication/__snapshots__/autentication-construct.test.ts.snap +++ b/tests/authentication/__snapshots__/autentication-construct.test.ts.snap @@ -3,11 +3,6 @@ exports[`snapshot test 1`] = ` { "Outputs": { - "AuthenticationConstructIdentityPoolIdE3D32668": { - "Value": { - "Ref": "AuthenticationConstructIdentityPool508C0E75", - }, - }, "AuthenticationConstructUserPoolId1F16F432": { "Value": { "Ref": "AuthenticationConstructUserPoolFE5ABE04", @@ -48,163 +43,6 @@ exports[`snapshot test 1`] = ` }, }, "Resources": { - "AuthenticationConstructIdentityPool508C0E75": { - "DependsOn": [ - "AuthenticationConstructUserPoolFE5ABE04", - "AuthenticationConstructUserPoolsmsRole65D4F248", - "AuthenticationConstructUserPoolUserPoolClientDCD6FB5B", - ], - "Properties": { - "AllowUnauthenticatedIdentities": false, - "CognitoIdentityProviders": [ - { - "ClientId": { - "Ref": "AuthenticationConstructUserPoolUserPoolClientDCD6FB5B", - }, - "ProviderName": { - "Fn::Join": [ - "", - [ - "cognito-idp.", - { - "Ref": "AWS::Region", - }, - ".", - { - "Ref": "AWS::URLSuffix", - }, - "/", - { - "Ref": "AuthenticationConstructUserPoolFE5ABE04", - }, - ], - ], - }, - "ServerSideTokenCheck": true, - }, - ], - }, - "Type": "AWS::Cognito::IdentityPool", - }, - "AuthenticationConstructIdentityPoolAuthenticatedRole8D31390C": { - "DependsOn": [ - "AuthenticationConstructUserPoolFE5ABE04", - "AuthenticationConstructUserPoolsmsRole65D4F248", - "AuthenticationConstructUserPoolUserPoolClientDCD6FB5B", - ], - "Properties": { - "AssumeRolePolicyDocument": { - "Statement": [ - { - "Action": "sts:AssumeRoleWithWebIdentity", - "Condition": { - "ForAnyValue:StringLike": { - "cognito-identity.amazonaws.com:amr": "authenticated", - }, - "StringEquals": { - "cognito-identity.amazonaws.com:aud": { - "Ref": "AuthenticationConstructIdentityPool508C0E75", - }, - }, - }, - "Effect": "Allow", - "Principal": { - "Federated": "cognito-identity.amazonaws.com", - }, - }, - ], - "Version": "2012-10-17", - }, - "Description": { - "Fn::Join": [ - "", - [ - "Default Authenticated Role for Identity Pool ", - { - "Fn::GetAtt": [ - "AuthenticationConstructIdentityPool508C0E75", - "Name", - ], - }, - ], - ], - }, - }, - "Type": "AWS::IAM::Role", - }, - "AuthenticationConstructIdentityPoolDefaultRoleAttachmentDC1BB441": { - "DependsOn": [ - "AuthenticationConstructUserPoolFE5ABE04", - "AuthenticationConstructUserPoolsmsRole65D4F248", - "AuthenticationConstructUserPoolUserPoolClientDCD6FB5B", - ], - "Properties": { - "IdentityPoolId": { - "Ref": "AuthenticationConstructIdentityPool508C0E75", - }, - "Roles": { - "authenticated": { - "Fn::GetAtt": [ - "AuthenticationConstructIdentityPoolAuthenticatedRole8D31390C", - "Arn", - ], - }, - "unauthenticated": { - "Fn::GetAtt": [ - "AuthenticationConstructIdentityPoolUnauthenticatedRoleCB2B4EA2", - "Arn", - ], - }, - }, - }, - "Type": "AWS::Cognito::IdentityPoolRoleAttachment", - }, - "AuthenticationConstructIdentityPoolUnauthenticatedRoleCB2B4EA2": { - "DependsOn": [ - "AuthenticationConstructUserPoolFE5ABE04", - "AuthenticationConstructUserPoolsmsRole65D4F248", - "AuthenticationConstructUserPoolUserPoolClientDCD6FB5B", - ], - "Properties": { - "AssumeRolePolicyDocument": { - "Statement": [ - { - "Action": "sts:AssumeRoleWithWebIdentity", - "Condition": { - "ForAnyValue:StringLike": { - "cognito-identity.amazonaws.com:amr": "unauthenticated", - }, - "StringEquals": { - "cognito-identity.amazonaws.com:aud": { - "Ref": "AuthenticationConstructIdentityPool508C0E75", - }, - }, - }, - "Effect": "Allow", - "Principal": { - "Federated": "cognito-identity.amazonaws.com", - }, - }, - ], - "Version": "2012-10-17", - }, - "Description": { - "Fn::Join": [ - "", - [ - "Default Unauthenticated Role for Identity Pool ", - { - "Fn::GetAtt": [ - "AuthenticationConstructIdentityPool508C0E75", - "Name", - ], - }, - ], - ], - }, - }, - "Type": "AWS::IAM::Role", - }, "AuthenticationConstructUserPoolFE5ABE04": { "DeletionPolicy": "Delete", "Metadata": { diff --git a/tests/chatbot-api/__snapshots__/chatbot-api-construct.test.ts.snap b/tests/chatbot-api/__snapshots__/chatbot-api-construct.test.ts.snap index cd9e324c3..f629865c9 100644 --- a/tests/chatbot-api/__snapshots__/chatbot-api-construct.test.ts.snap +++ b/tests/chatbot-api/__snapshots__/chatbot-api-construct.test.ts.snap @@ -297,11 +297,6 @@ exports[`snapshot test 1`] = ` }, }, "Outputs": { - "AuthenticationIdentityPoolIdB65295D6": { - "Value": { - "Ref": "AuthenticationIdentityPool98062B23", - }, - }, "AuthenticationUserPoolIdF0D106F7": { "Value": { "Ref": "AuthenticationUserPool28698864", @@ -413,163 +408,6 @@ exports[`snapshot test 1`] = ` }, }, "Resources": { - "AuthenticationIdentityPool98062B23": { - "DependsOn": [ - "AuthenticationUserPool28698864", - "AuthenticationUserPoolsmsRole5227CBC5", - "AuthenticationUserPoolUserPoolClient8AE1704E", - ], - "Properties": { - "AllowUnauthenticatedIdentities": false, - "CognitoIdentityProviders": [ - { - "ClientId": { - "Ref": "AuthenticationUserPoolUserPoolClient8AE1704E", - }, - "ProviderName": { - "Fn::Join": [ - "", - [ - "cognito-idp.", - { - "Ref": "AWS::Region", - }, - ".", - { - "Ref": "AWS::URLSuffix", - }, - "/", - { - "Ref": "AuthenticationUserPool28698864", - }, - ], - ], - }, - "ServerSideTokenCheck": true, - }, - ], - }, - "Type": "AWS::Cognito::IdentityPool", - }, - "AuthenticationIdentityPoolAuthenticatedRole04DA16FF": { - "DependsOn": [ - "AuthenticationUserPool28698864", - "AuthenticationUserPoolsmsRole5227CBC5", - "AuthenticationUserPoolUserPoolClient8AE1704E", - ], - "Properties": { - "AssumeRolePolicyDocument": { - "Statement": [ - { - "Action": "sts:AssumeRoleWithWebIdentity", - "Condition": { - "ForAnyValue:StringLike": { - "cognito-identity.amazonaws.com:amr": "authenticated", - }, - "StringEquals": { - "cognito-identity.amazonaws.com:aud": { - "Ref": "AuthenticationIdentityPool98062B23", - }, - }, - }, - "Effect": "Allow", - "Principal": { - "Federated": "cognito-identity.amazonaws.com", - }, - }, - ], - "Version": "2012-10-17", - }, - "Description": { - "Fn::Join": [ - "", - [ - "Default Authenticated Role for Identity Pool ", - { - "Fn::GetAtt": [ - "AuthenticationIdentityPool98062B23", - "Name", - ], - }, - ], - ], - }, - }, - "Type": "AWS::IAM::Role", - }, - "AuthenticationIdentityPoolDefaultRoleAttachment468E9E4F": { - "DependsOn": [ - "AuthenticationUserPool28698864", - "AuthenticationUserPoolsmsRole5227CBC5", - "AuthenticationUserPoolUserPoolClient8AE1704E", - ], - "Properties": { - "IdentityPoolId": { - "Ref": "AuthenticationIdentityPool98062B23", - }, - "Roles": { - "authenticated": { - "Fn::GetAtt": [ - "AuthenticationIdentityPoolAuthenticatedRole04DA16FF", - "Arn", - ], - }, - "unauthenticated": { - "Fn::GetAtt": [ - "AuthenticationIdentityPoolUnauthenticatedRoleC08D0BC1", - "Arn", - ], - }, - }, - }, - "Type": "AWS::Cognito::IdentityPoolRoleAttachment", - }, - "AuthenticationIdentityPoolUnauthenticatedRoleC08D0BC1": { - "DependsOn": [ - "AuthenticationUserPool28698864", - "AuthenticationUserPoolsmsRole5227CBC5", - "AuthenticationUserPoolUserPoolClient8AE1704E", - ], - "Properties": { - "AssumeRolePolicyDocument": { - "Statement": [ - { - "Action": "sts:AssumeRoleWithWebIdentity", - "Condition": { - "ForAnyValue:StringLike": { - "cognito-identity.amazonaws.com:amr": "unauthenticated", - }, - "StringEquals": { - "cognito-identity.amazonaws.com:aud": { - "Ref": "AuthenticationIdentityPool98062B23", - }, - }, - }, - "Effect": "Allow", - "Principal": { - "Federated": "cognito-identity.amazonaws.com", - }, - }, - ], - "Version": "2012-10-17", - }, - "Description": { - "Fn::Join": [ - "", - [ - "Default Unauthenticated Role for Identity Pool ", - { - "Fn::GetAtt": [ - "AuthenticationIdentityPool98062B23", - "Name", - ], - }, - ], - ], - }, - }, - "Type": "AWS::IAM::Role", - }, "AuthenticationUserPool28698864": { "DeletionPolicy": "Delete", "Metadata": { @@ -1545,7 +1383,7 @@ type EmbeddingModel @aws_cognito_user_pools { } input FileUploadInput { - workspaceId: String! + workspaceId: String fileName: String! } @@ -1777,6 +1615,8 @@ type Query { checkHealth: Boolean @aws_cognito_user_pools getUploadFileURL(input: FileUploadInput!): FileUploadResult @aws_cognito_user_pools + getFileURL(fileName: String!): String + @aws_cognito_user_pools listModels: [Model!]! @aws_cognito_user_pools listWorkspaces: [Workspace!]! @aws_cognito_user_pools getWorkspace(workspaceId: String!): Workspace @aws_cognito_user_pools @@ -2119,6 +1959,25 @@ schema { }, "Type": "AWS::AppSync::Resolver", }, + "ChatBotApiConstructChatbotApigetFileURLresolver3FBFB9AC": { + "DependsOn": [ + "ChatBotApiConstructChatbotApiproxyResolverFunction22AA16EE", + "ChatBotApiConstructChatbotApiSchema07900657", + ], + "Properties": { + "ApiId": { + "Fn::GetAtt": [ + "ChatBotApiConstructChatbotApi21E23C68", + "ApiId", + ], + }, + "DataSourceName": "proxyResolverFunction", + "FieldName": "getFileURL", + "Kind": "UNIT", + "TypeName": "Query", + }, + "Type": "AWS::AppSync::Resolver", + }, "ChatBotApiConstructChatbotApigetRSSPostsresolverFF07D9C6": { "DependsOn": [ "ChatBotApiConstructChatbotApiproxyResolverFunction22AA16EE", @@ -3350,6 +3209,9 @@ schema { "AURORA_DB_SECRET_ID": { "Ref": "RagEnginesAuroraPgVectorAuroraDatabaseSecretAttachmentA167E2A9", }, + "CHATBOT_FILES_BUCKET_NAME": { + "Ref": "ChatBotApiConstructChatBucketsFilesBucket66FE32D5", + }, "CONFIG_PARAMETER_NAME": { "Ref": "SharedConfig358B4A20", }, @@ -4032,6 +3894,43 @@ schema { }, ], }, + { + "Action": [ + "s3:GetObject*", + "s3:GetBucket*", + "s3:List*", + "s3:DeleteObject*", + "s3:PutObject", + "s3:PutObjectLegalHold", + "s3:PutObjectRetention", + "s3:PutObjectTagging", + "s3:PutObjectVersionTagging", + "s3:Abort*", + ], + "Effect": "Allow", + "Resource": [ + { + "Fn::GetAtt": [ + "ChatBotApiConstructChatBucketsFilesBucket66FE32D5", + "Arn", + ], + }, + { + "Fn::Join": [ + "", + [ + { + "Fn::GetAtt": [ + "ChatBotApiConstructChatBucketsFilesBucket66FE32D5", + "Arn", + ], + }, + "/*", + ], + ], + }, + ], + }, { "Action": [ "s3:GetObject*", diff --git a/tests/chatbot-api/functions/api-handler/routes/documents_test.py b/tests/chatbot-api/functions/api-handler/routes/documents_test.py index 70d22a8f8..5d7b286b4 100644 --- a/tests/chatbot-api/functions/api-handler/routes/documents_test.py +++ b/tests/chatbot-api/functions/api-handler/routes/documents_test.py @@ -34,7 +34,9 @@ def test_file_upload(mocker): - mocker.patch("genai_core.upload.generate_presigned_post", return_value="url") + mocker.patch( + "genai_core.presign.generate_workspace_presigned_post", return_value="url" + ) assert file_upload({"fileName": "fileName.txt", "workspaceId": "id"}) == "url" diff --git a/tests/chatbot-api/functions/api-handler/routes/sessions_test.py b/tests/chatbot-api/functions/api-handler/routes/sessions_test.py index 5fa47e471..8bcd278d5 100644 --- a/tests/chatbot-api/functions/api-handler/routes/sessions_test.py +++ b/tests/chatbot-api/functions/api-handler/routes/sessions_test.py @@ -1,6 +1,7 @@ from pydantic import ValidationError import pytest from genai_core.types import CommonError +from routes.sessions import get_file from routes.sessions import get_sessions from routes.sessions import get_session from routes.sessions import delete_user_sessions @@ -18,6 +19,12 @@ } +def test_get_file_url(mocker): + mocker.patch("genai_core.auth.get_user_id", return_value="userId") + mocker.patch("genai_core.presign.generate_user_presigned_get", return_value="url") + assert get_file("file") == "url" + + def test_get_sessions(mocker): mocker.patch("genai_core.auth.get_user_id", return_value="userId") mocker.patch("genai_core.sessions.list_sessions_by_user_id", return_value=[session])