Refactor prompt templates from gpt scraper to ScraperGenerator

Author: Dirk Breeuwer

gpt-scraper.py

@@ -3,8 +3,9 @@
 from langchain.llms import OpenAI
 import argparse
 from website_analysis.dom_analysis import HtmlLoader, UrlHtmlLoader
-from scraper_generation.scraper_generator import ScrapingCodeGenerator, CodeWriter 
+from scraper_generation.scraper_generator import ScrapingCodeGenerator, CodeWriter
 from data_extraction.data_extractor import CodeExecutor 
+from user_requirements_gathering import UserRequirementGatherer
 
 
 
@@ -19,48 +20,17 @@ def main():
     source_type = args.source_type
     USER_REQUIREMENTS = args.requirements
 
-
-
-    # Variables to define
-    MODEL_NAME = "text-davinci-003"
-    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.
-
-        USER REQUIREMENTS:
-        {requirements}
-
-        HTML CODE YOU NEED TO SCRAPE:
-        {html}
-
-        FINISH THE PYTHON CODE TO SCRAPE THE WEBSITE:
-
-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')
-        """
-
-    # Load environment variables
-    load_dotenv()
-
-    # Create the LLM
-    llm = OpenAI(model_name=MODEL_NAME, temperature=0)
-
-    # Create prompt template
-    prompt_template = PromptTemplate(input_variables=["requirements","html"], template=TEMPLATE)
-
     # create HtmlLoader or UrlHtmlLoader based on the source type
-    if source_type == 'url':
-        html_loader = UrlHtmlLoader(source)
-    else:  # source_type == 'file'
-        html_loader = HtmlLoader(source)
+    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)
 
     # Instantiate ScrapingCodeGenerator with the html_loader
-    code_generator = ScrapingCodeGenerator(html_loader, llm, prompt_template)
+    code_generator = ScrapingCodeGenerator(html_loader)
 
     # Generate scraping code
     scraping_code = code_generator.generate_scraping_code(USER_REQUIREMENTS)

scraper_generation/scraper_generator.py

@@ -3,6 +3,28 @@
 from langchain import PromptTemplate
 
 class ScrapingCodeGenerator:
+    MODEL_NAME = "text-davinci-003"
+    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.
+
+USER REQUIREMENTS:
+{requirements}
+
+HTML CODE YOU NEED TO SCRAPE:
+{html}
+
+FINISH THE PYTHON CODE TO SCRAPE THE WEBSITE:
+
+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')
+            """
+    
     static_code = """
 from bs4 import BeautifulSoup
 
@@ -11,11 +33,20 @@ class ScrapingCodeGenerator:
     response = f.read()
 
 html_soup = BeautifulSoup(response, 'html.parser')
-        """
-    def __init__(self, html_loader, llm, prompt_template):
+    """
+
+    def __init__(self, html_loader):
         self.html_loader = html_loader
-        self.llm = llm
-        self.prompt_template = prompt_template
+        self.llm = self.initialize_llm()
+        self.prompt_template = self.initialize_template()
+
+    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)
+
 
     def generate_scraping_code(self, user_requirements):
         """
@@ -26,11 +57,11 @@ def generate_scraping_code(self, user_requirements):
         generated_code = self.llm(formatted_prompt)
 
         full_scraping_code = f"""
-        {self.static_code}
-        {generated_code}
+{self.static_code}
+{generated_code}
         """
         return full_scraping_code
-    
+
 class CodeWriter:
     def __init__(self, file_name):
         self.file_name = file_name
@@ -40,5 +71,4 @@ def write(self, scraping_code):
         Writes the scraping code to a .py python file
         """
         with open(self.file_name, 'w') as file:
-            file.write(scraping_code)
-    
+            file.write(scraping_code)