diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index 4cd9fe6e..a0258fa9 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -304,6 +304,8 @@ }, "togetherai": {"Meta-Llama-3.1-70B-Instruct-Turbo": 128000}, "xai": { - "grok-1": 8192 + "grok-1": 8192, + "grok-3": 128000, + "grok-3-mini": 128000, }, } diff --git a/scrapegraphai/models/xai.py b/scrapegraphai/models/xai.py index 405f88b6..86969483 100644 --- a/scrapegraphai/models/xai.py +++ b/scrapegraphai/models/xai.py @@ -1,23 +1,22 @@ """ xAI Grok Module """ -from langchain_groq import ChatGroq as LangchainChatGroq +from langchain_openai import ChatOpenAI -class XAI(LangchainChatGroq): + +class XAI(ChatOpenAI): """ - Wrapper for the ChatGroq class from langchain_groq, for use with xAI models. - Handles API key mapping from generic 'api_key' to 'groq_api_key' and - maps 'model' to 'model_name'. + A wrapper for the ChatOpenAI class (xAI uses an OpenAI-compatible 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): - if "api_key" in llm_config and "groq_api_key" not in llm_config: - llm_config["groq_api_key"] = llm_config.pop("api_key") - - if "model" in llm_config and "model_name" not in llm_config: - llm_config["model_name"] = llm_config.pop("model") + if "api_key" in llm_config: + llm_config["openai_api_key"] = llm_config.pop("api_key") + llm_config["openai_api_base"] = "https://api.x.ai/v1" super().__init__(**llm_config) \ No newline at end of file