Service for interacting with ScrapeGraphAI API.
| 6 | logger = logging.getLogger(__name__) |
| 7 | |
| 8 | class ScrapingService: |
| 9 | """Service for interacting with ScrapeGraphAI API.""" |
| 10 | |
| 11 | def __init__(self, api_key: str): |
| 12 | self.api_key = api_key |
| 13 | |
| 14 | async def execute_pipeline( |
| 15 | self, |
| 16 | urls: List[str], |
| 17 | schema: Optional[Dict[str, Any]], |
| 18 | prompt: str |
| 19 | ) -> List[Dict]: |
| 20 | """Execute scraping for multiple URLs concurrently.""" |
| 21 | results = [] |
| 22 | |
| 23 | try: |
| 24 | async with AsyncClient(api_key=self.api_key) as client: |
| 25 | tasks = [] |
| 26 | |
| 27 | for url in urls: |
| 28 | # Create task for each URL using prompt-only extraction |
| 29 | task = client.smartscraper( |
| 30 | website_url=url, |
| 31 | user_prompt=prompt |
| 32 | ) |
| 33 | tasks.append(task) |
| 34 | |
| 35 | # Execute all tasks concurrently |
| 36 | raw_results = await asyncio.gather(*tasks, return_exceptions=True) |
| 37 | |
| 38 | # Process results |
| 39 | for i, result in enumerate(raw_results): |
| 40 | if isinstance(result, Exception): |
| 41 | results.append({ |
| 42 | "url": urls[i], |
| 43 | "success": False, |
| 44 | "data": None, |
| 45 | "error": str(result) |
| 46 | }) |
| 47 | logger.error(f"Scraping failed for {urls[i]}: {result}") |
| 48 | else: |
| 49 | # Handle the API response format |
| 50 | if isinstance(result, dict) and 'result' in result: |
| 51 | # Extract the actual data from the 'result' field |
| 52 | data = result.get('result', {}) |
| 53 | elif hasattr(result, 'model_dump'): |
| 54 | data = result.model_dump() |
| 55 | else: |
| 56 | data = result |
| 57 | |
| 58 | results.append({ |
| 59 | "url": urls[i], |
| 60 | "success": True, |
| 61 | "data": data, |
| 62 | "error": None |
| 63 | }) |
| 64 | logger.info(f"Successfully scraped {urls[i]}") |
| 65 |