community[patch]: Allow adding ARNs as model_id to support Amazon Bedrock custom models (#16800)

- **Description:** Adds an additional class variable to `BedrockBase`
called `provider` that allows sending a model provider such as amazon,
cohere, ai21, etc.
Up until now, the model provider is extracted from the `model_id` using
the first part before the `.`, such as `amazon` for
`amazon.titan-text-express-v1` (see [supported list of Bedrock model IDs
here](https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html)).
But for custom Bedrock models where the ARN of the provisioned
throughput must be supplied, the `model_id` is like
`arn:aws:bedrock:...` so the `model_id` cannot be extracted from this. A
model `provider` is required by the LangChain Bedrock class to perform
model-based processing. To allow the same processing to be performed for
custom-models of a specific base model type, passing this `provider`
argument can help solve the issues.
The alternative considered here was the use of
`provider.arn:aws:bedrock:...` which then requires ARN to be extracted
and passed separately when invoking the model. The proposed solution
here is simpler and also does not cause issues for current models
already using the Bedrock class.
  - **Issue:** N/A
  - **Dependencies:** N/A

---------

Co-authored-by: Piyush Jain <piyushjain@duck.com>
This commit is contained in:
Supreet Takkar
2024-02-05 17:28:03 -05:00
committed by GitHub
parent e022bfaa7d
commit ae33979813
2 changed files with 48 additions and 4 deletions

View File

@@ -111,7 +111,35 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### Custom models"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"custom_llm = Bedrock(\n",
" credentials_profile_name=\"bedrock-admin\",\n",
" provider=\"cohere\",\n",
" model_id=\"<Custom model ARN>\", # ARN like 'arn:aws:bedrock:...' obtained via provisioning the custom model\n",
" model_kwargs={\"temperature\": 1},\n",
" streaming=True,\n",
" callbacks=[StreamingStdOutCallbackHandler()],\n",
")\n",
"\n",
"conversation = ConversationChain(\n",
" llm=custom_llm, verbose=True, memory=ConversationBufferMemory()\n",
")\n",
"conversation.predict(input=\"What is the recipe of mayonnaise?\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Guardrails for Amazon Bedrock example \n",
"\n",
"## Guardrails for Amazon Bedrock (Preview) \n",
"[Guardrails for Amazon Bedrock](https://aws.amazon.com/bedrock/guardrails/) evaluates user inputs and model responses based on use case specific policies, and provides an additional layer of safeguards regardless of the underlying model. Guardrails can be applied across models, including Anthropic Claude, Meta Llama 2, Cohere Command, AI21 Labs Jurassic, and Amazon Titan Text, as well as fine-tuned models.\n",
@@ -139,7 +167,7 @@
" print(f\"Guardrails: {kwargs}\")\n",
"\n",
"\n",
"# guardrails for Amazon Bedrock with trace\n",
"# Guardrails for Amazon Bedrock with trace\n",
"llm = Bedrock(\n",
" credentials_profile_name=\"bedrock-admin\",\n",
" model_id=\"<Model_ID>\",\n",
@@ -166,7 +194,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.11.7"
}
},
"nbformat": 4,