178 lines
5.3 KiB
Python
178 lines
5.3 KiB
Python
"""
|
|
Firecrawl service for website scraping and enrichment.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Optional, Dict, Any, List
|
|
|
|
import httpx
|
|
|
|
from ..config import get_settings
|
|
from ..models.schemas import FirecrawlResult
|
|
from ..utils.retry import retry_with_backoff
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FirecrawlService:
|
|
"""Service for scraping websites using Firecrawl."""
|
|
|
|
BASE_URL = "https://api.firecrawl.dev/v1"
|
|
|
|
def __init__(self):
|
|
"""Initialize Firecrawl client."""
|
|
settings = get_settings()
|
|
self.api_key = settings.firecrawl_api_key
|
|
self.headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
@retry_with_backoff(max_attempts=3, exceptions=(httpx.HTTPError,))
|
|
def scrape_url(
|
|
self,
|
|
url: str,
|
|
formats: Optional[List[str]] = None,
|
|
include_images: bool = True,
|
|
) -> FirecrawlResult:
|
|
"""
|
|
Scrape a single URL.
|
|
|
|
Args:
|
|
url: URL to scrape
|
|
formats: Output formats (default: markdown)
|
|
include_images: Whether to extract image URLs
|
|
|
|
Returns:
|
|
Scraped content
|
|
"""
|
|
if formats is None:
|
|
formats = ["markdown"]
|
|
|
|
endpoint = f"{self.BASE_URL}/scrape"
|
|
|
|
payload = {
|
|
"url": url,
|
|
"formats": formats,
|
|
}
|
|
|
|
with httpx.Client(timeout=60.0) as client:
|
|
response = client.post(endpoint, json=payload, headers=self.headers)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
if not data.get("success"):
|
|
raise RuntimeError(f"Firecrawl scrape failed: {data.get('error')}")
|
|
|
|
result_data = data.get("data", {})
|
|
|
|
# Extract images from metadata if available
|
|
images = []
|
|
if include_images:
|
|
metadata = result_data.get("metadata", {})
|
|
# Try to get og:image and other image sources
|
|
if metadata.get("ogImage"):
|
|
images.append(metadata["ogImage"])
|
|
if metadata.get("logo"):
|
|
images.append(metadata["logo"])
|
|
|
|
return FirecrawlResult(
|
|
markdown=result_data.get("markdown"),
|
|
metadata=result_data.get("metadata"),
|
|
images=images,
|
|
)
|
|
|
|
def extract_company_info(
|
|
self, url: str
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Extract structured company information from a website.
|
|
|
|
Args:
|
|
url: Company website URL
|
|
|
|
Returns:
|
|
Dictionary with extracted information
|
|
"""
|
|
try:
|
|
result = self.scrape_url(url, formats=["markdown"], include_images=True)
|
|
|
|
info = {
|
|
"markdown": result.markdown,
|
|
"images": result.images or [],
|
|
"logo_url": None,
|
|
"email": None,
|
|
"title": None,
|
|
"description": None,
|
|
}
|
|
|
|
# Extract metadata
|
|
if result.metadata:
|
|
info["title"] = result.metadata.get("title")
|
|
info["description"] = result.metadata.get("description")
|
|
if result.metadata.get("ogImage"):
|
|
info["logo_url"] = result.metadata["ogImage"]
|
|
|
|
# Try to find email in markdown content
|
|
if result.markdown:
|
|
import re
|
|
email_pattern = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"
|
|
emails = re.findall(email_pattern, result.markdown)
|
|
if emails:
|
|
# Filter out common non-company emails
|
|
company_emails = [
|
|
e for e in emails
|
|
if not any(
|
|
domain in e.lower()
|
|
for domain in ["example.com", "test.com", "email.com"]
|
|
)
|
|
]
|
|
if company_emails:
|
|
info["email"] = company_emails[0]
|
|
|
|
return info
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to scrape {url}: {e}")
|
|
return {
|
|
"markdown": None,
|
|
"images": [],
|
|
"logo_url": None,
|
|
"email": None,
|
|
"title": None,
|
|
"description": None,
|
|
"error": str(e),
|
|
}
|
|
|
|
def batch_scrape(
|
|
self, urls: List[str], batch_size: int = 5
|
|
) -> Dict[str, FirecrawlResult]:
|
|
"""
|
|
Scrape multiple URLs in batches.
|
|
|
|
Args:
|
|
urls: List of URLs to scrape
|
|
batch_size: Number of concurrent scrapes
|
|
|
|
Returns:
|
|
Dictionary mapping URLs to results
|
|
"""
|
|
results = {}
|
|
|
|
for i in range(0, len(urls), batch_size):
|
|
batch = urls[i : i + batch_size]
|
|
logger.info(f"Scraping batch {i // batch_size + 1} ({len(batch)} URLs)")
|
|
|
|
for url in batch:
|
|
try:
|
|
results[url] = self.scrape_url(url)
|
|
except Exception as e:
|
|
logger.error(f"Failed to scrape {url}: {e}")
|
|
results[url] = FirecrawlResult(
|
|
markdown=None,
|
|
metadata={"error": str(e)},
|
|
images=[],
|
|
)
|
|
|
|
return results
|