refactoring of loggers

VinciGit00 · VinciGit00 · commit 05890835f522 · 2024-05-15T10:54:53.000+02:00
diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py
@@ -10,7 +10,7 @@
 from ..helpers import models_tokens
 from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Ollama, OpenAI, Anthropic
 from langchain_google_genai.embeddings import GoogleGenerativeAIEmbeddings
-
+from ..utils.logging import get_logger
 
 class AbstractGraph(ABC):
     """
@@ -61,6 +61,7 @@ def __init__(self, prompt: str, config: dict, source: Optional[str] = None):
         self.headless = True if config is None else config.get(
             "headless", True)
         self.loader_kwargs = config.get("loader_kwargs", {})
+        self.logger = get_logger("graph")
 
         common_params = {"headless": self.headless,
                          "verbose": self.verbose,
@@ -79,7 +80,7 @@ def set_common_params(self, params: dict, overwrite=False):
 
         for node in self.graph.nodes:
             node.update_config(params, overwrite)
-
+        
     def _set_model_token(self, llm):
 
         if 'Azure' in str(type(llm)):
diff --git a/scrapegraphai/nodes/fetch_node.py b/scrapegraphai/nodes/fetch_node.py
@@ -74,8 +74,7 @@ def execute(self, state):
                     necessary information to perform the operation is missing.
         """
         if self.verbose:
-            logger = get_logger("fetch node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
@@ -129,7 +128,7 @@ def execute(self, state):
                 cleanedup_html = cleanup_html(response.text, source)
                 compressed_document = [Document(page_content=cleanedup_html)]
             else:	
-                logger.warning(f"Failed to retrieve contents from the webpage at url: {source}")
+                self.logger.warning(f"Failed to retrieve contents from the webpage at url: {source}")
 
         else:
             loader_kwargs = {}
diff --git a/scrapegraphai/nodes/generate_answer_csv_node.py b/scrapegraphai/nodes/generate_answer_csv_node.py
@@ -1,4 +1,5 @@
 """
+gg
 Module for generating the answer node
 """
 # Imports from standard library
@@ -73,8 +74,7 @@ def execute(self, state):
         """
 
         if self.verbose:
-            logger = get_logger("generate_answer csv node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py
@@ -60,8 +60,7 @@ def execute(self, state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("generate answer node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
diff --git a/scrapegraphai/nodes/generate_answer_pdf_node.py b/scrapegraphai/nodes/generate_answer_pdf_node.py
@@ -73,8 +73,7 @@ def execute(self, state):
         """
 
         if self.verbose:
-            logger = get_logger("generate answer pdf node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
diff --git a/scrapegraphai/nodes/generate_scraper_node.py b/scrapegraphai/nodes/generate_scraper_node.py
@@ -64,8 +64,7 @@ def execute(self, state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("generate scraper node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
diff --git a/scrapegraphai/nodes/get_probable_tags_node.py b/scrapegraphai/nodes/get_probable_tags_node.py
@@ -51,8 +51,7 @@ def execute(self, state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("get probable tags node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
diff --git a/scrapegraphai/nodes/graph_iterator_node.py b/scrapegraphai/nodes/graph_iterator_node.py
@@ -60,7 +60,8 @@ def execute(self, state: dict) -> dict:
         batchsize = self.node_config.get("batchsize", _default_batchsize)
 
         if self.verbose:
-            print(f"--- Executing {self.node_name} Node with batchsize {batchsize} ---")
+            self.logger.info(f"--- Executing {self.node_name} Node with batchsize {batchsize} ---")
+
 
         try:
             eventloop = asyncio.get_event_loop()
diff --git a/scrapegraphai/nodes/image_to_text_node.py b/scrapegraphai/nodes/image_to_text_node.py
@@ -43,8 +43,7 @@ def execute(self, state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("image to text node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            self.logger.info(f"--- Executing {self.node_name} Node ---")
             
         input_keys = self.get_input_keys(state)
         input_data = [state[key] for key in input_keys]
diff --git a/scrapegraphai/nodes/merge_answers_node.py b/scrapegraphai/nodes/merge_answers_node.py
@@ -54,8 +54,7 @@ def execute(self, state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("fetch node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            self.ogger.info(f"--- Executing {self.node_name} Node ---")
 
         # Interpret input keys based on the provided input expression
         input_keys = self.get_input_keys(state)
diff --git a/scrapegraphai/nodes/parse_node.py b/scrapegraphai/nodes/parse_node.py
@@ -49,8 +49,7 @@ def execute(self,  state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("fetch node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
diff --git a/scrapegraphai/nodes/rag_node.py b/scrapegraphai/nodes/rag_node.py
@@ -56,10 +56,9 @@ def execute(self, state: dict) -> dict:
             KeyError: If the input keys are not found in the state, indicating that the
                         necessary information for compressing the content is missing.
         """
-        logger = get_logger("rag node")
 
         if self.verbose:
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
@@ -82,7 +81,7 @@ def execute(self, state: dict) -> dict:
             chunked_docs.append(doc)
 
         if self.verbose:
-            logger.info("--- (updated chunks metadata) ---")
+           self.logger.info("--- (updated chunks metadata) ---")
 
         # check if embedder_model is provided, if not use llm_model
         self.embedder_model = self.embedder_model if self.embedder_model else self.llm_model
@@ -110,7 +109,7 @@ def execute(self, state: dict) -> dict:
         compressed_docs = compression_retriever.invoke(user_prompt)
 
         if self.verbose:
-            print("--- (tokens compressed and vector stored) ---")
+            self.logger.info("--- (tokens compressed and vector stored) ---")
 
         state.update({self.output[0]: compressed_docs})
         return state
diff --git a/scrapegraphai/nodes/robots_node.py b/scrapegraphai/nodes/robots_node.py
@@ -122,17 +122,17 @@ def execute(self, state: dict) -> dict:
 
             if "no" in is_scrapable:
                 if self.verbose:
-                    logger.warning("\033[31m(Scraping this website is not allowed)\033[0m")
+                    self.logger.warning("\033[31m(Scraping this website is not allowed)\033[0m")
                     
                 if not self.force_scraping:
                     raise ValueError(
                         'The website you selected is not scrapable')
                 else:
                     if self.verbose:
-                        logger.warning("\033[33m(WARNING: Scraping this website is not allowed but you decided to force it)\033[0m")
+                        self.logger.warning("\033[33m(WARNING: Scraping this website is not allowed but you decided to force it)\033[0m")
             else:
                 if self.verbose:
-                    logger.warning("\033[32m(Scraping this website is allowed)\033[0m")
+                    self.logger.warning("\033[32m(Scraping this website is allowed)\033[0m")
 
         state.update({self.output[0]: is_scrapable})
         return state
diff --git a/scrapegraphai/nodes/search_internet_node.py b/scrapegraphai/nodes/search_internet_node.py
@@ -54,10 +54,9 @@ def execute(self, state: dict) -> dict:
             KeyError: If the input keys are not found in the state, indicating that the
                         necessary information for generating the answer is missing.
         """
-        logger = get_logger("search interne node")
 
         if self.verbose:
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            self.logger.info(f"--- Executing {self.node_name} Node ---")
 
         input_keys = self.get_input_keys(state)
 
@@ -89,7 +88,7 @@ def execute(self, state: dict) -> dict:
         search_query = search_answer.invoke({"user_prompt": user_prompt})[0]
 
         if self.verbose:
-            logger.info(f"Search Query: {search_query}")
+            self.logger.info(f"Search Query: {search_query}")
 
 
         answer = search_on_web(
diff --git a/scrapegraphai/nodes/search_link_node.py b/scrapegraphai/nodes/search_link_node.py
@@ -59,8 +59,7 @@ def execute(self, state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("search link node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
diff --git a/scrapegraphai/nodes/text_to_speech_node.py b/scrapegraphai/nodes/text_to_speech_node.py
@@ -45,8 +45,7 @@ def execute(self, state: dict) -> dict:
         """
 
         if self.verbose:
-            logger = get_logger("text to speach node")
-            logger.info(f"--- Executing {self.node_name} Node ---")
+            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)
