Skip to content

Graph Builder failed import bugfix #600

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions scrapegraphai/builders/graph_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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):
Expand Down