diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index 5a81ecaa..f06713b2 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -379,6 +379,10 @@ "mixtral-moe-8x22B-instruct": 65536, "mixtral-moe-8x7B-instruct": 65536, }, + "futureinfra": { + "openai/gpt-4o-mini": 128000, + "openai/gpt-4o": 128000, + }, "clod": { "open-mistral-7b": 32000, "Llama-3.1-70b": 128000, diff --git a/scrapegraphai/models/__init__.py b/scrapegraphai/models/__init__.py index db01cd2e..c24c819f 100644 --- a/scrapegraphai/models/__init__.py +++ b/scrapegraphai/models/__init__.py @@ -6,6 +6,7 @@ from .cheaperinference import CheaperInference from .clod import CLoD from .deepseek import DeepSeek +from .futureinfra import FutureInfra from .minimax import MiniMax from .nvidia import Nvidia from .oneapi import OneApi @@ -13,4 +14,4 @@ from .openai_tts import OpenAITextToSpeech from .xai import XAI -__all__ = ["AtlasCloud", "CheaperInference", "DeepSeek", "MiniMax", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD", "XAI", "Nvidia"] +__all__ = ["AtlasCloud", "CheaperInference", "DeepSeek", "FutureInfra", "MiniMax", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD", "XAI", "Nvidia"] diff --git a/scrapegraphai/models/futureinfra.py b/scrapegraphai/models/futureinfra.py new file mode 100644 index 00000000..971d7e10 --- /dev/null +++ b/scrapegraphai/models/futureinfra.py @@ -0,0 +1,23 @@ +""" +FutureInfra Module +""" + +from langchain_openai import ChatOpenAI + + +class FutureInfra(ChatOpenAI): + """ + A wrapper for ChatOpenAI configured for FutureInfra's OpenAI-compatible + AI API router. Model ids use the ``provider/model`` format, + e.g. ``openai/gpt-4o-mini``. + + 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://futureinfra.ai/v1/ai" + + super().__init__(**llm_config) diff --git a/tests/test_futureinfra_model.py b/tests/test_futureinfra_model.py new file mode 100644 index 00000000..6856ce28 --- /dev/null +++ b/tests/test_futureinfra_model.py @@ -0,0 +1,48 @@ +"""Tests for FutureInfra model configuration.""" + +import importlib.util +import os + + +def test_futureinfra_model_sets_openai_compatible_base_url(): + """FutureInfra should map api_key and set its base URL.""" + spec = importlib.util.spec_from_file_location( + "futureinfra", + os.path.join( + os.path.dirname(__file__), + "..", + "scrapegraphai", + "models", + "futureinfra.py", + ), + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + model = module.FutureInfra( + api_key="test-key", + model="openai/gpt-4o-mini", + ) + + assert str(model.openai_api_base).rstrip("/") == "https://futureinfra.ai/v1/ai" + assert model.openai_api_key.get_secret_value() == "test-key" + + +def test_futureinfra_models_in_token_list(): + """FutureInfra defaults should be listed with their 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) + + futureinfra_models = module.models_tokens["futureinfra"] + assert futureinfra_models["openai/gpt-4o-mini"] == 128000 + assert futureinfra_models["openai/gpt-4o"] == 128000