mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-07 15:18:29 +00:00
lint fixes
This commit is contained in:
@@ -1,60 +1,137 @@
|
||||
import logging
|
||||
from typing import Any, Dict, Literal, Optional, Type
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from crewai.tools import BaseTool
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
logger = logging.getLogger(__file__)
|
||||
|
||||
|
||||
class SpiderToolSchema(BaseModel):
|
||||
url: str = Field(description="Website URL")
|
||||
params: Optional[Dict[str, Any]] = Field(
|
||||
description="Set additional params. Options include:\n"
|
||||
"- `limit`: Optional[int] - The maximum number of pages allowed to crawl per website. Remove the value or set it to `0` to crawl all pages.\n"
|
||||
"- `depth`: Optional[int] - The crawl limit for maximum depth. If `0`, no limit will be applied.\n"
|
||||
"- `metadata`: Optional[bool] - Boolean to include metadata or not. Defaults to `False` unless set to `True`. If the user wants metadata, include params.metadata = True.\n"
|
||||
"- `query_selector`: Optional[str] - The CSS query selector to use when extracting content from the markup.\n"
|
||||
"""Input schema for SpiderTool."""
|
||||
|
||||
website_url: str = Field(
|
||||
..., description="Mandatory website URL to scrape or crawl"
|
||||
)
|
||||
mode: Literal["scrape", "crawl"] = Field(
|
||||
default="scrape",
|
||||
description="Mode, the only two allowed modes are `scrape` or `crawl`. Use `scrape` to scrape a single page and `crawl` to crawl the entire website following subpages. These modes are the only allowed values even when ANY params is set.",
|
||||
description="The mode of the SpiderTool. The only two allowed modes are `scrape` or `crawl`. Crawl mode will follow up to 5 links and return their content in markdown format.",
|
||||
)
|
||||
|
||||
|
||||
class SpiderTool(BaseTool):
|
||||
name: str = "Spider scrape & crawl tool"
|
||||
description: str = "Scrape & Crawl any url and return LLM-ready data."
|
||||
args_schema: Type[BaseModel] = SpiderToolSchema
|
||||
api_key: Optional[str] = None
|
||||
spider: Optional[Any] = None
|
||||
"""Tool for scraping and crawling websites."""
|
||||
|
||||
DEFAULT_CRAWL_LIMIT: int = 5
|
||||
DEFAULT_RETURN_FORMAT: str = "markdown"
|
||||
|
||||
name: str = "SpiderTool"
|
||||
description: str = (
|
||||
"A tool to scrape or crawl a website and return LLM-ready content."
|
||||
)
|
||||
args_schema: Type[BaseModel] = SpiderToolSchema
|
||||
custom_params: Optional[Dict[str, Any]] = None
|
||||
website_url: Optional[str] = None
|
||||
api_key: Optional[str] = None
|
||||
spider: Any = None
|
||||
log_failures: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
website_url: Optional[str] = None,
|
||||
custom_params: Optional[Dict[str, Any]] = None,
|
||||
log_failures: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize SpiderTool for web scraping and crawling.
|
||||
|
||||
Args:
|
||||
api_key (Optional[str]): Spider API key for authentication. Required for production use.
|
||||
website_url (Optional[str]): Default website URL to scrape/crawl. Can be overridden during execution.
|
||||
custom_params (Optional[Dict[str, Any]]): Additional parameters to pass to Spider API.
|
||||
These override any parameters set by the LLM.
|
||||
log_failures (bool): If True, logs errors. Defaults to True.
|
||||
**kwargs: Additional arguments passed to BaseTool.
|
||||
|
||||
Raises:
|
||||
ImportError: If spider-client package is not installed.
|
||||
RuntimeError: If Spider client initialization fails.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
if website_url is not None:
|
||||
self.website_url = website_url
|
||||
|
||||
self.log_failures = log_failures
|
||||
self.custom_params = custom_params
|
||||
|
||||
try:
|
||||
from spider import Spider # type: ignore
|
||||
|
||||
self.spider = Spider(api_key=api_key)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"`spider-client` package not found, please run `pip install spider-client`"
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to initialize Spider client: {str(e)}")
|
||||
|
||||
self.spider = Spider(api_key=api_key)
|
||||
def _validate_url(self, url: str) -> bool:
|
||||
"""Validate URL format.
|
||||
|
||||
Args:
|
||||
url (str): URL to validate.
|
||||
Returns:
|
||||
bool: True if valid URL.
|
||||
"""
|
||||
try:
|
||||
result = urlparse(url)
|
||||
return all([result.scheme, result.netloc])
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _run(
|
||||
self,
|
||||
url: str,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
mode: Optional[Literal["scrape", "crawl"]] = "scrape",
|
||||
):
|
||||
if mode not in ["scrape", "crawl"]:
|
||||
website_url: str,
|
||||
mode: Literal["scrape", "crawl"] = "scrape",
|
||||
) -> str:
|
||||
params = {}
|
||||
url = website_url or self.website_url
|
||||
|
||||
if not self._validate_url(url):
|
||||
raise ValueError("Invalid URL format")
|
||||
|
||||
if not url:
|
||||
raise ValueError(
|
||||
"Unknown mode in `mode` parameter, `scrape` or `crawl` are the allowed modes"
|
||||
"Website URL must be provided either during initialization or execution"
|
||||
)
|
||||
|
||||
# Ensure 'return_format': 'markdown' is always included
|
||||
if params:
|
||||
params["return_format"] = "markdown"
|
||||
else:
|
||||
params = {"return_format": "markdown"}
|
||||
if mode not in ["scrape", "crawl"]:
|
||||
raise ValueError("Mode must be either 'scrape' or 'crawl'")
|
||||
|
||||
action = self.spider.scrape_url if mode == "scrape" else self.spider.crawl_url
|
||||
spider_docs = action(url=url, params=params)
|
||||
params["request"] = "smart"
|
||||
params["filter_output_svg"] = True
|
||||
params["return_format"] = self.DEFAULT_RETURN_FORMAT
|
||||
|
||||
return spider_docs
|
||||
if mode == "crawl":
|
||||
params["limit"] = self.DEFAULT_CRAWL_LIMIT
|
||||
|
||||
# Update params with custom params if provided.
|
||||
# This will override any params passed by LLM.
|
||||
if self.custom_params:
|
||||
params.update(self.custom_params)
|
||||
|
||||
try:
|
||||
action = (
|
||||
self.spider.scrape_url if mode == "scrape" else self.spider.crawl_url
|
||||
)
|
||||
return action(url=url, params=params)
|
||||
|
||||
except Exception as e:
|
||||
if self.log_failures:
|
||||
logger.error(f"Error fetching data from {url}, exception: {e}")
|
||||
return None
|
||||
else:
|
||||
raise e
|
||||
|
||||
Reference in New Issue
Block a user