Merge pull request #17 from Querent-ai/setup_webscaper_collector

modyfing interfaces and setting up scraper

Author: saraswatpuneet

querent/collectors/collector_factory.py

@@ -3,7 +3,7 @@
 from enum import Enum
 from querent.collectors.collector_base import Collector
 from querent.collectors.collector_errors import CollectorResolverError, CollectorErrorKind
-from querent.config.collector_config import CollectorBackend
+from querent.config.collector_config import CollectConfig, CollectorBackend
 
 
 class CollectorFactory(ABC):
@@ -12,7 +12,7 @@ def backend(self) -> CollectorBackend:
         pass
 
     @abstractmethod
-    async def resolve(self, uri: str) -> Optional[CollectorBackend]:
+    async def resolve(self, uri: str, config: CollectConfig) -> Optional[CollectorBackend]:
         pass
 
 
@@ -21,7 +21,7 @@ def __init__(self, backend: CollectorBackend, message: str):
         self.backend = backend
         self.message = message
 
-    async def resolve(self, uri: str) -> Optional[Collector]:
+    async def resolve(self, uri: str, config: CollectConfig) -> Optional[Collector]:
         raise CollectorResolverError(
             CollectorErrorKind.NotSupported, self.backend, self.message
         )

querent/collectors/collector_resolver.py

@@ -1,6 +1,7 @@
 from typing import Optional
 from querent.collectors.fs.fs_collector import FSCollectorFactory
-from querent.config.collector_config import CollectorBackend
+from querent.collectors.webscaper.web_scraper_collector import WebScraperFactory
+from querent.config.collector_config import CollectConfig, CollectorBackend
 from querent.collectors.collector_base import Collector
 from querent.collectors.collector_errors import CollectorResolverError, CollectorErrorKind
 from querent.common.uri import Protocol, Uri
@@ -9,15 +10,16 @@ class CollectorResolver:
     def __init__(self):
         self.collector_factories = {
             CollectorBackend.LocalFile: FSCollectorFactory(),
+            CollectorBackend.WebScraper: WebScraperFactory(),
             # Add other collector factories as needed
         }
 
-    def resolve(self, uri: Uri) -> Optional[Collector]:
+    def resolve(self, uri: Uri, config: CollectConfig) -> Optional[Collector]:
         backend = self._determine_backend(uri.protocol)
 
         if backend in self.collector_factories:
             factory = self.collector_factories[backend]
-            return factory.resolve(uri)
+            return factory.resolve(uri, config)
         else:
             raise CollectorResolverError(
                 CollectorErrorKind.NotSupported, backend, "Unsupported backend"

querent/collectors/fs/fs_collector.py

@@ -40,15 +40,15 @@ async def walk_files(self, root: Path) -> AsyncGenerator[Path, None]:
         for item in root.iterdir():
             if item.is_file():
                 yield item
-
-
+            elif item.is_dir():
+                async for file_path in self.walk_files(item):
+                    yield file_path
 class FSCollectorFactory(CollectorFactory):
     def __init__(self):
         pass
 
     def backend(self) -> CollectorBackend:
         return CollectorBackend.LocalFile
 
-    def resolve(self, uri: Uri) -> Collector:
-        config = FSCollectorConfig(root_path=uri.path)
+    def resolve(self, uri: Uri, config: FSCollectorConfig) -> Collector:
         return FSCollector(config)

querent/collectors/webscaper/web_scraper_collector.py

@@ -0,0 +1,34 @@
+from querent.collectors.collector_base import Collector
+from querent.collectors.collector_factory import CollectorFactory
+from querent.collectors.collector_result import CollectorResult
+from querent.config.collector_config import CollectorBackend, WebScraperConfig
+from querent.tools.web_page_extractor import WebpageExtractor
+
+class WebScraperCollector(Collector):
+    def __init__(self, config: WebScraperConfig):
+        self.website_url = config.website_url
+
+    async def connect(self):
+        pass  # Any setup logic before scraping
+
+    async def disconnect(self):
+        pass  # Any cleanup logic after scraping
+
+    async def poll(self):
+        content = await self.scrape_website(self.website_url)
+        yield CollectorResult(content)
+
+    async def scrape_website(self, website_url: str):
+        content = WebpageExtractor().extract_with_bs4(website_url)
+        max_length = len(' '.join(content.split(" ")[:600]))
+        return content[:max_length]
+
+class WebScraperFactory(CollectorFactory):
+    def __init__(self):
+        pass
+    
+    def backend(self) -> CollectorBackend:
+        return CollectorBackend.WebScraper
+
+    def resolve(self, config: WebScraperConfig) -> Collector:
+        return WebScraperCollector(config)

querent/config/collector_config.py

@@ -1,6 +1,6 @@
 from enum import Enum
 from typing import Optional
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class CollectorBackend(str, Enum):
@@ -35,6 +35,10 @@ class GcsCollectConfig(BaseModel):
     access_key: str
     secret_key: str
 
+class WebScraperConfig(BaseModel):
+    website_url: str = Field(
+        ..., description="The URL of the website to scrape."
+    )
 
 class CollectConfigWrapper(BaseModel):
     backend: CollectorBackend
@@ -50,5 +54,9 @@ def from_collect_config(cls, collect_config: CollectConfig):
             return cls(backend=CollectorBackend.S3, config=S3CollectConfig())
         elif collect_config.backend == CollectorBackend.Gcs:
             return cls(backend=CollectorBackend.Gcs, config=GcsCollectConfig())
+        elif collect_config.backend == CollectorBackend.WebScraper:
+            return cls(
+                backend=CollectorBackend.WebScraper, config=WebScraperConfig()
+            )
         else:
             raise ValueError(f"Unsupported collector backend: {collect_config.backend}")

querent/tools/__init__.py

@@ -0,0 +1,166 @@
+from io import BytesIO
+from PyPDF2 import PdfFileReader
+from PyPDF2 import PdfReader
+import requests
+import re
+from requests.exceptions import RequestException
+from bs4 import BeautifulSoup
+from newspaper import Article, ArticleException, Config
+from requests_html import HTMLSession
+import time
+import random
+from lxml import html
+from querent.lib.logger import logger
+
+USER_AGENTS = [
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:66.0) Gecko/20100101 Firefox/66.0",
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_5) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/12.1.1 Safari/605.1.15",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/64.0.3282.140 Safari/537.36 Edge/17.0",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:72.0) Gecko/20100101 Firefox/72.0",
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/78.0.3904.97 Safari/537.36",
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/13.0.3 Safari/605.1.15",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:80.0) Gecko/20100101 Firefox/80.0",
+    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/80.0.3987.87 Safari/537.36",
+    "Mozilla/5.0 (iPhone; CPU iPhone OS 13_3 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/13.0.4 Mobile/15E148 Safari/604.1",
+    "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:73.0) Gecko/20100101 Firefox/73.0",
+    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.129 Safari/537.36"
+]
+
+class WebpageExtractor:
+
+    def __init__(self, num_extracts=3):
+        """
+        Initialize the WebpageExtractor class.
+        """
+        self.num_extracts = num_extracts
+
+    def extract_with_3k(self, url):
+        """
+        Extract the text from a webpage using the 3k method.
+
+        Args:
+            url (str): The URL of the webpage to extract from.
+
+        Returns:
+            str: The extracted text.
+        """
+        try:
+            if url.lower().endswith(".pdf"):
+                response = requests.get(url)
+                response.raise_for_status()
+
+                with BytesIO(response.content) as pdf_data:
+                    reader = PdfReader(pdf_data)
+                    content = " ".join([reader.getPage(i).extract_text() for i in range(reader.getNumPages())])
+
+            else:
+                config = Config()
+                config.browser_user_agent = random.choice(USER_AGENTS)
+                config.request_timeout = 10
+                session = HTMLSession()
+
+                response = session.get(url)
+                response.html.render(timeout=config.request_timeout)
+                html_content = response.html.html
+
+                article = Article(url, config=config)
+                article.set_html(html_content)
+                article.parse()
+                content = article.text.replace('\t', ' ').replace('\n', ' ').strip()
+
+            return content[:1500]
+
+        except ArticleException as ae:
+            logger.error(f"Error while extracting text from HTML (newspaper3k): {str(ae)}")
+            return f"Error while extracting text from HTML (newspaper3k): {str(ae)}"
+
+        except RequestException as re:
+            logger.error(f"Error while making the request to the URL (newspaper3k): {str(re)}")
+            return f"Error while making the request to the URL (newspaper3k): {str(re)}"
+
+        except Exception as e:
+            logger.error(f"Unknown error while extracting text from HTML (newspaper3k): {str(e)}")
+            return ""
+
+    def extract_with_bs4(self, url):
+        """
+        Extract the text from a webpage using the BeautifulSoup4 method.
+
+        Args:
+            url (str): The URL of the webpage to extract from.
+
+        Returns:
+            str: The extracted text.
+        """
+        headers = {
+            "User-Agent": random.choice(USER_AGENTS)
+        }
+
+        try:
+            response = requests.get(url, headers=headers, timeout=10)
+            if response.status_code == 200:
+                soup = BeautifulSoup(response.text, 'html.parser')
+                for tag in soup(['script', 'style', 'nav', 'footer', 'head', 'link', 'meta', 'noscript']):
+                    tag.decompose()
+
+                main_content_areas = soup.find_all(['main', 'article', 'section', 'div'])
+                if main_content_areas:
+                    main_content = max(main_content_areas, key=lambda x: len(x.text))
+                    content_tags = ['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']
+                    content = ' '.join([tag.text.strip() for tag in main_content.find_all(content_tags)])
+                else:
+                    content = ' '.join([tag.text.strip() for tag in soup.find_all(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6'])])
+
+                content = re.sub(r'\t', ' ', content)
+                content = re.sub(r'\s+', ' ', content)
+                return content
+            elif response.status_code == 404:
+                return f"Error: 404. Url is invalid or does not exist. Try with valid url..."
+            else:
+                logger.error(f"Error while extracting text from HTML (bs4): {response.status_code}")
+                return f"Error while extracting text from HTML (bs4): {response.status_code}"
+
+        except Exception as e:
+            logger.error(f"Unknown error while extracting text from HTML (bs4): {str(e)}")
+            return ""
+
+    def extract_with_lxml(self, url):
+        """
+        Extract the text from a webpage using the lxml method.
+
+        Args:
+            url (str): The URL of the webpage to extract from.
+
+        Returns:
+            str: The extracted text.
+        """
+        try:
+            config = Config()
+            config.browser_user_agent = random.choice(USER_AGENTS)
+            config.request_timeout = 10
+            session = HTMLSession()
+
+            response = session.get(url)
+            response.html.render(timeout=config.request_timeout)
+            html_content = response.html.html
+
+            tree = html.fromstring(html_content)
+            paragraphs = tree.cssselect('p, h1, h2, h3, h4, h5, h6')
+            content = ' '.join([para.text_content() for para in paragraphs if para.text_content()])
+            content = content.replace('\t', ' ').replace('\n', ' ').strip()
+
+            return content
+
+        except ArticleException as ae:
+            logger.error("Error while extracting text from HTML (lxml): {str(ae)}")
+            return ""
+
+        except RequestException as re:
+            logger.error(f"Error while making the request to the URL (lxml): {str(re)}")
+            return ""
+
+        except Exception as e:
+            logger.error(f"Unknown error while extracting text from HTML (lxml): {str(e)}")
+            return ""
+    

querent/tools/web_page_extractor.py

@@ -6,7 +6,7 @@
 import pytest
 
 from querent.common.uri import Uri
-from querent.config.collector_config import CollectorBackend
+from querent.config.collector_config import CollectorBackend, FSCollectorConfig
 
 
 @pytest.fixture
@@ -19,7 +19,8 @@ def temp_dir():
 def test_fs_collector(temp_dir):
     uri = Uri("file://" + temp_dir)
     resolver = CollectorResolver()
-    collector = resolver.resolve(uri)
+    fileConfig = FSCollectorConfig(root_path=uri.path)
+    collector = resolver.resolve(uri, fileConfig)
     assert collector is not None
 
 
@@ -35,7 +36,8 @@ def test_add_files_read_via_collector(temp_dir):
         file.write(b"test_add_files_read_via_collector")
     uri = Uri("file://" + temp_dir)
     resolver = CollectorResolver()
-    collector = resolver.resolve(uri)
+    fileConfig = FSCollectorConfig(root_path=uri.path)
+    collector = resolver.resolve(uri, fileConfig)
     assert collector is not None
 
     async def poll_and_print():