Skip to content

Commit 203ee2c

Browse files
committed
removed unused imports
1 parent 597f2ac commit 203ee2c

File tree

2 files changed

+3
-5
lines changed

2 files changed

+3
-5
lines changed

scrapegraphai/builders/graph_builder.py

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

55
from langchain_core.prompts import ChatPromptTemplate
66
from langchain.chains import create_extraction_chain
7-
from ..models import OpenAI, Gemini
7+
from ..models import Gemini
88
from ..helpers import nodes_metadata, graph_schema
9-
from ..models.ernie import Ernie
10-
9+
from langchain_openai import ChatOpenAI
1110

1211
class GraphBuilder:
1312
"""
@@ -71,7 +70,7 @@ def _create_llm(self, llm_config: dict):
7170

7271
# select the model based on the model name
7372
if "gpt-" in llm_params["model"]:
74-
return OpenAI(llm_params)
73+
return ChatOpenAI(llm_params)
7574
elif "gemini" in llm_params["model"]:
7675
return Gemini(llm_params)
7776
elif "ernie" in llm_params["model"]:

scrapegraphai/graphs/abstract_graph.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,6 @@ def _create_llm(self, llm_config: dict) -> object:
135135
raise KeyError("model_tokens not specified") from exc
136136
return llm_params["model_instance"]
137137

138-
# Instantiate the language model based on the model name (models that use the common interface)
139138
def handle_model(model_name, provider, token_key, default_token=8192):
140139
try:
141140
self.model_token = models_tokens[provider][token_key]

0 commit comments

Comments
 (0)