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 via an API utilizing Amazon Bedrock. Amazon Bedrock is a completely 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 you should 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 person enter and group information. A typical problem in multi-tenant choices, equivalent to software program as a service (SaaS) merchandise, is tenant isolation. Tenant isolation makes certain 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 software 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 forestall cross-tenant entry. In different eventualities, you may need pooled sources (equivalent to one database desk containing rows from completely 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 steadiness the fashions’ tradeoffs.

Isolating tenants in a pooled mannequin is achieved by utilizing tenant context data in numerous software elements. The tenant context might be injected by an authoritative supply, such because the id supplier (IdP) throughout the authentication of a person. Integrity of the tenant context have to be preserved all through the system to forestall malicious customers from performing on behalf of a tenant that they shouldn’t have entry to, leading to doubtlessly delicate information being disclosed or modified.

FMs act on unstructured information and reply in a probabilistic vogue. 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 a substitute, tenant context needs to be securely handed between deterministic elements of an software, which might in flip devour 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 show this utilizing a pattern multi-tenant e-commerce software that gives a service for varied 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 data, equivalent 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 stage for return coverage data and the person stage 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 completely different elements this resolution is utilizing.

1. A tenant person indicators in to an id supplier equivalent to Amazon Cognito. They get a JSON Net Token (JWT), which they use for API requests. The JWT incorporates claims such because the person ID (or topic, sub), which identifies the tenant person, and the tenantId, which defines which tenant the person belongs to.
2. The tenant person inputs their query into the shopper software. The shopper software sends the query to a GraphQL API endpoint offered by AWS AppSync, within the type of a GraphQL mutation. You’ll be able to study extra about this sample within the weblog put up Construct a Actual-time, WebSockets API for Amazon Bedrock. The shopper software authenticates to AWS AppSync utilizing the JWT from Amazon Cognito. The person is permitted utilizing the Cognito Person Swimming pools integration.
3. The GraphQL mutation invokes utilizing the EventBridge resolver. The occasion triggers an AWS Lambda operate utilizing an EventBridge rule.
4. The Lambda operate calls the Amazon Bedrock InvokeAgent API. This operate 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, person 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 person request by utilizing the reasoning capabilities of the related giant 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 person ID for additional processing steps.
6. This Lambda operate makes use of the offered tenant particular scoped credentials and tenant ID to fetch data from Amazon DynamoDB. Tenant configuration information is saved in a single, shared desk, whereas person information is cut 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 person primarily based on the offered information.
7. The agent’s response is revealed 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 person after it’s obtained from the server.

Be aware 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 via 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 just securely cross tenant and person context between its completely different layers. Don’t depend on probabilistic elements of your stack, equivalent to an LLM, to precisely transmit safety data.

How tenant and person information is remoted

This part describes how person 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 person and tenant ID, is made obtainable to downstream elements.

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 person ID (saved within the sub area 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 = person?.userId + "." + uuidv4();

2. The frontend makes use of the AWS Amplify SDK, which takes care of authenticating the GraqhQL request. That is completed 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 person 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 beneath in a resolver hooked up to a subscription, person 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 person id.

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

Be aware that the authorization is checked in opposition to the cryptographically signed JWT from the Amazon Cognito person pool. Therefore, even when a malicious person might tamper with the token regionally to alter 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 operate will get the person 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 operate assumes the beneath IAM position that has permissions scoped right down to a particular tenant and generates tenant particular scoped credentials. This position 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],
		situations: {
			"ForAllValues:StringEquals": {
				"dynamodb:LeadingKeys": [
					"${aws:PrincipalTag/TenantId}"
				]}}}),
        new PolicyStatement({
actions: ["dynamodb:Query"], sources: ["arn:aws:dynamodb:*:*:table/${aws:PrincipalTag/TenantId}-orders"] }) ]

5. This id data and tenant particular scoped credentials are handed to the agent via 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 throughout the agent activity to grant the agent entry to solely the database tables and rows for the precise tenant person. The duty creates a DynamoDB shopper utilizing the tenant-scoped credentials. Utilizing the scoped shopper, it appears up the right order desk identify 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 operate, 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 individual AWS account.

Stipulations

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

Allow giant language mannequin

An agent wants a big language mannequin (LLM) to purpose about the easiest way to fulfil a person 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 with a standing of Entry granted:

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

Deploy pattern software

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

When you’ve got by no means used the CDK within the present account and Area (us-east-1), you could bootstrap the setting utilizing the next command:

Utilizing your native command line interface, problem 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’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/...

✨  Whole 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. Be aware down the passwords for the three customers from the CDK output, labelled MultiTenantAiAssistantStack.tenantXPassword. You’ll need them within the subsequent part. Should you come again to this walkthrough later, you possibly can retrieve these values from the file cdk/cdk-output.json generated by the CDK. Be aware that these are solely preliminary passwords and should be modified on first sign-in of every person.

You may have now efficiently deployed the stack referred to 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 robotically pulls data from the AWS CDK output, so that you don’t have to configure it manually.

1. Challenge 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

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 person.
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 may have accomplished the setup of the structure proven in Determine 1 in your individual setting. You can begin exploring the online software by your self or observe the steps steered beneath.

1. Beneath Enter your Immediate, enter the next query logged in as tenant1-user:
   What's your return coverage?
   It is best to obtain a response that you would be able to return objects 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?
   It is best to obtain a response that you haven’t positioned any orders but.

Determine 5: Pattern software screenshot

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

[ERROR] ClientError: An error occurred (AccessDeniedException) when calling the Question operation: Person: *****/agentTaskLambda isn't approved 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 now have 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 just created.
3. Challenge 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’ll be able to create a brand new person in Amazon Cognito and check in to see that no information from different customers might be accessed. The implementation is detailed in Determine 2.

Clear up

To keep away from incurring future costs, 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 put up, 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, at all times implement tenant isolation (leverage IAM the place ever attainable).
2. Securely cross tenant and person context between deterministic elements of your software, 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 cross alongside tenant context.
4. Implement isolation at completely different layers of your software to confirm that customers can solely entry information and sources related to their respective tenant and person context.

By following these rules, you possibly can construct AI-powered purposes that present a customized 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 software demonstrated on this weblog put up is only one approach to strategy multi-tenant AI assistants. Relying in your particular necessities, you may have to 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. When you’ve got any questions, you can begin a brand new thread on AWS re:Submit 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 professional at AWS. He’s a technologist that helps clients in Europe succeed and innovate by fixing enterprise challenges utilizing AWS Cloud companies. Apart from working as a Options Architect, Florian is a passionate mountaineer and has climbed a few of the highest mountains throughout Europe.