Merge pull request #600 from alexljenkins/bugfix/graph_builder-import-error

Author: VinciGit00

scrapegraphai/builders/graph_builder.py

@@ -4,10 +4,12 @@
 
 from langchain_core.prompts import ChatPromptTemplate
 from langchain.chains import create_extraction_chain
-from ..models import Gemini
-from ..helpers import nodes_metadata, graph_schema
+from langchain_community.chat_models import ErnieBotChat
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_openai import ChatOpenAI
 
+from ..helpers import nodes_metadata, graph_schema
+
 class GraphBuilder:
     """
     GraphBuilder is a dynamic tool for constructing web scraping graphs based on user prompts. 
@@ -72,9 +74,9 @@ def _create_llm(self, llm_config: dict):
         if "gpt-" in llm_params["model"]:
             return ChatOpenAI(llm_params)
         elif "gemini" in llm_params["model"]:
-            return Gemini(llm_params)
+            return ChatGoogleGenerativeAI(llm_params)
         elif "ernie" in llm_params["model"]:
-            return Ernie(llm_params)
+            return ErnieBotChat(llm_params)
         raise ValueError("Model not supported")
 
     def _generate_nodes_description(self):