doc list not empty (#21208)

Make sure the doc list is not empty, and set Metadata: true in param, to
enable the user to disable metadata for slightly faster crawls.
This commit is contained in:
WilliamEspegren 2024-05-20 17:24:06 +02:00 committed by GitHub
parent 8da35fba7f
commit 30bca57aae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -18,7 +18,7 @@ class SpiderLoader(BaseLoader):
*, *,
api_key: Optional[str] = None, api_key: Optional[str] = None,
mode: Literal["scrape", "crawl"] = "scrape", mode: Literal["scrape", "crawl"] = "scrape",
params: Optional[dict] = {"return_format": "markdown"}, params: Optional[dict] = None,
): ):
"""Initialize with API key and URL. """Initialize with API key and URL.
@ -31,6 +31,12 @@ class SpiderLoader(BaseLoader):
crawling following subpages). crawling following subpages).
params: Additional parameters for the Spider API. params: Additional parameters for the Spider API.
""" """
if params is None:
params = {
"return_format": "markdown",
"metadata": True,
} # Using the metadata param slightly slows down the output
try: try:
from spider import Spider from spider import Spider
except ImportError: except ImportError:
@ -41,13 +47,6 @@ class SpiderLoader(BaseLoader):
raise ValueError( raise ValueError(
f"Unrecognized mode '{mode}'. Expected one of 'scrape', 'crawl'." f"Unrecognized mode '{mode}'. Expected one of 'scrape', 'crawl'."
) )
# If `params` is `None`, initialize it as an empty dictionary
if params is None:
params = {}
# Add a default value for 'metadata' if it's not already present
if "metadata" not in params:
params["metadata"] = True
# Use the environment variable if the API key isn't provided # Use the environment variable if the API key isn't provided
api_key = api_key or get_from_env("api_key", "SPIDER_API_KEY") api_key = api_key or get_from_env("api_key", "SPIDER_API_KEY")
@ -79,6 +78,7 @@ class SpiderLoader(BaseLoader):
# Ensure metadata is also not None # Ensure metadata is also not None
metadata = doc[0].get("metadata", {}) metadata = doc[0].get("metadata", {})
if page_content is not None:
yield Document(page_content=page_content, metadata=metadata) yield Document(page_content=page_content, metadata=metadata)
if self.mode == "crawl": if self.mode == "crawl":
# Ensure page_content is also not None # Ensure page_content is also not None