feat: add grok integration

Author: VinciGit00

scrapegraphai/graphs/abstract_graph.py

@@ -13,7 +13,7 @@
 from pydantic import BaseModel
 
 from ..helpers import models_tokens
-from ..models import CLoD, DeepSeek, OneApi
+from ..models import CLoD, DeepSeek, OneApi, XAI
 from ..utils.logging import set_verbosity_info, set_verbosity_warning
 
 
@@ -163,6 +163,7 @@ def _create_llm(self, llm_config: dict) -> object:
  "fireworks",
  "clod",
  "togetherai",
+ "xai",
  }
 
  if "/" in llm_params["model"]:
@@ -217,6 +218,7 @@ def _create_llm(self, llm_config: dict) -> object:
  "deepseek",
  "togetherai",
  "clod",
+ "xai",
  }:
  if llm_params["model_provider"] == "bedrock":
  llm_params["model_kwargs"] = {
@@ -242,6 +244,9 @@ def _create_llm(self, llm_config: dict) -> object:
  elif model_provider == "oneapi":
  return OneApi(**llm_params)
 
+ elif model_provider == "xai":
+ return XAI(**llm_params)
+
  elif model_provider == "togetherai":
  try:
  from langchain_together import ChatTogether

scrapegraphai/helpers/models_tokens.py

@@ -150,7 +150,7 @@
  "llama3-70b-8192": 8192,
  "mixtral-8x7b-32768": 32768,
  "gemma-7b-it": 8192,
- "claude-3-haiku-20240307'": 8192,
+ "claude-3-haiku-20240307": 8192,
  },
  "toghetherai": {
  "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": 128000,
@@ -303,4 +303,7 @@
  "grok-2-latest": 128000,
  },
  "togetherai": {"Meta-Llama-3.1-70B-Instruct-Turbo": 128000},
+ "xai": {
+ "grok-1": 8192
+ },
 }

scrapegraphai/models/__init__.py

@@ -7,5 +7,6 @@
 from .oneapi import OneApi
 from .openai_itt import OpenAIImageToText
 from .openai_tts import OpenAITextToSpeech
+from .xai import XAI
 
-__all__ = ["DeepSeek", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD"]
+__all__ = ["DeepSeek", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD", "XAI"]

scrapegraphai/models/xai.py

@@ -0,0 +1,23 @@
+"""
+xAI Grok Module
+"""
+from langchain_groq import ChatGroq as LangchainChatGroq
+
+class XAI(LangchainChatGroq):
+ """
+ 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'.
+
+ 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")
+
+ super().__init__(**llm_config)