集成 LangChain Bedrock LLM 和 Elasticsearch
了解如何将 LangChain Bedrock LLM 组件与 Elasticsearch 集成。本笔记本演示了如何生成嵌入、查询向量存储以及构建检索 QA 链。
在 Colab 中打开下载笔记本This workbook demonstrates how to work with Langchain Amazon Bedrock. Amazon Bedrock is a managed service that makes foundation models from leading AI startup and Amazon's own Titan models available through APIs.
Install packages and import modules
# install packages
!python3 -m pip install -qU langchain langchain-elasticsearch langchain_community boto3 tiktoken
# import modules
from getpass import getpass
from urllib.request import urlopen
from langchain_elasticsearch import ElasticsearchStore
from langchain_community.embeddings.bedrock import BedrockEmbeddings
from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
import boto3
import jsonNote: boto3 is part of AWS SDK for Python and is required to use Bedrock LLM
Init Bedrock client
To authorize in AWS service we can use ~/.aws/config file with configuring credentials or pass AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_REGION to boto3 module.
We're using second approach for our example.
default_region = "us-east-1"
AWS_ACCESS_KEY = getpass("AWS Acces key: ")
AWS_SECRET_KEY = getpass("AWS Secret key: ")
AWS_REGION = input(f"AWS Region [default: {default_region}]: ") or default_region
bedrock_client = boto3.client(
service_name="bedrock-runtime",
region_name=AWS_REGION,
aws_access_key_id=AWS_ACCESS_KEY,
aws_secret_access_key=AWS_SECRET_KEY,
)Connect to Elasticsearch
ℹ️ We're using an Elastic Cloud deployment of Elasticsearch for this notebook. If you don't have an Elastic Cloud deployment, sign up here for a free trial.
We'll use the Cloud ID to identify our deployment, because we are using Elastic Cloud deployment. To find the Cloud ID for your deployment, go to https://cloud.elastic.co/deployments and select your deployment.
We will use ElasticsearchStore to connect to our elastic cloud deployment. This would help create and index data easily. In the ElasticsearchStore instance, will set embedding to BedrockEmbeddings to embed the texts and elasticsearch index name that will be used in this example.
# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id
ELASTIC_CLOUD_ID = getpass("Elastic Cloud ID: ")
# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key
ELASTIC_API_KEY = getpass("Elastic Api Key: ")
bedrock_embedding = BedrockEmbeddings(client=bedrock_client)
vector_store = ElasticsearchStore(
es_cloud_id=ELASTIC_CLOUD_ID,
es_api_key=ELASTIC_API_KEY,
index_name="workplace_index",
embedding=bedrock_embedding,
)Download the dataset
Let's download the sample dataset and deserialize the document.
url = "https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/example-apps/chatbot-rag-app/data/data.json"
response = urlopen(url)
workplace_docs = json.loads(response.read())Split Documents into Passages
We’ll chunk documents into passages in order to improve the retrieval specificity and to ensure that we can provide multiple passages within the context window of the final question answering prompt.
Here we are chunking documents into 500 token passages with an overlap of 0 tokens.
Here we are using a simple splitter but Langchain offers more advanced splitters to reduce the chance of context being lost.
from langchain.text_splitter import RecursiveCharacterTextSplitter
metadata = []
content = []
for doc in workplace_docs:
content.append(doc["content"])
metadata.append(
{
"name": doc["name"],
"summary": doc["summary"],
"rolePermissions": doc["rolePermissions"],
}
)
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
chunk_size=512, chunk_overlap=256
)
docs = text_splitter.create_documents(content, metadatas=metadata)Index data into elasticsearch
Next, we will index data to elasticsearch using ElasticsearchStore.from_documents. We will use Cloud ID, Password and Index name values set in the Create cloud deployment step.
documents = vector_store.from_documents(
docs,
es_cloud_id=ELASTIC_CLOUD_ID,
es_api_key=ELASTIC_API_KEY,
index_name="workplace_index",
embedding=bedrock_embedding,
)Init Bedrock LLM
Next, we will initialize Bedrock LLM. In the Bedrock instance, will pass bedrock_client and specific model_id: amazon.titan-text-express-v1, ai21.j2-ultra-v1, anthropic.claude-v2, cohere.command-text-v14 or etc. You can see list of available base models on Amazon Bedrock User Guide
default_model_id = "amazon.titan-text-express-v1"
AWS_MODEL_ID = input(f"AWS model [default: {default_model_id}]: ") or default_model_id
llm = Bedrock(client=bedrock_client, model_id=AWS_MODEL_ID)Asking a question
Now that we have the passages stored in Elasticsearch and llm is initialized, we can now ask a question to get the relevant passages.
retriever = vector_store.as_retriever()
qa = RetrievalQA.from_llm(llm=llm, retriever=retriever, return_source_documents=True)
questions = [
"What is the nasa sales team?",
"What is our work from home policy?",
"Does the company own my personal project?",
"What job openings do we have?",
"How does compensation work?",
]
question = questions[1]
print(f"Question: {question}")
ans = qa({"query": question})
print("\033[92m ---- Answer ---- \033[0m")
print(ans["result"] + "\n")
print("\033[94m ---- Sources ---- \033[0m")
for doc in ans["source_documents"]:
print("Name: " + doc.metadata["name"])
print("Content: " + doc.page_content)
print("-------")