diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index 17f85583..d0234c17 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -335,6 +335,10 @@ "deepseek-ai/DeepSeek-V2-Chat": 131072, "claude-3-haiku": 200000, }, + "atlascloud": { + "deepseek-ai/deepseek-v4-pro": 1048576, + "qwen/qwen3.5-flash": 1000000, + }, "deepseek": { "deepseek-chat": 128000, "deepseek-coder": 128000, diff --git a/scrapegraphai/models/__init__.py b/scrapegraphai/models/__init__.py index d2157d37..1cbf0e06 100644 --- a/scrapegraphai/models/__init__.py +++ b/scrapegraphai/models/__init__.py @@ -2,6 +2,7 @@ This module contains the model definitions used in the ScrapeGraphAI application. """ +from .atlascloud import AtlasCloud from .clod import CLoD from .deepseek import DeepSeek from .minimax import MiniMax @@ -11,4 +12,4 @@ from .openai_tts import OpenAITextToSpeech from .xai import XAI -__all__ = ["DeepSeek", "MiniMax", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD", "XAI", "Nvidia"] +__all__ = ["AtlasCloud", "DeepSeek", "MiniMax", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD", "XAI", "Nvidia"] diff --git a/scrapegraphai/models/atlascloud.py b/scrapegraphai/models/atlascloud.py new file mode 100644 index 00000000..e3626f95 --- /dev/null +++ b/scrapegraphai/models/atlascloud.py @@ -0,0 +1,22 @@ +""" +Atlas Cloud Module +""" + +from langchain_openai import ChatOpenAI + + +class AtlasCloud(ChatOpenAI): + """ + A wrapper for ChatOpenAI configured for Atlas Cloud's OpenAI-compatible + LLM API. + + Args: + llm_config (dict): Configuration parameters for the language model. + """ + + def __init__(self, **llm_config): + if "api_key" in llm_config: + llm_config["openai_api_key"] = llm_config.pop("api_key") + llm_config["openai_api_base"] = "https://api.atlascloud.ai/v1" + + super().__init__(**llm_config) diff --git a/tests/test_atlascloud_model.py b/tests/test_atlascloud_model.py new file mode 100644 index 00000000..78b42edf --- /dev/null +++ b/tests/test_atlascloud_model.py @@ -0,0 +1,48 @@ +"""Tests for Atlas Cloud model configuration.""" + +import importlib.util +import os + + +def test_atlascloud_model_sets_openai_compatible_base_url(): + """AtlasCloud should map api_key and set the Atlas Cloud base URL.""" + spec = importlib.util.spec_from_file_location( + "atlascloud", + os.path.join( + os.path.dirname(__file__), + "..", + "scrapegraphai", + "models", + "atlascloud.py", + ), + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + model = module.AtlasCloud( + api_key="test-key", + model="qwen/qwen3.5-flash", + ) + + assert str(model.openai_api_base).rstrip("/") == "https://api.atlascloud.ai/v1" + assert model.openai_api_key.get_secret_value() == "test-key" + + +def test_atlascloud_models_in_token_list(): + """Atlas Cloud defaults should be listed with current context lengths.""" + spec = importlib.util.spec_from_file_location( + "models_tokens", + os.path.join( + os.path.dirname(__file__), + "..", + "scrapegraphai", + "helpers", + "models_tokens.py", + ), + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + atlascloud_models = module.models_tokens["atlascloud"] + assert atlascloud_models["deepseek-ai/deepseek-v4-pro"] == 1048576 + assert atlascloud_models["qwen/qwen3.5-flash"] == 1000000