Spaces:
Sleeping
Sleeping
| 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") | |