Spaces:
Sleeping
Sleeping
File size: 2,926 Bytes
c96b98a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | from typing import Dict, List, Optional, Literal
from phi.document.base import Document
from phi.document.reader.base import Reader
from phi.utils.log import logger
from firecrawl import FirecrawlApp
class FirecrawlReader(Reader):
api_key: Optional[str] = None
params: Optional[Dict] = None
mode: Literal["scrape", "crawl"] = "scrape"
def scrape(self, url: str) -> List[Document]:
"""
Scrapes a website and returns a list of documents.
Args:
url: The URL of the website to scrape
Returns:
A list of documents
"""
logger.debug(f"Scraping: {url}")
app = FirecrawlApp(api_key=self.api_key)
scraped_data = app.scrape_url(url, params=self.params)
# print(scraped_data)
content = scraped_data.get("markdown", "")
# Debug logging
logger.debug(f"Received content type: {type(content)}")
logger.debug(f"Content empty: {not bool(content)}")
# Ensure content is a string
if content is None:
content = "" # or you could use metadata to create a meaningful message
logger.warning(f"No content received for URL: {url}")
documents = []
if self.chunk and content: # Only chunk if there's content
documents.extend(self.chunk_document(Document(name=url, id=url, content=content)))
else:
documents.append(Document(name=url, id=url, content=content))
return documents
def crawl(self, url: str) -> List[Document]:
"""
Crawls a website and returns a list of documents.
Args:
url: The URL of the website to crawl
Returns:
A list of documents
"""
logger.debug(f"Crawling: {url}")
app = FirecrawlApp(api_key=self.api_key)
crawl_result = app.crawl_url(url, params=self.params)
documents = []
# Extract data from crawl results
results_data = crawl_result.get("data", [])
for result in results_data:
# Get markdown content, default to empty string if not found
content = result.get("markdown", "")
if content: # Only create document if content exists
if self.chunk:
documents.extend(self.chunk_document(Document(name=url, id=url, content=content)))
else:
documents.append(Document(name=url, id=url, content=content))
return documents
def read(self, url: str) -> List[Document]:
"""
Args:
url: The URL of the website to scrape
Returns:
A list of documents
"""
if self.mode == "scrape":
return self.scrape(url)
elif self.mode == "crawl":
return self.crawl(url)
else:
raise NotImplementedError(f"Mode {self.mode} not implemented")
|