Fix: generated scraper to include source as an argument

Author: Dirk Breeuwer

gpt-scraper.py

@@ -5,8 +5,6 @@
 from website_analysis.dom_analysis import HtmlLoader, UrlHtmlLoader
 from scraper_generation.scraper_generator import ScrapingCodeGenerator, CodeWriter
 from data_extraction.data_extractor import CodeExecutor 
-from user_requirements_gathering import UserRequirementGatherer
-
 
 
 def main():
@@ -20,7 +18,7 @@ def main():
     source_type = args.source_type
     USER_REQUIREMENTS = args.requirements
 
-    # create HtmlLoader or UrlHtmlLoader based on the source type
+    # Create HtmlLoader or UrlHtmlLoader based on the source type
     def create_html_loader(source, source_type):
         if source_type == 'url':
             return UrlHtmlLoader(source)
@@ -30,7 +28,7 @@ def create_html_loader(source, source_type):
     html_loader = create_html_loader(source, source_type)
 
     # Instantiate ScrapingCodeGenerator with the html_loader
-    code_generator = ScrapingCodeGenerator(html_loader)
+    code_generator = ScrapingCodeGenerator(html_loader, source=source, source_type=source_type)
 
     # Generate scraping code
     scraping_code = code_generator.generate_scraping_code(USER_REQUIREMENTS)
@@ -51,8 +49,4 @@ def create_html_loader(source, source_type):
     main()
 
 
-
-
-
-
 # python3 gpt-scraper.py --source-type "file" --source "./results/denver.html" --requirements "Extract the average monthly temperature in denver"

scraper_generation/scraper_generator.py

@@ -4,7 +4,23 @@
 
 class ScrapingCodeGenerator:
     MODEL_NAME = "text-davinci-003"
-    TEMPLATE = """
+    SCRAPING_CODE = f"""
+from bs4 import BeautifulSoup
+from website_analysis.dom_analysis import HtmlLoader, UrlHtmlLoader
+
+# Create HtmlLoader or UrlHtmlLoader based on the source type
+def create_html_loader(source, source_type):
+    if source_type == 'url':
+        return UrlHtmlLoader(source)
+    else:  # source_type == 'file'
+        return HtmlLoader(source)
+
+html_loader = create_html_loader("{{source}}", "{{source_type}}")
+response = html_loader.load()
+
+html_soup = BeautifulSoup(response, 'html.parser')
+    """
+    PROMPT_TEMPLATE = """
 You are an expert website analyzer for a web scraping process.
 Take the user requirements and convert it into clean python code to scrape the website.
 
@@ -25,27 +41,20 @@ class ScrapingCodeGenerator:
 html_soup = BeautifulSoup(response, 'html.parser')
             """
 
-    static_code = """
-from bs4 import BeautifulSoup
 
-# Get the URL of the website
-with open('./results/denver.html') as f:
-    response = f.read()
-
-html_soup = BeautifulSoup(response, 'html.parser')
-    """
 
-    def __init__(self, html_loader):
+    def __init__(self, html_loader, source, source_type):
         self.html_loader = html_loader
         self.llm = self.initialize_llm()
         self.prompt_template = self.initialize_template()
+        self.scraping_code = self.SCRAPING_CODE.format(source=source, source_type=source_type)
 
     def initialize_llm(self):
         load_dotenv()
         return OpenAI(model_name=self.MODEL_NAME, temperature=0)
 
     def initialize_template(self):
-        return PromptTemplate(input_variables=["requirements","html"], template=self.TEMPLATE)
+        return PromptTemplate(input_variables=["requirements","html"], template=self.PROMPT_TEMPLATE)
 
 
     def generate_scraping_code(self, user_requirements):
@@ -57,7 +66,7 @@ def generate_scraping_code(self, user_requirements):
         generated_code = self.llm(formatted_prompt)
 
         full_scraping_code = f"""
-{self.static_code}
+{self.scraping_code}
 {generated_code}
         """
         return full_scraping_code