Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion scrapegraphai/helpers/models_tokens.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
}
19 changes: 9 additions & 10 deletions scrapegraphai/models/xai.py
Original file line number Diff line number Diff line change
@@ -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)
Loading