Updated the Bedrock rag template (#12462)

Updates the bedrock rag template.
- Removes pinecone and replaces with FAISS as the vector store
- Fixes the environment variables, setting defaults
- Adds a `main.py` test file quick sanity testing
- Updates README.md with correct instructions
This commit is contained in:
Piyush Jain
2023-10-27 17:02:28 -07:00
committed by GitHub
parent 5c2243ee91
commit 5545de0466
5 changed files with 294 additions and 1030 deletions

View File

@@ -5,56 +5,31 @@ 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
from langchain.vectorstores import FAISS
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()
# Get region and profile from env
region = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
profile = os.environ.get("AWS_PROFILE", "default")
# 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",
region_name=region,
credentials_profile_name=profile,
model_kwargs={'max_tokens_to_sample':200}
)
bedrock_embeddings = BedrockEmbeddings(
model_id="amazon.titan-embed-text-v1"
)
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)
# Add to vectorDB
vectorstore = FAISS.from_texts(
["harrison worked at kensho"],
embedding=bedrock_embeddings
)
retriever = vectorstore.as_retriever()
# Get retriever from vectorstore
retriever = vectorstore.as_retriever()
# RAG prompt
@@ -64,10 +39,11 @@ Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
# RAG
chain = (
RunnableParallel({"context": retriever, "question": RunnablePassthrough()})
| prompt
| model
| StrOutputParser()
)
)