feat: web scraping connector task

cognee-mcp/pyproject.toml

@@ -37,5 +37,4 @@ dev = [
 allow-direct-references = true
 
 [project.scripts]
-cognee = "src:main"
-cognee-mcp = "src:main_mcp"
+cognee-mcp = "src:main"

cognee/api/v1/add/add.py

@@ -1,5 +1,6 @@
 from uuid import UUID
-from typing import Union, BinaryIO, List, Optional
+import os
+from typing import Union, BinaryIO, List, Optional, Dict, Any
 
 from cognee.modules.users.models import User
 from cognee.modules.pipelines import Task, run_pipeline
@@ -11,6 +12,12 @@
 )
 from cognee.modules.engine.operations.setup import setup
 from cognee.tasks.ingestion import ingest_data, resolve_data_directories
+from cognee.tasks.web_scraper.config import TavilyConfig, SoupCrawlerConfig
+from cognee.context_global_variables import (
+    tavily_config as tavily,
+    soup_crawler_config as soup_crawler,
+)
+from urllib.parse import urlparse
 
 
 async def add(
@@ -23,12 +30,15 @@ async def add(
     dataset_id: Optional[UUID] = None,
     preferred_loaders: List[str] = None,
     incremental_loading: bool = True,
+    extraction_rules: Optional[Dict[str, Any]] = None,
+    tavily_config: Optional[TavilyConfig] = None,
+    soup_crawler_config: Optional[SoupCrawlerConfig] = None,
 ):
     """
     Add data to Cognee for knowledge graph processing.
 
     This is the first step in the Cognee workflow - it ingests raw data and prepares it
-    for processing. The function accepts various data formats including text, files, and
+    for processing. The function accepts various data formats including text, files, urls and
     binary streams, then stores them in a specified dataset for further processing.
 
     Prerequisites:
@@ -68,6 +78,7 @@ async def add(
             - S3 path: "s3://my-bucket/documents/file.pdf"
             - List of mixed types: ["text content", "/path/file.pdf", "file://doc.txt", file_handle]
             - Binary file object: open("file.txt", "rb")
+            - url: A web link url (https or http)
         dataset_name: Name of the dataset to store data in. Defaults to "main_dataset".
                     Create separate datasets to organize different knowledge domains.
         user: User object for authentication and permissions. Uses default user if None.
@@ -78,6 +89,9 @@ async def add(
         vector_db_config: Optional configuration for vector database (for custom setups).
         graph_db_config: Optional configuration for graph database (for custom setups).
         dataset_id: Optional specific dataset UUID to use instead of dataset_name.
+        extraction_rules: Optional dictionary of rules (e.g., CSS selectors, XPath) for extracting specific content from web pages using BeautifulSoup
+        tavily_config: Optional configuration for Tavily API, including API key and extraction settings
+        soup_crawler_config: Optional configuration for BeautifulSoup crawler, specifying concurrency, crawl delay, and extraction rules.
 
     Returns:
         PipelineRunInfo: Information about the ingestion pipeline execution including:
@@ -126,6 +140,21 @@ async def add(
 
         # Add a single file
         await cognee.add("/home/user/documents/analysis.pdf")
+
+        # Add a single url and bs4 extract ingestion method
+        extraction_rules = {
+            "title": "h1",
+            "description": "p",
+            "more_info": "a[href*='more-info']"
+        }
+        await cognee.add("https://example.com",extraction_rules=extraction_rules)
+
+        # Add a single url and tavily extract ingestion method
+        Make sure to set    TAVILY_API_KEY = YOUR_TAVILY_API_KEY as a environment variable
+        await cognee.add("https://example.com")
+
+        # Add multiple urls
+        await cognee.add(["https://example.com","https://books.toscrape.com"])
         ```
 
     Environment Variables:
@@ -139,11 +168,38 @@ async def add(
         - DEFAULT_USER_PASSWORD: Custom default user password
         - VECTOR_DB_PROVIDER: "lancedb" (default), "chromadb", "pgvector"
         - GRAPH_DATABASE_PROVIDER: "kuzu" (default), "neo4j"
+        - TAVILY_API_KEY: YOUR_TAVILY_API_KEY
 
     """
+
+    if not soup_crawler_config and extraction_rules:
+        soup_crawler_config = SoupCrawlerConfig(extraction_rules=extraction_rules)
+    if not tavily_config and os.getenv("TAVILY_API_KEY"):
+        tavily_config = TavilyConfig(api_key=os.getenv("TAVILY_API_KEY"))
+
+    soup_crawler.set(soup_crawler_config)
+    tavily.set(tavily_config)
+
+    http_schemes = {"http", "https"}
+
+    def _is_http_url(item: Union[str, BinaryIO]) -> bool:
+        return isinstance(item, str) and urlparse(item).scheme in http_schemes
+
+    if _is_http_url(data):
+        node_set = ["web_content"] if not node_set else node_set + ["web_content"]
+    elif isinstance(data, list) and any(_is_http_url(item) for item in data):
+        node_set = ["web_content"] if not node_set else node_set + ["web_content"]
+
     tasks = [
         Task(resolve_data_directories, include_subdirectories=True),
-        Task(ingest_data, dataset_name, user, node_set, dataset_id, preferred_loaders),
+        Task(
+            ingest_data,
+            dataset_name,
+            user,
+            node_set,
+            dataset_id,
+            preferred_loaders,
+        ),
     ]
 
     await setup()

cognee/context_global_variables.py

@@ -12,6 +12,8 @@
 #       for different async tasks, threads and processes
 vector_db_config = ContextVar("vector_db_config", default=None)
 graph_db_config = ContextVar("graph_db_config", default=None)
+soup_crawler_config = ContextVar("soup_crawler_config", default=None)
+tavily_config = ContextVar("tavily_config", default=None)
 
 
 async def set_database_global_context_variables(dataset: Union[str, UUID], user_id: UUID):

cognee/tasks/ingestion/save_data_item_to_storage.py

@@ -7,6 +7,7 @@
 from cognee.modules.ingestion import save_data_to_file
 from cognee.shared.logging_utils import get_logger
 from pydantic_settings import BaseSettings, SettingsConfigDict
+from cognee.context_global_variables import tavily_config, soup_crawler_config
 
 logger = get_logger()
 
@@ -17,6 +18,13 @@ class SaveDataSettings(BaseSettings):
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
 
+class HTMLContent(str):
+    def __new__(cls, value: str):
+        if not ("<" in value and ">" in value):
+            raise ValueError("Not valid HTML-like content")
+        return super().__new__(cls, value)
+
+
 settings = SaveDataSettings()
 
 
@@ -48,6 +56,39 @@ async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str
         # data is s3 file path
         if parsed_url.scheme == "s3":
             return data_item
+        elif parsed_url.scheme == "http" or parsed_url.scheme == "https":
+            # Validate URL by sending a HEAD request
+            try:
+                from cognee.tasks.web_scraper import fetch_page_content
+
+                tavily = tavily_config.get()
+                soup_crawler = soup_crawler_config.get()
+                preferred_tool = "beautifulsoup" if soup_crawler else "tavily"
+                if preferred_tool == "tavily" and tavily is None:
+                    raise IngestionError(
+                        message="TavilyConfig must be set on the ingestion context when fetching HTTP URLs without a SoupCrawlerConfig."
+                    )
+                if preferred_tool == "beautifulsoup" and soup_crawler is None:
+                    raise IngestionError(
+                        message="SoupCrawlerConfig must be set on the ingestion context when using the BeautifulSoup scraper."
+                    )
+
+                data = await fetch_page_content(
+                    data_item,
+                    preferred_tool=preferred_tool,
+                    tavily_config=tavily,
+                    soup_crawler_config=soup_crawler,
+                )
+                content = ""
+                for key, value in data.items():
+                    content += f"{key}:\n{value}\n\n"
+                return await save_data_to_file(content)
+            except IngestionError:
+                raise
+            except Exception as e:
+                raise IngestionError(
+                    message=f"Error ingesting webpage results of url {data_item}: {str(e)}"
+                )
 
         # data is local file path
         elif parsed_url.scheme == "file":

cognee/tasks/web_scraper/__init__.py

@@ -0,0 +1,18 @@
+"""Web scraping module for cognee.
+
+This module provides tools for scraping web content, managing scraping jobs, and storing
+data in a graph database. It includes classes and functions for crawling web pages using
+BeautifulSoup or Tavily, defining data models, and handling scraping configurations.
+"""
+
+from .bs4_crawler import BeautifulSoupCrawler
+from .utils import fetch_page_content
+from .web_scraper_task import cron_web_scraper_task, web_scraper_task
+
+
+__all__ = [
+    "BeautifulSoupCrawler",
+    "fetch_page_content",
+    "cron_web_scraper_task",
+    "web_scraper_task",
+]