Merge pull request #229 from JGalego/feat/custom-aws-creds

feat: Custom AWS credentials
This commit is contained in:
Marco Vinciguerra 2024-05-13 20:14:28 +02:00 committed by GitHub
commit 28ab8da1fd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 2 deletions

View File

@ -190,12 +190,13 @@ class AbstractGraph(ABC):
elif "bedrock" in llm_params["model"]:
llm_params["model"] = llm_params["model"].split("/")[-1]
model_id = llm_params["model"]
client = llm_params.get('client', None)
try:
self.model_token = models_tokens["bedrock"][llm_params["model"]]
except KeyError as exc:
raise KeyError("Model not supported") from exc
return Bedrock({
"client": client,
"model_id": model_id,
"model_kwargs": {
"temperature": llm_params["temperature"],
@ -289,11 +290,12 @@ class AbstractGraph(ABC):
return GoogleGenerativeAIEmbeddings(model=embedder_config["model"])
elif "bedrock" in embedder_config["model"]:
embedder_config["model"] = embedder_config["model"].split("/")[-1]
client = embedder_config.get('client', None)
try:
models_tokens["bedrock"][embedder_config["model"]]
except KeyError as exc:
raise KeyError("Model not supported") from exc
return BedrockEmbeddings(client=None, model_id=embedder_config["model"])
return BedrockEmbeddings(client=client, model_id=embedder_config["model"])
else:
raise ValueError(
"Model provided by the configuration not supported")

View File

@ -0,0 +1,17 @@
"""
Anthropic Module
"""
from langchain_anthropic import ChatAnthropic
class Anthropic(ChatAnthropic):
"""
A wrapper for the ChatAnthropic class that provides default configuration
and could be extended with additional methods if needed.
Args:
llm_config (dict): Configuration parameters for the language model.
"""
def __init__(self, llm_config: dict):
super().__init__(**llm_config)