update versions, using langchain community, added loguru

Author: omerh

ask-bedrock-with-rag.py

@@ -1,18 +1,20 @@
-import coloredlogs
-import logging
 import argparse
 from utils import opensearch, secret
-from langchain.embeddings import BedrockEmbeddings
-from langchain.vectorstores import OpenSearchVectorSearch
+from langchain_community.embeddings import BedrockEmbeddings
+from langchain_community.vectorstores import OpenSearchVectorSearch
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
 from langchain.llms.bedrock import Bedrock
 import boto3
+from loguru import logger
+import sys
+import os
 
 
-coloredlogs.install(fmt='%(asctime)s %(levelname)s %(message)s', datefmt='%H:%M:%S', level='INFO')
-logging.basicConfig(level=logging.INFO) 
-logger = logging.getLogger(__name__)
+# logger
+logger.remove()
+logger.add(sys.stdout, level=os.getenv("LOG_LEVEL", "INFO"))
+
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -58,7 +60,7 @@ def create_bedrock_llm(bedrock_client, model_version_id):
 
 
 def main():
-    logging.info("Starting")
+    logger.info("Starting")
     args, _ = parse_args()
     region = args.region
     index_name = args.index
@@ -78,7 +80,7 @@ def main():
         question = args.ask
     else:
         question = "what is the meaning of <3?"
-        logging.info(f"No question provided, using default question {question}")
+        logger.info(f"No question provided, using default question {question}")
 
     prompt_template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. don't include harmful content
 
@@ -90,23 +92,23 @@ def main():
         template=prompt_template, input_variables=["context", "question"]
     )
 
-    logging.info(f"Starting the chain with KNN similarity using OpenSearch, Bedrock FM {bedrock_model_id}, and Bedrock embeddings with {bedrock_embedding_model_id}")
+    logger.info(f"Starting the chain with KNN similarity using OpenSearch, Bedrock FM {bedrock_model_id}, and Bedrock embeddings with {bedrock_embedding_model_id}")
     qa = RetrievalQA.from_chain_type(llm=bedrock_llm, 
                                      chain_type="stuff", 
                                      retriever=opensearch_vector_search_client.as_retriever(),
                                      return_source_documents=True,
                                      chain_type_kwargs={"prompt": PROMPT, "verbose": True},
                                      verbose=True)
 
-    response = qa(question, return_only_outputs=False)
+    response = qa.invoke(question, return_only_outputs=False)
 
-    logging.info("This are the similar documents from OpenSearch based on the provided query")
+    logger.info("This are the similar documents from OpenSearch based on the provided query")
     source_documents = response.get('source_documents')
     for d in source_documents:
-        logging.info(f"With the following similar content from OpenSearch:\n{d.page_content}\n")
-        logging.info(f"Text: {d.metadata['text']}")
+        logger.info(f"With the following similar content from OpenSearch:\n{d.page_content}\n")
+        logger.info(f"Text: {d.metadata['text']}")
 
-    logging.info(f"The answer from Bedrock {bedrock_model_id} is: {response.get('result')}")
+    logger.info(f"The answer from Bedrock {bedrock_model_id} is: {response.get('result')}")
 
 
 if __name__ == "__main__":

load-data-to-opensearch.py

@@ -1,13 +1,15 @@
-import logging
-import coloredlogs
 import json
 import argparse
 import boto3
 from utils import dataset, secret, opensearch
+from loguru import logger
+import sys
+import os
 
-coloredlogs.install(fmt='%(asctime)s %(levelname)s %(message)s', datefmt='%H:%M:%S', level='INFO')
-logging.basicConfig(level=logging.INFO) 
-logger = logging.getLogger(__name__)
+
+# logger
+logger.remove()
+logger.add(sys.stdout, level=os.getenv("LOG_LEVEL", "INFO"))
 
 
 def parse_args():
@@ -42,7 +44,7 @@ def create_vector_embedding_with_bedrock(text, name, bedrock_client):
 
 
 def main():
-    logging.info("Starting")
+    logger.info("Starting")
 
     dataset_url = "https://huggingface.co/datasets/sentence-transformers/embedding-training-data/resolve/main/gooaq_pairs.jsonl.gz"
     early_stop_record_count = 100
@@ -52,29 +54,29 @@ def main():
     name = args.index
 
     # Prepare OpenSearch index with vector embeddings index mapping
-    logging.info(f"recreating opensearch index: {args.recreate}, using early stop: {args.early_stop} to insert only {early_stop_record_count} records")
-    logging.info("Preparing OpenSearch Index")
+    logger.info(f"recreating opensearch index: {args.recreate}, using early stop: {args.early_stop} to insert only {early_stop_record_count} records")
+    logger.info("Preparing OpenSearch Index")
     opensearch_password = secret.get_secret(name, region)
     opensearch_client =  opensearch.get_opensearch_cluster_client(name, opensearch_password, region)
 
     # Check if to delete OpenSearch index with the argument passed to the script --recreate 1
     if args.recreate:
         response = opensearch.delete_opensearch_index(opensearch_client, name)
         if response:
-            logging.info("OpenSearch index successfully deleted")
+            logger.info("OpenSearch index successfully deleted")
 
-    logging.info(f"Checking if index {name} exists in OpenSearch cluster")
+    logger.info(f"Checking if index {name} exists in OpenSearch cluster")
     exists = opensearch.check_opensearch_index(opensearch_client, name)    
     if not exists:
-        logging.info("Creating OpenSearch index")
+        logger.info("Creating OpenSearch index")
         success = opensearch.create_index(opensearch_client, name)
         if success:
-            logging.info("Creating OpenSearch index mapping")
+            logger.info("Creating OpenSearch index mapping")
             success = opensearch.create_index_mapping(opensearch_client, name)
-            logging.info(f"OpenSearch Index mapping created")
+            logger.info(f"OpenSearch Index mapping created")
 
     # Download sample dataset from HuggingFace 
-    logging.info("Downloading dataset from HuggingFace")        
+    logger.info("Downloading dataset from HuggingFace")        
     compressed_file_path = dataset.download_dataset(dataset_url)
     if compressed_file_path is not None:
         file_path = dataset.decompress_dataset(compressed_file_path)
@@ -86,7 +88,7 @@ def main():
 
     # Vector embedding using Amazon Bedrock Titan text embedding
     all_json_records = []
-    logging.info(f"Creating embeddings for records")
+    logger.info(f"Creating embeddings for records")
 
     # using the arg --early-stop
     i = 0
@@ -96,24 +98,24 @@ def main():
             if i > early_stop_record_count:
                 # Bulk put all records to OpenSearch
                 success, failed = opensearch.put_bulk_in_opensearch(all_json_records, opensearch_client)
-                logging.info(f"Documents saved {success}, documents failed to save {failed}")
+                logger.info(f"Documents saved {success}, documents failed to save {failed}")
                 break
         records_with_embedding = create_vector_embedding_with_bedrock(record, name, bedrock_client)
-        logging.info(f"Embedding for record {i} created")
+        logger.info(f"Embedding for record {i} created")
         all_json_records.append(records_with_embedding)
         if i % 500 == 0 or i == len(all_records)-1:
             # Bulk put all records to OpenSearch
             success, failed = opensearch.put_bulk_in_opensearch(all_json_records, opensearch_client)
             all_json_records = []
-            logging.info(f"Documents saved {success}, documents failed to save {failed}")
+            logger.info(f"Documents saved {success}, documents failed to save {failed}")
 
-    logging.info("Finished creating records using Amazon Bedrock Titan text embedding")
+    logger.info("Finished creating records using Amazon Bedrock Titan text embedding")
 
-    logging.info("Cleaning up")
+    logger.info("Cleaning up")
     dataset.delete_file(compressed_file_path)
     dataset.delete_file(file_path)
 
-    logging.info("Finished")
+    logger.info("Finished")
 
 if __name__ == "__main__":
     main()

requirements.txt

@@ -1,5 +1,7 @@
-boto3>=1.28.57
-langchain>=0.0.312
-coloredlogs==15.0.1
-jq==1.6.0
-opensearch-py==2.3.1
+boto3>=1.34.79
+langchain==0.1.14
+langchain-community==0.0.31
+coloredlogs>=15.0.1
+jq==1.7.0
+opensearch-py==2.5.0
+loguru==0.7.2

utils/dataset.py

@@ -1,47 +1,53 @@
-import logging
 import requests
 import gzip
 import json
 import tempfile
+from loguru import logger
+import sys
 import os
 
 
+# logger
+logger.remove()
+logger.add(sys.stdout, level=os.getenv("LOG_LEVEL", "INFO"))
+
+
 def download_dataset(url):
     # download the dataset and store it to tmp dir
     try:
-        logging.info("Downloading dataset")
+        logger.info("Downloading dataset")
         response = requests.get(url)
         if response.status_code == 200:
             temp = tempfile.NamedTemporaryFile(delete=False)
             temp.write(response.content)
             temp.close()
             return temp.name
         else:
-            logging.error("Failed to download dataset")
+            logger.error("Failed to download dataset")
             return None
     except Exception as e:
-        logging.error(e)
+        logger.error(e)
         return None
 
 
 
 def decompress_dataset(file_path):
     # decompress the dataset
     temp_fd, temp_path = tempfile.mkstemp()
-    logging.info(f"Decompressing dataset {file_path} to new file {temp_path}") 
+    logger.info(f"Decompressing dataset {file_path} to new file {temp_path}") 
     try:    
         with gzip.open(file_path, 'rb') as compressed:
             with open(temp_path, 'wb') as decompressed:
                 decompressed.write(compressed.read())
-        logging.info("Decompression complete") 
+        logger.info("Decompression complete") 
         return temp_path
     except Exception as e:
-        logging.error(e)
+        logger.error(e)
         return None
 
 
 def prep_for_put(file_path):
-    logging.info(f"Loading file {file_path}")
+    logger.info(f"Loading file {file_path}")
     all_records = []
     with open(file_path, 'r') as f:
         for line in f:
@@ -52,8 +58,8 @@ def prep_for_put(file_path):
 
 
 def delete_file(file_path):
-    logging.info(f"Deleting file {file_path}")
+    logger.info(f"Deleting file {file_path}")
     try:
         os.remove(file_path)
     except Exception as e:
-        logging.error(e)
+        logger.error(e)

utils/opensearch.py

@@ -1,7 +1,13 @@
 import boto3
 from opensearchpy import OpenSearch, RequestsHttpConnection
 from opensearchpy.helpers import bulk
-import logging
+from loguru import logger
+import sys
+import os
+
+# logger
+logger.remove()
+logger.add(sys.stdout, level=os.getenv("LOG_LEVEL", "INFO"))
 
 
 def get_opensearch_cluster_client(name, password, region):
@@ -30,7 +36,7 @@ def get_opensearch_endpoint(name, region):
 
 
 def put_bulk_in_opensearch(list, client):
-    logging.info(f"Putting {len(list)} documents in OpenSearch")
+    logger.info(f"Putting {len(list)} documents in OpenSearch")
     success, failed = bulk(client, list)
     return success, failed
 
@@ -70,12 +76,12 @@ def create_index_mapping(opensearch_client, index_name):
 
 
 def delete_opensearch_index(opensearch_client, index_name):
-    logging.info(f"Trying to delete index {index_name}")
+    logger.info(f"Trying to delete index {index_name}")
     try:
         response = opensearch_client.indices.delete(index=index_name)
-        logging.info(f"Index {index_name} deleted")
+        logger.info(f"Index {index_name} deleted")
         return response['acknowledged']
     except Exception as e:
-        logging.info(f"Index {index_name} not found, nothing to delete")
+        logger.info(f"Index {index_name} not found, nothing to delete")
         return True
 

utils/secret.py

@@ -1,5 +1,5 @@
 import boto3
-import logging
+
 
 def get_secret(secret_prefix, region):
     client = boto3.client('secretsmanager', region_name=region)