AWS Bedrock RAG template (#12450)

templates/rag-aws-bedrock/rag_aws_bedrock/__init__.py

@@ -0,0 +1,3 @@
+from rag_aws_bedrock.chain import chain
+
+__all__ = ["chain"]

templates/rag-aws-bedrock/rag_aws_bedrock/chain.py

@@ -0,0 +1,73 @@
+import os
+
+from langchain.embeddings import BedrockEmbeddings
+from langchain.llms.bedrock import Bedrock
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema.output_parser import StrOutputParser
+from langchain.schema.runnable import RunnableParallel, RunnablePassthrough
+from langchain.vectorstores import Pinecone
+from utils import bedrock
+
+if os.environ.get("PINECONE_API_KEY", None) is None:
+    raise Exception("Missing `PINECONE_API_KEY` environment variable.")
+
+if os.environ.get("PINECONE_ENVIRONMENT", None) is None:
+    raise Exception("Missing `PINECONE_ENVIRONMENT` environment variable.")
+
+if os.environ.get("AWS_DEFAULT_REGION", None) is None:
+    raise Exception("Missing `AWS_DEFAULT_REGION` environment variable.")
+
+if os.environ.get("AWS_PROFILE", None) is None:
+    raise Exception("Missing `AWS_PROFILE` environment variable.")
+
+if os.environ.get("BEDROCK_ASSUME_ROLE", None) is None:
+    raise Exception("Missing `BEDROCK_ASSUME_ROLE` environment variable.")
+
+PINECONE_INDEX_NAME = os.environ.get("PINECONE_INDEX", "langchain-test")
+
+### Ingest code - you may need to run this the first time
+# Load
+# from langchain.document_loaders import WebBaseLoader
+# loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
+# data = loader.load()
+
+# # Split
+# from langchain.text_splitter import RecursiveCharacterTextSplitter
+# text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
+# all_splits = text_splitter.split_documents(data)
+
+# # Add to vectorDB
+# vectorstore = Pinecone.from_documents(
+#     documents=all_splits, embedding=OpenAIEmbeddings(), index_name=PINECONE_INDEX_NAME
+# )
+# retriever = vectorstore.as_retriever()
+
+# Set LLM and embeddings
+boto3_bedrock = bedrock.get_bedrock_client(
+    assumed_role=os.environ.get("BEDROCK_ASSUME_ROLE", None),
+    region=os.environ.get("AWS_DEFAULT_REGION", None)
+)
+model = Bedrock(model_id="anthropic.claude-v2", 
+                client=boto3_bedrock, 
+                model_kwargs={'max_tokens_to_sample':200})
+bedrock_embeddings = BedrockEmbeddings(model_id="amazon.titan-embed-text-v1", 
+                                       client=boto3_bedrock)
+
+# Set vectostore
+vectorstore = Pinecone.from_existing_index(PINECONE_INDEX_NAME, bedrock_embeddings)
+retriever = vectorstore.as_retriever()
+
+# RAG prompt
+template = """Answer the question based only on the following context:
+{context}
+Question: {question}
+"""
+prompt = ChatPromptTemplate.from_template(template)
+
+# RAG
+chain = (
+    RunnableParallel({"context": retriever, "question": RunnablePassthrough()})
+    | prompt
+    | model
+    | StrOutputParser()
+)