Skip to content

Commit bda30a9

Browse files
author
ajenkins
committed
swapped failing imports (local Gemini and non-imported Ernie) for langchain chat models in graph_builder.py
1 parent 4eccc76 commit bda30a9

File tree

1 file changed

+6
-4
lines changed

1 file changed

+6
-4
lines changed

scrapegraphai/builders/graph_builder.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@
44

55
from langchain_core.prompts import ChatPromptTemplate
66
from langchain.chains import create_extraction_chain
7-
from ..models import Gemini
8-
from ..helpers import nodes_metadata, graph_schema
7+
from langchain_community.chat_models import ErnieBotChat
8+
from langchain_google_genai import ChatGoogleGenerativeAI
99
from langchain_openai import ChatOpenAI
1010

11+
from ..helpers import nodes_metadata, graph_schema
12+
1113
class GraphBuilder:
1214
"""
1315
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):
7274
if "gpt-" in llm_params["model"]:
7375
return ChatOpenAI(llm_params)
7476
elif "gemini" in llm_params["model"]:
75-
return Gemini(llm_params)
77+
return ChatGoogleGenerativeAI(llm_params)
7678
elif "ernie" in llm_params["model"]:
77-
return Ernie(llm_params)
79+
return ErnieBotChat(llm_params)
7880
raise ValueError("Model not supported")
7981

8082
def _generate_nodes_description(self):

0 commit comments

Comments
 (0)