Update agent.py

agent.py

@@ -4,6 +4,7 @@ import hashlib
 import logging
 import datetime
 import csv
+import threading
 from urllib.parse import urlparse
 from selenium import webdriver
 from selenium.webdriver.chrome.service import Service
@@ -11,67 +12,130 @@ from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
-from selenium.common.exceptions import TimeoutException, NoSuchElementException, StaleElementReferenceException
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from transformers import pipeline
+from selenium.common.exceptions import (
+    TimeoutException,
+    NoSuchElementException,
+    StaleElementReferenceException,
+)
+from webdriver_manager.chrome import ChromeDriverManager  # Added import
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import feedparser
+import gradio as gr
+import xml.etree.ElementTree as ET
 
 # Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
 # Define constants
 DEFAULT_FILE_PATH = "scraped_data"
-PURPOSE = f"You go to Culvers sites, you continuously seek changes on them since your last observation. Anything new that gets logged and dumped into csv, stored in your log folder at user/app/scraped_data."
+PURPOSE = (
+    "You go to Culvers sites, you continuously seek changes on them since your last observation. "
+    "Anything new that gets logged and dumped into csv, stored in your log folder at user/app/scraped_data."
+)
 HISTORY = []
 CURRENT_TASK = None
+STOP_THREADS = False  # Flag to stop scraping threads
 
 # Function to monitor URLs for changes
 def monitor_urls(storage_location, urls, scrape_interval, content_type, selector=None):
-    global HISTORY
-    previous_hashes = {url: "" for url in urls} 
+    global HISTORY, STOP_THREADS
+    previous_hashes = {url: "" for url in urls}
+
+    options = Options()
+    options.add_argument("--headless")  # Run Chrome in headless mode
+    options.add_argument("--no-sandbox")
+    options.add_argument("--disable-dev-shm-usage")
 
     try:
-        with webdriver.Chrome(service=Service(webdriver.ChromeDriverManager().install()), options=Options()) as driver:
-            while True:
-                for url in urls:
-                    try:
-                        driver.get(url)
-                        WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.TAG_NAME, 'body')))  # Wait for basic page load
-                        time.sleep(2)  # Additional wait for dynamic content
-
-                        if content_type == "text":
-                            current_content = driver.page_source
-                        elif content_type == "media":
-                            if selector:
-                                try:
-                                    elements = WebDriverWait(driver, 5).until(EC.presence_of_all_elements_located((By.CSS_SELECTOR, selector)))
-                                    current_content = [element.get_attribute('src') for element in elements]
-                                except TimeoutException:
-                                    logging.warning(f"Timeout waiting for media elements with selector '{selector}' on {url}")
-                                    current_content = []
-                            else:
-                                current_content = driver.find_elements(By.TAG_NAME, "img")
-                        else:
-                            current_content = driver.page_source
-
-                        current_hash = hashlib.md5(str(current_content).encode('utf-8')).hexdigest()
-                        if current_hash != previous_hashes[url]:
-                            previous_hashes[url] = current_hash
-                            date_time_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                            HISTORY.append(f"Change detected at {url} on {date_time_str}")
-                            with open(os.path.join(storage_location, f"{urlparse(url).hostname}_changes.csv"), "a", newline="") as csvfile:
-                                csv_writer = csv.DictWriter(csvfile, fieldnames=["date", "time", "url", "change"])
-                                csv_writer.writerow({"date": date_time_str.split()[0], "time": date_time_str.split()[1], "url": url, "change": "Content changed"})
-                                logging.info(f"Change detected at {url} on {date_time_str}")
-                    except (NoSuchElementException, StaleElementReferenceException, Exception) as e:
-                        logging.error(f"Error accessing {url}: {e}")
-                time.sleep(scrape_interval * 60)  # Check every scrape_interval minutes
+        driver = webdriver.Chrome(
+            service=Service(ChromeDriverManager().install()), options=options
+        )
     except Exception as e:
-        logging.error(f"Error starting ChromeDriver: {e}")
+        logging.error(f"Error initializing ChromeDriver: {e}")
+        return
+
+    try:
+        while not STOP_THREADS:
+            for url in urls:
+                try:
+                    driver.get(url)
+                    WebDriverWait(driver, 10).until(
+                        EC.presence_of_element_located((By.TAG_NAME, "body"))
+                    )  # Wait for basic page load
+                    time.sleep(2)  # Additional wait for dynamic content
+
+                    if content_type == "text":
+                        current_content = driver.page_source
+                    elif content_type == "media":
+                        if selector:
+                            try:
+                                elements = WebDriverWait(driver, 5).until(
+                                    EC.presence_of_all_elements_located(
+                                        (By.CSS_SELECTOR, selector)
+                                    )
+                                )
+                                current_content = [
+                                    element.get_attribute("src") for element in elements
+                                ]
+                            except TimeoutException:
+                                logging.warning(
+                                    f"Timeout waiting for media elements with selector '{selector}' on {url}"
+                                )
+                                current_content = []
+                        else:
+                            elements = driver.find_elements(By.TAG_NAME, "img")
+                            current_content = [element.get_attribute("src") for element in elements]
+                    else:
+                        current_content = driver.page_source
+
+                    current_hash = hashlib.md5(
+                        str(current_content).encode("utf-8")
+                    ).hexdigest()
+                    if current_hash != previous_hashes[url]:
+                        previous_hashes[url] = current_hash
+                        date_time_str = datetime.datetime.now().strftime(
+                            "%Y-%m-%d %H:%M:%S"
+                        )
+                        HISTORY.append(f"Change detected at {url} on {date_time_str}")
+                        csv_file_path = os.path.join(
+                            storage_location, f"{urlparse(url).hostname}_changes.csv"
+                        )
+                        os.makedirs(storage_location, exist_ok=True)
+                        file_exists = os.path.isfile(csv_file_path)
+                        with open(csv_file_path, "a", newline="", encoding="utf-8") as csvfile:
+                            fieldnames = ["date", "time", "url", "change"]
+                            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+                            if not file_exists:
+                                writer.writeheader()
+                            writer.writerow(
+                                {
+                                    "date": date_time_str.split()[0],
+                                    "time": date_time_str.split()[1],
+                                    "url": url,
+                                    "change": "Content changed",
+                                }
+                            )
+                        logging.info(f"Change detected at {url} on {date_time_str}")
+                except (
+                    NoSuchElementException,
+                    StaleElementReferenceException,
+                    TimeoutException,
+                    Exception,
+                ) as e:
+                    logging.error(f"Error accessing {url}: {e}")
+            time.sleep(scrape_interval * 60)  # Check every scrape_interval minutes
+    finally:
+        driver.quit()
+        logging.info("ChromeDriver session ended.")
 
 # Function to start scraping
 def start_scraping(storage_location, urls, scrape_interval, content_type, selector=None):
-    global CURRENT_TASK, HISTORY
+    global CURRENT_TASK, HISTORY, STOP_THREADS
+
+    if STOP_THREADS:
+        STOP_THREADS = False  # Reset the flag if previously stopped
 
     CURRENT_TASK = f"Monitoring URLs: {', '.join(urls)}"
     HISTORY.append(f"Task started: {CURRENT_TASK}")
@@ -84,123 +148,316 @@ def start_scraping(storage_location, urls, scrape_interval, content_type, select
 
         # Log the initial observation
         try:
-            with webdriver.Chrome(service=Service(webdriver.ChromeDriverManager().install()), options=Options()) as driver:
-                driver.get(url)
-                WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.TAG_NAME, 'body')))  # Wait for basic page load
-                time.sleep(2)  # Additional wait for dynamic content
-
-                if content_type == "text":
-                    initial_content = driver.page_source
-                elif content_type == "media":
-                    if selector:
-                        try:
-                            elements = WebDriverWait(driver, 5).until(EC.presence_of_all_elements_located((By.CSS_SELECTOR, selector)))
-                            initial_content = [element.get_attribute('src') for element in elements]
-                        except TimeoutException:
-                            logging.warning(f"Timeout waiting for media elements with selector '{selector}' on {url}")
-                            initial_content = []
-                    else:
-                        initial_content = driver.find_elements(By.TAG_NAME, "img")
+            options = Options()
+            options.add_argument("--headless")  # Run Chrome in headless mode
+            options.add_argument("--no-sandbox")
+            options.add_argument("--disable-dev-shm-usage")
+
+            driver = webdriver.Chrome(
+                service=Service(ChromeDriverManager().install()), options=options
+            )
+            driver.get(url)
+            WebDriverWait(driver, 10).until(
+                EC.presence_of_element_located((By.TAG_NAME, "body"))
+            )  # Wait for basic page load
+            time.sleep(2)  # Additional wait for dynamic content
+
+            if content_type == "text":
+                initial_content = driver.page_source
+            elif content_type == "media":
+                if selector:
+                    try:
+                        elements = WebDriverWait(driver, 5).until(
+                            EC.presence_of_all_elements_located(
+                                (By.CSS_SELECTOR, selector)
+                            )
+                        )
+                        initial_content = [
+                            element.get_attribute("src") for element in elements
+                        ]
+                    except TimeoutException:
+                        logging.warning(
+                            f"Timeout waiting for media elements with selector '{selector}' on {url}"
+                        )
+                        initial_content = []
                 else:
-                    initial_content = driver.page_source
+                    elements = driver.find_elements(By.TAG_NAME, "img")
+                    initial_content = [element.get_attribute("src") for element in elements]
+            else:
+                initial_content = driver.page_source
 
-                initial_hash = hashlib.md5(str(initial_content).encode('utf-8')).hexdigest()
-                HISTORY.append(f"Initial observation at {url}: {initial_hash}")
-                with open(os.path.join(folder_path, f"{hostname}_initial_observation.txt"), "w") as file:
-                    file.write(f"Initial observation at {url}: {initial_hash}")
-        except (NoSuchElementException, StaleElementReferenceException, Exception) as e:
+            initial_hash = hashlib.md5(
+                str(initial_content).encode("utf-8")
+            ).hexdigest()
+            HISTORY.append(f"Initial observation at {url}: {initial_hash}")
+            initial_observation_path = os.path.join(
+                folder_path, f"{hostname}_initial_observation.txt"
+            )
+            with open(initial_observation_path, "w", encoding="utf-8") as file:
+                file.write(f"Initial observation at {url}: {initial_hash}")
+            logging.info(f"Initial observation logged for {url}")
+        except (
+            NoSuchElementException,
+            StaleElementReferenceException,
+            TimeoutException,
+            Exception,
+        ) as e:
             HISTORY.append(f"Error accessing {url}: {e}")
+            logging.error(f"Error accessing {url}: {e}")
+        finally:
+            driver.quit()
 
-    # Monitor the URLs
-    monitor_urls(storage_location, urls, scrape_interval, content_type, selector)
+    # Start a new thread for monitoring URLs
+    monitor_thread = threading.Thread(
+        target=monitor_urls,
+        args=(storage_location, urls, scrape_interval, content_type, selector),
+        daemon=True,
+    )
+    monitor_thread.start()
+    logging.info("Started scraping thread.")
 
     return f"Started scraping {', '.join(urls)} every {scrape_interval} minutes."
 
+# Function to stop scraping
+def stop_scraping():
+    global STOP_THREADS
+    STOP_THREADS = True
+    HISTORY.append("Scraping stopped by user.")
+    logging.info("Scraping stop signal sent.")
+    return "Scraping has been stopped."
+
 # Function to display CSV content
 def display_csv(storage_location, url):
     hostname = urlparse(url).hostname
-    folder_path = os.path.join(storage_location, hostname)
-    csv_path = os.path.join(folder_path, f"{hostname}_changes.csv")
+    csv_path = os.path.join(storage_location, f"{hostname}_changes.csv")
     if os.path.exists(csv_path):
-        with open(csv_path, "r") as file:
-            return file.read()
+        try:
+            with open(csv_path, "r", encoding="utf-8") as file:
+                content = file.read()
+            return content
+        except Exception as e:
+            logging.error(f"Error reading CSV file for {url}: {e}")
+            return f"Error reading CSV file for {url}: {e}"
     else:
         return "No data available."
 
 # Function to generate RSS feed for a given URL
 def generate_rss_feed(storage_location, url):
     hostname = urlparse(url).hostname
-    folder_path = os.path.join(storage_location, hostname)
-    csv_path = os.path.join(folder_path, f"{hostname}_changes.csv")
+    csv_path = os.path.join(storage_location, f"{hostname}_changes.csv")
     if os.path.exists(csv_path):
-        with open(csv_path, "r") as file:
-            reader = csv.DictReader(file)
-            feed = feedparser.parse(f"rss.xml")  # Create a new feed object
-            feed.feed.title = f"Changes for {hostname}"
-            feed.feed.link = url
-            feed.feed.description = "Recent changes detected on the website."
-            feed.entries = []
-            for row in reader:
-                feed.entries.append({
-                    "title": f"Change detected at {row['url']}",
-                    "link": row['url'],
-                    "description": f"Content changed on {row['date']} at {row['time']}",
-                    "published": datetime.datetime.strptime(f"{row['date']} {row['time']}", "%Y-%m-%d %H:%M:%S").isoformat(),
-                })
-            return feed.entries
+        try:
+            # Parse the CSV file
+            with open(csv_path, "r", encoding="utf-8") as file:
+                reader = csv.DictReader(file)
+                changes = list(reader)
+
+            # Create the root RSS element
+            rss = ET.Element("rss", version="2.0")
+            channel = ET.SubElement(rss, "channel")
+
+            # Add channel elements
+            title = ET.SubElement(channel, "title")
+            title.text = f"RSS Feed for {hostname}"
+
+            link = ET.SubElement(channel, "link")
+            link.text = url
+
+            description = ET.SubElement(channel, "description")
+            description.text = "Recent changes detected on the website."
+
+            # Add items to the feed
+            for change in changes[-10:]:  # Last 10 changes
+                item = ET.SubElement(channel, "item")
+
+                item_title = ET.SubElement(item, "title")
+                item_title.text = f"Change detected at {change['url']}"
+
+                item_link = ET.SubElement(item, "link")
+                item_link.text = change["url"]
+
+                item_description = ET.SubElement(item, "description")
+                item_description.text = f"Content changed on {change['date']} at {change['time']}"
+
+                pub_date = ET.SubElement(item, "pubDate")
+                pub_date.text = datetime.datetime.strptime(
+                    f"{change['date']} {change['time']}", "%Y-%m-%d %H:%M:%S"
+                ).strftime("%a, %d %b %Y %H:%M:%S +0000")
+
+            # Generate the XML string
+            rss_feed = ET.tostring(rss, encoding="utf-8")
+            return rss_feed.decode("utf-8")
+        except Exception as e:
+            logging.error(f"Error generating RSS feed for {url}: {e}")
+            return f"Error generating RSS feed for {url}: {e}"
     else:
         return "No data available."
 
 # Function to define the chat response function using the Mistral model
 def respond(message, history, system_message, max_tokens, temperature, top_p):
-    model = AutoModelForSeq2SeqLM.from_pretrained_model("mistralai/Mixtral-8x7B-Instruct-v0.1")
-    tokenizer = AutoTokenizer.from_pretrained_model("mistralai/Mixtral-8x7B-Instruct-v0.1")
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    response = pipe(f"User: {message}\nHistory: {history}\nSystem: {system_message}", max_length=max_tokens, temperature=temperature, top_p=top_p)[0]
-    return response
+    # Load the model and tokenizer once
+    if not hasattr(respond, "pipe"):
+        try:
+            model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+            respond.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            respond.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+            respond.pipe = pipeline(
+                "text-generation",
+                model=respond.model,
+                tokenizer=respond.tokenizer,
+                device=0 if torch.cuda.is_available() else -1,
+            )
+            logging.info("Model loaded successfully.")
+        except Exception as e:
+            logging.error(f"Error loading model: {e}")
+            return "Error loading the response model."
+
+    try:
+        prompt = (
+            f"System: {system_message}\n"
+            f"History: {history}\n"
+            f"User: {message}\n"
+            f"Assistant:"
+        )
+        response = respond.pipe(
+            prompt, max_length=max_tokens, temperature=temperature, top_p=top_p
+        )[0]["generated_text"]
+        return response
+    except Exception as e:
+        logging.error(f"Error generating response: {e}")
+        return "Error generating response."
 
 # Define the Gradio interface
 def create_interface():
     with gr.Blocks() as demo:
+        gr.Markdown("# All-in-One Scraper, Database, and RSS Feeder")
+
         with gr.Row():
             with gr.Column():
-                message = gr.Textbox(label="Message")
-                system_message = gr.Textbox(value="You are a helpful assistant.", label="System message")
-                max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
-                temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
-                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
-                storage_location = gr.Textbox(value="scraped_data", label="Storage Location")
-                urls = gr.Textbox(label="URLs (comma separated)")
-                scrape_interval = gr.Slider(minimum=1, maximum=60, value=5, step=1, label="Scrape Interval (minutes)")
-                content_type = gr.Radio(choices=["text", "media", "both"], value="text", label="Content Type")
+                storage_location = gr.Textbox(
+                    value=DEFAULT_FILE_PATH, label="Storage Location"
+                )
+                urls = gr.Textbox(
+                    label="URLs (comma separated)",
+                    placeholder="https://example.com, https://anotherexample.com",
+                )
+                scrape_interval = gr.Slider(
+                    minimum=1,
+                    maximum=60,
+                    value=5,
+                    step=1,
+                    label="Scrape Interval (minutes)",
+                )
+                content_type = gr.Radio(
+                    choices=["text", "media", "both"],
+                    value="text",
+                    label="Content Type",
+                )
+                selector = gr.Textbox(
+                    label="CSS Selector for Media (Optional)",
+                    placeholder="e.g., img.main-image",
+                )
                 start_button = gr.Button("Start Scraping")
                 stop_button = gr.Button("Stop Scraping")
-                csv_output = gr.Textbox(label="CSV Output", interactive=False)
+                csv_output = gr.Textbox(
+                    label="CSV Output", interactive=False, lines=2
+                )
 
             with gr.Column():
                 chat_history = gr.Chatbot(label="Chat History")
-                response_box = gr.Textbox(label="Response")
+                with gr.Row():
+                    message = gr.Textbox(label="Message", placeholder="Type your message here...")
+                system_message = gr.Textbox(
+                    value="You are a helpful assistant.", label="System message"
+                )
+                max_tokens = gr.Slider(
+                    minimum=1,
+                    maximum=2048,
+                    value=512,
+                    step=1,
+                    label="Max new tokens",
+                )
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=4.0,
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature",
+                )
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.95,
+                    step=0.05,
+                    label="Top-p (nucleus sampling)",
+                )
+                response_box = gr.Textbox(label="Response", interactive=False, lines=2)
 
-# Function to stop scraping
-def stop_scraping(stop_scraping_flag):
-    stop_scraping_flag[0] = True
-    return "Scraping stopped."
-
-        # Add a button to display the CSV content for a selected URL
         with gr.Row():
-            selected_url = gr.Textbox(label="Select URL for CSV Content")
+            selected_url_csv = gr.Textbox(
+                label="Select URL for CSV Content",
+                placeholder="https://example.com",
+            )
             csv_button = gr.Button("Display CSV Content")
-            csv_output = gr.Textbox(label="CSV Content Output", interactive=False)
-
-        csv_button.click(display_csv, inputs=[storage_location, selected_url], outputs=csv_output)
+            csv_content_output = gr.Textbox(
+                label="CSV Content Output", interactive=False, lines=10
+            )
 
-        # Add a button to display the RSS feed for a selected URL
         with gr.Row():
-            selected_url = gr.Textbox(label="Select URL for RSS Feed")
+            selected_url_rss = gr.Textbox(
+                label="Select URL for RSS Feed",
+                placeholder="https://example.com",
+            )
             rss_button = gr.Button("Generate RSS Feed")
-            rss_output = gr.Textbox(label="RSS Feed Output", interactive=False)
+            rss_output = gr.Textbox(
+                label="RSS Feed Output", interactive=False, lines=20
+            )
+
+        # Connect buttons to their respective functions
+        start_button.click(
+            fn=start_scraping,
+            inputs=[
+                storage_location,
+                gr.Textbox.value,
+                scrape_interval,
+                content_type,
+                selector,
+            ],
+            outputs=csv_output,
+        )
+
+        stop_button.click(fn=stop_scraping, outputs=csv_output)
+
+        csv_button.click(
+            fn=display_csv,
+            inputs=[storage_location, selected_url_csv],
+            outputs=csv_content_output,
+        )
+
+        rss_button.click(
+            fn=generate_rss_feed,
+            inputs=[storage_location, selected_url_rss],
+            outputs=rss_output,
+        )
+
+        # Connect message submission to the chat interface
+        def update_chat(message, history, system_message, max_tokens, temperature, top_p):
+            response = respond(message, history, system_message, max_tokens, temperature, top_p)
+            history.append((message, response))
+            return history, response
 
-        rss_button.click(generate_rss_feed, inputs=[storage_location, selected_url], outputs=rss_output)
+        message.submit(
+            update_chat,
+            inputs=[
+                message,
+                chat_history,
+                system_message,
+                max_tokens,
+                temperature,
+                top_p,
+            ],
+            outputs=[chat_history, response_box],
+        )
 
     return demo