Merge pull request #69 from VinciGit00/refactor-script-generator

PeriniM · web-flow · commit 4c0118c00890 · 2024-04-19T09:34:56.000+02:00
Refactor script generator
diff --git a/.gitignore b/.gitignore
@@ -31,3 +31,7 @@ examples/graph_examples/ScrapeGraphAI_generated_graph
 examples/**/*.csv
 examples/**/*.json
 main.py
+poetry.lock
+
+# lock files
+*.lock
diff --git a/examples/gemini/script_generator_gemini.py b/examples/gemini/script_generator_gemini.py
@@ -21,6 +21,7 @@
         "api_key": gemini_key,
         "model": "gpt-3.5-turbo",
     },
+    "library": "beautifoulsoup"
 }
 
 # ************************************************
diff --git a/examples/local_models/Docker/script_generator_docker.py b/examples/local_models/Docker/script_generator_docker.py
@@ -18,7 +18,8 @@
     "embeddings": {
         "model": "ollama/nomic-embed-text",
         "temperature": 0,
-    }
+    },
+    "library": "beautifoulsoup"
 }
 
 # ************************************************
diff --git a/examples/local_models/Ollama/script_generator_ollama.py b/examples/local_models/Ollama/script_generator_ollama.py
@@ -1,4 +1,4 @@
-""" 
+"""
 Basic example of scraping pipeline using ScriptCreatorGraph
 """
 from scrapegraphai.graphs import ScriptCreatorGraph
@@ -11,15 +11,15 @@
     "llm": {
         "model": "ollama/mistral",
         "temperature": 0,
-        "format": "json",  # Ollama needs the format to be specified explicitly
         # "model_tokens": 2000, # set context length arbitrarily,
         "base_url": "http://localhost:11434",  # set ollama URL arbitrarily
     },
     "embeddings": {
         "model": "ollama/nomic-embed-text",
         "temperature": 0,
         "base_url": "http://localhost:11434",  # set ollama URL arbitrarily
-    }
+    },
+    "library": "beautifoulsoup"
 }
 
 # ************************************************
diff --git a/examples/local_models/Ollama/smart_scraper_ollama.py b/examples/local_models/Ollama/smart_scraper_ollama.py
@@ -10,7 +10,7 @@
 graph_config = {
     "llm": {
         "model": "ollama/mistral",
-        "temperature": 0,
+        "temperature": 1,
         "format": "json",  # Ollama needs the format to be specified explicitly
         # "model_tokens": 2000, # set context length arbitrarily,
         "base_url": "http://localhost:11434",  # set ollama URL arbitrarily
diff --git a/examples/openai/script_generator_openai.py b/examples/openai/script_generator_openai.py
@@ -20,6 +20,7 @@
         "api_key": openai_key,
         "model": "gpt-3.5-turbo",
     },
+    "library": "beautifoulsoup"
 }
 
 # ************************************************
diff --git a/manual deployment/commit_and_push.sh b/manual deployment/commit_and_push.sh
@@ -21,7 +21,7 @@ cd ..
 commit_message="$1"
 
 # Run Pylint on the specified Python files
-pylint pylint scrapegraphai/**/*.py scrapegraphai/*.py 
+pylint pylint scrapegraphai/**/*.py scrapegraphai/*.py tests/*.py
 #Make the pull
 git pull
 
diff --git a/poetry.lock b/poetry.lock
diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py
@@ -6,6 +6,7 @@
 from ..models import OpenAI, Gemini, Ollama, AzureOpenAI, HuggingFace
 from ..helpers import models_tokens
 
+
 class AbstractGraph(ABC):
     """
     Abstract class representing a generic graph-based tool.
@@ -22,7 +23,6 @@ def __init__(self, prompt: str, config: dict, source: Optional[str] = None):
         self.embedder_model = None if "embeddings" not in config else self._create_llm(
             config["embeddings"])
         self.graph = self._create_graph()
-        
         self.final_state = None
         self.execution_info = None
 
@@ -88,7 +88,7 @@ def get_execution_info(self):
         Returns the execution information of the graph.
         """
         return self.execution_info
-    
+
     @abstractmethod
     def _create_graph(self):
         """
diff --git a/scrapegraphai/graphs/script_creator_graph.py b/scrapegraphai/graphs/script_creator_graph.py
@@ -21,6 +21,8 @@ def __init__(self, prompt: str, source: str, config: dict):
         """
         Initializes the ScriptCreatorGraph with a prompt, source, and configuration.
         """
+        self.library = config['library']
+
         super().__init__(prompt, config, source)
 
         self.input_key = "url" if source.startswith("http") else "local_dir"
@@ -50,6 +52,8 @@ def _create_graph(self):
             input="user_prompt & (relevant_chunks | parsed_doc | doc)",
             output=["answer"],
             node_config={"llm": self.llm_model},
+            library=self.library,
+            website=self.source
         )
 
         return BaseGraph(
diff --git a/scrapegraphai/nodes/fetch_node.py b/scrapegraphai/nodes/fetch_node.py
@@ -72,15 +72,16 @@ def execute(self, state):
 
         # if it is a local directory
         if not source.startswith("http"):
-            compressedDocument = [Document(page_content=remover(source), metadata={
+            compressed_document = [Document(page_content=remover(source), metadata={
                 "source": "local_dir"
             })]
 
         # if it is a URL
         else:
             loader = AsyncHtmlLoader(source)
             document = loader.load()
-            compressedDocument = [Document(page_content=remover(str(document)))]
+            compressed_document = [
+                Document(page_content=remover(str(document)))]
 
-        state.update({self.output[0]: compressedDocument})
+        state.update({self.output[0]: compressed_document})
         return state
diff --git a/scrapegraphai/nodes/generate_scraper_node.py b/scrapegraphai/nodes/generate_scraper_node.py
@@ -7,7 +7,7 @@
 
 # Imports from Langchain
 from langchain.prompts import PromptTemplate
-from langchain_core.output_parsers import JsonOutputParser
+from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableParallel
 
 # Imports from the library
@@ -40,7 +40,7 @@ class GenerateScraperNode(BaseNode):
     """
 
     def __init__(self, input: str, output: List[str], node_config: dict,
-                 node_name: str = "GenerateAnswer"):
+                 library: str, website: str, node_name: str = "GenerateAnswer"):
         """
         Initializes the GenerateScraperNode with a language model client and a node name.
         Args:
@@ -49,6 +49,8 @@ def __init__(self, input: str, output: List[str], node_config: dict,
         """
         super().__init__(node_name, "node", input, output, 2, node_config)
         self.llm_model = node_config["llm"]
+        self.library = library
+        self.source = website
 
     def execute(self, state):
         """
@@ -80,29 +82,36 @@ def execute(self, state):
         user_prompt = input_data[0]
         doc = input_data[1]
 
-        output_parser = JsonOutputParser()
-        format_instructions = output_parser.get_format_instructions()
+        output_parser = StrOutputParser()
 
         template_chunks = """
         PROMPT:
         You are a website scraper script creator and you have just scraped the
         following content from a website.
-        Write the code in python with the Beautiful Soup library to extract the informations requested by the task.\n  \n
+        Write the code in python for extracting the informations requested by the task.\n 
+        The python library to use is specified in the instructions \n
         The website is big so I am giving you one chunk at the time to be merged later with the other chunks.\n
         CONTENT OF {chunk_id}: {context}. 
         Ignore all the context sentences that ask you not to extract information from the html code
-        INSTRUCTIONS: {format_instructions}
+        The output should be just pyton code without any comment and should implement the main, the HTML code
+        should do a get to the website and use the library request for making the GET. 
+        LIBRARY: {library}.
+        SOURCE: {source}
+        The output should be just pyton code without any comment and should implement the main.
         QUESTION: {question}
         """
         template_no_chunks = """
         PROMPT:
         You are a website scraper script creator and you have just scraped the
         following content from a website.
-        Write the code in python with the Beautiful Soup library to extract the informations requested by the task.\n  \n
+        Write the code in python for extracting the informations requested by the task.\n 
+        The python library to use is specified in the instructions \n
         The website is big so I am giving you one chunk at the time to be merged later with the other chunks.\n
-        CONTENT OF {chunk_id}: {context}. 
         Ignore all the context sentences that ask you not to extract information from the html code
-        INSTRUCTIONS: {format_instructions}
+        The output should be just pyton code without any comment and should implement the main, the HTML code
+        should do a get to the website and use the library request for making the GET. 
+        LIBRARY: {library}
+        SOURCE: {source}
         QUESTION: {question}
         """
 
@@ -130,8 +139,10 @@ def execute(self, state):
                 template=template,
                 input_variables=["question"],
                 partial_variables={"context": chunk.page_content,
-                                    "chunk_id": i + 1,
-                                    "format_instructions": format_instructions},
+                                   "chunk_id": i + 1,
+                                   "library": self.library,
+                                   "source": self.source
+                                   },
             )
             # Dynamically name the chains based on their index
             chain_name = f"chunk{i+1}"
@@ -148,7 +159,6 @@ def execute(self, state):
             merge_prompt = PromptTemplate(
                 template=template_merge,
                 input_variables=["context", "question"],
-                partial_variables={"format_instructions": format_instructions},
             )
             merge_chain = merge_prompt | self.llm_model | output_parser
             answer = merge_chain.invoke(
diff --git a/scrapegraphai/utils/remover.py b/scrapegraphai/utils/remover.py
@@ -31,17 +31,8 @@ def remover(html_content: str) -> str:
     # Body Extraction (if it exists)
     body_content = soup.find('body')
     if body_content:
-        # Remove some attributes from tags
-        """ tagsToRemove = ['style', 'rel', 'width',
-                        'height', 'target', 'media',
-                        'onerror', 'onload', 'onclick']
-        for tag in body_content.find_all():
-            for attr in tagsToRemove:
-                if tag.has_attr(attr):
-                    del tag.attrs[attr] """
-
         # Minify the HTML within the body tag
         minimized_body = minify(str(body_content))
         return "Title: " + title + ", Body: " + minimized_body
-    else:
-        return "Title: " + title + ", Body: No body content found"
+
+    return "Title: " + title + ", Body: No body content found"
diff --git a/tests/script_generator_test.py b/tests/script_generator_test.py
@@ -1,3 +1,6 @@
+""" 
+Module for making the tests for ScriptGeneratorGraph
+"""
 import pytest
 from scrapegraphai.graphs import ScriptCreatorGraph
 from scrapegraphai.utils import prettify_exec_info
@@ -11,6 +14,7 @@ def graph_config():
             "temperature": 0,
             "format": "json",
             "base_url": "http://localhost:11434",
+            "library": "beautifoulsoup",
         },
         "embeddings": {
             "model": "ollama/nomic-embed-text",

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@`
`21`	`21`	`"api_key": gemini_key,`
`22`	`22`	`"model": "gpt-3.5-turbo",`
`23`	`23`	`},`
	`24`	`+ "library": "beautifoulsoup"`
`24`	`25`	`}`
`25`	`26`
`26`	`27`	`# ************************************************`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,8 @@`
`18`	`18`	`"embeddings": {`
`19`	`19`	`"model": "ollama/nomic-embed-text",`
`20`	`20`	`"temperature": 0,`
`21`		`- }`
	`21`	`+ },`
	`22`	`+ "library": "beautifoulsoup"`
`22`	`23`	`}`
`23`	`24`
`24`	`25`	`# ************************************************`
Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@`
`20`	`20`	`"api_key": openai_key,`
`21`	`21`	`"model": "gpt-3.5-turbo",`
`22`	`22`	`},`
	`23`	`+ "library": "beautifoulsoup"`
`23`	`24`	`}`
`24`	`25`
`25`	`26`	`# ************************************************`