From 786af992f8fbdadfdc3d2d6a06c0cfd81289f8f2 Mon Sep 17 00:00:00 2001
From: Federico Aguzzi <62149513+f-aguzzi@users.noreply.github.com>
Date: Wed, 7 Aug 2024 15:15:35 +0200
Subject: [PATCH] fix(AbstractGraph): LangChain warnings handling, Mistral
 tokens

---
 scrapegraphai/graphs/abstract_graph.py | 5 ++++-
 scrapegraphai/helpers/models_tokens.py | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py
index d5fa2c47..83b532bc 100644
--- a/scrapegraphai/graphs/abstract_graph.py
+++ b/scrapegraphai/graphs/abstract_graph.py
@@ -5,6 +5,7 @@ AbstractGraph Module
 from abc import ABC, abstractmethod
 from typing import Optional
 import uuid
+import warnings
 from pydantic import BaseModel
 
 from langchain_community.chat_models import ErnieBotChat
@@ -144,7 +145,9 @@ class AbstractGraph(ABC):
                 self.model_token = default_token
             llm_params["model_provider"] = provider
             llm_params["model"] = model_name
-            return init_chat_model(**llm_params)
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                return init_chat_model(**llm_params)
 
         if "azure" in llm_params["model"]:
             model_name = llm_params["model"].split("/")[-1]
diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py
index 7cf0111c..e32838f1 100644
--- a/scrapegraphai/helpers/models_tokens.py
+++ b/scrapegraphai/helpers/models_tokens.py
@@ -145,7 +145,7 @@ models_tokens = {
         "cohere.embed-english-v3": 512,
         "cohere.embed-multilingual-v3": 512
     },
-    "mistral": {
+    "mistralai": {
         "mistral-large-latest": 128000,
         "open-mistral-nemo": 128000,
         "codestral-latest": 32000,