Implementing tenant isolation utilizing Brokers for Amazon Bedrock in a multi-tenant surroundings

The variety of generative synthetic intelligence (AI) options is rising inside software program choices, particularly after market-leading foundational fashions (FMs) grew to become consumable via an API utilizing Amazon Bedrock. Amazon Bedrock is a totally managed service that provides a selection of high-performing basis fashions from main AI firms like AI21 Labs, Anthropic, Cohere, Meta, Stability AI, and Amazon via a single API, together with a broad set of capabilities it’s good to construct generative AI purposes with safety, privateness, and accountable AI.

Brokers for Amazon Bedrock allows software program builders to finish actions and duties primarily based on consumer enter and group information. A typical problem in multi-tenant choices, corresponding to software program as a service (SaaS) merchandise, is tenant isolation. Tenant isolation makes positive every tenant can entry solely their very own sources—even when all tenants run on shared infrastructure.

You possibly can isolate tenants in an software utilizing totally different multi-tenant structure patterns. In some instances, isolation will be achieved by having whole stacks of sources devoted to 1 tenant (silo mannequin) with coarse-grained insurance policies to stop cross-tenant entry. In different situations, you may need pooled sources (corresponding to one database desk containing rows from totally different tenants) that require fine-grained insurance policies to manage entry. Oftentimes, Amazon Net Companies (AWS) clients design their purposes utilizing a mixture of each fashions to stability the fashions’ tradeoffs.

Isolating tenants in a pooled mannequin is achieved through the use of tenant context info in numerous software parts. The tenant context will be injected by an authoritative supply, such because the id supplier (IdP) in the course of the authentication of a consumer. Integrity of the tenant context should be preserved all through the system to stop malicious customers from appearing on behalf of a tenant that they shouldn’t have entry to, leading to probably delicate information being disclosed or modified.

FMs act on unstructured information and reply in a probabilistic trend. These properties make FMs unfit to deal with tenant context securely. For instance, FMs are inclined to immediate injection, which can be utilized by malicious actors to alter the tenant context. As an alternative, tenant context ought to be securely handed between deterministic parts of an software, which may in flip eat FM capabilities, giving the FM solely info that’s already scoped all the way down to the particular tenant.

On this weblog publish, you’ll learn to implement tenant isolation utilizing Amazon Bedrock brokers inside a multi-tenant surroundings. We’ll display this utilizing a pattern multi-tenant e-commerce software that gives a service for numerous tenants to create on-line shops. This software makes use of Amazon Bedrock brokers to develop an AI assistant or chatbot able to offering tenant-specific info, corresponding to return insurance policies and user-specific info like order counts and standing updates. This structure showcases how you should utilize pooled Amazon Bedrock brokers and implement tenant isolation at each the tenant degree for return coverage info and the consumer degree for user-related information, offering a safe and customized expertise for every tenant and their customers.

Structure overview

Determine 1: Structure of the pattern AI assistant software

Let’s discover the totally different parts this resolution is utilizing.

1. A tenant consumer indicators in to an id supplier corresponding to Amazon Cognito. They get a JSON Net Token (JWT), which they use for API requests. The JWT incorporates claims such because the consumer ID (or topic, sub), which identifies the tenant consumer, and the tenantId, which defines which tenant the consumer belongs to.
2. The tenant consumer inputs their query into the shopper software. The shopper software sends the query to a GraphQL API endpoint supplied by AWS AppSync, within the type of a GraphQL mutation. You possibly can be taught extra about this sample within the weblog publish Construct a Actual-time, WebSockets API for Amazon Bedrock. The shopper software authenticates to AWS AppSync utilizing the JWT from Amazon Cognito. The consumer is allowed utilizing the Cognito Person Swimming pools integration.
3. The GraphQL mutation invokes utilizing the EventBridge resolver. The occasion triggers an AWS Lambda perform utilizing an EventBridge rule.
4. The Lambda perform calls the Amazon Bedrock InvokeAgent API. This perform makes use of a tenant isolation coverage to scope the permissions and generates tenant particular scoped credentials. Extra about this may be learn within the weblog Constructing a Multi-Tenant SaaS Resolution Utilizing AWS Serverless Companies. Then, it sends the tenant ID, consumer ID and tenant particular scoped credentials to this API utilizing the sessionAttributes parameter from the agent’s sessionState.
5. The Amazon Bedrock agent determines what it must do to fulfill the consumer request through the use of the reasoning capabilities of the related massive language mannequin (LLM). A wide range of LLMs can be utilized, and for this resolution we used Anthropic Claude 3 Sonnet. It passes the sessionAttributes object to an motion group decided to assist with the request, thereby securely forwarding tenant and consumer ID for additional processing steps.
6. This Lambda perform makes use of the supplied tenant particular scoped credentials and tenant ID to fetch info from Amazon DynamoDB. Tenant configuration information is saved in a single, shared desk, whereas consumer information is break up in a single desk per tenant. After the right information is fetched, it’s returned to the agent. The agent interacts with the LLM for the second time to formulate a natural-language reply to the consumer primarily based on the supplied information.
7. The agent’s response is printed as one other GraphQL mutation via AWS AppSync.
8. The shopper listens to the response utilizing a GraphQL subscription. It renders the response to the consumer after it’s obtained from the server.

Notice that every part on this pattern structure will be modified to suit into your pre-existing structure and information within the group. For instance, you may select to make use of a WebSocket implementation via Amazon API Gateway as an alternative of utilizing GraphQL or implement a synchronous request and response sample. Whichever know-how stack you select to make use of, confirm that you simply securely cross tenant and consumer context between its totally different layers. Don’t depend on probabilistic parts of your stack, corresponding to an LLM, to precisely transmit safety info.

How tenant and consumer information is remoted

This part describes how consumer and tenant information is remoted when a request is processed all through the system. Every step is mentioned in additional element following the diagram. For every immediate within the UI, the frontend sends the immediate as a mutation request to the AWS AppSync API and listens for the response via a subscription, as defined in step 8 of Determine 1 proven above. The subscription is required to obtain the reply from the immediate, because the agent is invoked asynchronously. Each the request and response are authenticated utilizing Amazon Cognito, and the request’s context, together with consumer and tenant ID, is made out there to downstream parts.

Determine 2: Person and tenant information isolation

1. For every immediate created within the pattern UI, a novel ID(answerId) is generated. The answerId is required to correlate the enter immediate with the reply from the agent. It makes use of the Cognito consumer ID (saved within the sub subject within the JWT and accessible as userId within the AWS Amplify SDK) as a prefix to allow fine-grained permissions. That is defined in additional depth in step 3. The answerId is generated within the web page.tsx file:

const answerId = consumer?.userId + "." + uuidv4();

2. The frontend makes use of the AWS Amplify SDK, which takes care of authenticating the GraqhQL request. That is carried out for the immediate request (a GraphQL mutation request) and for the response (a GraphQL subscription which listens to a solution to the immediate). The authentication mode is about within the tsx file. Amplify makes use of the Amazon Cognito consumer pool it has been configured with. Additionally, the beforehand generated answerId is used as a novel identifier for the request.

await shopper.graphql({
	authMode: "userPool",
    ...
    variables: {
      answerId,
      ...
    },
  });

3. The frontend sends the GraphQL mutation request and the response is obtained by the subscription. To correlate the mutation request and response within the subscription, the answerId, generated in Step1, is used. By working the code under in a resolver connected to a subscription, consumer isolation is enforced. Customers can’t subscribe to arbitrary mutations and obtain their response. The code verifies that that the userId within the mutation request matches the userId within the response obtained by the subscription. The ctx variable is populated by AWS AppSync with the request’s payload and metadata such because the consumer id.

if (!ctx.args.answerId.startsWith(ctx.id.sub + ".")) {
  util.unauthorized()
}

Notice that the authorization is checked towards the cryptographically signed JWT from the Amazon Cognito consumer pool. Therefore, even when a malicious consumer may tamper with the token domestically to alter the userId, the authorization verify would nonetheless fail.

4. The userId and tenantId (from the AWS AppSync context) is handed on to Amazon EventBridge and to AWS Lambda, which invokes the Agent. The Lambda perform will get the consumer info from the occasion object in file invokeAgent/index.py:

tenant_id = occasion["detail"]["identity"]["claims"]["custom:tenantId"]
user_id = occasion["detail"]["identity"]["claims"]["sub"]

The Lambda perform assumes the under IAM function that has permissions scoped all the way down to a particular tenant and generates tenant particular scoped credentials. This function solely grants entry to DynamoDB objects which has the given tenant ID because the main key.

statements: [
	new PolicyStatement({
		actions: ["dynamodb:Query"],
		sources: [tenantConfigurationTable.tableArn],
		circumstances: {
			"ForAllValues:StringEquals": {
				"dynamodb:LeadingKeys": [
					"${aws:PrincipalTag/TenantId}"
				]}}}),
        new PolicyStatement({
actions: ["dynamodb:Query"], sources: ["arn:aws:dynamodb:*:*:table/${aws:PrincipalTag/TenantId}-orders"] }) ]

5. This id info and tenant particular scoped credentials are handed to the agent via sessionAttributes within the Amazon Bedrock InvokeAgent API name as proven under.

response = shopper.invoke_agent(
    ...
sessionState={
"sessionAttributes": {
		"tenantId": tenant_id,
		"userId": user_id,
		"accessKeyId": credentials["accessKeyId"],
		"secretAccessKey":credentials["secretAccessKey"],
		"sessionToken": credentials["sessionToken"],
},)

6. The sessionAttributes are used inside the agent activity to grant the agent entry to solely the database tables and rows for the particular tenant consumer. The duty creates a DynamoDB shopper utilizing the tenant-scoped credentials. Utilizing the scoped shopper, it seems up the right order desk title within the tenant configuration and queries the order desk for information:

tenant_id = occasion["sessionAttributes"]["tenantId"]
user_id = occasion["sessionAttributes"]["userId"]
access_key_id = occasion["sessionAttributes"]["accessKeyId"]
secret_access_key = occasion["sessionAttributes"]["secretAccessKey"]
session_token = occasion["sessionAttributes"]["sessionToken"]

dynamodb = boto3.useful resource(
        "dynamodb",
        aws_access_key_id=occasion["sessionAttributes"]["accessKeyId"],
        aws_secret_access_key=occasion["sessionAttributes"]["secretAccessKey"],
        aws_session_token=occasion["sessionAttributes"]["sessionToken"],
    )
tenant_config_table_name = os.getenv("TENANT_CONFIG_TABLE_NAME")
tenant_config_table = dynamodb.Desk(tenant_config_table_name)

orders_table_name = tenant_config_table.question(
    KeyConditionExpression=Key("tenantId").eq(tenant_id)
)["Items"][0]["ordersTableName"]
...
orders_table.question(KeyConditionExpression=Key("userId").eq(user_id))[
    "Items"
]

When modifying / debugging this perform, just remember to don’t log any credentials or the entire occasion object.

Walkthrough

On this part, you’ll arrange the pattern AI assistant described within the earlier sections in your personal AWS account.

Conditions

For this walkthrough, it’s best to have the next conditions:

Allow massive language mannequin

An agent wants a big language mannequin (LLM) to purpose about the easiest way to fulfil a consumer request and formulate natural-language solutions. Observe the Amazon Bedrock mannequin entry documentation to allow Anthropic Claude 3 Sonnet mannequin entry within the us-east-1 (N. Virginia) Area. After enabling the LLM, you will note the next display screen with a standing of Entry granted:

Determine 3: You’ve got now enabled Anthropic Claude 3 Sonnet in Amazon Bedrock on your AWS account.

Deploy pattern software

We ready many of the pattern software’s infrastructure as an AWS Cloud Growth Package (AWS CDK) mission.

If in case you have by no means used the CDK within the present account and Area (us-east-1), you should bootstrap the surroundings utilizing the next command:

Utilizing your native command line interface, problem the next instructions to clone the mission repository and deploy the CDK mission to your AWS account:

git clone https://github.com/aws-samples/multi-tenant-ai-assistant
cd multi-tenant-ai-assistant/cdk
npm set up
cdk deploy 
cd ..

This takes about 3 minutes, after which it’s best to see output much like the next:

✅ MultiTenantAiAssistantStack

✨  Deployment time: 132.24s

Outputs:
MultiTenantAiAssistantStack.appClientId = ...
MultiTenantAiAssistantStack.graphqlEndpoint = https://...
MultiTenantAiAssistantStack.tenant1Password = Preliminary-...
MultiTenantAiAssistantStack.tenant2Password = Preliminary-...
MultiTenantAiAssistantStack.tenant3Password = Preliminary-...
MultiTenantAiAssistantStack.userPoolId = us-east-1_...
Stack ARN:
arn:aws:cloudformation:us-east-1:...:stack/MultiTenantAiAssistantStack/...

✨  Complete time: 179.54s

Along with the AWS sources proven in Figure1, this AWS CDK stack provisions three customers, every for a separate tenant, into your AWS account. Notice down the passwords for the three customers from the CDK output, labelled MultiTenantAiAssistantStack.tenantXPassword. You have to them within the subsequent part. For those who come again to this walkthrough later, you possibly can retrieve these values from the file cdk/cdk-output.json generated by the CDK. Notice that these are solely preliminary passwords and have to be modified on first sign-in of every consumer.

You’ve got now efficiently deployed the stack known as MultiTenantAiAssistantStack.

Begin the frontend and check in

Now that the backend is deployed and configured, you can begin the frontend in your native machine, which is inbuilt JavaScript utilizing React. The frontend mechanically pulls info from the AWS CDK output, so that you don’t have to configure it manually.

1. Concern the next instructions to put in dependencies and begin the native webserver:

   cd frontend
   npm set up
   npm run dev

Open the frontend software by visiting localhost:3000 in your browser. It is best to see a sign-in web page:

Determine 4: Signal-in display screen

2. For Username, enter tenant1-user. For Password, enter the password you will have beforehand retrieved from CDK output.
3. Set a brand new password for the consumer.
4. On the web page Account restoration requires verified contact info, select Skip.

You’re now signed in and might begin interacting with the agent.

Work together with the agent

You’ve got accomplished the setup of the structure proven in Determine 1 in your personal surroundings. You can begin exploring the online software by your self or observe the steps advised under.

1. Underneath Enter your Immediate, enter the next query logged in as tenant1-user:
   What's your return coverage?
   It is best to obtain a response which you could return objects for 10 days. Tenant 2 has a return coverage of 20 days, tenant 3 of 30 days.
2. Underneath Enter your Immediate, enter the next query:
   Which orders did I place?
   It is best to obtain a response that you haven’t positioned any orders but.

Determine 5: Pattern software screenshot

You’ve got now verified the performance of the applying. You may also attempt to entry information from one other consumer, and you’ll not get a solution as a result of scoped IAM coverage. For instance, you possibly can modify the agent and hardcode a tenant ID (corresponding to tenant2). Within the UI, check in because the tenant1 consumer and you will note that with the generated tenant1 scoped credentials you won’t be able to entry tenant2 sources and you’ll get an AccessDeniedException. You may also see the error within the CloudWatch Logs for the AgentTask Lambda perform:

[ERROR] ClientError: An error occurred (AccessDeniedException) when calling the Question operation: Person: *****/agentTaskLambda shouldn't be licensed to carry out: dynamodb:Question on useful resource: TABLE  as a result of no identity-based coverage permits the dynamodb:Question motion

Add take a look at information

To simplify the method of including orders to your database, we’ve written a bash script that inserts entries into the order tables.

1. In your CLI, from the repository root folder, problem this command so as to add an order for tenant1-user:
   ./manage-orders.sh tenant1-user add
2. Return to the online software and problem the next immediate:
   Which orders did I place?
   The agent ought to now reply with the order that you simply created.
3. Concern the next command to delete the orders for tenant1-user:
   ./manage-orders.sh tenant1-user clear

Repeat steps 1 via 3 with a number of orders. You possibly can create a brand new consumer in Amazon Cognito and check in to see that no information from different customers will be accessed. The implementation is detailed in Determine 2.

Clear up

To keep away from incurring future fees, delete the sources created throughout this walkthrough. From the cdk folder of the repository, run the next command:

cdk destroy

Conclusion

Enabling safe multi-tenant capabilities in AI assistants is essential for sustaining information privateness and stopping unauthorized entry. By following the strategy outlined on this weblog publish, you possibly can create an AI assistant that isolates tenants whereas utilizing the ability of huge language fashions.

The important thing factors to recollect are:

1. When constructing multi-tenant SaaS purposes, all the time implement tenant isolation (leverage IAM the place ever doable).
2. Securely cross tenant and consumer context between deterministic parts of your software, with out counting on an AI mannequin to deal with this delicate info.
3. Use Brokers for Amazon Bedrock to assist construct an AI assistant that may securely cross alongside tenant context.
4. Implement isolation at totally different layers of your software to confirm that customers can solely entry information and sources related to their respective tenant and consumer context.

By following these ideas, you possibly can construct AI-powered purposes that present a personalised expertise to customers whereas sustaining strict isolation and safety. As AI capabilities proceed to advance, it’s important to design architectures that use these applied sciences responsibly and securely.

Bear in mind, the pattern software demonstrated on this weblog publish is only one option to strategy multi-tenant AI assistants. Relying in your particular necessities, you may have to adapt the structure or use totally different AWS providers.

To proceed studying about generative AI patterns on AWS, go to the AWS Machine Studying Weblog. To discover SaaS on AWS, begin by visiting our SaaS touchdown web page. If in case you have any questions, you can begin a brand new thread on AWS re:Put up or attain out to AWS Assist.

Concerning the authors

Ulrich Hinze is a Options Architect at AWS. He companions with software program firms to architect and implement cloud-based options on AWS. Earlier than becoming a member of AWS, he labored for AWS clients and companions in software program engineering, consulting, and structure roles for 8+ years.

Florian Mair is a Senior Options Architect and information streaming skilled at AWS. He’s a technologist that helps clients in Europe succeed and innovate by fixing enterprise challenges utilizing AWS Cloud providers. In addition to working as a Options Architect, Florian is a passionate mountaineer and has climbed a few of the highest mountains throughout Europe.