Skip to content

Commit

Permalink
Community: Adds ability to pass a Config to the boto3 client used by …
Browse files Browse the repository at this point in the history
…Bedrock (#15029)

# Description  
This PR adds the ability to pass a `botocore.config.Config` instance to
the boto3 client instantiated by the Bedrock LLM.

Currently, the Bedrock LLM doesn't support a way to pass a Config, which
means that some settings (e.g., timeouts and retry configuration)
require instantiating a new boto3 client with a Config and then
replacing the LLM's client:

```python
llm = Bedrock(
        region_name='us-west-2',
        model_id="anthropic.claude-v2",
        model_kwargs={'max_tokens_to_sample': 4096, 'temperature': 0},
)

llm.client = boto_client('bedrock-runtime', region_name='us-west-2', config=Config({'read_timeout': 300}))
```

# Issue
N/A

# Dependencies
N/A
  • Loading branch information
blanehoneycutt-addepar authored Dec 22, 2023
1 parent dc71fcf commit 3fc1b35
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion libs/community/langchain_community/llms/bedrock.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
from __future__ import annotations

import json
import warnings
from abc import ABC
from typing import Any, Dict, Iterator, List, Mapping, Optional
from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Mapping, Optional

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
Expand All @@ -15,6 +17,9 @@
get_token_ids_anthropic,
)

if TYPE_CHECKING:
from botocore.config import Config

HUMAN_PROMPT = "\n\nHuman:"
ASSISTANT_PROMPT = "\n\nAssistant:"
ALTERNATION_ERROR = (
Expand Down Expand Up @@ -163,6 +168,9 @@ class BedrockBase(BaseModel, ABC):
See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
"""

config: Optional[Config] = None
"""An optional botocore.config.Config instance to pass to the client."""

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"""
Expand Down Expand Up @@ -212,6 +220,8 @@ def validate_environment(cls, values: Dict) -> Dict:
client_params["region_name"] = values["region_name"]
if values["endpoint_url"]:
client_params["endpoint_url"] = values["endpoint_url"]
if values["config"]:
client_params["config"] = values["config"]

values["client"] = session.client("bedrock-runtime", **client_params)

Expand Down

0 comments on commit 3fc1b35

Please sign in to comment.