From 562fdfc8f9d7cd1b4fcf98c29081a75065aaf4f8 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 31 May 2023 07:17:01 -0700 Subject: [PATCH] Bedrock llm and embeddings (#5464) # Bedrock LLM and Embeddings This PR adds a new LLM and an Embeddings class for the [Bedrock](https://aws.amazon.com/bedrock) service. The PR also includes example notebooks for using the LLM class in a conversation chain and embeddings usage in creating an embedding for a query and document. **Note**: AWS is doing a private release of the Bedrock service on 05/31/2023; users need to request access and added to an allowlist in order to start using the Bedrock models and embeddings. Please use the [Bedrock Home Page](https://aws.amazon.com/bedrock) to request access and to learn more about the models available in Bedrock. --- docs/integrations/bedrock.md | 24 +++ .../models/llms/integrations/bedrock.ipynb | 86 ++++++++ .../text_embedding/examples/bedrock.ipynb | 75 +++++++ langchain/embeddings/__init__.py | 2 + langchain/embeddings/bedrock.py | 157 +++++++++++++++ langchain/llms/__init__.py | 2 + langchain/llms/bedrock.py | 187 ++++++++++++++++++ 7 files changed, 533 insertions(+) create mode 100644 docs/integrations/bedrock.md create mode 100644 docs/modules/models/llms/integrations/bedrock.ipynb create mode 100644 docs/modules/models/text_embedding/examples/bedrock.ipynb create mode 100644 langchain/embeddings/bedrock.py create mode 100644 langchain/llms/bedrock.py diff --git a/docs/integrations/bedrock.md b/docs/integrations/bedrock.md new file mode 100644 index 00000000..9dabac7a --- /dev/null +++ b/docs/integrations/bedrock.md @@ -0,0 +1,24 @@ +# Amazon Bedrock + +>[Amazon Bedrock](https://aws.amazon.com/bedrock/) is a fully managed service that makes FMs from leading AI startups and Amazon available via an API, so you can choose from a wide range of FMs to find the model that is best suited for your use case. + +## Installation and Setup + +```bash +pip install boto3 +``` + +## LLM + +See a [usage example](../modules/models/llms/integrations/bedrock.ipynb). + +```python +from langchain import Bedrock +``` + +## Text Embedding Models + +See a [usage example](../modules/models/text_embedding/examples/bedrock.ipynb). +```python +from langchain.embeddings import BedrockEmbeddings +``` diff --git a/docs/modules/models/llms/integrations/bedrock.ipynb b/docs/modules/models/llms/integrations/bedrock.ipynb new file mode 100644 index 00000000..d2d5e54f --- /dev/null +++ b/docs/modules/models/llms/integrations/bedrock.ipynb @@ -0,0 +1,86 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Amazon Bedrock" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[Amazon Bedrock](https://aws.amazon.com/bedrock/) is a fully managed service that makes FMs from leading AI startups and Amazon available via an API, so you can choose from a wide range of FMs to find the model that is best suited for your use case" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%pip install boto3" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from langchain.llms.bedrock import Bedrock\n", + "\n", + "llm = Bedrock(credentials_profile_name=\"bedrock-admin\", model_id=\"amazon.titan-tg1-large\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Using in a conversation chain" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.chains import ConversationChain\n", + "from langchain.memory import ConversationBufferMemory\n", + "\n", + "conversation = ConversationChain(\n", + " llm=llm,\n", + " verbose=True,\n", + " memory=ConversationBufferMemory()\n", + ")\n", + "\n", + "conversation.predict(input=\"Hi there!\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/modules/models/text_embedding/examples/bedrock.ipynb b/docs/modules/models/text_embedding/examples/bedrock.ipynb new file mode 100644 index 00000000..ae161a52 --- /dev/null +++ b/docs/modules/models/text_embedding/examples/bedrock.ipynb @@ -0,0 +1,75 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "75e378f5-55d7-44b6-8e2e-6d7b8b171ec4", + "metadata": {}, + "source": [ + "# Bedrock Embeddings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2dbe40fa-7c0b-4bcb-a712-230bf613a42f", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install boto3" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "282239c8-e03a-4abc-86c1-ca6120231a20", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.embeddings import BedrockEmbeddings\n", + "\n", + "embeddings = BedrockEmbeddings(credentials_profile_name=\"bedrock-admin\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19a46868-4bed-40cd-89ca-9813fbfda9cb", + "metadata": {}, + "outputs": [], + "source": [ + "embeddings.embed_query(\"This is a content of the document\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cf0349c4-6408-4342-8691-69276a388784", + "metadata": {}, + "outputs": [], + "source": [ + "embeddings.embed_documents([\"This is a content of the document\"])" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/langchain/embeddings/__init__.py b/langchain/embeddings/__init__.py index dc6b8672..6261b3dc 100644 --- a/langchain/embeddings/__init__.py +++ b/langchain/embeddings/__init__.py @@ -6,6 +6,7 @@ from langchain.embeddings.aleph_alpha import ( AlephAlphaAsymmetricSemanticEmbedding, AlephAlphaSymmetricSemanticEmbedding, ) +from langchain.embeddings.bedrock import BedrockEmbeddings from langchain.embeddings.cohere import CohereEmbeddings from langchain.embeddings.elasticsearch import ElasticsearchEmbeddings from langchain.embeddings.fake import FakeEmbeddings @@ -56,6 +57,7 @@ __all__ = [ "GooglePalmEmbeddings", "MiniMaxEmbeddings", "VertexAIEmbeddings", + "BedrockEmbeddings", ] diff --git a/langchain/embeddings/bedrock.py b/langchain/embeddings/bedrock.py new file mode 100644 index 00000000..e884b139 --- /dev/null +++ b/langchain/embeddings/bedrock.py @@ -0,0 +1,157 @@ +import json +import os +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Extra, root_validator + +from langchain.embeddings.base import Embeddings + + +class BedrockEmbeddings(BaseModel, Embeddings): + """Embeddings provider to invoke Bedrock embedding models. + + To authenticate, the AWS client uses the following methods to + automatically load credentials: + https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html + + If a specific credential profile should be used, you must pass + the name of the profile from the ~/.aws/credentials file that is to be used. + + Make sure the credentials / roles used have the required policies to + access the Bedrock service. + """ + + """ + Example: + .. code-block:: python + + from langchain.bedrock_embeddings import BedrockEmbeddings + + region_name ="us-east-1" + credentials_profile_name = "default" + model_id = "amazon.titan-e1t-medium" + + be = BedrockEmbeddings( + credentials_profile_name=credentials_profile_name, + region_name=region_name, + model_id=model_id + ) + """ + + client: Any #: :meta private: + + region_name: Optional[str] = None + """The aws region e.g., `us-west-2`. Fallsback to AWS_DEFAULT_REGION env variable + or region specified in ~/.aws/config in case it is not provided here. + """ + + credentials_profile_name: Optional[str] = None + """The name of the profile in the ~/.aws/credentials or ~/.aws/config files, which + has either access keys or role information specified. + If not specified, the default credential profile or, if on an EC2 instance, + credentials from IMDS will be used. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html + """ + + model_id: str = "amazon.titan-e1t-medium" + """Id of the model to call, e.g., amazon.titan-e1t-medium, this is + equivalent to the modelId property in the list-foundation-models api""" + + model_kwargs: Optional[Dict] = None + """Key word arguments to pass to the model.""" + + class Config: + """Configuration for this pydantic object.""" + + extra = Extra.forbid + + @root_validator() + def validate_environment(cls, values: Dict) -> Dict: + """Validate that AWS credentials to and python package exists in environment.""" + try: + import boto3 + + if values["credentials_profile_name"] is not None: + session = boto3.Session(profile_name=values["credentials_profile_name"]) + else: + # use default credentials + session = boto3.Session() + + client_params = {} + if values["region_name"]: + client_params["region_name"] = values["region_name"] + + values["client"] = session.client("bedrock", **client_params) + + except ImportError: + raise ModuleNotFoundError( + "Could not import boto3 python package. " + "Please install it with `pip install boto3`." + ) + except Exception as e: + raise ValueError( + "Could not load credentials to authenticate with AWS client. " + "Please check that credentials in the specified " + "profile name are valid." + ) from e + + return values + + def _embedding_func(self, text: str) -> List[float]: + """Call out to Bedrock embedding endpoint.""" + # replace newlines, which can negatively affect performance. + text = text.replace(os.linesep, " ") + _model_kwargs = self.model_kwargs or {} + + input_body = {**_model_kwargs} + input_body["inputText"] = text + body = json.dumps(input_body) + content_type = "application/json" + accepts = "application/json" + + embeddings = [] + try: + response = self.client.invoke_model( + body=body, + modelId=self.model_id, + accept=accepts, + contentType=content_type, + ) + response_body = json.loads(response.get("body").read()) + embeddings = response_body.get("embedding") + except Exception as e: + raise ValueError(f"Error raised by inference endpoint: {e}") + + return embeddings + + def embed_documents( + self, texts: List[str], chunk_size: int = 1 + ) -> List[List[float]]: + """Compute doc embeddings using a Bedrock model. + + Args: + texts: The list of texts to embed. + chunk_size: Bedrock currently only allows single string + inputs, so chunk size is always 1. This input is here + only for compatibility with the embeddings interface. + + + Returns: + List of embeddings, one for each text. + """ + results = [] + for text in texts: + response = self._embedding_func(text) + results.append(response) + return results + + def embed_query(self, text: str) -> List[float]: + """Compute query embeddings using a Bedrock model. + + Args: + text: The text to embed. + + Returns: + Embeddings for the text. + """ + return self._embedding_func(text) diff --git a/langchain/llms/__init__.py b/langchain/llms/__init__.py index 3e62a59d..9dc97c56 100644 --- a/langchain/llms/__init__.py +++ b/langchain/llms/__init__.py @@ -8,6 +8,7 @@ from langchain.llms.anyscale import Anyscale from langchain.llms.bananadev import Banana from langchain.llms.base import BaseLLM from langchain.llms.beam import Beam +from langchain.llms.bedrock import Bedrock from langchain.llms.cerebriumai import CerebriumAI from langchain.llms.cohere import Cohere from langchain.llms.ctransformers import CTransformers @@ -48,6 +49,7 @@ __all__ = [ "Anyscale", "Banana", "Beam", + "Bedrock", "CerebriumAI", "Cohere", "CTransformers", diff --git a/langchain/llms/bedrock.py b/langchain/llms/bedrock.py new file mode 100644 index 00000000..a59a81d2 --- /dev/null +++ b/langchain/llms/bedrock.py @@ -0,0 +1,187 @@ +import json +from typing import Any, Dict, List, Mapping, Optional + +from pydantic import Extra, root_validator + +from langchain.callbacks.manager import CallbackManagerForLLMRun +from langchain.llms.base import LLM +from langchain.llms.utils import enforce_stop_tokens + + +class LLMInputOutputAdapter: + """Adapter class to prepare the inputs from Langchain to a format + that LLM model expects. Also, provides helper function to extract + the generated text from the model response.""" + + @classmethod + def prepare_input( + cls, provider: str, prompt: str, model_kwargs: Dict[str, Any] + ) -> Dict[str, Any]: + input_body = {**model_kwargs} + if provider == "anthropic" or provider == "ai21": + input_body["prompt"] = prompt + else: + input_body["inputText"] = prompt + + if provider == "anthropic" and "max_tokens_to_sample" not in input_body: + input_body["max_tokens_to_sample"] = 50 + + return input_body + + @classmethod + def prepare_output(cls, provider: str, response: Any) -> str: + if provider == "anthropic": + response_body = json.loads(response.get("body").read().decode()) + return response_body.get("completion") + else: + response_body = json.loads(response.get("body").read()) + + if provider == "ai21": + return response_body.get("completions")[0].get("data").get("text") + else: + return response_body.get("results")[0].get("outputText") + + +class Bedrock(LLM): + """LLM provider to invoke Bedrock models. + + To authenticate, the AWS client uses the following methods to + automatically load credentials: + https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html + + If a specific credential profile should be used, you must pass + the name of the profile from the ~/.aws/credentials file that is to be used. + + Make sure the credentials / roles used have the required policies to + access the Bedrock service. + """ + + """ + Example: + .. code-block:: python + + from bedrock_langchain.bedrock_llm import BedrockLLM + + llm = BedrockLLM( + credentials_profile_name="default", + model_id="amazon.titan-tg1-large" + ) + + """ + + client: Any #: :meta private: + + region_name: Optional[str] = None + """The aws region e.g., `us-west-2`. Fallsback to AWS_DEFAULT_REGION env variable + or region specified in ~/.aws/config in case it is not provided here. + """ + + credentials_profile_name: Optional[str] = None + """The name of the profile in the ~/.aws/credentials or ~/.aws/config files, which + has either access keys or role information specified. + If not specified, the default credential profile or, if on an EC2 instance, + credentials from IMDS will be used. + See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html + """ + + model_id: str + """Id of the model to call, e.g., amazon.titan-tg1-large, this is + equivalent to the modelId property in the list-foundation-models api""" + + model_kwargs: Optional[Dict] = None + """Key word arguments to pass to the model.""" + + class Config: + """Configuration for this pydantic object.""" + + extra = Extra.forbid + + @root_validator() + def validate_environment(cls, values: Dict) -> Dict: + """Validate that AWS credentials to and python package exists in environment.""" + try: + import boto3 + + if values["credentials_profile_name"] is not None: + session = boto3.Session(profile_name=values["credentials_profile_name"]) + else: + # use default credentials + session = boto3.Session() + + client_params = {} + if values["region_name"]: + client_params["region_name"] = values["region_name"] + + values["client"] = session.client("bedrock", **client_params) + + except ImportError: + raise ModuleNotFoundError( + "Could not import boto3 python package. " + "Please install it with `pip install boto3`." + ) + except Exception as e: + raise ValueError( + "Could not load credentials to authenticate with AWS client. " + "Please check that credentials in the specified " + "profile name are valid." + ) from e + + return values + + @property + def _identifying_params(self) -> Mapping[str, Any]: + """Get the identifying parameters.""" + _model_kwargs = self.model_kwargs or {} + return { + **{"model_kwargs": _model_kwargs}, + } + + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "amazon_bedrock" + + def _call( + self, + prompt: str, + stop: Optional[List[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + ) -> str: + """Call out to Bedrock service model. + + Args: + prompt: The prompt to pass into the model. + stop: Optional list of stop words to use when generating. + + Returns: + The string generated by the model. + + Example: + .. code-block:: python + + response = se("Tell me a joke.") + """ + _model_kwargs = self.model_kwargs or {} + + provider = self.model_id.split(".")[0] + + input_body = LLMInputOutputAdapter.prepare_input( + provider, prompt, _model_kwargs + ) + body = json.dumps(input_body) + accept = "application/json" + contentType = "application/json" + + try: + response = self.client.invoke_model( + body=body, modelId=self.model_id, accept=accept, contentType=contentType + ) + text = LLMInputOutputAdapter.prepare_output(provider, response) + + except Exception as e: + raise ValueError(f"Error raised by bedrock service: {e}") + + if stop is not None: + text = enforce_stop_tokens(text, stop) + + return text