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

The variety of generative synthetic intelligence (AI) options is rising inside software program choices, particularly after market-leading foundational fashions (FMs) turned consumable by means of an API utilizing Amazon Bedrock. Amazon Bedrock is a totally managed service that provides a selection of high-performing basis fashions from main AI corporations like AI21 Labs, Anthropic, Cohere, Meta, Stability AI, and Amazon by means of a single API, together with a broad set of capabilities you could construct generative AI functions with safety, privateness, and accountable AI.

Brokers for Amazon Bedrock permits software program builders to finish actions and duties based mostly on consumer enter and group knowledge. A typical problem in multi-tenant choices, akin 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’ll be able to isolate tenants in an utility utilizing completely different multi-tenant structure patterns. In some circumstances, isolation might be achieved by having total stacks of sources devoted to 1 tenant (silo mannequin) with coarse-grained insurance policies to stop cross-tenant entry. In different eventualities, you might need pooled sources (akin to one database desk containing rows from completely different tenants) that require fine-grained insurance policies to regulate entry. Oftentimes, Amazon Internet Providers (AWS) clients design their functions utilizing a mixture of each fashions to steadiness the fashions’ tradeoffs.

Isolating tenants in a pooled mannequin is achieved by utilizing tenant context data in several utility parts. The tenant context might be injected by an authoritative supply, such because the identification supplier (IdP) through the authentication of a consumer. Integrity of the tenant context have to 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 doubtlessly delicate knowledge being disclosed or modified.

FMs act on unstructured knowledge and reply in a probabilistic trend. These properties make FMs unfit to deal with tenant context securely. For instance, FMs are prone to prompt injection, which can be utilized by malicious actors to vary the tenant context. As an alternative, tenant context must be securely handed between deterministic parts of an utility, which may in flip eat FM capabilities, giving the FM solely data that’s already scoped right down to the precise tenant.

On this weblog put up, you’ll discover ways to implement tenant isolation utilizing Amazon Bedrock brokers inside a multi-tenant setting. We’ll exhibit this utilizing a pattern multi-tenant e-commerce utility that gives a service for varied tenants to create on-line shops. This utility makes use of Amazon Bedrock brokers to develop an AI assistant or chatbot able to offering tenant-specific data, akin to return insurance policies and user-specific data 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 data and the consumer degree for user-related knowledge, offering a safe and customized expertise for every tenant and their customers.

Structure overview

Determine 1: Structure of the pattern AI assistant utility

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

1. A tenant consumer indicators in to an identification supplier akin to Amazon Cognito. They get a JSON Internet Token (JWT), which they use for API requests. The JWT comprises 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 utility. The shopper utility sends the query to a GraphQL API endpoint supplied by AWS AppSync, within the type of a GraphQL mutation. You’ll be able to be taught extra about this sample within the weblog put up Construct a Actual-time, WebSockets API for Amazon Bedrock. The shopper utility authenticates to AWS AppSync utilizing the JWT from Amazon Cognito. The consumer is allowed utilizing the Cognito Consumer 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 Answer Utilizing AWS Serverless Providers. 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 by utilizing 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 data from Amazon DynamoDB. Tenant configuration knowledge is saved in a single, shared desk, whereas consumer knowledge is break up in a single desk per tenant. After the right knowledge 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 based mostly on the supplied knowledge.
7. The agent’s response is revealed as one other GraphQL mutation by means of AWS AppSync.
8. The shopper listens to the response utilizing a GraphQL subscription. It renders the response to the consumer after it’s acquired from the server.

Notice that every part on this pattern structure might 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 by means of Amazon API Gateway as a substitute 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 go tenant and consumer context between its completely different layers. Don’t depend on probabilistic parts of your stack, akin to an LLM, to precisely transmit safety data.

How tenant and consumer knowledge is remoted

This part describes how consumer and tenant knowledge 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 by means of 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 obtainable to downstream parts.

Determine 2: Consumer and tenant knowledge 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 accomplished 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 acquired by the subscription. To correlate the mutation request and response within the subscription, the answerId, generated in Step1, is used. By operating the code beneath in a resolver hooked up 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 acquired by the subscription. The ctx variable is populated by AWS AppSync with the request’s payload and metadata such because the consumer identification.

if (!ctx.args.answerId.startsWith(ctx.identification.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 regionally to vary the userId, the authorization test 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 data 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 beneath IAM function that has permissions scoped right down to a particular tenant and generates tenant particular scoped credentials. This function solely grants entry to DynamoDB gadgets 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 identification data and tenant particular scoped credentials are handed to the agent by means of sessionAttributes within the Amazon Bedrock InvokeAgent API name as proven beneath.

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 process to grant the agent entry to solely the database tables and rows for the precise tenant consumer. The duty creates a DynamoDB shopper utilizing the tenant-scoped credentials. Utilizing the scoped shopper, it seems to be up the right order desk title within the tenant configuration and queries the order desk for knowledge:

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.

Stipulations

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

Allow massive language mannequin

An agent wants a big language mannequin (LLM) to cause about the easiest way to fulfil a consumer request and formulate natural-language solutions. Comply with 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 with a standing of Entry granted:

Determine 3: You will have now enabled Anthropic Claude 3 Sonnet in Amazon Bedrock in your AWS account.

Deploy pattern utility

We ready many of the pattern utility’s infrastructure as an AWS Cloud Improvement Package (AWS CDK) undertaking.

You probably have by no means used the CDK within the present account and Area (us-east-1), you have to bootstrap the setting utilizing the next command:

Utilizing your native command line interface, challenge the next instructions to clone the undertaking repository and deploy the CDK undertaking 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 is best to see output just 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 will have them within the subsequent part. When you come again to this walkthrough later, you’ll be able to 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 will have now efficiently deployed the stack referred to as MultiTenantAiAssistantStack.

Begin the frontend and sign up

Now that the backend is deployed and configured, you can begin the frontend in your native machine, which is in-built JavaScript utilizing React. The frontend routinely pulls data from the AWS CDK output, so that you don’t must configure it manually.

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

   cd frontend
   npm set up
   npm run dev

Open the frontend utility by visiting localhost:3000 in your browser. You must see a sign-in web page:

Determine 4: Signal-in display

2. For Username, enter tenant1-user. For Password, enter the password you may 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 data, select Skip.

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

Work together with the agent

You will have accomplished the setup of the structure proven in Determine 1 in your personal setting. You can begin exploring the net utility by your self or comply with the steps instructed beneath.

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

Determine 5: Pattern utility screenshot

You will have now verified the performance of the applying. It’s also possible to attempt to entry knowledge from one other consumer, and you’ll not get a solution because of the scoped IAM coverage. For instance, you’ll be able to modify the agent and hardcode a tenant ID (akin to tenant2). Within the UI, sign up 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’re going to get an AccessDeniedException. It’s also possible to see the error within the CloudWatch Logs for the AgentTask Lambda perform:

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

Add check knowledge

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, challenge this command so as to add an order for tenant1-user:
   ./manage-orders.sh tenant1-user add
2. Return to the net utility and challenge the next immediate:
   Which orders did I place?
   The agent ought to now reply with the order that you simply created.
3. Difficulty the next command to delete the orders for tenant1-user:
   ./manage-orders.sh tenant1-user clear

Repeat steps 1 by means of 3 with a number of orders. You’ll be able to create a brand new consumer in Amazon Cognito and sign up to see that no knowledge from different customers might be accessed. The implementation is detailed in Determine 2.

Clear up

To keep away from incurring future expenses, 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 knowledge privateness and stopping unauthorized entry. By following the strategy outlined on this weblog put up, you’ll be able to create an AI assistant that isolates tenants whereas utilizing the ability of enormous language fashions.

The important thing factors to recollect are:

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

By following these ideas, you’ll be able to construct AI-powered functions 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.

Keep in mind, the pattern utility demonstrated on this weblog put up is only one technique to strategy multi-tenant AI assistants. Relying in your particular necessities, you may must adapt the structure or use completely different AWS companies.

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. You probably have any questions, you can begin a brand new thread on AWS re:Post or attain out to AWS Assist.

Concerning the authors

Ulrich Hinze is a Options Architect at AWS. He companions with software program corporations 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 knowledge streaming professional at AWS. He’s a technologist that helps clients in Europe succeed and innovate by fixing enterprise challenges utilizing AWS Cloud companies. In addition to working as a Options Architect, Florian is a passionate mountaineer and has climbed a number of the highest mountains throughout Europe.