mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-09 15:03:21 +00:00
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:
@@ -229,9 +229,17 @@ class BedrockBase(BaseModel, ABC):
|
||||
config: Optional[Config] = None
|
||||
"""An optional botocore.config.Config instance to pass to the client."""
|
||||
|
||||
provider: Optional[str] = None
|
||||
"""The model provider, e.g., amazon, cohere, ai21, etc. When not supplied, provider
|
||||
is extracted from the first part of the model_id e.g. 'amazon' in
|
||||
'amazon.titan-text-express-v1'. This value should be provided for model ids that do
|
||||
not have the provider in them, e.g., custom and provisioned models that have an ARN
|
||||
associated with them."""
|
||||
|
||||
model_id: str
|
||||
"""Id of the model to call, e.g., amazon.titan-text-express-v1, this is
|
||||
equivalent to the modelId property in the list-foundation-models api"""
|
||||
equivalent to the modelId property in the list-foundation-models api. For custom and
|
||||
provisioned models, an ARN value is expected."""
|
||||
|
||||
model_kwargs: Optional[Dict] = None
|
||||
"""Keyword arguments to pass to the model."""
|
||||
@@ -353,6 +361,14 @@ class BedrockBase(BaseModel, ABC):
|
||||
}
|
||||
|
||||
def _get_provider(self) -> str:
|
||||
if self.provider:
|
||||
return self.provider
|
||||
if self.model_id.startswith("arn"):
|
||||
raise ValueError(
|
||||
"Model provider should be supplied when passing a model ARN as "
|
||||
"model_id"
|
||||
)
|
||||
|
||||
return self.model_id.split(".")[0]
|
||||
|
||||
@property
|
||||
|
Reference in New Issue
Block a user