refactoring of the code

Author: VinciGit00

scrapegraphai/graphs/abstract_graph.py

@@ -7,8 +7,6 @@
 import uuid
 import warnings
 from pydantic import BaseModel
-from langchain_community.chat_models import ErnieBotChat
-from langchain_nvidia_ai_endpoints import ChatNVIDIA
 from langchain.chat_models import init_chat_model
 from ..helpers import models_tokens
 from ..models import (
@@ -147,8 +145,7 @@ def handle_model(model_name, provider, token_key, default_token=8192):
                 warnings.simplefilter("ignore")
                 return init_chat_model(**llm_params)
 
-        known_models = ["chatgpt","gpt","openai", "azure_openai", "google_genai", "ollama", "oneapi", "nvidia", "groq", "google_vertexai", "bedrock", "mistralai", "hugging_face", "deepseek", "ernie", "fireworks"]
-
+        known_models = {"chatgpt","gpt","openai", "azure_openai", "google_genai", "ollama", "oneapi", "nvidia", "groq", "google_vertexai", "bedrock", "mistralai", "hugging_face", "deepseek", "ernie", "fireworks"}
         if llm_params["model"].split("/")[0] not in known_models and llm_params["model"].split("-")[0] not in known_models:
             raise ValueError(f"Model '{llm_params['model']}' is not supported")
 
@@ -198,6 +195,8 @@ def handle_model(model_name, provider, token_key, default_token=8192):
                 return DeepSeek(llm_params)
 
             elif "ernie" in llm_params["model"]:
+                from langchain_community.chat_models import ErnieBotChat
+
                 try:
                     self.model_token = models_tokens["ernie"][llm_params["model"]]
                 except KeyError:
@@ -215,6 +214,8 @@ def handle_model(model_name, provider, token_key, default_token=8192):
                 return OneApi(llm_params)
 
             elif "nvidia" in llm_params["model"]:
+                from langchain_nvidia_ai_endpoints import ChatNVIDIA
+
                 try:
                     self.model_token = models_tokens["nvidia"][llm_params["model"].split("/")[-1]]
                     llm_params["model"] = "/".join(llm_params["model"].split("/")[1:])

scrapegraphai/nodes/generate_answer_csv_node.py

@@ -9,7 +9,8 @@
 from tqdm import tqdm
 from ..utils.logging import get_logger
 from .base_node import BaseNode
-from ..prompts.generate_answer_node_csv_prompts import TEMPLATE_CHUKS_CSV, TEMPLATE_NO_CHUKS_CSV, TEMPLATE_MERGE_CSV
+from ..prompts.generate_answer_node_csv_prompts import (TEMPLATE_CHUKS_CSV,
+                                                        TEMPLATE_NO_CHUKS_CSV, TEMPLATE_MERGE_CSV)
 
 class GenerateAnswerCSVNode(BaseNode):
     """
@@ -95,22 +96,22 @@ def execute(self, state):
         else:
             output_parser = JsonOutputParser()
 
-        TEMPLATE_NO_CHUKS_CSV_prompt = TEMPLATE_NO_CHUKS_CSV
-        TEMPLATE_CHUKS_CSV_prompt = TEMPLATE_CHUKS_CSV
-        TEMPLATE_MERGE_CSV_prompt  = TEMPLATE_MERGE_CSV
+        TEMPLATE_NO_CHUKS_CSV_PROMPT = TEMPLATE_NO_CHUKS_CSV
+        TEMPLATE_CHUKS_CSV_PROMPT = TEMPLATE_CHUKS_CSV
+        TEMPLATE_MERGE_CSV_PROMPT  = TEMPLATE_MERGE_CSV
 
         if self.additional_info is not None:
-            TEMPLATE_NO_CHUKS_CSV_prompt = self.additional_info + TEMPLATE_NO_CHUKS_CSV
-            TEMPLATE_CHUKS_CSV_prompt = self.additional_info + TEMPLATE_CHUKS_CSV
-            TEMPLATE_MERGE_CSV_prompt = self.additional_info + TEMPLATE_MERGE_CSV
+            TEMPLATE_NO_CHUKS_CSV_PROMPT = self.additional_info + TEMPLATE_NO_CHUKS_CSV
+            TEMPLATE_CHUKS_CSV_PROMPT = self.additional_info + TEMPLATE_CHUKS_CSV
+            TEMPLATE_MERGE_CSV_PROMPT = self.additional_info + TEMPLATE_MERGE_CSV
 
         format_instructions = output_parser.get_format_instructions()
 
         chains_dict = {}
 
         if len(doc) == 1:
             prompt = PromptTemplate(
-                template=TEMPLATE_NO_CHUKS_CSV_prompt,
+                template=TEMPLATE_NO_CHUKS_CSV_PROMPT,
                 input_variables=["question"],
                 partial_variables={
                     "context": doc,
@@ -127,7 +128,7 @@ def execute(self, state):
             tqdm(doc, desc="Processing chunks", disable=not self.verbose)
         ):
             prompt = PromptTemplate(
-                    template=TEMPLATE_CHUKS_CSV_prompt,
+                    template=TEMPLATE_CHUKS_CSV_PROMPT,
                     input_variables=["question"],
                     partial_variables={
                         "context": chunk,
@@ -144,7 +145,7 @@ def execute(self, state):
         batch_results =  async_runner.invoke({"question": user_prompt})
 
         merge_prompt = PromptTemplate(
-                template = TEMPLATE_MERGE_CSV_prompt,
+                template = TEMPLATE_MERGE_CSV_PROMPT,
                 input_variables=["context", "question"],
                 partial_variables={"format_instructions": format_instructions},
             )
@@ -153,4 +154,4 @@ def execute(self, state):
         answer = merge_chain.invoke({"context": batch_results, "question": user_prompt})
 
         state.update({self.output[0]: answer})
-        return state
+        return state

scrapegraphai/nodes/generate_scraper_node.py

@@ -67,10 +67,8 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-        # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
 
         user_prompt = input_data[0]

scrapegraphai/nodes/get_probable_tags_node.py

@@ -58,10 +58,8 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-        # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
 
         user_prompt = input_data[0]
@@ -88,10 +86,8 @@ def execute(self, state: dict) -> dict:
             },
         )
 
-        # Execute the chain to get probable tags
         tag_answer = tag_prompt | self.llm_model | output_parser
         probable_tags = tag_answer.invoke({"question": user_prompt})
 
-        # Update the dictionary with probable tags
         state.update({self.output[0]: probable_tags})
         return state

scrapegraphai/nodes/graph_iterator_node.py

@@ -103,7 +103,6 @@ async def _async_execute(self, state: dict, batchsize: int) -> dict:
         if graph_instance is None:
             raise ValueError("graph instance is required for concurrent execution")
 
-        # Assign depth level to the graph
         if "graph_depth" in graph_instance.config:
             graph_instance.config["graph_depth"] += 1
         else:
@@ -113,14 +112,12 @@ async def _async_execute(self, state: dict, batchsize: int) -> dict:
 
         participants = []
 
-        # semaphore to limit the number of concurrent tasks
         semaphore = asyncio.Semaphore(batchsize)
 
         async def _async_run(graph):
             async with semaphore:
                 return await asyncio.to_thread(graph.run)
 
-        # creates a deepcopy of the graph instance for each endpoint
         for url in urls:
             instance = copy.copy(graph_instance)
             instance.source = url

scrapegraphai/nodes/merge_answers_node.py

@@ -56,21 +56,17 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-        # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
 
         user_prompt = input_data[0]
         answers = input_data[1]
 
-        # merge the answers in one string
         answers_str = ""
         for i, answer in enumerate(answers):
             answers_str += f"CONTENT WEBSITE {i+1}: {answer}\n"
 
-        # Initialize the output parser
         if self.node_config.get("schema", None) is not None:
             output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
         else:
@@ -90,6 +86,5 @@ def execute(self, state: dict) -> dict:
         merge_chain = prompt_template | self.llm_model | output_parser
         answer = merge_chain.invoke({"user_prompt": user_prompt})
 
-        # Update the state with the generated answer
         state.update({self.output[0]: answer})
         return state

scrapegraphai/nodes/parse_node.py

@@ -59,13 +59,11 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-        # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
-        # Parse the document
         docs_transformed = input_data[0]
+
         if self.parse_html:
             docs_transformed = Html2TextTransformer().transform_documents(input_data[0])
             docs_transformed = docs_transformed[0]
@@ -77,7 +75,6 @@ def execute(self, state: dict) -> dict:
         else:
             docs_transformed = docs_transformed[0]
 
-            # Adapt the chunk size, leaving room for the reply, the prompt and the schema
             chunk_size = self.node_config.get("chunk_size", 4096)
             chunk_size = min(chunk_size - 500, int(chunk_size * 0.9))
 

scrapegraphai/nodes/rag_node.py

@@ -80,10 +80,8 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-        # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
 
         user_prompt = input_data[0]
@@ -102,7 +100,6 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info("--- (updated chunks metadata) ---")
 
-        # check if embedder_model is provided, if not use llm_model
         if self.embedder_model is not None:
             embeddings = self.embedder_model
         elif 'embeddings' in self.node_config:
@@ -144,23 +141,17 @@ def execute(self, state: dict) -> dict:
         pipeline_compressor = DocumentCompressorPipeline(
             transformers=[redundant_filter, relevant_filter]
         )
-        # redundant + relevant filter compressor
         compression_retriever = ContextualCompressionRetriever(
             base_compressor=pipeline_compressor, base_retriever=retriever
         )
 
-        # relevant filter compressor only
-        # compression_retriever = ContextualCompressionRetriever(
-        #     base_compressor=relevant_filter, base_retriever=retriever
-        # )
-
         compressed_docs = compression_retriever.invoke(user_prompt)
 
         self.logger.info("--- (tokens compressed and vector stored) ---")
 
         state.update({self.output[0]: compressed_docs})
         return state
-    
+
 
     def _create_default_embedder(self, llm_config=None) -> object:
         """
@@ -223,7 +214,6 @@ def _create_embedder(self, embedder_config: dict) -> object:
         embedder_params = {**embedder_config}
         if "model_instance" in embedder_config:
             return embedder_params["model_instance"]
-        # Instantiate the embedding model based on the model name
         if "openai" in embedder_params["model"]:
             return OpenAIEmbeddings(api_key=embedder_params["api_key"])
         if "azure" in embedder_params["model"]:

scrapegraphai/nodes/robots_node.py

@@ -75,10 +75,8 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-        # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
 
         source = input_data[0]

scrapegraphai/nodes/search_internet_node.py

@@ -67,7 +67,6 @@ def execute(self, state: dict) -> dict:
 
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
 
         user_prompt = input_data[0]
@@ -79,10 +78,8 @@ def execute(self, state: dict) -> dict:
             input_variables=["user_prompt"],
         )
 
-        # Execute the chain to get the search query
         search_answer = search_prompt | self.llm_model | output_parser
-        
-        # Ollama: Use no json format when creating the search query
+
         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]
@@ -96,9 +93,7 @@ def execute(self, state: dict) -> dict:
                                search_engine=self.search_engine)
 
         if len(answer) == 0:
-            # raise an exception if no answer is found
             raise ValueError("Zero results found for the search query.")
 
-        # Update the state with the generated answer
         state.update({self.output[0]: answer})
         return state

scrapegraphai/nodes/search_link_node.py

@@ -49,7 +49,6 @@ def __init__(
             self.filter_config = {**default_filters.filter_dict, **provided_filter_config}
             self.filter_links = True
         else:
-            # Skip filtering if not enabled
             self.filter_config = None
             self.filter_links = False
 
@@ -58,42 +57,38 @@ def __init__(
 
     def _is_same_domain(self, url, domain):
         if not self.filter_links or not self.filter_config.get("diff_domain_filter", True):
-            return True  # Skip the domain filter if not enabled
+            return True
         parsed_url = urlparse(url)
         parsed_domain = urlparse(domain)
         return parsed_url.netloc == parsed_domain.netloc
 
     def _is_image_url(self, url):
         if not self.filter_links:
-            return False  # Skip image filtering if filtering is not enabled
-        
+            return False
         image_extensions = self.filter_config.get("img_exts", [])
         return any(url.lower().endswith(ext) for ext in image_extensions)
 
     def _is_language_url(self, url):
         if not self.filter_links:
-            return False  # Skip language filtering if filtering is not enabled
+            return False
 
         lang_indicators = self.filter_config.get("lang_indicators", [])
         parsed_url = urlparse(url)
         query_params = parse_qs(parsed_url.query)
 
-        # Check if the URL path or query string indicates a language-specific version
         return any(indicator in parsed_url.path.lower() or indicator in query_params for indicator in lang_indicators)
-
     def _is_potentially_irrelevant(self, url):
         if not self.filter_links:
             return False  # Skip irrelevant URL filtering if filtering is not enabled
 
         irrelevant_keywords = self.filter_config.get("irrelevant_keywords", [])
         return any(keyword in url.lower() for keyword in irrelevant_keywords)
 
-    
+
     def execute(self, state: dict) -> dict:
         """
-        Filter out relevant links from the webpage that are relavant to prompt. Out of the filtered links, also
-        ensure that all links are navigable.
-
+        Filter out relevant links from the webpage that are relavant to prompt. 
+        Out of the filtered links, also ensure that all links are navigable.
         Args:
             state (dict): The current state of the graph. The input keys will be used to fetch the
                             correct data types from the state.
@@ -108,7 +103,6 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-
         parsed_content_chunks = state.get("doc")
         source_url = state.get("url") or state.get("local_dir")
         output_parser = JsonOutputParser()
@@ -148,7 +142,7 @@ def execute(self, state: dict) -> dict:
             except Exception as e:
                 # Fallback approach: Using the LLM to extract links
                 self.logger.error(f"Error extracting links: {e}. Falling back to LLM.")
-                
+
                 merge_prompt = PromptTemplate(
                     template=TEMPLATE_RELEVANT_LINKS,
                     input_variables=["content", "user_prompt"],

scrapegraphai/nodes/search_node_with_context.py

@@ -58,10 +58,8 @@ def execute(self, state: dict) -> dict:
 
         self.logger.info(f"--- Executing {self.node_name} Node ---")
 
-        # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
 
-        # Fetching data from the state based on the input keys
         input_data = [state[key] for key in input_keys]
 
         doc = input_data[1]
@@ -71,7 +69,6 @@ def execute(self, state: dict) -> dict:
 
         result = []
 
-        # Use tqdm to add progress bar
         for i, chunk in enumerate(
             tqdm(doc, desc="Processing chunks", disable=not self.verbose)
         ):