Merge pull request #12 from aws-samples/multi-tenant

Multi tenant

Author: omerh

README.md

@@ -61,6 +61,8 @@ AI21 Labs:
     >>- `--early-stop` to load only 100 embedded documents into OpenSearch
     >>- `--index` to use a different index than the default **rag**
     >>- `--region` in case you are not using the default **us-east-1**
+    >>- `--multi-tenant` to use multi tenancy, will load data with tenant IDs (1-5)
+    
 
 3. Now that we have embedded text, into our OpenSearch cluster, we can start querying our LLM model Titan text in Amazon Bedrock with RAG
 
@@ -72,6 +74,8 @@ AI21 Labs:
     >>- `--index` to use a different index than the default **rag**
     >>- `--region` in case you are not using the default **us-east-1**
     >>- `--bedrock-model-id` to choose different models than Anthropic's Claude v2
+    >>- `--tenant-id` to filter only a specific tenant ID
+
 
 ### Cleanup
 

ask-bedrock-with-rag.py

@@ -22,6 +22,7 @@ def parse_args():
     parser.add_argument("--ask", type=str, default="What is the meaning of <3?")
     parser.add_argument("--index", type=str, default="rag")
     parser.add_argument("--region", type=str, default="us-east-1")
+    parser.add_argument("--tenant-id", type=str, default=None)
     parser.add_argument("--bedrock-model-id", type=str, default="anthropic.claude-3-sonnet-20240229-v1:0")
     parser.add_argument("--bedrock-embedding-model-id", type=str, default="amazon.titan-embed-text-v1")
 
@@ -68,6 +69,7 @@ def main():
     bedrock_model_id = args.bedrock_model_id
     bedrock_embedding_model_id = args.bedrock_embedding_model_id
     question = args.ask
+    tenant_id = args.tenant_id
     logger.info(f"Question provided: {question}")
 
     # Creating all clients for chain
@@ -87,23 +89,30 @@ def main():
     Answer:""")
 
     docs_chain = create_stuff_documents_chain(bedrock_llm, prompt)
+
+    search_kwargs = {}
+    if tenant_id: 
+        search_kwargs["filter"] = {
+            "term": {
+                "tenant_id": tenant_id
+        }
+    }
+
     retrieval_chain = create_retrieval_chain(
-        retriever=opensearch_vector_search_client.as_retriever(),
+        retriever=opensearch_vector_search_client.as_retriever(search_kwargs=search_kwargs),
         combine_docs_chain = docs_chain
     )
 
     logger.info(f"Invoking the chain with KNN similarity using OpenSearch, Bedrock FM {bedrock_model_id}, and Bedrock embeddings with {bedrock_embedding_model_id}")
     response = retrieval_chain.invoke({"input": question})
 
-    print("")
     logger.info("These are the similar documents from OpenSearch based on the provided query:")
     source_documents = response.get('context')
     for d in source_documents:
-        print("")
         logger.info(f"Text: {d.page_content}")
 
     print("")
-    logger.info(f"The answer from Bedrock {bedrock_model_id} is: {response.get('answer')}")
+    logger.info(f"The answer from Bedrock!!!!! {bedrock_model_id} is: {response.get('answer')}")
 
 
 if __name__ == "__main__":

load-data-to-opensearch.py

@@ -5,6 +5,8 @@
 from loguru import logger
 import sys
 import os
+import random
+
 
 
 # logger
@@ -18,6 +20,7 @@ def parse_args():
     parser.add_argument("--early-stop", type=bool, default=0)
     parser.add_argument("--index", type=str, default="rag")
     parser.add_argument("--region", type=str, default="us-east-1")
+    parser.add_argument("--multi-tenant", type=bool, default=0)
 
     return parser.parse_known_args()
 
@@ -33,16 +36,28 @@ def create_vector_embedding_with_bedrock(text, name, bedrock_client):
     modelId = "amazon.titan-embed-text-v1"
     accept = "application/json"
     contentType = "application/json"
+    args, _ = parse_args()
+    multi_tenant = args.multi_tenant
 
     response = bedrock_client.invoke_model(
         body=body, modelId=modelId, accept=accept, contentType=contentType
     )
     response_body = json.loads(response.get("body").read())
 
     embedding = response_body.get("embedding")
-    return {"_index": name, "text": text, "vector_field": embedding}
 
-            
+    document = {
+        "_index": name,
+        "text": text,
+        "vector_field": embedding
+    }
+    
+
+    if multi_tenant == 1:
+        document["tenant_id"] = random.randint(1, 5)
+
+    return document
+
 def main():
     logger.info("Starting")
 
@@ -52,9 +67,13 @@ def main():
     args, _ = parse_args()
     region = args.region
     name = args.index
+    multi_tenant = args.multi_tenant
+
 
     # Prepare OpenSearch index with vector embeddings index mapping
-    logger.info(f"recreating opensearch index: {args.recreate}, using early stop: {args.early_stop} to insert only {early_stop_record_count} records")
+    logger.info(f"Recreating opensearch index: {args.recreate}, using early stop: {args.early_stop} to insert only {early_stop_record_count} records")
+    if multi_tenant:
+        logger.info("Using multi tenant mode")    
     logger.info("Preparing OpenSearch Index")
     opensearch_password = secret.get_secret(name, region)
     opensearch_client =  opensearch.get_opensearch_cluster_client(name, opensearch_password, region)

requirements.txt

@@ -1,6 +1,6 @@
-boto3>=1.34.79
+boto3>=1.35.73
 langchain==0.1.14
-langchain-community>=0.0.31
+langchain-community==0.0.36
 langchain-core==0.1.50
 coloredlogs>=15.0.1
 jq==1.7.0