forked from Archives/langchain
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. <!-- For a quicker response, figure out the right person to tag with @ @hwchase17 - project lead Tracing / Callbacks - @agola11 Async - @agola11 DataLoaders - @eyurtsev Models - @hwchase17 - @agola11 Agents / Tools / Toolkits - @vowelparrot VectorStores / Retrievers / Memory - @dev2049 -->searx_updates
parent
5ce74b5958
commit
562fdfc8f9
@ -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
|
||||
```
|
@ -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
|
||||
}
|
@ -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
|
||||
}
|
@ -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)
|
@ -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
|
Loading…
Reference in New Issue