Skip to content

Commit 2c5f934

Browse files
committed
refactor: remove LangChain wrappers
1 parent 9e795f4 commit 2c5f934

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

scrapegraphai/nodes/generate_answer_node.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,10 @@
66
from langchain.prompts import PromptTemplate
77
from langchain_core.output_parsers import JsonOutputParser
88
from langchain_core.runnables import RunnableParallel
9+
from langchain_openai import ChatOpenAI
10+
from langchain_community.chat_models import ChatOllama
911
from tqdm import tqdm
1012
from ..utils.logging import get_logger
11-
from ..models import Ollama, OpenAI
1213
from .base_node import BaseNode
1314
from ..helpers import template_chunks, template_no_chunks, template_merge, template_chunks_md, template_no_chunks_md, template_merge_md
1415

@@ -41,7 +42,7 @@ def __init__(
4142

4243
self.llm_model = node_config["llm_model"]
4344

44-
if isinstance(node_config["llm_model"], Ollama):
45+
if isinstance(node_config["llm_model"], ChatOllama):
4546
self.llm_model.format="json"
4647

4748
self.verbose = (
@@ -93,7 +94,7 @@ def execute(self, state: dict) -> dict:
9394

9495
format_instructions = output_parser.get_format_instructions()
9596

96-
if isinstance(self.llm_model, OpenAI) and not self.script_creator or self.force and not self.script_creator or self.is_md_scraper:
97+
if isinstance(self.llm_model, ChatOpenAI) and not self.script_creator or self.force and not self.script_creator or self.is_md_scraper:
9798
template_no_chunks_prompt = template_no_chunks_md
9899
template_chunks_prompt = template_chunks_md
99100
template_merge_prompt = template_merge_md

0 commit comments

Comments
 (0)