First release

Author: omerh

.gitignore

@@ -0,0 +1,7 @@
+.terraform/
+*.tfstate*
+
+dependencies/
+
+.venv/
+__pycache__/

CODE_OF_CONDUCT.md

@@ -1,4 +1,5 @@
-## Code of Conduct
+# Code of Conduct
+
 This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
 For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
-[email protected] with any additional questions or comments.
+[[email protected]](mailto:[email protected]) with any additional questions or comments.

CONTRIBUTING.md

@@ -6,7 +6,6 @@ documentation, we greatly value feedback and contributions from our community.
 Please read through this document before submitting any issues or pull requests to ensure we have all the necessary
 information to effectively respond to your bug report or contribution.
 
-
 ## Reporting Bugs/Feature Requests
 
 We welcome you to use the GitHub issue tracker to report bugs or suggest features.
@@ -19,8 +18,8 @@ reported the issue. Please try to include as much information as you can. Detail
 * Any modifications you've made relevant to the bug
 * Anything unusual about your environment or deployment
 
-
 ## Contributing via Pull Requests
+
 Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that:
 
 1. You are working against the latest source on the *main* branch.
@@ -39,20 +38,19 @@ To send us a pull request, please:
 GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and
 [creating a pull request](https://help.github.com/articles/creating-a-pull-request/).
 
-
 ## Finding contributions to work on
-Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start.
 
+Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start.
 
 ## Code of Conduct
+
 This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
 For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
-[email protected] with any additional questions or comments.
-
+[[email protected]](mailto:[email protected]) with any additional questions or comments.
 
 ## Security issue notifications
-If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue.
 
+If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue.
 
 ## Licensing
 

README.md

@@ -1,11 +1,61 @@
-## My Project
+# RAG using LangChain with Amazon Bedrock Titan text, and embedding, using OpenSearch vector engine
 
-TODO: Fill this README out!
+This sample repository provides a sample code for using RAG (Retrieval augmented generation) method relaying on [Amazon Bedrock](https://aws.amazon.com/bedrock/) [Titan text embedding](https://aws.amazon.com/bedrock/titan/) LLM (Large Language Model), for creating text embedding that will be stored in [Amazon OpenSearch](https://aws.amazon.com/opensearch-service/) with [vector engine support](https://aws.amazon.com/about-aws/whats-new/2023/07/vector-engine-amazon-opensearch-serverless-preview/) for assisting with the prompt engineering task for more accurate response from LLMs.
 
-Be sure to:
+After we successfully loaded embeddings into OpenSearch, we will then start querying our LLM, by using [LangChain](https://www.langchain.com/). We will ask questions, retrieving similar embedding for a more accurate prompt.
 
-* Change the title in this README
-* Edit your repository description on GitHub
+## Prerequisites
+
+1. This was tested on Python 3.11.4
+2. It is advise to work on a clean environment, use `virtualenv` or any other virtual environment packages.
+
+    ```bash
+    pip install virtualenv
+    python -m virtualenv venv
+    source ./venv/bin/activate
+    ```
+
+3. Run `./download-beta-sdk.sh` to download the beta SDK for using Amazon Bedrock
+4. Install requirements `pip install -r requirements.txt`
+5. Install [terraform](https://developer.hashicorp.com/terraform/downloads?product_intent=terraform) to create the OpenSearch cluster
+
+    ```bash
+    brew tap hashicorp/tap
+    brew install hashicorp/tap/terraform
+    ```
+
+## Steps for using this sample code
+
+1. In the first step we will launch an OpenSearch cluster using Terraform.
+
+    ```bash
+    cd ./terraform
+    terraform init
+    terraform apply -auto-approve
+    ```
+
+    >>This cluster configuration is for testing proposes only, as it's endpoint is public for simplifying the use of this sample code.
+
+2. Now that we have a running OpenSearch cluster with vector engine support we will start uploading our data that will help us with prompt engineering. For this sample, we will use a data source from [Hugging Face](https://huggingface.co) [embedding-training-data](https://huggingface.co/datasets/sentence-transformers/embedding-training-data) [gooaq_pairs](https://huggingface.co/datasets/sentence-transformers/embedding-training-data/resolve/main/gooaq_pairs.jsonl.gz), we will download it, and invoke Titan embedding to get a text embedding, that we will store in OpenSearch for next steps.
+
+    ```bash
+    python load-data-to-opensearch.py --recreate 1 --early-stop 1
+    ```
+
+    >>Optional arguments: `--recreate` for recreating the index in OpenSearch, and `--early-stop` to load only 100 embedded documents into OpenSearch
+
+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
+
+    ```bash
+    python ask-titan-with-rag.py --ask "your question here"
+    ```
+
+### Cleanup
+
+```bash
+cd ./terraform
+terraform destroy # When prompt for confirmation, type yes, and press enter.
+```
 
 ## Security
 
@@ -14,4 +64,3 @@ See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more inform
 ## License
 
 This library is licensed under the MIT-0 License. See the LICENSE file.
-

ask-titan-with-rag.py

@@ -0,0 +1,113 @@
+import coloredlogs
+import logging
+import argparse
+import sys
+import os
+from utils.bedrock import get_bedrock_client
+from utils import bedrock, opensearch, secret, iam
+from langchain.embeddings import BedrockEmbeddings
+from langchain.vectorstores import OpenSearchVectorSearch
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain.llms.bedrock import Bedrock
+
+
+coloredlogs.install(fmt='%(asctime)s %(levelname)s %(message)s', datefmt='%H:%M:%S', level='INFO')
+logging.basicConfig(level=logging.INFO) 
+logger = logging.getLogger(__name__)
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ask", type=str, default="What is <3?")
+    
+    return parser.parse_known_args()
+
+
+def get_bedrock_client(region, account_id):
+    module_path = "."
+    sys.path.append(os.path.abspath(module_path))
+    os.environ['AWS_DEFAULT_REGION'] = region
+
+    boto3_bedrock = bedrock.get_bedrock_client(
+        assumed_role=f'arn:aws:iam::{account_id}:role/bedrock',
+        region=region, 
+        )
+    return boto3_bedrock
+
+def create_langchain_vector_embedding_using_bedrock(bedrock_client):
+    bedrock_embeddings_client = BedrockEmbeddings(
+        client=bedrock_client,
+        model_id="amazon.titan-e1t-medium")
+    return bedrock_embeddings_client
+    
+
+def create_opensearch_vector_search_client(index_name, opensearch_password, bedrock_embeddings_client, opensearch_endpoint, _is_aoss=False):
+    docsearch = OpenSearchVectorSearch(
+        index_name=index_name,
+        embedding_function=bedrock_embeddings_client,
+        opensearch_url=f"https://{opensearch_endpoint}",
+        http_auth=(index_name, opensearch_password),
+        is_aoss=_is_aoss
+    )
+    return docsearch
+
+
+def create_bedrock_llm(bedrock_client):
+    bedrock_llm = Bedrock(
+        model_id="amazon.titan-tg1-large", 
+        client=bedrock_client,
+        model_kwargs={'temperature': 0}
+        )
+    return bedrock_llm
+    
+
+def main():
+    logging.info("Starting")
+    # vars
+    region = "us-west-2"
+    index_name = 'rag'
+    args, _ = parse_args()
+    
+    # Creating all clients for chain
+    account_id = iam.get_account_id()
+    bedrock_client = get_bedrock_client(region, account_id)
+    bedrock_llm = create_bedrock_llm(bedrock_client)
+    bedrock_embeddings_client = create_langchain_vector_embedding_using_bedrock(bedrock_client)
+    opensearch_endpoint = opensearch.get_opensearch_endpoint(index_name, region)
+    opensearch_password = secret.get_secret(index_name, region)
+    opensearch_vector_search_client = create_opensearch_vector_search_client(index_name, opensearch_password, bedrock_embeddings_client, opensearch_endpoint)
+    
+    # LangChain prompt template
+    question = "what is the meaning of <3?"
+    
+    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
+
+    {context}
+
+    Question: {question}
+    Answer:"""
+    PROMPT = PromptTemplate(
+        template=prompt_template, input_variables=["context", "question"]
+    )
+    
+    logging.info(f"Starting the chain with KNN similarity using OpenSearch, and than passing to Bedrock Titan FM")
+    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)
+    
+    logging.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"vector: {d.metadata}")
+    
+    logging.info(f"\nThe answer from Titan: {response.get('result')}")
+    
+
+if __name__ == "__main__":
+    main()

download-beta-sdk.sh

@@ -0,0 +1,14 @@
+#!/bin/sh
+
+echo "Creating directory"
+mkdir -p ./dependencies && \
+cd ./dependencies && \
+echo "Downloading dependencies"
+curl -sS https://d2eo22ngex1n9g.cloudfront.net/Documentation/SDK/bedrock-python-sdk.zip > sdk.zip && \
+echo "Unpacking dependencies"
+unzip -o -q sdk.zip && \
+rm sdk.zip
+
+pip install --force-reinstall --no-cache-dir awscli-1.29.21-py3-none-any.whl
+pip install --force-reinstall --no-cache-dir boto3-1.28.21-py3-none-any.whl 
+pip install --force-reinstall --no-cache-dir botocore-1.31.21-py3-none-any.whl 

load-data-to-opensearch.py

@@ -0,0 +1,123 @@
+import logging
+import coloredlogs
+import json
+import argparse
+from utils.bedrock import get_bedrock_client
+import sys
+import os
+from utils import bedrock, dataset, secret, opensearch, iam
+
+coloredlogs.install(fmt='%(asctime)s %(levelname)s %(message)s', datefmt='%H:%M:%S', level='INFO')
+logging.basicConfig(level=logging.INFO) 
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--recreate", type=bool, default=0)
+    parser.add_argument("--early-stop", type=bool, default=0)
+    
+    return parser.parse_known_args()
+
+
+
+def get_bedrock_client(region, account_id):
+    module_path = "."
+    sys.path.append(os.path.abspath(module_path))
+    os.environ['AWS_DEFAULT_REGION'] = region
+
+    boto3_bedrock = bedrock.get_bedrock_client(
+        assumed_role=f'arn:aws:iam::{account_id}:role/bedrock',
+        region=region, 
+        )
+    return boto3_bedrock
+
+
+def create_vector_embedding_with_bedrock(text, name, bedrock_client):
+    payload = {"inputText": f"{text}"}
+    body = json.dumps(payload)
+    modelId = "amazon.titan-e1t-medium"
+    accept = "application/json"
+    contentType = "application/json"
+
+    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}
+
+            
+def main():
+    logging.info("Starting")
+    # vars
+    region = "us-west-2"
+    name = 'rag'
+    dataset_url = "https://huggingface.co/datasets/sentence-transformers/embedding-training-data/resolve/main/gooaq_pairs.jsonl.gz"
+    early_stop_record_count = 100
+    args, _ = parse_args()
+    
+    # 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")
+    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")
+    
+    logging.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")
+        success = opensearch.create_index(opensearch_client, name)
+        if success:
+            logging.info("Creating OpenSearch index mapping")
+            success = opensearch.create_index_mapping(opensearch_client, name)
+            logging.info(f"OpenSearch Index mapping created")
+    
+    # Download sample dataset from HuggingFace 
+    logging.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)
+        if file_path is not None:
+            all_records = dataset.prep_for_put(file_path)
+    
+    # Initialize bedrock client
+    account_id = iam.get_account_id()
+    bedrock_client = get_bedrock_client(region, account_id)
+    
+    # Vector embedding using Amazon Bedrock Titan text embedding
+    all_json_records = []
+    logging.info(f"Creating embeddings for {len(all_records)} records")
+    
+    # using the arg --early-stop
+    i = 0
+    for record in all_records:
+        i += 1
+        if args.early_stop:
+            if i > early_stop_record_count:
+                break
+        records_with_embedding = create_vector_embedding_with_bedrock(record, name, bedrock_client)
+        logging.info(f"Embedding for record {i} created")
+        all_json_records.append(records_with_embedding)
+    
+    logging.info("Finished creating records using Amazon Bedrock Titan text embedding")
+    
+    # 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}")
+    
+    logging.info("Cleaning up")
+    dataset.delete_file(compressed_file_path)
+    dataset.delete_file(file_path)
+    
+    logging.info("Finished")
+        
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,4 @@
+langchain==0.0.266
+coloredlogs==15.0.1
+jq==1.4.1
+opensearch-py==2.3.0