Skip to content

Commit 859c5d5

Browse files
committed
Refactored to include custom AWS client for bedrock; Added missing Anthropic class
1 parent f3d44c0 commit 859c5d5

File tree

2 files changed

+21
-2
lines changed

2 files changed

+21
-2
lines changed

scrapegraphai/graphs/abstract_graph.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -190,12 +190,13 @@ def _create_llm(self, llm_config: dict, chat=False) -> object:
190190
elif "bedrock" in llm_params["model"]:
191191
llm_params["model"] = llm_params["model"].split("/")[-1]
192192
model_id = llm_params["model"]
193-
193+
client = llm_params.get('client', None)
194194
try:
195195
self.model_token = models_tokens["bedrock"][llm_params["model"]]
196196
except KeyError as exc:
197197
raise KeyError("Model not supported") from exc
198198
return Bedrock({
199+
"client": client,
199200
"model_id": model_id,
200201
"model_kwargs": {
201202
"temperature": llm_params["temperature"],
@@ -289,11 +290,12 @@ def _create_embedder(self, embedder_config: dict) -> object:
289290
return GoogleGenerativeAIEmbeddings(model=embedder_config["model"])
290291
elif "bedrock" in embedder_config["model"]:
291292
embedder_config["model"] = embedder_config["model"].split("/")[-1]
293+
client = embedder_config.get('client', None)
292294
try:
293295
models_tokens["bedrock"][embedder_config["model"]]
294296
except KeyError as exc:
295297
raise KeyError("Model not supported") from exc
296-
return BedrockEmbeddings(client=None, model_id=embedder_config["model"])
298+
return BedrockEmbeddings(client=client, model_id=embedder_config["model"])
297299
else:
298300
raise ValueError(
299301
"Model provided by the configuration not supported")

scrapegraphai/models/anthropic.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
"""
2+
Anthropic Module
3+
"""
4+
from langchain_anthropic import ChatAnthropic
5+
6+
7+
class Anthropic(ChatAnthropic):
8+
"""
9+
A wrapper for the ChatAnthropic class that provides default configuration
10+
and could be extended with additional methods if needed.
11+
12+
Args:
13+
llm_config (dict): Configuration parameters for the language model.
14+
"""
15+
16+
def __init__(self, llm_config: dict):
17+
super().__init__(**llm_config)

0 commit comments

Comments
 (0)