diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index a71508a4..f58f04a2 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -7,7 +7,7 @@ from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings from langchain_community.embeddings import HuggingFaceHubEmbeddings, OllamaEmbeddings, BedrockEmbeddings from langchain_google_genai import GoogleGenerativeAIEmbeddings from ..helpers import models_tokens -from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic, Claude +from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic, Claude, DeepSeek class AbstractGraph(ABC): @@ -200,6 +200,12 @@ class AbstractGraph(ABC): elif "claude-3-" in llm_params["model"]: self.model_token = models_tokens["claude"]["claude3"] return Anthropic(llm_params) + elif "deepseek" in llm_params["model"]: + try: + self.model_token = models_tokens["deepseek"][llm_params["model"]] + except KeyError as exc: + raise KeyError("Model not supported") from exc + return DeepSeek(llm_params) else: raise ValueError( "Model provided by the configuration not supported") diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index 7b28f9af..ea330bee 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -102,6 +102,12 @@ models_tokens = { "cognitivecomputations/dolphin-2.9-llama3-8b-gguf": 8192, "cognitivecomputations/dolphin-2.8-mistral-7b-v02": 32768, "cognitivecomputations/dolphin-2.5-mixtral-8x7b": 32768, - "TheBloke/dolphin-2.7-mixtral-8x7b-GGUF": 32768 + "TheBloke/dolphin-2.7-mixtral-8x7b-GGUF": 32768, + "deepseek-ai/DeepSeek-V2": 131072, + "deepseek-ai/DeepSeek-V2-Chat": 131072 + }, + "deepseek": { + "deepseek-chat": 32768, + "deepseek-coder": 16384 } } diff --git a/scrapegraphai/models/__init__.py b/scrapegraphai/models/__init__.py index a6cab89d..9c049fcd 100644 --- a/scrapegraphai/models/__init__.py +++ b/scrapegraphai/models/__init__.py @@ -13,3 +13,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/deepseek.py b/scrapegraphai/models/deepseek.py new file mode 100644 index 00000000..523fe667 --- /dev/null +++ b/scrapegraphai/models/deepseek.py @@ -0,0 +1,18 @@ +""" +DeepSeek Module +""" +from langchain_openai import ChatOpenAI + + +class DeepSeek(ChatOpenAI): + """ + A wrapper for the ChatOpenAI class (DeepSeek uses an OpenAI-like API) 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)