Skip to content

Commit

Permalink
Merge pull request #600 from alexljenkins/bugfix/graph_builder-import…
Browse files Browse the repository at this point in the history
…-error
  • Loading branch information
VinciGit00 authored Aug 28, 2024
2 parents 4eccc76 + bda30a9 commit edfe45e
Showing 1 changed file with 6 additions and 4 deletions.
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

0 comments on commit edfe45e

Please sign in to comment.