From bda30a9ecbaa65668a3f21fbc9a5c50dcd21576c Mon Sep 17 00:00:00 2001 From: ajenkins Date: Wed, 28 Aug 2024 16:48:00 +1000 Subject: [PATCH] swapped failing imports (local Gemini and non-imported Ernie) for langchain chat models in graph_builder.py --- scrapegraphai/builders/graph_builder.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/scrapegraphai/builders/graph_builder.py b/scrapegraphai/builders/graph_builder.py index 303f1457..69ebe492 100644 --- a/scrapegraphai/builders/graph_builder.py +++ b/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):