This repo provides a set of utility classes to work with Langchain. It currently has a retriever class KendraIndexRetriever
for working with a Kendra index and sample scripts to execute the QA chain for SageMaker, Open AI and Anthropic providers.
Clone the repository
git clone https://github.com/aws-samples/amazon-kendra-langchain-extensions.git
Move to the repo dir
cd amazon-kendra-langchain-extensions
Install the classes
pip install .
Usage with SageMaker Endpoint for Flan-T-XXL
from aws_langchain.kendra_index_retriever import KendraIndexRetriever
from langchain.chains import RetrievalQA
from langchain import OpenAI
from langchain.prompts import PromptTemplate
from langchain import SagemakerEndpoint
from langchain.llms.sagemaker_endpoint import ContentHandlerBase
import json
class ContentHandler(ContentHandlerBase):
content_type = "application/json"
accepts = "application/json"
def transform_input(self, prompt: str, model_kwargs: dict) -> bytes:
input_str = json.dumps({"inputs": prompt, "parameters": model_kwargs})
return input_str.encode('utf-8')
def transform_output(self, output: bytes) -> str:
response_json = json.loads(output.read().decode("utf-8"))
return response_json[0]["generated_text"]
content_handler = ContentHandler()
llm=SagemakerEndpoint(
endpoint_name=endpoint_name,
region_name="us-east-1",
model_kwargs={"temperature":1e-10, "max_length": 500},
content_handler=content_handler
)
retriever = KendraIndexRetriever(kendraindex=kendra_index_id,
awsregion=region,
return_source_documents=True
)
prompt_template = """
The following is a friendly conversation between a human and an AI.
The AI is talkative and provides lots of specific details from its context.
If the AI does not know the answer to a question, it truthfully says it
does not know.
{context}
Instruction: Based on the above documents, provide a detailed answer for, {question} Answer "don't know" if not present in the document. Solution:
"""
PROMPT = PromptTemplate(
template=prompt_template, input_variables=["context", "question"]
)
chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(
llm,
chain_type="stuff",
retriever=retriever,
chain_type_kwargs=chain_type_kwargs,
return_source_documents=True
)
result = qa("What's SageMaker?")
print(result['answer'])
If you wish to create a sample Kendra index and index sample data and experiment with the index using the sample applications you can deploy the CloudFormation template samples/kendra-docs-index.yaml
For executing sample chains, install the optional dependencies
pip install ".[samples]"
Ensure that the environment variables are set for the aws region, kendra index id and the provider/model used by the sample.
For example, for running the kendra_chat_flan_xl.py
sample, these environment variables must be set: AWS_REGION, KENDRA_INDEX_ID
and FLAN_XL_ENDPOINT.
You can use commands as below to set the environment variables.
export AWS_REGION="<YOUR-AWS-REGION>"
export KENDRA_INDEX_ID="<YOUR-KENDRA-INDEX-ID>"
export FLAN_XL_ENDPOINT="<YOUR-SAGEMAKER-ENDPOINT-FOR-FLAN-T-XL>"
export FLAN_XXL_ENDPOINT="<YOUR-SAGEMAKER-ENDPOINT-FOR-FLAN-T-XXL>"
export OPENAI_API_KEY="<YOUR-OPEN-AI-API-KEY>"
export ANTHROPIC_API_KEY="<YOUR-ANTHROPIC-API-KEY>"
The samples directory is bundled with an app.py
file that can be run as a web app using streamlit.
cd samples
streamlit run app.py anthropic
The above command will run the kendra_chat_anthropic
as the LLM chain. In order to run a different chain, pass a different provider, for example for running the open_ai
chain run this command streamlit run app.py openai
.
python samples/<sample-file-name.py>
pip uninstall aws-langchain
Create your GitHub branch and make a pull request. See CONTRIBUTING for more information.
Install in editable mode, this will make sure your changes are synced in local python env
pip install -e ".[dev]"
This library is licensed under the MIT-0 License. See the LICENSE file.