From 39be38f292f7e8ccd40aa7af5d69db880b6971b5 Mon Sep 17 00:00:00 2001 From: JGalego Date: Mon, 13 May 2024 12:35:42 +0100 Subject: [PATCH] Fixed anthropic/bedrock conflict; Removed duplicate class Claude; Updated BedrockEmbeddings class --- scrapegraphai/graphs/abstract_graph.py | 9 +++++---- scrapegraphai/models/__init__.py | 1 - scrapegraphai/models/claude.py | 19 ------------------- 3 files changed, 5 insertions(+), 24 deletions(-) delete mode 100644 scrapegraphai/models/claude.py diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index f58f04a2..b96503ba 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -3,11 +3,12 @@ """ from abc import ABC, abstractmethod from typing import Optional +from langchain_aws import BedrockEmbeddings from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings -from langchain_community.embeddings import HuggingFaceHubEmbeddings, OllamaEmbeddings, BedrockEmbeddings +from langchain_community.embeddings import HuggingFaceHubEmbeddings, OllamaEmbeddings from langchain_google_genai import GoogleGenerativeAIEmbeddings from ..helpers import models_tokens -from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic, Claude, DeepSeek +from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic, DeepSeek class AbstractGraph(ABC): @@ -145,12 +146,12 @@ def _create_llm(self, llm_config: dict, chat=False) -> object: except KeyError as exc: raise KeyError("Model not supported") from exc return Gemini(llm_params) - elif "claude" in llm_params["model"]: + elif llm_params["model"].startswith("claude"): try: self.model_token = models_tokens["claude"][llm_params["model"]] except KeyError as exc: raise KeyError("Model not supported") from exc - return Claude(llm_params) + return Anthropic(llm_params) elif "ollama" in llm_params["model"]: llm_params["model"] = llm_params["model"].split("/")[-1] diff --git a/scrapegraphai/models/__init__.py b/scrapegraphai/models/__init__.py index 9c049fcd..7e7d5e18 100644 --- a/scrapegraphai/models/__init__.py +++ b/scrapegraphai/models/__init__.py @@ -12,5 +12,4 @@ from .groq import Groq from .bedrock import Bedrock from .anthropic import Anthropic -from .claude import Claude from .deepseek import DeepSeek diff --git a/scrapegraphai/models/claude.py b/scrapegraphai/models/claude.py deleted file mode 100644 index b0030fc4..00000000 --- a/scrapegraphai/models/claude.py +++ /dev/null @@ -1,19 +0,0 @@ -""" -Claude model -""" -from langchain_anthropic import ChatAnthropic - - -class Claude(ChatAnthropic): - """Class for wrapping bedrock module""" - - def __init__(self, llm_config: dict): - """ - A wrapper for the Claude class that provides default configuration - and could be extended with additional methods if needed. - - Args: - llm_config (dict): Configuration parameters for the language model. - """ - # Initialize the superclass (ChatAnthropic) with provided config parameters - super().__init__(**llm_config)