new site

directory/services/firecrawl.py

@@ -0,0 +1,177 @@
+"""
+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