refactoring

Co-Authored-By: Federico Aguzzi <[email protected]>

Author: VinciGit00

scrapegraphai/nodes/fetch_node.py

@@ -4,7 +4,7 @@
 
 import json
 from typing import List, Optional
-
+from langchain_openai import ChatOpenAI
 import pandas as pd
 import requests
 from langchain_community.document_loaders import PyPDFLoader
@@ -14,7 +14,6 @@
 from ..utils.convert_to_md import convert_to_md
 from ..utils.logging import get_logger
 from .base_node import BaseNode
-from ..models import OpenAI
 
 
 class FetchNode(BaseNode):
@@ -165,7 +164,7 @@ def execute(self, state):
 
             parsed_content = source
 
-            if  isinstance(self.llm_model, OpenAI) and not self.script_creator or self.force and not self.script_creator:
+            if  isinstance(self.llm_model, ChatOpenAI) and not self.script_creator or self.force and not self.script_creator:
                 parsed_content = convert_to_md(source)
 
             compressed_document = [
@@ -184,7 +183,7 @@ def execute(self, state):
                 if not self.cut:
                     parsed_content = cleanup_html(response, source)
 
-                if  (isinstance(self.llm_model, OpenAI) and not self.script_creator) or (self.force and not self.script_creator):
+                if  (isinstance(self.llm_model, ChatOpenAI) and not self.script_creator) or (self.force and not self.script_creator):
                     parsed_content = convert_to_md(source, input_data[0])
                 compressed_document = [Document(page_content=parsed_content)]
             else:
@@ -206,7 +205,7 @@ def execute(self, state):
                 raise ValueError("No HTML body content found in the document fetched by ChromiumLoader.")
             parsed_content = document[0].page_content
 
-            if  isinstance(self.llm_model, OpenAI) and not self.script_creator or self.force and not self.script_creator and not self.openai_md_enabled:
+            if  isinstance(self.llm_model, ChatOpenAI) and not self.script_creator or self.force and not self.script_creator and not self.openai_md_enabled:
 
                 parsed_content = convert_to_md(document[0].page_content, input_data[0])
 

scrapegraphai/nodes/generate_answer_node.py

@@ -9,6 +9,7 @@
 from langchain_openai import ChatOpenAI
 from langchain_community.chat_models import ChatOllama
 from tqdm import tqdm
+from langchain_openai import ChatOpenAI
 from ..utils.logging import get_logger
 from .base_node import BaseNode
 from ..helpers import template_chunks, template_no_chunks, template_merge, template_chunks_md, template_no_chunks_md, template_merge_md

scrapegraphai/nodes/search_internet_node.py

@@ -6,12 +6,10 @@
 
 from langchain.output_parsers import CommaSeparatedListOutputParser
 from langchain.prompts import PromptTemplate
-
+from langchain_community.chat_models import ChatOllama
 from ..utils.logging import get_logger
 from ..utils.research_web import search_on_web
 from .base_node import BaseNode
-from ..models import Ollama
-
 
 class SearchInternetNode(BaseNode):
     """
@@ -97,7 +95,7 @@ def execute(self, state: dict) -> dict:
         search_answer = search_prompt | self.llm_model | output_parser
 
         # Ollama: Use no json format when creating the search query
-        if isinstance(self.llm_model, Ollama) and self.llm_model.format == 'json':
+        if isinstance(self.llm_model, ChatOllama) and self.llm_model.format == 'json':
             self.llm_model.format = None
             search_query = search_answer.invoke({"user_prompt": user_prompt})[0]
             self.llm_model.format = 'json'