AWS Langchain

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.

Installing

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

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'])

Creating an Amazon Kendra index with test data

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

Running samples

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>"

Running samples from the streamlit app

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.

Running samples from the command line

python samples/<sample-file-name.py>

Uninstall

pip uninstall aws-langchain

Contributing

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]"

License

This library is licensed under the MIT-0 License. See the LICENSE file.