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