add multi tenant support

Author: yoavfish

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,31 @@ 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("")
+        print (f"tenant_id={tenant_id}")
         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,8 +1,8 @@
-boto3>=1.34.79
-langchain==0.1.14
-langchain-community>=0.0.31
-langchain-core==0.1.50
-coloredlogs>=15.0.1
-jq==1.7.0
-opensearch-py==2.5.0
-loguru==0.7.2
+boto3
+langchain
+langchain-community
+langchain-core
+coloredlogs
+jq
+opensearch-py
+loguru